Refactor identifying function calls from pg_catalog

BemiHQ · Jan 6, 2025 · 27dcf0f · 27dcf0f
1 parent ef143d9
commit 27dcf0f
Show file tree

Hide file tree

Showing 4 changed files with 42 additions and 99 deletions.
diff --git a/src/query_parser_select.go b/src/query_parser_select.go
@@ -37,21 +37,7 @@ func (parser *QueryParserSelect) NestedFunctionCalls(functionCall *pgQuery.FuncC
 }
 
 func (parser *QueryParserSelect) SchemaFunction(functionCall *pgQuery.FuncCall) PgSchemaFunction {
-	if len(functionCall.Funcname) == 1 {
-		return PgSchemaFunction{
-			Schema:   "",
-			Function: functionCall.Funcname[0].GetString_().Sval,
-		}
-	}
-
-	if len(functionCall.Funcname) == 2 {
-		return PgSchemaFunction{
-			Schema:   functionCall.Funcname[0].GetString_().Sval,
-			Function: functionCall.Funcname[1].GetString_().Sval,
-		}
-	}
-
-	return PgSchemaFunction{}
+	return parser.utils.SchemaFunction(functionCall)
 }
 
 // quote_ident(str) -> concat("\""+str+"\"")

diff --git a/src/query_parser_table.go b/src/query_parser_table.go
@@ -175,27 +175,19 @@ func (parser *QueryParserTable) MakeIcebergTableNode(tablePath string, qSchemaTa
 	return parser.utils.MakeSubselectFromNode(qSchemaTable.Table, []*pgQuery.Node{selectStarNode}, node, qSchemaTable.Alias)
 }
 
-// pg_catalog.pg_get_keywords()
-func (parser *QueryParserTable) IsPgGetKeywordsFunction(node *pgQuery.Node) bool {
+func (parser *QueryParserTable) SchemaFunction(node *pgQuery.Node) PgSchemaFunction {
 	for _, funcNode := range node.GetRangeFunction().Functions {
 		for _, funcItemNode := range funcNode.GetList().Items {
 			funcCallNode := funcItemNode.GetFuncCall()
 			if funcCallNode == nil {
 				continue
 			}
-			if len(funcCallNode.Funcname) != 2 {
-				continue
-			}
 
-			schema := funcCallNode.Funcname[0].GetString_().Sval
-			function := funcCallNode.Funcname[1].GetString_().Sval
-			if schema == PG_SCHEMA_PG_CATALOG && function == PG_FUNCTION_PG_GET_KEYWORDS {
-				return true
-			}
+			return parser.utils.SchemaFunction(funcCallNode)
 		}
 	}
 
-	return false
+	return PgSchemaFunction{}
 }
 
 // pg_catalog.pg_get_keywords() -> VALUES(values...) t(columns...)
@@ -255,35 +247,7 @@ func (parser *QueryParserTable) MakeArrayUpperNode(funcCallNode *pgQuery.FuncCal
 	).GetFuncCall()
 }
 
