Repository: xataio/pgstream
Pull request #149 (2025-02-18 09:51:07)

Add pg_dump/pg_restore schema snapshot generator
This PR adds an implementation of the snapshot generator that generates a postgres schema snapshot using `pg_dump` and `pg_restore` postgres utilities.


Patch:
diff --git a/pkg/snapshot/generator/postgres/data/pg_snapshot_generator.go b/pkg/snapshot/generator/postgres/data/pg_snapshot_generator.go
index 36978fd..0c17a24 100644
--- a/pkg/snapshot/generator/postgres/data/pg_snapshot_generator.go
+++ b/pkg/snapshot/generator/postgres/data/pg_snapshot_generator.go
@@ -70,7 +70,7 @@ func NewSnapshotGenerator(ctx context.Context, cfg *Config, processRow snapshot.
 func WithLogger(logger loglib.Logger) Option {
 	return func(sg *SnapshotGenerator) {
 		sg.logger = loglib.NewLogger(logger).WithFields(loglib.Fields{
-			loglib.ModuleField: "postgres_snapshot_generator",
+			loglib.ModuleField: "postgres_data_snapshot_generator",
 		})
 	}
 }
diff --git a/pkg/snapshot/generator/postgres/schema/pgdumprestore/snapshot_pg_dump_restore_generator.go b/pkg/snapshot/generator/postgres/schema/pgdumprestore/snapshot_pg_dump_restore_generator.go
new file mode 100644
index 0000000..49f1802
--- /dev/null
+++ b/pkg/snapshot/generator/postgres/schema/pgdumprestore/snapshot_pg_dump_restore_generator.go
@@ -0,0 +1,119 @@
+// SPDX-License-Identifier: Apache-2.0
+
+package pgdumprestore
+
+import (
+	"context"
+	"fmt"
+
+	pglib "github.com/xataio/pgstream/internal/postgres"
+	loglib "github.com/xataio/pgstream/pkg/log"
+	"github.com/xataio/pgstream/pkg/snapshot"
+)
+
+// SnapshotGenerator generates postgres schema snapshots using pg_dump and
+// pg_restore
+type SnapshotGenerator struct {
+	sourceURL   string
+	targetURL   string
+	pgDumpFn    pgdumpFn
+	pgRestoreFn pgrestoreFn
+	targetConn  pglib.Querier
+	logger      loglib.Logger
+}
+
+type Config struct {
+	SourcePGURL string
+	TargetPGURL string
+}
+
+type (
+	pgdumpFn    func(pglib.PGDumpOptions) ([]byte, error)
+	pgrestoreFn func(pglib.PGRestoreOptions, []byte) (string, error)
+)
+
+type Option func(s *SnapshotGenerator)
+
+const publicSchema = "public"
+
+// NewSnapshotGenerator will return a postgres schema snapshot generator that
+// uses pg_dump and pg_restore to sync the schema of two postgres databases
+func NewSnapshotGenerator(ctx context.Context, c *Config, opts ...Option) (*SnapshotGenerator, error) {
+	targetConn, err := pglib.NewConnPool(ctx, c.TargetPGURL)
+	if err != nil {
+		return nil, err
+	}
+	sg := &SnapshotGenerator{
+		sourceURL:   c.SourcePGURL,
+		targetURL:   c.TargetPGURL,
+		pgDumpFn:    pglib.RunPGDump,
+		pgRestoreFn: pglib.RunPGRestore,
+		targetConn:  targetConn,
+		logger:      loglib.NewNoopLogger(),
+	}
+
+	for _, opt := range opts {
+		opt(sg)
+	}
+
+	return sg, nil
+}
+
+func WithLogger(logger loglib.Logger) Option {
+	return func(sg *SnapshotGenerator) {
+		sg.logger = loglib.NewLogger(logger).WithFields(loglib.Fields{
+			loglib.ModuleField: "postgres_schema_snapshot_generator",
+		})
+	}
+}
+
+func (s *SnapshotGenerator) CreateSnapshot(ctx context.Context, ss *snapshot.Snapshot) error {
+	s.logger.Info("creating schema snapshot", loglib.Fields{"schema": ss.SchemaName, "tables": ss.TableNames})
+	dump, err := s.pgDumpFn(s.pgdumpOptions(ss))
+	if err != nil {
+		return err
+	}
+
+	// if we use table filtering in the pg_dump command, the schema creation
+	// will not be dumped, so it needs to be created explicitly (except for
+	// public schema)
+	if len(ss.TableNames) > 0 && ss.SchemaName != publicSchema {
+		_, err = s.targetConn.Exec(ctx, fmt.Sprintf("CREATE SCHEMA IF NOT EXISTS %s", ss.SchemaName))
+		if err != nil {
+			return err
+		}
+	}
+
+	_, err = s.pgRestoreFn(s.pgrestoreOptions(), dump)
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
+func (s *SnapshotGenerator) Close() error {
+	return s.targetConn.Close(context.Background())
+}
+
+func (s *SnapshotGenerator) pgdumpOptions(ss *snapshot.Snapshot) pglib.PGDumpOptions {
+	opts := pglib.PGDumpOptions{
+		ConnectionString: s.sourceURL,
+		Format:           "c",
+		SchemaOnly:       true,
+		Schemas:          []string{ss.SchemaName},
+	}
+
+	for _, table := range ss.TableNames {
+		opts.Tables = append(opts.Tables, ss.SchemaName+"."+table)
+	}
+
+	return opts
+}
+
+func (s *SnapshotGenerator) pgrestoreOptions() pglib.PGRestoreOptions {
+	return pglib.PGRestoreOptions{
+		ConnectionString: s.targetURL,
+		SchemaOnly:       true,
+	}
+}


Test patch:
diff --git a/pkg/snapshot/generator/postgres/schema/pgdumprestore/snapshot_pg_dump_restore_generator_test.go b/pkg/snapshot/generator/postgres/schema/pgdumprestore/snapshot_pg_dump_restore_generator_test.go
new file mode 100644
index 0000000..35b7b95
--- /dev/null
+++ b/pkg/snapshot/generator/postgres/schema/pgdumprestore/snapshot_pg_dump_restore_generator_test.go
@@ -0,0 +1,146 @@
+// SPDX-License-Identifier: Apache-2.0
+
+package pgdumprestore
+
+import (
+	"context"
+	"errors"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+	pglib "github.com/xataio/pgstream/internal/postgres"
+	"github.com/xataio/pgstream/internal/postgres/mocks"
+	"github.com/xataio/pgstream/pkg/log"
+	"github.com/xataio/pgstream/pkg/snapshot"
+)
+
+func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
+	t.Parallel()
+
+	testDump := []byte("test dump")
+	testSchema := "test_schema"
+	testTable := "test_table"
+	errTest := errors.New("oh noes")
+
+	tests := []struct {
+		name        string
+		snapshot    *snapshot.Snapshot
+		conn        pglib.Querier
+		pgdumpFn    pgdumpFn
+		pgrestoreFn pgrestoreFn
+
+		wantErr error
+	}{
+		{
+			name: "ok",
+			snapshot: &snapshot.Snapshot{
+				SchemaName: testSchema,
+				TableNames: []string{testTable},
+			},
+			conn: &mocks.Querier{
+				ExecFn: func(ctx context.Context, i uint, query string, args ...any) (pglib.CommandTag, error) {
+					require.Equal(t, "CREATE SCHEMA IF NOT EXISTS "+testSchema, query)
+					return pglib.CommandTag{}, nil
+				},
+			},
+			pgdumpFn: func(po pglib.PGDumpOptions) ([]byte, error) {
+				require.Equal(t, pglib.PGDumpOptions{
+					ConnectionString: "source-url",
+					Format:           "c",
+					SchemaOnly:       true,
+					Schemas:          []string{testSchema},
+					Tables:           []string{testSchema + "." + testTable},
+				}, po)
+				return testDump, nil
+			},
+			pgrestoreFn: func(po pglib.PGRestoreOptions, dump []byte) (string, error) {
+				require.Equal(t, pglib.PGRestoreOptions{
+					ConnectionString: "target-url",
+					SchemaOnly:       true,
+				}, po)
+				require.Equal(t, testDump, dump)
+				return "", nil
+			},
+
+			wantErr: nil,
+		},
+		{
+			name: "error - performing pgdump",
+			snapshot: &snapshot.Snapshot{
+				SchemaName: testSchema,
+				TableNames: []string{testTable},
+			},
+			conn: &mocks.Querier{
+				ExecFn: func(ctx context.Context, i uint, query string, args ...any) (pglib.CommandTag, error) {
+					return pglib.CommandTag{}, errors.New("ExecFn: should not be called")
+				},
+			},
+			pgdumpFn: func(po pglib.PGDumpOptions) ([]byte, error) {
+				return nil, errTest
+			},
+			pgrestoreFn: func(po pglib.PGRestoreOptions, dump []byte) (string, error) {
+				return "", errors.New("pgrestoreFn: should not be called")
+			},
+
+			wantErr: errTest,
+		},
+		{
+			name: "error - performing pgrestore",
+			snapshot: &snapshot.Snapshot{
+				SchemaName: publicSchema,
+				TableNames: []string{testTable},
+			},
+			conn: &mocks.Querier{
+				ExecFn: func(ctx context.Context, i uint, query string, args ...any) (pglib.CommandTag, error) {
+					return pglib.CommandTag{}, errors.New("ExecFn: should not be called")
+				},
+			},
+			pgdumpFn: func(po pglib.PGDumpOptions) ([]byte, error) {
+				return testDump, nil
+			},
+			pgrestoreFn: func(po pglib.PGRestoreOptions, dump []byte) (string, error) {
+				return "", errTest
+			},
+
+			wantErr: errTest,
+		},
+		{
+			name: "error - creating schema",
+			snapshot: &snapshot.Snapshot{
+				SchemaName: testSchema,
+				TableNames: []string{testTable},
+			},
+			conn: &mocks.Querier{
+				ExecFn: func(ctx context.Context, i uint, query string, args ...any) (pglib.CommandTag, error) {
+					return pglib.CommandTag{}, errTest
+				},
+			},
+			pgdumpFn: func(po pglib.PGDumpOptions) ([]byte, error) {
+				return testDump, nil
+			},
+			pgrestoreFn: func(po pglib.PGRestoreOptions, dump []byte) (string, error) {
+				return "", errors.New("pgrestoreFn: should not be called")
+			},
+
+			wantErr: errTest,
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+
+			sg := SnapshotGenerator{
+				sourceURL:   "source-url",
+				targetURL:   "target-url",
+				targetConn:  tc.conn,
+				pgDumpFn:    tc.pgdumpFn,
+				pgRestoreFn: tc.pgrestoreFn,
+				logger:      log.NewNoopLogger(),
+			}
+
+			err := sg.CreateSnapshot(context.Background(), tc.snapshot)
+			require.ErrorIs(t, err, tc.wantErr)
+		})
+	}
+}