Remap functions in CASE -> THEN statements

BemiHQ · Feb 20, 2025 · 42a4902 · 42a4902
1 parent e453ccb
commit 42a4902
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 10 deletions.
diff --git a/src/query_handler_test.go b/src/query_handler_test.go
@@ -832,11 +832,11 @@ func TestHandleQuery(t *testing.T) {
 			"types":       {Uint32ToString(pgtype.Int8OID), Uint32ToString(pgtype.TextOID), Uint32ToString(pgtype.BoolOID), Uint32ToString(pgtype.BoolOID)},
 			"values":      {},
 		},
-		// "SELECT CASE WHEN TRUE THEN pg_catalog.pg_is_in_recovery() END AS CASE": {
-		// 	"description": {"case"},
-		// 	"types":       {Uint32ToString(pgtype.BoolOID)},
-		// 	"values":      {"f"},
-		// },
+		"SELECT CASE WHEN TRUE THEN pg_catalog.pg_is_in_recovery() END AS CASE": {
+			"description": {"case"},
+			"types":       {Uint32ToString(pgtype.BoolOID)},
+			"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)},

diff --git a/src/query_remapper.go b/src/query_remapper.go
@@ -184,12 +184,12 @@ func (remapper *QueryRemapper) remapCaseExpression(caseExpr *pgQuery.CaseExpr, i
 						whenClause.Expr = remapper.remapperSelect.parserSelect.ConvertAnyToIn(aExpr)
 					}
 					if subLink := aExpr.Lexpr.GetSubLink(); subLink != nil {
-						remapper.traceTreeTraversal("CASE WHEN left", indentLevel+1)
+						remapper.traceTreeTraversal("CASE->WHEN left", indentLevel+1)
 						subSelect := subLink.Subselect.GetSelectStmt()
 						remapper.remapSelectStatement(subSelect, indentLevel+1)
 					}
 					if subLink := aExpr.Rexpr.GetSubLink(); subLink != nil {
-						remapper.traceTreeTraversal("CASE WHEN right", indentLevel+1)
+						remapper.traceTreeTraversal("CASE->WHEN right", indentLevel+1)
 						subSelect := subLink.Subselect.GetSelectStmt()
 						remapper.remapSelectStatement(subSelect, indentLevel+1)
 					}
@@ -202,12 +202,16 @@ func (remapper *QueryRemapper) remapCaseExpression(caseExpr *pgQuery.CaseExpr, i
 				if typeName != "" {
 					whenClause.Result = remapper.parserTypeCast.MakeCaseTypeCastNode(whenClause.Result, typeName)
 				}
-
 				if subLink := whenClause.Result.GetSubLink(); subLink != nil {
-					remapper.traceTreeTraversal("CASE THEN", indentLevel+1)
+					remapper.traceTreeTraversal("CASE->THEN", indentLevel+1)
 					subSelect := subLink.Subselect.GetSelectStmt()
 					remapper.remapSelectStatement(subSelect, indentLevel+1)
 				}
+				if functionCall := whenClause.Result.GetFuncCall(); functionCall != nil {
+					remapper.traceTreeTraversal("CASE->THEN function", indentLevel+1)
+					remapper.remapperFunction.RemapFunctionCall(functionCall)
+					remapper.remapperFunction.RemapNestedFunctionCalls(functionCall) // recursion
+				}
 			}
 		}
 	}
@@ -220,7 +224,7 @@ func (remapper *QueryRemapper) remapCaseExpression(caseExpr *pgQuery.CaseExpr, i
 		}
 
 		if subLink := caseExpr.Defresult.GetSubLink(); subLink != nil {
-			remapper.traceTreeTraversal("CASE ELSE", indentLevel+1)
+			remapper.traceTreeTraversal("CASE->ELSE", indentLevel+1)
 			subSelect := subLink.Subselect.GetSelectStmt()
 			remapper.remapSelectStatement(subSelect, indentLevel+1)
 		}