Identify system PG tables without pg_catalog schema

BemiHQ · Jan 6, 2025 · fdf83e5 · fdf83e5
1 parent b2df4ff
commit fdf83e5
Show file tree

Hide file tree

Showing 4 changed files with 148 additions and 144 deletions.
diff --git a/src/query_parser_table.go b/src/query_parser_table.go
@@ -119,12 +119,6 @@ func (parser *QueryParserTable) MakePgUserNode(user string, alias string) *pgQue
 	return parser.utils.MakeSubselectWithRowsNode(PG_TABLE_PG_USER, columns, [][]string{rowValues}, alias)
 }
 
-// System pg_* tables
-func (parser *QueryParserTable) IsTableFromPgCatalog(qSchemaTable QuerySchemaTable) bool {
-	return parser.isPgCatalogSchema(qSchemaTable) &&
-		(PG_SYSTEM_TABLES.Contains(qSchemaTable.Table) || PG_SYSTEM_VIEWS.Contains(qSchemaTable.Table))
-}
-
 // Other information_schema.* tables
 func (parser *QueryParserTable) IsTableFromInformationSchema(qSchemaTable QuerySchemaTable) bool {
 	return qSchemaTable.Schema == PG_SCHEMA_INFORMATION_SCHEMA
@@ -444,124 +438,6 @@ func (parser *QueryParserTable) isPgCatalogSchema(qSchemaTable QuerySchemaTable)
 	return qSchemaTable.Schema == PG_SCHEMA_PG_CATALOG || qSchemaTable.Schema == ""
 }
 
