Support WHEN ANY in CASE expressions

BemiHQ · Jan 11, 2025 · 07c8c02 · 07c8c02
1 parent a642cc0
commit 07c8c02
Show file tree

Hide file tree

Showing 3 changed files with 52 additions and 0 deletions.
diff --git a/src/parser_select.go b/src/parser_select.go
@@ -1,6 +1,8 @@
 package main
 
 import (
+	"strings"
+
 	pgQuery "github.com/pganalyze/pg_query_go/v5"
 )
 
@@ -23,3 +25,45 @@ func (parser *ParserSelect) SetDefaultTargetName(targetNode *pgQuery.Node, name
 		target.Name = name
 	}
 }
+
+// = ANY({schema_information}) -> IN (schema_information)
+func (parser *ParserSelect) ConvertAnyToIn(aExpr *pgQuery.A_Expr) *pgQuery.Node {
+	arrayStr := aExpr.Rexpr.GetAConst().GetSval().Sval
+	arrayStr = strings.Trim(arrayStr, "{}")
+	values := strings.Split(arrayStr, ",")
+
+	items := make([]*pgQuery.Node, len(values))
+	for i, value := range values {
+		value = strings.Trim(value, " ")
+		items[i] = &pgQuery.Node{
+			Node: &pgQuery.Node_AConst{
+				AConst: &pgQuery.A_Const{
+					Val: &pgQuery.A_Const_Sval{
+						Sval: &pgQuery.String{
+							Sval: value,
+						},
+					},
+					Location: 0,
+				},
+			},
+		}
+	}
+
+	return &pgQuery.Node{
+		Node: &pgQuery.Node_AExpr{
+			AExpr: &pgQuery.A_Expr{
+				Kind:  pgQuery.A_Expr_Kind_AEXPR_IN,
+				Name:  []*pgQuery.Node{{Node: &pgQuery.Node_String_{String_: &pgQuery.String{Sval: "="}}}},
+				Lexpr: aExpr.Lexpr,
+				Rexpr: &pgQuery.Node{
+					Node: &pgQuery.Node_List{
+						List: &pgQuery.List{
+							Items: items,
+						},
+					},
+				},
+				Location: aExpr.Location,
+			},
+		},
+	}
+}
diff --git a/src/query_handler_test.go b/src/query_handler_test.go
@@ -744,6 +744,11 @@ func TestHandleQuery(t *testing.T) {
 			"types":       {Uint32ToString(pgtype.TextOID)},
 			"values":      {"f"},
 		},
+		"SELECT CASE WHEN nsp.nspname = ANY('{information_schema}') THEN false ELSE true END AS db_support FROM pg_catalog.pg_namespace nsp WHERE nsp.oid = 1268::OID;": {
+			"description": {"db_support"},
+			"types":       {Uint32ToString(pgtype.BoolOID)},
+			"values":      {"true"},
+		},
 
 		// WHERE pg functions
 		"SELECT gss_authenticated, encrypted FROM (SELECT false, false, false, false, false WHERE false) t(pid, gss_authenticated, principal, encrypted, credentials_delegated) WHERE pid = pg_backend_pid()": {

diff --git a/src/query_remapper.go b/src/query_remapper.go
@@ -203,6 +203,9 @@ func (remapper *QueryRemapper) remapCaseExpressions(selectStatement *pgQuery.Sel
 				if whenClause := when.GetCaseWhen(); whenClause != nil {
 					if whenClause.Expr != nil {
 						if aExpr := whenClause.Expr.GetAExpr(); aExpr != nil {
+							if aExpr.Kind == pgQuery.A_Expr_Kind_AEXPR_OP_ANY {
+								whenClause.Expr = remapper.remapperSelect.parserSelect.ConvertAnyToIn(aExpr)
+							}
 							if subLink := aExpr.Lexpr.GetSubLink(); subLink != nil {
 								remapper.traceTreeTraversal("CASE WHEN left", indentLevel+1)
 								subSelect := subLink.Subselect.GetSelectStmt()