-// pg_show_all_settings()
-func (parser *QueryParserTable) IsPgShowAllSettingsFunction(node *pgQuery.Node) bool {
-	for _, funcNode := range node.GetRangeFunction().Functions {
-		for _, funcItemNode := range funcNode.GetList().Items {
-			funcCallNode := funcItemNode.GetFuncCall()
-			if funcCallNode == nil {
-				continue
-			}
-
-			if len(funcCallNode.Funcname) == 1 {
-				function := funcCallNode.Funcname[0].GetString_().Sval
-				if function == PG_FUNCTION_PG_SHOW_ALL_SETTINGS {
-					return true
-				}
-			}
-
-			if len(funcCallNode.Funcname) == 2 {
-				schema := funcCallNode.Funcname[0].GetString_().Sval
-				function := funcCallNode.Funcname[1].GetString_().Sval
-				if schema == PG_SCHEMA_PG_CATALOG && function == PG_FUNCTION_PG_SHOW_ALL_SETTINGS {
-					return true
-				}
-			}
-		}
-	}
-	return false
-}
-
-// pg_show_all_settings() -> duckdb_settings() mapped to pg format
+// pg_catalog.pg_show_all_settings() -> duckdb_settings() mapped to pg format
 func (parser *QueryParserTable) MakePgShowAllSettingsNode(node *pgQuery.Node) *pgQuery.Node {
 	targetList := []*pgQuery.Node{
 		pgQuery.MakeResTargetNodeWithNameAndVal(
@@ -392,35 +356,7 @@ func (parser *QueryParserTable) MakePgShowAllSettingsNode(node *pgQuery.Node) *p
 	return parser.utils.MakeSubselectFromNode(PG_FUNCTION_PG_SHOW_ALL_SETTINGS, targetList, fromNode, alias)
 }
 
-// pg_is_in_recovery()
-func (parser *QueryParserTable) IsPgIsInRecoveryFunction(node *pgQuery.Node) bool {
-	for _, funcNode := range node.GetRangeFunction().Functions {
-		for _, funcItemNode := range funcNode.GetList().Items {
-			funcCallNode := funcItemNode.GetFuncCall()
-			if funcCallNode == nil {
-				continue
-			}
-
-			if len(funcCallNode.Funcname) == 1 {
-				function := funcCallNode.Funcname[0].GetString_().Sval
-				if function == PG_FUNCTION_PG_IS_IN_RECOVERY {
-					return true
-				}
-			}
-
-			if len(funcCallNode.Funcname) == 2 {
-				schema := funcCallNode.Funcname[0].GetString_().Sval
-				function := funcCallNode.Funcname[1].GetString_().Sval
-				if schema == PG_SCHEMA_PG_CATALOG && function == PG_FUNCTION_PG_IS_IN_RECOVERY {
-					return true
-				}
-			}
-		}
-	}
-	return false
-}
-
-// pg_is_in_recovery() -> 'f'::bool
+// pg_catalog.pg_is_in_recovery() -> 'f'::bool
 func (parser *QueryParserTable) MakePgIsInRecoveryNode(node *pgQuery.Node) *pgQuery.Node {
 	var alias string
 	if node.GetAlias() != nil {
@@ -435,10 +371,6 @@ func (parser *QueryParserTable) MakePgIsInRecoveryNode(node *pgQuery.Node) *pgQu
 	)
 }
 
-func (parser *QueryParserTable) isPgCatalogSchema(qSchemaTable QuerySchemaTable) bool {
-	return qSchemaTable.Schema == PG_SCHEMA_PG_CATALOG || qSchemaTable.Schema == ""
-}
-
 var PG_SHADOW_VALUE_BY_COLUMN = NewOrderedMap([][]string{
 	{"usename", "bemidb"},
 	{"usesysid", "10"},

diff --git a/src/query_parser_utils.go b/src/query_parser_utils.go
@@ -15,6 +15,23 @@ func NewQueryParserUtils(config *Config) *QueryParserUtils {
 	return &QueryParserUtils{config: config}
 }
 
+func (utils *QueryParserUtils) SchemaFunction(functionCall *pgQuery.FuncCall) PgSchemaFunction {
+	switch len(functionCall.Funcname) {
+	case 1:
+		return PgSchemaFunction{
+			Schema:   "",
+			Function: functionCall.Funcname[0].GetString_().Sval,
+		}
+	case 2:
+		return PgSchemaFunction{
+			Schema:   functionCall.Funcname[0].GetString_().Sval,
+			Function: functionCall.Funcname[1].GetString_().Sval,
+		}
+	default:
+		panic("Invalid function call")
+	}
+}
+
 func (utils *QueryParserUtils) MakeSubselectWithRowsNode(tableName string, columns []string, rowsValues [][]string, alias string) *pgQuery.Node {
 	parserType := NewQueryParserType(utils.config)
 

diff --git a/src/select_remapper_table.go b/src/select_remapper_table.go
@@ -154,19 +154,23 @@ func (remapper *SelectRemapperTable) RemapTable(node *pgQuery.Node) *pgQuery.Nod
 func (remapper *SelectRemapperTable) RemapTableFunction(node *pgQuery.Node) *pgQuery.Node {
 	parser := remapper.parserTable
 
-	// pg_catalog.pg_get_keywords() -> hard-coded keywords
-	if parser.IsPgGetKeywordsFunction(node) {
-		return parser.MakePgGetKeywordsNode(node)
-	}
+	schemaFunction := parser.SchemaFunction(node)
 
-	// pg_show_all_settings() -> duckdb_settings()
-	if parser.IsPgShowAllSettingsFunction(node) {
-		return parser.MakePgShowAllSettingsNode(node)
-	}
+	if remapper.isFunctionFromPgCatalog(schemaFunction) {
+		switch {
+
+		// pg_catalog.pg_get_keywords() -> hard-coded keywords
+		case schemaFunction.Function == PG_FUNCTION_PG_GET_KEYWORDS:
+			return parser.MakePgGetKeywordsNode(node)
+
+		// pg_catalog.pg_show_all_settings() -> duckdb_settings()
+		case schemaFunction.Function == PG_FUNCTION_PG_SHOW_ALL_SETTINGS:
+			return parser.MakePgShowAllSettingsNode(node)
 
-	// pg_is_in_recovery() -> 'f'::bool
-	if parser.IsPgIsInRecoveryFunction(node) {
-		return parser.MakePgIsInRecoveryNode(node)
+		// pg_catalog.pg_is_in_recovery() -> 'f'::bool
+		case schemaFunction.Function == PG_FUNCTION_PG_IS_IN_RECOVERY:
+			return parser.MakePgIsInRecoveryNode(node)
+		}
 	}
 
 	return node
@@ -239,6 +243,10 @@ func (remapper *SelectRemapperTable) isTableFromPgCatalog(qSchemaTable QuerySche
 			!remapper.icebergSchemaTableExists(qSchemaTable.ToIcebergSchemaTable()))
 }
 
+func (remapper *SelectRemapperTable) isFunctionFromPgCatalog(schemaFunction PgSchemaFunction) bool {
+	return schemaFunction.Schema == PG_SCHEMA_PG_CATALOG || schemaFunction.Schema == ""
+}
+
 var PG_INHERITS_COLUMNS = []string{
 	"inhrelid",
 	"inhparent",