-var PG_SYSTEM_TABLES = NewSet([]string{
-	"pg_aggregate",
-	"pg_am",
-	"pg_amop",
-	"pg_amproc",
-	"pg_attrdef",
-	"pg_attribute",
-	"pg_auth_members",
-	"pg_authid",
-	"pg_cast",
-	"pg_class",
-	"pg_collation",
-	"pg_constraint",
-	"pg_conversion",
-	"pg_database",
-	"pg_db_role_setting",
-	"pg_default_acl",
-	"pg_depend",
-	"pg_description",
-	"pg_enum",
-	"pg_event_trigger",
-	"pg_extension",
-	"pg_foreign_data_wrapper",
-	"pg_foreign_server",
-	"pg_foreign_table",
-	"pg_index",
-	"pg_inherits",
-	"pg_init_privs",
-	"pg_language",
-	"pg_largeobject",
-	"pg_largeobject_metadata",
-	"pg_matviews",
-	"pg_namespace",
-	"pg_opclass",
-	"pg_operator",
-	"pg_opfamily",
-	"pg_parameter_acl",
-	"pg_partitioned_table",
-	"pg_policy",
-	"pg_proc",
-	"pg_publication",
-	"pg_publication_namespace",
-	"pg_publication_rel",
-	"pg_user",
-	"pg_range",
-	"pg_replication_origin",
-	"pg_replication_slots",
-	"pg_rewrite",
-	"pg_roles",
-	"pg_seclabel",
-	"pg_sequence",
-	"pg_shadow",
-	"pg_shdepend",
-	"pg_shdescription",
-	"pg_shseclabel",
-	"pg_statistic",
-	"pg_statistic_ext",
-	"pg_statistic_ext_data",
-	"pg_subscription",
-	"pg_subscription_rel",
-	"pg_tablespace",
-	"pg_transform",
-	"pg_trigger",
-	"pg_ts_config",
-	"pg_ts_config_map",
-	"pg_ts_dict",
-	"pg_ts_parser",
-	"pg_ts_template",
-	"pg_type",
-	"pg_user_mapping",
-})
-
-var PG_SYSTEM_VIEWS = NewSet([]string{
-	"pg_stat_activity",
-	"pg_stat_replication",
-	"pg_stat_wal_receiver",
-	"pg_stat_recovery_prefetch",
-	"pg_stat_subscription",
-	"pg_stat_ssl",
-	"pg_stat_gssapi",
-	"pg_stat_progress_analyze",
-	"pg_stat_progress_create_index",
-	"pg_stat_progress_vacuum",
-	"pg_stat_progress_cluster",
-	"pg_stat_progress_basebackup",
-	"pg_stat_progress_copy",
-	"pg_stat_archiver",
-	"pg_stat_bgwriter",
-	"pg_stat_checkpointer",
-	"pg_stat_database",
-	"pg_stat_database_conflicts",
-	"pg_stat_io",
-	"pg_stat_replication_slots",
-	"pg_stat_slru",
-	"pg_stat_subscription_stats",
-	"pg_stat_wal",
-	"pg_stat_all_tables",
-	"pg_stat_sys_tables",
-	"pg_stat_user_tables",
-	"pg_stat_xact_all_tables",
-	"pg_stat_xact_sys_tables",
-	"pg_stat_xact_user_tables",
-	"pg_stat_all_indexes",
-	"pg_stat_sys_indexes",
-	"pg_stat_user_indexes",
-	"pg_stat_user_functions",
-	"pg_stat_xact_user_functions",
-	"pg_statio_all_tables",
-	"pg_statio_sys_tables",
-	"pg_statio_user_tables",
-	"pg_statio_all_indexes",
-	"pg_statio_sys_indexes",
-	"pg_statio_user_indexes",
-	"pg_statio_all_sequences",
-	"pg_statio_sys_sequences",
-	"pg_statio_user_sequences",
-})
-
 var PG_SHADOW_VALUE_BY_COLUMN = NewOrderedMap([][]string{
 	{"usename", "bemidb"},
 	{"usesysid", "10"},

diff --git a/src/select_remapper.go b/src/select_remapper.go
@@ -162,7 +162,7 @@ func (selectRemapper *SelectRemapper) remapSelectStatement(selectStatement *pgQu
 				selectRemapper.traceTreeTraversal("FROM table", indentLevel)
 				selectStatement.FromClause[i] = selectRemapper.remapperTable.RemapTable(fromNode)
 				qSchemaTable := selectRemapper.parserTable.NodeToQuerySchemaTable(fromNode)
-				selectStatement = selectRemapper.remapperWhere.RemapWhereClauseForTable(qSchemaTable, selectStatement)
+				selectStatement = selectRemapper.remapperTable.RemapWhereClauseForTable(qSchemaTable, selectStatement)
 			} else if fromNode.GetRangeSubselect() != nil {
 				// FROM (SELECT ...)
 				selectRemapper.traceTreeTraversal("FROM subselect", indentLevel)
@@ -367,7 +367,7 @@ func (selectRemapper *SelectRemapper) remapJoinExpressions(selectStatement *pgQu
 		// WHERE
 		selectRemapper.traceTreeTraversal("WHERE left", indentLevel+1)
 		qSchemaTable := selectRemapper.parserTable.NodeToQuerySchemaTable(leftJoinNode)
-		selectStatement = selectRemapper.remapperWhere.RemapWhereClauseForTable(qSchemaTable, selectStatement)
+		selectStatement = selectRemapper.remapperTable.RemapWhereClauseForTable(qSchemaTable, selectStatement)
 		// TABLE
 		selectRemapper.traceTreeTraversal("TABLE left", indentLevel+1)
 		leftJoinNode = selectRemapper.remapperTable.RemapTable(leftJoinNode)
@@ -385,7 +385,7 @@ func (selectRemapper *SelectRemapper) remapJoinExpressions(selectStatement *pgQu
 		// WHERE
 		selectRemapper.traceTreeTraversal("WHERE right", indentLevel+1)
 		qSchemaTable := selectRemapper.parserTable.NodeToQuerySchemaTable(rightJoinNode)
-		selectStatement = selectRemapper.remapperWhere.RemapWhereClauseForTable(qSchemaTable, selectStatement)
+		selectStatement = selectRemapper.remapperTable.RemapWhereClauseForTable(qSchemaTable, selectStatement)
 		// TABLE
 		selectRemapper.traceTreeTraversal("TABLE right", indentLevel+1)
 		rightJoinNode = selectRemapper.remapperTable.RemapTable(rightJoinNode)

diff --git a/src/select_remapper_table.go b/src/select_remapper_table.go
@@ -30,6 +30,7 @@ const (
 
 type SelectRemapperTable struct {
 	parserTable         *QueryParserTable
+	parserWhere         *QueryParserWhere
 	icebergSchemaTables []IcebergSchemaTable
 	icebergReader       *IcebergReader
 	duckdb              *Duckdb
@@ -39,6 +40,7 @@ type SelectRemapperTable struct {
 func NewSelectRemapperTable(config *Config, icebergReader *IcebergReader, duckdb *Duckdb) *SelectRemapperTable {
 	remapper := &SelectRemapperTable{
 		parserTable:   NewQueryParserTable(config),
+		parserWhere:   NewQueryParserWhere(config),
 		icebergReader: icebergReader,
 		duckdb:        duckdb,
 		config:        config,
@@ -53,7 +55,7 @@ func (remapper *SelectRemapperTable) RemapTable(node *pgQuery.Node) *pgQuery.Nod
 	qSchemaTable := parser.NodeToQuerySchemaTable(node)
 
 	// pg_catalog.pg_* system tables
-	if parser.IsTableFromPgCatalog(qSchemaTable) {
+	if remapper.isTableFromPgCatalog(qSchemaTable) {
 		switch qSchemaTable.Table {
 		case PG_TABLE_PG_SHADOW:
 			// pg_catalog.pg_shadow -> return hard-coded credentials
@@ -178,6 +180,22 @@ func (remapper *SelectRemapperTable) RemapNestedTableFunction(funcCallNode *pgQu
 	return funcCallNode
 }
 
+func (remapper *SelectRemapperTable) RemapWhereClauseForTable(qSchemaTable QuerySchemaTable, selectStatement *pgQuery.SelectStmt) *pgQuery.SelectStmt {
+	if remapper.isTableFromPgCatalog(qSchemaTable) {
+		switch qSchemaTable.Table {
+		case PG_TABLE_PG_NAMESPACE:
+			// FROM pg_catalog.pg_namespace -> FROM pg_catalog.pg_namespace WHERE nspname != 'main'
+			withoutMainSchemaWhereCondition := remapper.parserWhere.MakeExpressionNode("nspname", "!=", "main")
+			return remapper.parserWhere.AppendWhereCondition(selectStatement, withoutMainSchemaWhereCondition)
+		case PG_TABLE_PG_STATIO_USER_TABLES:
+			// FROM pg_catalog.pg_statio_user_tables -> FROM pg_catalog.pg_statio_user_tables WHERE false
+			falseWhereCondition := remapper.parserWhere.MakeFalseConditionNode()
+			return remapper.parserWhere.OverrideWhereCondition(selectStatement, falseWhereCondition)
+		}
+	}
+	return selectStatement
+}
+
 func (remapper *SelectRemapperTable) overrideTable(node *pgQuery.Node, fromClause *pgQuery.Node) *pgQuery.Node {
 	node = fromClause
 	return node
@@ -204,6 +222,14 @@ func (remapper *SelectRemapperTable) icebergSchemaTableExists(schemaTable Iceber
 	return false
 }
 
+// System pg_* tables
+func (remapper *SelectRemapperTable) isTableFromPgCatalog(qSchemaTable QuerySchemaTable) bool {
+	return qSchemaTable.Schema == PG_SCHEMA_PG_CATALOG ||
+		(qSchemaTable.Schema == "" &&
+			(PG_SYSTEM_TABLES.Contains(qSchemaTable.Table) || PG_SYSTEM_VIEWS.Contains(qSchemaTable.Table)) &&
+			!remapper.icebergSchemaTableExists(qSchemaTable.ToIcebergSchemaTable()))
+}
+
 var PG_INHERITS_COLUMNS = []string{
 	"inhrelid",
 	"inhparent",
@@ -300,3 +326,121 @@ var PG_MATVIEWS_COLUMNS = []string{
 	"ispopulated",
 	"definition",
 }
+
+var PG_SYSTEM_TABLES = NewSet([]string{
+	"pg_aggregate",
+	"pg_am",
+	"pg_amop",
+	"pg_amproc",
+	"pg_attrdef",
+	"pg_attribute",
+	"pg_auth_members",
+	"pg_authid",
+	"pg_cast",
+	"pg_class",
+	"pg_collation",
+	"pg_constraint",
+	"pg_conversion",
+	"pg_database",
+	"pg_db_role_setting",
+	"pg_default_acl",
+	"pg_depend",
+	"pg_description",
+	"pg_enum",
+	"pg_event_trigger",
+	"pg_extension",
+	"pg_foreign_data_wrapper",
+	"pg_foreign_server",
+	"pg_foreign_table",
+	"pg_index",
+	"pg_inherits",
+	"pg_init_privs",
+	"pg_language",
+	"pg_largeobject",
+	"pg_largeobject_metadata",
+	"pg_matviews",
+	"pg_namespace",
+	"pg_opclass",
+	"pg_operator",
+	"pg_opfamily",
+	"pg_parameter_acl",
+	"pg_partitioned_table",
+	"pg_policy",
+	"pg_proc",
+	"pg_publication",
+	"pg_publication_namespace",
+	"pg_publication_rel",
+	"pg_user",
+	"pg_range",
+	"pg_replication_origin",
+	"pg_replication_slots",
+	"pg_rewrite",
+	"pg_roles",
+	"pg_seclabel",
+	"pg_sequence",
+	"pg_shadow",
+	"pg_shdepend",
+	"pg_shdescription",
+	"pg_shseclabel",
+	"pg_statistic",
+	"pg_statistic_ext",
+	"pg_statistic_ext_data",
+	"pg_subscription",
+	"pg_subscription_rel",
+	"pg_tablespace",
+	"pg_transform",
+	"pg_trigger",
+	"pg_ts_config",
+	"pg_ts_config_map",
+	"pg_ts_dict",
+	"pg_ts_parser",
+	"pg_ts_template",
+	"pg_type",
+	"pg_user_mapping",
+})
+
+var PG_SYSTEM_VIEWS = NewSet([]string{
+	"pg_stat_activity",
+	"pg_stat_replication",
+	"pg_stat_wal_receiver",
+	"pg_stat_recovery_prefetch",
+	"pg_stat_subscription",
+	"pg_stat_ssl",
+	"pg_stat_gssapi",
+	"pg_stat_progress_analyze",
+	"pg_stat_progress_create_index",
+	"pg_stat_progress_vacuum",
+	"pg_stat_progress_cluster",
+	"pg_stat_progress_basebackup",
+	"pg_stat_progress_copy",
+	"pg_stat_archiver",
+	"pg_stat_bgwriter",
+	"pg_stat_checkpointer",
+	"pg_stat_database",
+	"pg_stat_database_conflicts",
+	"pg_stat_io",
+	"pg_stat_replication_slots",
+	"pg_stat_slru",
+	"pg_stat_subscription_stats",
+	"pg_stat_wal",
+	"pg_stat_all_tables",
+	"pg_stat_sys_tables",
+	"pg_stat_user_tables",
+	"pg_stat_xact_all_tables",
+	"pg_stat_xact_sys_tables",
+	"pg_stat_xact_user_tables",
+	"pg_stat_all_indexes",
+	"pg_stat_sys_indexes",
+	"pg_stat_user_indexes",
+	"pg_stat_user_functions",
+	"pg_stat_xact_user_functions",
+	"pg_statio_all_tables",
+	"pg_statio_sys_tables",
+	"pg_statio_user_tables",
+	"pg_statio_all_indexes",
+	"pg_statio_sys_indexes",
+	"pg_statio_user_indexes",
+	"pg_statio_all_sequences",
+	"pg_statio_sys_sequences",
+	"pg_statio_user_sequences",
+})
diff --git a/src/select_remapper_where.go b/src/select_remapper_where.go
@@ -37,19 +37,3 @@ func (remapper *SelectRemapperWhere) RemapWhereExpressions(selectStatement *pgQu
 
 	return selectStatement
 }
-
-func (remapper *SelectRemapperWhere) RemapWhereClauseForTable(qSchemaTable QuerySchemaTable, selectStatement *pgQuery.SelectStmt) *pgQuery.SelectStmt {
-	if remapper.parserTable.IsTableFromPgCatalog(qSchemaTable) {
-		switch qSchemaTable.Table {
-		case PG_TABLE_PG_NAMESPACE:
-			// FROM pg_catalog.pg_namespace -> FROM pg_catalog.pg_namespace WHERE nspname != 'main'
-			withoutMainSchemaWhereCondition := remapper.parserWhere.MakeExpressionNode("nspname", "!=", "main")
-			return remapper.parserWhere.AppendWhereCondition(selectStatement, withoutMainSchemaWhereCondition)
-		case PG_TABLE_PG_STATIO_USER_TABLES:
-			// FROM pg_catalog.pg_statio_user_tables -> FROM pg_catalog.pg_statio_user_tables WHERE false
-			falseWhereCondition := remapper.parserWhere.MakeFalseConditionNode()
-			return remapper.parserWhere.OverrideWhereCondition(selectStatement, falseWhereCondition)
-		}
-	}
-	return selectStatement
-}