Add support for type casting 'value'::regclass::oid

BemiHQ · Feb 7, 2025 · 01ad4b5 · 01ad4b5
1 parent 112dd74
commit 01ad4b5
Show file tree

Hide file tree

Showing 4 changed files with 118 additions and 5 deletions.
diff --git a/src/parser_type_cast.go b/src/parser_type_cast.go
@@ -66,6 +66,98 @@ func (parser *ParserTypeCast) MakeListValueFromArray(node *pgQuery.Node) *pgQuer
 	}
 }
 
+// SELECT c.oid
+// FROM pg_class c
+// JOIN pg_namespace n ON n.oid = c.relnamespace
+// WHERE n.nspname = 'schema' AND c.relname = 'table'
+func (parser *ParserTypeCast) MakeSubselectOidBySchemaTable(argumentNode *pgQuery.Node) *pgQuery.Node {
+	targetNode := pgQuery.MakeResTargetNodeWithVal(
+		pgQuery.MakeColumnRefNode([]*pgQuery.Node{
+			pgQuery.MakeStrNode("c"),
+			pgQuery.MakeStrNode("oid"),
+		}, 0),
+		0,
+	)
+
+	joinNode := pgQuery.MakeJoinExprNode(
+		pgQuery.JoinType_JOIN_INNER,
+		pgQuery.MakeFullRangeVarNode("", "pg_class", "c", 0),
+		pgQuery.MakeFullRangeVarNode("", "pg_namespace", "n", 0),
+		pgQuery.MakeAExprNode(
+			pgQuery.A_Expr_Kind_AEXPR_OP,
+			[]*pgQuery.Node{
+				pgQuery.MakeStrNode("="),
+			},
+			pgQuery.MakeColumnRefNode([]*pgQuery.Node{
+				pgQuery.MakeStrNode("n"),
+				pgQuery.MakeStrNode("oid"),
+			}, 0),
+			pgQuery.MakeColumnRefNode([]*pgQuery.Node{
+				pgQuery.MakeStrNode("c"),
+				pgQuery.MakeStrNode("relnamespace"),
+			}, 0),
+			0,
+		),
+	)
+
+	value := argumentNode.GetAConst().GetSval().Sval
+	parts := strings.Split(value, ".")
+	schema := PG_SCHEMA_PUBLIC
+	if len(parts) > 1 {
+		schema = parts[0]
+	}
+	table := parts[len(parts)-1]
+
+	whereNode := pgQuery.MakeBoolExprNode(
+		pgQuery.BoolExprType_AND_EXPR,
+		[]*pgQuery.Node{
+			pgQuery.MakeAExprNode(
+				pgQuery.A_Expr_Kind_AEXPR_OP,
+				[]*pgQuery.Node{
+					pgQuery.MakeStrNode("="),
+				},
+				pgQuery.MakeColumnRefNode([]*pgQuery.Node{
+					pgQuery.MakeStrNode("n"),
+					pgQuery.MakeStrNode("nspname"),
+				}, 0),
+				pgQuery.MakeAConstStrNode(schema, 0),
+				0,
+			),
+			pgQuery.MakeAExprNode(
+				pgQuery.A_Expr_Kind_AEXPR_OP,
+				[]*pgQuery.Node{
+					pgQuery.MakeStrNode("="),
+				},
+				pgQuery.MakeColumnRefNode([]*pgQuery.Node{
+					pgQuery.MakeStrNode("c"),
+					pgQuery.MakeStrNode("relname"),
+				}, 0),
+				pgQuery.MakeAConstStrNode(table, 0),
+				0,
+			),
+		},
+		0,
+	)
+
+	return &pgQuery.Node{
+		Node: &pgQuery.Node_SubLink{
+			SubLink: &pgQuery.SubLink{
+				SubLinkType: pgQuery.SubLinkType_EXPR_SUBLINK,
+				Subselect: &pgQuery.Node{
+					Node: &pgQuery.Node_SelectStmt{
+						SelectStmt: &pgQuery.SelectStmt{
+							TargetList:  []*pgQuery.Node{targetNode},
+							FromClause:  []*pgQuery.Node{joinNode},
+							WhereClause: whereNode,
+						},
+					},
+				},
+			},
+		},
+	}
+
+}
+
 func (parser *ParserTypeCast) inferNodeType(node *pgQuery.Node) string {
 	if typeCast := node.GetTypeCast(); typeCast != nil {
 		return typeCast.TypeName.Names[0].GetString_().Sval

diff --git a/src/query_handler_test.go b/src/query_handler_test.go
@@ -643,7 +643,12 @@ func TestHandleQuery(t *testing.T) {
 			"values":      {""},
 		},
 
-		// Typecasts
+		// Type casts
+		"SELECT 'public.test_table'::regclass::oid AS oid": {
+			"description": {"oid"},
+			"types":       {Uint32ToString(pgtype.OIDOID)},
+			"values":      {"1270"},
+		},
 		"SELECT objoid, classoid, objsubid, description FROM pg_description WHERE classoid = 'pg_class'::regclass": {
 			"description": {"objoid", "classoid", "objsubid", "description"},
 			"types":       {Uint32ToString(pgtype.OIDOID), Uint32ToString(pgtype.TextOID), Uint32ToString(pgtype.Int4OID), Uint32ToString(pgtype.TextOID)},

diff --git a/src/query_remapper.go b/src/query_remapper.go
@@ -60,7 +60,7 @@ func (remapper *QueryRemapper) RemapStatements(statements []*pgQuery.RawStmt) ([
 	}
 
 	for i, stmt := range statements {
-		LogTrace(remapper.config, "Remapping statement", i+1)
+		LogTrace(remapper.config, "Remapping statement #"+IntToString(i+1))
 
 		node := stmt.Stmt
 
@@ -115,7 +115,7 @@ func (remapper *QueryRemapper) remapSetStatement(stmt *pgQuery.RawStmt) *pgQuery
 }
 
 func (remapper *QueryRemapper) remapSelectStatement(selectStatement *pgQuery.SelectStmt, indentLevel int) *pgQuery.SelectStmt {
-	// Target Sublinks's
+	// Nested SELECT (SELECT ...) ...
 	for _, target := range selectStatement.TargetList {
 		if subLink := target.GetResTarget().Val.GetSubLink(); subLink != nil {
 			remapper.traceTreeTraversal("Target SubLink", indentLevel)

diff --git a/src/query_remapper_type_cast.go b/src/query_remapper_type_cast.go
@@ -19,7 +19,7 @@ func NewQueryRemapperTypeCast(config *Config) *QueryRemapperTypeCast {
 	return remapper
 }
 
-// value::type -> value
+// value::type or CAST(value AS type)
 func (remapper *QueryRemapperTypeCast) RemapTypeCast(node *pgQuery.Node) *pgQuery.Node {
 	typeCast := remapper.parserTypeCast.TypeCast(node)
 	if typeCast == nil {
@@ -29,12 +29,28 @@ func (remapper *QueryRemapperTypeCast) RemapTypeCast(node *pgQuery.Node) *pgQuer
 	typeName := remapper.parserTypeCast.TypeName(typeCast)
 	switch typeName {
 	case "regclass":
+		// 'schema.table'::regclass -> 'schema.table'
 		return typeCast.Arg
 	case "text":
+		// '{a,b,c}'::text[] -> ARRAY['a', 'b', 'c']
 		return remapper.parserTypeCast.MakeListValueFromArray(typeCast.Arg)
 	case "regproc":
-		functionNameParts := strings.Split(remapper.parserTypeCast.ArgStringValue(typeCast), ".") // pg_catalog.func_name
+		// 'schema.function_name'::regproc -> 'function_name'
+		functionNameParts := strings.Split(remapper.parserTypeCast.ArgStringValue(typeCast), ".")
 		return pgQuery.MakeAConstStrNode(functionNameParts[len(functionNameParts)-1], 0)
+	case "oid":
+		// 'schema.table'::regclass::oid -> SELECT c.oid FROM pg_class c JOIN pg_namespace n ON n.oid = c.relnamespace WHERE n.nspname = 'schema' AND c.relname = 'table'
+		nestedNode := typeCast.Arg
+		nestedTypeCast := remapper.parserTypeCast.TypeCast(nestedNode)
+		if nestedTypeCast == nil {
+			return node
+		}
+		nestedTypeName := remapper.parserTypeCast.TypeName(nestedTypeCast)
+		if nestedTypeName != "regclass" {
+			return node
+		}
+
+		return remapper.parserTypeCast.MakeSubselectOidBySchemaTable(nestedTypeCast.Arg)
 	}
 
 	return node