From f1074a23424515be923bb4dc7940939ff8bd9167 Mon Sep 17 00:00:00 2001 From: Arjun Lall Date: Thu, 19 Dec 2024 21:55:13 -0500 Subject: [PATCH] Allow combination expression remaps --- src/query_handler_test.go | 4 ++++ src/select_remapper.go | 12 +----------- 2 files changed, 5 insertions(+), 11 deletions(-) diff --git a/src/query_handler_test.go b/src/query_handler_test.go index 8efb898..4275c77 100644 --- 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 index bf1e103..21a4a6a 100644 --- 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