Repository: xataio/pgstream
Pull request #412 (2025-06-20 08:46:42)

Add schema wildcard supp for transformer validator
Adding the wildcard schema support for transformer validator.


Patch:
diff --git a/pkg/wal/processor/transformer/wal_postgres_transformer_parser.go b/pkg/wal/processor/transformer/wal_postgres_transformer_parser.go
index f8b1670..87ccf13 100644
--- a/pkg/wal/processor/transformer/wal_postgres_transformer_parser.go
+++ b/pkg/wal/processor/transformer/wal_postgres_transformer_parser.go
@@ -140,7 +140,19 @@ func (v *PostgresTransformerParser) getRequiredTablesList(ctx context.Context) (
 			return nil, err
 		}
 		if schemaName == wildcard {
-			return nil, fmt.Errorf("wildcard schema name is not supported: %s", table)
+			if tableName != wildcard {
+				return nil, fmt.Errorf("wildcard schema must be used with wildcard table, got: %q", tableName)
+			}
+
+			// if schemaName is wildcard, fetch all schemas
+			allSchemas, err := v.getAllSchemaNames(ctx)
+			if err != nil {
+				return nil, fmt.Errorf("fetching all schemas for wildcard: %w", err)
+			}
+			for _, schema := range allSchemas {
+				v.requiredTables = append(v.requiredTables, schema+"."+wildcard)
+			}
+			continue
 		}
 
 		if tableName != wildcard {
@@ -195,6 +207,30 @@ func (v *PostgresTransformerParser) getAllSchemaTables(ctx context.Context, sche
 	return tableNames, nil
 }
 
+func (v *PostgresTransformerParser) getAllSchemaNames(ctx context.Context) ([]string, error) {
+	const query = "SELECT nspname FROM pg_catalog.pg_namespace WHERE nspname NOT IN ('pg_catalog', 'information_schema', 'pg_toast', 'pgstream')"
+	rows, err := v.conn.Query(ctx, query)
+	if err != nil {
+		return nil, fmt.Errorf("discovering all schemas for wildcard: %w", err)
+	}
+	defer rows.Close()
+
+	schemas := []string{}
+	for rows.Next() {
+		var schemaName string
+		if err := rows.Scan(&schemaName); err != nil {
+			return nil, fmt.Errorf("scanning schema name: %w", err)
+		}
+		schemas = append(schemas, schemaName)
+	}
+
+	if err := rows.Err(); err != nil {
+		return nil, err
+	}
+
+	return schemas, nil
+}
+
 func parseTableName(qualifiedTableName string) (string, string, error) {
 	parts := strings.Split(qualifiedTableName, ".")
 	switch len(parts) {


Test patch:
diff --git a/pkg/wal/processor/transformer/wal_postgres_transformer_parser_test.go b/pkg/wal/processor/transformer/wal_postgres_transformer_parser_test.go
index be16cca..55c71cf 100644
--- a/pkg/wal/processor/transformer/wal_postgres_transformer_parser_test.go
+++ b/pkg/wal/processor/transformer/wal_postgres_transformer_parser_test.go
@@ -52,6 +52,19 @@ func TestPostgresTransformerParser_ParseAndValidate(t *testing.T) {
 					},
 					ErrFn: func() error { return nil },
 				}, nil
+			case "SELECT nspname FROM pg_catalog.pg_namespace WHERE nspname NOT IN ('pg_catalog', 'information_schema', 'pg_toast', 'pgstream')":
+				return &pgmocks.Rows{
+					CloseFn: func() {},
+					NextFn:  func(i uint) bool { return i == 1 },
+					ScanFn: func(dest ...any) error {
+						require.Len(t, dest, 1)
+						schemaName, ok := dest[0].(*string)
+						require.True(t, ok)
+						*schemaName = "public"
+						return nil
+					},
+					ErrFn: func() error { return nil },
+				}, nil
 			default:
 				return nil, fmt.Errorf("unexpected query: %s", query)
 			}
@@ -114,6 +127,60 @@ func TestPostgresTransformerParser_ParseAndValidate(t *testing.T) {
 			wantTransformersFor: []string{"name"},
 			wantErr:             nil,
 		},
+		{
+			name: "ok - with wildcard table",
+			transformerRules: []TableRules{
+				{
+					Schema:         "public",
+					Table:          "test",
+					ValidationMode: "strict",
+					ColumnRules: map[string]TransformerRules{
+						"id": {
+							Name: "greenmask_integer",
+						},
+						"name": {
+							Name: "string",
+						},
+					},
+				},
+			},
+			validator: PostgresTransformerParser{
+				conn:           testQuerier,
+				builder:        builder.NewTransformerBuilder(),
+				pgtypeMap:      pgtype.NewMap(),
+				requiredTables: []string{"*"},
+			},
+
+			wantTransformersFor: []string{"id", "name"},
+			wantErr:             nil,
+		},
+		{
+			name: "ok - with wildcard schema and table",
+			transformerRules: []TableRules{
+				{
+					Schema:         "public",
+					Table:          "test",
+					ValidationMode: "strict",
+					ColumnRules: map[string]TransformerRules{
+						"id": {
+							Name: "greenmask_integer",
+						},
+						"name": {
+							Name: "string",
+						},
+					},
+				},
+			},
+			validator: PostgresTransformerParser{
+				conn:           testQuerier,
+				builder:        builder.NewTransformerBuilder(),
+				pgtypeMap:      pgtype.NewMap(),
+				requiredTables: []string{"*.*"},
+			},
+
+			wantTransformersFor: []string{"id", "name"},
+			wantErr:             nil,
+		},
 		{
 			name: "error - missing column for strict validation",
 			transformerRules: []TableRules{
@@ -229,7 +296,7 @@ func TestPostgresTransformerParser_ParseAndValidate(t *testing.T) {
 			wantErr: errInvalidTableName,
 		},
 		{
-			name:             "error - wildcard schema name",
+			name:             "error - wildcard schema name with non-wildcard table name",
 			transformerRules: []TableRules{},
 			validator: PostgresTransformerParser{
 				conn:           testQuerier,
@@ -237,7 +304,7 @@ func TestPostgresTransformerParser_ParseAndValidate(t *testing.T) {
 				pgtypeMap:      pgtype.NewMap(),
 				requiredTables: []string{"*.test"},
 			},
-			wantErr: fmt.Errorf("wildcard schema name is not supported: *.test"),
+			wantErr: fmt.Errorf("wildcard schema must be used with wildcard table, got: \"test\""),
 		},
 	}
 	for _, tc := range tests {