Add support for pg_index, fix type casting to array types

BemiHQ · Feb 12, 2025 · b0c9e83 · b0c9e83
1 parent 8e6bc71
commit b0c9e83
Show file tree

Hide file tree

Showing 6 changed files with 57 additions and 8 deletions.
diff --git a/src/parser_table.go b/src/parser_table.go
@@ -128,6 +128,28 @@ func (parser *ParserTable) MakePgStatUserTablesNode(schemaTables []IcebergSchema
 	return parser.utils.MakeSubselectWithRowsNode(PG_TABLE_PG_STAT_USER_TABLES, tableDef, rowsValues, alias)
 }
 
+// pg_index -> returns (SELECT *, FALSE AS indnullsnotdistinct FROM pg_index)
+func (parser *ParserTable) MakePgIndexNode(qSchemaTable QuerySchemaTable) *pgQuery.Node {
+	targetList := []*pgQuery.Node{
+		pgQuery.MakeResTargetNodeWithVal(
+			pgQuery.MakeColumnRefNode(
+				[]*pgQuery.Node{pgQuery.MakeAStarNode()},
+				0,
+			),
+			0,
+		),
+		pgQuery.MakeResTargetNodeWithNameAndVal(
+			"indnullsnotdistinct",
+			parser.utils.MakeAConstBoolNode(false),
+			0,
+		),
+	}
+
+	fromNode := pgQuery.MakeSimpleRangeVarNode(PG_TABLE_PG_INDEX, 0)
+
+	return parser.utils.MakeSubselectFromNode(qSchemaTable.Table, targetList, fromNode, qSchemaTable.Alias)
+}
+
 // Other information_schema.* tables
 func (parser *ParserTable) IsTableFromInformationSchema(qSchemaTable QuerySchemaTable) bool {
 	return qSchemaTable.Schema == PG_SCHEMA_INFORMATION_SCHEMA

diff --git a/src/parser_utils.go b/src/parser_utils.go
@@ -1,6 +1,8 @@
 package main
 
 import (
+	"strings"
+
 	pgQuery "github.com/pganalyze/pg_query_go/v5"
 )
 
@@ -169,16 +171,30 @@ func (utils *ParserUtils) makeTypedConstNode(val string, pgType string) *pgQuery
 }
 
 func (utils *ParserUtils) MakeTypeCastNode(arg *pgQuery.Node, typeName string) *pgQuery.Node {
+	var typeNameNode *pgQuery.TypeName
+
+	if strings.HasSuffix(typeName, "[]") {
+		typeNameNode = &pgQuery.TypeName{
+			Names: []*pgQuery.Node{
+				pgQuery.MakeStrNode(strings.TrimSuffix(typeName, "[]")),
+			},
+			ArrayBounds: []*pgQuery.Node{
+				pgQuery.MakeIntNode(-1),
+			},
+		}
+	} else {
+		typeNameNode = &pgQuery.TypeName{
+			Names: []*pgQuery.Node{
+				pgQuery.MakeStrNode(typeName),
+			},
+		}
+	}
+
 	return &pgQuery.Node{
 		Node: &pgQuery.Node_TypeCast{
 			TypeCast: &pgQuery.TypeCast{
-				Arg: arg,
-				TypeName: &pgQuery.TypeName{
-					Names: []*pgQuery.Node{
-						pgQuery.MakeStrNode(typeName),
-					},
-					Location: 0,
-				},
+				Arg:      arg,
+				TypeName: typeNameNode,
 			},
 		},
 	}

diff --git a/src/pg_constants.go b/src/pg_constants.go
@@ -24,6 +24,7 @@ const (
 	PG_TABLE_PG_COLLATION          = "pg_collation"
 	PG_TABLE_PG_DATABASE           = "pg_database"
 	PG_TABLE_PG_EXTENSION          = "pg_extension"
+	PG_TABLE_PG_INDEX              = "pg_index"
 	PG_TABLE_PG_INHERITS           = "pg_inherits"
 	PG_TABLE_PG_MATVIEWS           = "pg_matviews"
 	PG_TABLE_PG_NAMESPACE          = "pg_namespace"

diff --git a/src/query_handler_test.go b/src/query_handler_test.go
@@ -249,6 +249,10 @@ func TestHandleQuery(t *testing.T) {
 			"types":       {Uint32ToString(pgtype.TextArrayOID)},
 			"values":      {},
 		},
+		"SELECT indnullsnotdistinct FROM pg_index": {
+			"description": {"indnullsnotdistinct"},
+			"types":       {Uint32ToString(pgtype.BoolOID)},
+		},
 
 		// Information schema
 		"SELECT * FROM information_schema.tables": {

diff --git a/src/query_remapper_select.go b/src/query_remapper_select.go
@@ -22,13 +22,15 @@ func (remapper *QueryRemapperSelect) RemapFunctionToConstant(functionCall *pgQue
 	return remapper.parserFunction.RemapToConstant(functionCall)
 }
 
-// SELECT [PG_FUNCTION()]
+// SELECT ...
 func (remapper *QueryRemapperSelect) RemapSelect(targetNode *pgQuery.Node) *pgQuery.Node {
+	// PG_FUNCTION().value
 	newTargetNode := remapper.remappedInderectionFunctionCall(targetNode)
 	if newTargetNode != nil {
 		return newTargetNode
 	}
 
+	// PG_FUNCTION()
 	functionCall := remapper.parserFunction.FunctionCall(targetNode)
 	if functionCall == nil {
 		return targetNode

diff --git a/src/query_remapper_table.go b/src/query_remapper_table.go
@@ -113,6 +113,10 @@ func (remapper *QueryRemapperTable) RemapTable(node *pgQuery.Node) *pgQuery.Node
 		case PG_TABLE_PG_OPCLASS:
 			return parser.MakeEmptyTableNode(PG_TABLE_PG_OPCLASS, PG_OPCLASS_DEFINITION, qSchemaTable.Alias)
 
+		// pg_index -> returns (SELECT *, FALSE AS indnullsnotdistinct FROM pg_index)
+		case PG_TABLE_PG_INDEX:
+			return parser.MakePgIndexNode(qSchemaTable)
+
 		// pg_catalog.pg_* other system tables -> return as is
 		default:
 			return node