Repository: xataio/pgstream
Pull request #104 (2024-12-11 12:25:07)

Add table parser in snapshot generator
This PR adds a table parser to the postgres snapshot generator. This allows to translate the wildcard table into the tables associated to the request schema without having to provide all of them explicitly. 


Patch:
diff --git a/pkg/snapshot/generator/data/postgres/pg_schema_table_parser.go b/pkg/snapshot/generator/data/postgres/pg_schema_table_parser.go
new file mode 100644
index 0000000..a636a22
--- /dev/null
+++ b/pkg/snapshot/generator/data/postgres/pg_schema_table_parser.go
@@ -0,0 +1,58 @@
+// SPDX-License-Identifier: Apache-2.0
+
+package postgres
+
+import (
+	"context"
+	"fmt"
+	"slices"
+
+	pglib "github.com/xataio/pgstream/internal/postgres"
+	"github.com/xataio/pgstream/pkg/snapshot"
+)
+
+type schemaTableParser struct {
+	conn pglib.Querier
+}
+
+func newSchemaTableParser(conn pglib.Querier) *schemaTableParser {
+	return &schemaTableParser{
+		conn: conn,
+	}
+}
+
+func (p *schemaTableParser) parseSnapshotTables(ctx context.Context, snapshot *snapshot.Snapshot) error {
+	if slices.Contains(snapshot.TableNames, "*") {
+		var err error
+		snapshot.TableNames, err = p.discoverAllSchemaTables(ctx, snapshot.SchemaName)
+		if err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
+func (p *schemaTableParser) discoverAllSchemaTables(ctx context.Context, schema string) ([]string, error) {
+	const query = "SELECT tablename FROM pg_tables WHERE schemaname=$1"
+	rows, err := p.conn.Query(ctx, query, schema)
+	if err != nil {
+		return nil, fmt.Errorf("discovering all tables for schema %s: %w", schema, err)
+	}
+	defer rows.Close()
+
+	tableNames := []string{}
+	for rows.Next() {
+		var tableName string
+		if err := rows.Scan(&tableName); err != nil {
+			return nil, fmt.Errorf("scanning table name: %w", err)
+		}
+		tableNames = append(tableNames, tableName)
+	}
+
+	if err := rows.Err(); err != nil {
+		return nil, err
+	}
+
+	return tableNames, nil
+}
diff --git a/pkg/snapshot/generator/data/postgres/pg_snapshot_generator.go b/pkg/snapshot/generator/data/postgres/pg_snapshot_generator.go
index 3377eab..a5769df 100644
--- a/pkg/snapshot/generator/data/postgres/pg_snapshot_generator.go
+++ b/pkg/snapshot/generator/data/postgres/pg_snapshot_generator.go
@@ -17,9 +17,10 @@ import (
 )
 
 type SnapshotGenerator struct {
-	logger loglib.Logger
-	conn   pglib.Querier
-	mapper mapper
+	logger      loglib.Logger
+	conn        pglib.Querier
+	mapper      mapper
+	tableParser tableParser
 
 	schemaWorkers uint
 	tableWorkers  uint
@@ -38,6 +39,8 @@ type pageRange struct {
 	end   uint
 }
 
+type tableParser func(ctx context.Context, snapshot *snapshot.Snapshot) error
+
 type Option func(sg *SnapshotGenerator)
 
 func NewSnapshotGenerator(ctx context.Context, cfg *Config, processRow snapshot.RowProcessor, opts ...Option) (*SnapshotGenerator, error) {
@@ -54,6 +57,7 @@ func NewSnapshotGenerator(ctx context.Context, cfg *Config, processRow snapshot.
 		batchPageSize: cfg.batchPageSize(),
 		tableWorkers:  cfg.tableWorkers(),
 		schemaWorkers: cfg.schemaWorkers(),
+		tableParser:   newSchemaTableParser(conn).parseSnapshotTables,
 	}
 
 	for _, opt := range opts {
@@ -72,6 +76,9 @@ func WithLogger(logger loglib.Logger) Option {
 }
 
 func (sg *SnapshotGenerator) CreateSnapshot(ctx context.Context, ss *snapshot.Snapshot) error {
+	if err := sg.tableParser(ctx, ss); err != nil {
+		return err
+	}
 	return sg.conn.ExecInTxWithOptions(ctx, func(tx pglib.Tx) error {
 		snapshotID, err := sg.exportSnapshot(ctx, tx)
 		if err != nil {


Test patch:
diff --git a/pkg/snapshot/generator/data/postgres/pg_schema_table_parser_test.go b/pkg/snapshot/generator/data/postgres/pg_schema_table_parser_test.go
new file mode 100644
index 0000000..1d30e10
--- /dev/null
+++ b/pkg/snapshot/generator/data/postgres/pg_schema_table_parser_test.go
@@ -0,0 +1,172 @@
+// SPDX-License-Identifier: Apache-2.0
+
+package postgres
+
+import (
+	"context"
+	"errors"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+	pglib "github.com/xataio/pgstream/internal/postgres"
+	pgmocks "github.com/xataio/pgstream/internal/postgres/mocks"
+	"github.com/xataio/pgstream/pkg/snapshot"
+)
+
+func TestSchemaTableParser_parseSnapshotTables(t *testing.T) {
+	t.Parallel()
+
+	testSchema := "test-schema"
+	errTest := errors.New("oh noes")
+
+	tests := []struct {
+		name     string
+		conn     *pgmocks.Querier
+		snapshot *snapshot.Snapshot
+
+		wantTables []string
+		wantErr    error
+	}{
+		{
+			name: "ok - no wildcard",
+			conn: &pgmocks.Querier{},
+			snapshot: &snapshot.Snapshot{
+				SchemaName: testSchema,
+				TableNames: []string{"table-a", "table-b", "table-c"},
+			},
+
+			wantTables: []string{"table-a", "table-b", "table-c"},
+			wantErr:    nil,
+		},
+		{
+			name: "ok - only wildcard",
+			conn: &pgmocks.Querier{
+				QueryFn: func(ctx context.Context, query string, args ...any) (pglib.Rows, error) {
+					require.Equal(t, "SELECT tablename FROM pg_tables WHERE schemaname=$1", query)
+					require.Equal(t, []any{testSchema}, args)
+					return &pgmocks.Rows{
+						ScanFn: func(dest ...any) error {
+							require.Len(t, dest, 1)
+							table, ok := dest[0].(*string)
+							require.True(t, ok)
+							*table = "table-1"
+							return nil
+						},
+						NextFn:  func(i uint) bool { return i == 1 },
+						CloseFn: func() {},
+						ErrFn:   func() error { return nil },
+					}, nil
+				},
+			},
+			snapshot: &snapshot.Snapshot{
+				SchemaName: testSchema,
+				TableNames: []string{"*"},
+			},
+
+			wantTables: []string{"table-1"},
+			wantErr:    nil,
+		},
+		{
+			name: "ok - with wildcard",
+			conn: &pgmocks.Querier{
+				QueryFn: func(ctx context.Context, query string, args ...any) (pglib.Rows, error) {
+					require.Equal(t, "SELECT tablename FROM pg_tables WHERE schemaname=$1", query)
+					require.Equal(t, []any{testSchema}, args)
+					return &pgmocks.Rows{
+						ScanFn: func(dest ...any) error {
+							require.Len(t, dest, 1)
+							table, ok := dest[0].(*string)
+							require.True(t, ok)
+							*table = "table-1"
+							return nil
+						},
+						NextFn:  func(i uint) bool { return i == 1 },
+						CloseFn: func() {},
+						ErrFn:   func() error { return nil },
+					}, nil
+				},
+			},
+			snapshot: &snapshot.Snapshot{
+				SchemaName: testSchema,
+				TableNames: []string{"table-a", "table-b", "*"},
+			},
+
+			wantTables: []string{"table-1"},
+			wantErr:    nil,
+		},
+		{
+			name: "error - querying schema tables",
+			conn: &pgmocks.Querier{
+				QueryFn: func(ctx context.Context, query string, args ...any) (pglib.Rows, error) {
+					require.Equal(t, "SELECT tablename FROM pg_tables WHERE schemaname=$1", query)
+					require.Equal(t, []any{testSchema}, args)
+					return nil, errTest
+				},
+			},
+			snapshot: &snapshot.Snapshot{
+				SchemaName: testSchema,
+				TableNames: []string{"table-a", "table-b", "*"},
+			},
+
+			wantTables: nil,
+			wantErr:    errTest,
+		},
+		{
+			name: "error - scanning row",
+			conn: &pgmocks.Querier{
+				QueryFn: func(ctx context.Context, query string, args ...any) (pglib.Rows, error) {
+					require.Equal(t, "SELECT tablename FROM pg_tables WHERE schemaname=$1", query)
+					require.Equal(t, []any{testSchema}, args)
+					return &pgmocks.Rows{
+						ScanFn: func(dest ...any) error {
+							return errTest
+						},
+						NextFn:  func(i uint) bool { return i == 1 },
+						CloseFn: func() {},
+						ErrFn:   func() error { return nil },
+					}, nil
+				},
+			},
+			snapshot: &snapshot.Snapshot{
+				SchemaName: testSchema,
+				TableNames: []string{"table-a", "table-b", "*"},
+			},
+
+			wantTables: nil,
+			wantErr:    errTest,
+		},
+		{
+			name: "error - rows error",
+			conn: &pgmocks.Querier{
+				QueryFn: func(ctx context.Context, query string, args ...any) (pglib.Rows, error) {
+					require.Equal(t, "SELECT tablename FROM pg_tables WHERE schemaname=$1", query)
+					require.Equal(t, []any{testSchema}, args)
+					return &pgmocks.Rows{
+						NextFn:  func(i uint) bool { return i == 0 },
+						CloseFn: func() {},
+						ErrFn:   func() error { return errTest },
+					}, nil
+				},
+			},
+			snapshot: &snapshot.Snapshot{
+				SchemaName: testSchema,
+				TableNames: []string{"table-a", "table-b", "*"},
+			},
+
+			wantTables: nil,
+			wantErr:    errTest,
+		},
+	}
+
+	for _, tc := range tests {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+
+			p := newSchemaTableParser(tc.conn)
+			err := p.parseSnapshotTables(context.Background(), tc.snapshot)
+			require.ErrorIs(t, err, tc.wantErr)
+			require.Equal(t, tc.wantTables, tc.snapshot.TableNames)
+		})
+	}
+}
diff --git a/pkg/snapshot/generator/data/postgres/pg_snapshot_generator_test.go b/pkg/snapshot/generator/data/postgres/pg_snapshot_generator_test.go
index cb036b0..4953183 100644
--- a/pkg/snapshot/generator/data/postgres/pg_snapshot_generator_test.go
+++ b/pkg/snapshot/generator/data/postgres/pg_snapshot_generator_test.go
@@ -874,6 +874,7 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 				schemaWorkers: 1,
 				tableWorkers:  1,
 				batchPageSize: 10,
+				tableParser:   func(ctx context.Context, snapshot *snapshot.Snapshot) error { return nil },
 			}
 
 			if tc.schemaWorkers != 0 {