Allow combination expression remaps

BemiHQ · Dec 20, 2024 · f1074a2 · f1074a2
1 parent 85474c2
commit f1074a2
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 11 deletions.
diff --git a/src/query_handler_test.go b/src/query_handler_test.go
@@ -499,6 +499,10 @@ func TestHandleQuery(t *testing.T) {
 			"description": {"type"},
 			"values":      {""},
 		},
+		"SELECT roles.oid AS id, roles.rolname AS name, roles.rolsuper AS is_superuser, CASE WHEN roles.rolsuper THEN true ELSE false END AS can_create_role FROM pg_catalog.pg_roles roles WHERE rolname = current_user": {
+			"description": {"id", "name", "is_superuser", "can_create_role"},
+			"values":      {},
+		},
 		// 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()": {
 			"description": {"gss_authenticated", "encrypted"},

diff --git a/src/select_remapper.go b/src/select_remapper.go
@@ -62,7 +62,7 @@ func (selectRemapper *SelectRemapper) remapSelectStatement(selectStatement *pgQu
 	// CASE
 	if hasCaseExpr := selectRemapper.hasCaseExpressions(selectStatement); hasCaseExpr {
 		selectRemapper.traceTreeTraversal("CASE expressions", indentLevel)
-		return selectRemapper.remapCaseExpressions(selectStatement, indentLevel)
+		selectRemapper.remapCaseExpressions(selectStatement, indentLevel)
 	}
 
 	// UNION
@@ -74,16 +74,11 @@ func (selectRemapper *SelectRemapper) remapSelectStatement(selectStatement *pgQu
 		selectRemapper.traceTreeTraversal("UNION right", indentLevel)
 		rightSelectStatement := selectStatement.Rarg
 		rightSelectStatement = selectRemapper.remapSelectStatement(rightSelectStatement, indentLevel+1) // self-recursion
-
-		return selectStatement
 	}
 
 	// JOIN
 	if len(selectStatement.FromClause) > 0 && selectStatement.FromClause[0].GetJoinExpr() != nil {
-		// SELECT
-		selectStatement = selectRemapper.remapSelect(selectStatement, indentLevel)                                                         // recursive
 		selectStatement.FromClause[0] = selectRemapper.remapJoinExpressions(selectStatement, selectStatement.FromClause[0], indentLevel+1) // recursive with self-recursion
-		return selectStatement
 	}
 
 	// WHERE
@@ -93,10 +88,6 @@ func (selectRemapper *SelectRemapper) remapSelectStatement(selectStatement *pgQu
 
 	// FROM
 	if len(selectStatement.FromClause) > 0 {
-		// SELECT
-		selectStatement = selectRemapper.remapSelect(selectStatement, indentLevel) // recursive
-
-		// FROM
 		for i, fromNode := range selectStatement.FromClause {
 			if fromNode.GetRangeVar() != nil {
 				// WHERE
@@ -118,7 +109,6 @@ func (selectRemapper *SelectRemapper) remapSelectStatement(selectStatement *pgQu
 				selectStatement.FromClause[i] = selectRemapper.remapTableFunction(fromNode, indentLevel+1) // recursive
 			}
 		}
-		return selectStatement
 	}
 
 	selectStatement = selectRemapper.remapSelect(selectStatement, indentLevel) // recursive