Repository: xataio/pgstream (go)

Issue:
[Snapshots] `pg_dump/pg_restore` schema snapshot not working well with replication
When using the `pg_dump`/`pg_restore` schema snapshot generator (which is the default when replicating to postgres), the replication of schema changes is not working well, since the schema view is not captured in the `pgstream.schema_log`, which means the first schema change that gets replicated will be ignored, since there's no record of the schema in the `pgstream.schema_log` table. 

Fix the behaviour so that replication of schema changes works from the get go when using `pg_dump`/`pg_restore` for schema snapshoting.

Pull request description:
[Snapshots] Fix DDL replication with pg_dump/pg_restore schema snapshot
This PR updates the pg_dump/pg_restore schema snapshot generator to ensure the schema is populated on the `schema_log` table if it exists, since it needs to be initialised for replication to work properly with schema changes.

It will only update the view if the table exists. The snapshot mode supports no pgstream state when using pg2pg snapshots without replication.

Closes #199 

Patch:
diff --git a/cmd/config.go b/cmd/config.go
index 9c52c34..c0d04b4 100644
--- a/cmd/config.go
+++ b/cmd/config.go
@@ -115,13 +115,14 @@ func parseSnapshotConfig(pgURL, prefix string) *snapshotbuilder.SnapshotListener
 			Tables:          viper.GetStringSlice(fmt.Sprintf("%s_SNAPSHOT_TABLES", prefix)),
 			SnapshotWorkers: viper.GetUint(fmt.Sprintf("%s_SNAPSHOT_WORKERS", prefix)),
 		},
-		Schema: parseSchemaSnapshotConfig(pgURL),
+		Schema: parseSchemaSnapshotConfig(prefix, pgURL),
 	}
 }
 
-func parseSchemaSnapshotConfig(pgurl string) snapshotbuilder.SchemaSnapshotConfig {
+func parseSchemaSnapshotConfig(prefix, pgurl string) snapshotbuilder.SchemaSnapshotConfig {
+	useSchemaLog := viper.GetBool(fmt.Sprintf("%s_SNAPSHOT_USE_SCHEMALOG", prefix))
 	pgTargetURL := viper.GetString("PGSTREAM_POSTGRES_WRITER_TARGET_URL")
-	if pgTargetURL != "" {
+	if pgTargetURL != "" && !useSchemaLog {
 		return snapshotbuilder.SchemaSnapshotConfig{
 			DumpRestore: &pgdumprestore.Config{
 				SourcePGURL: pgurl,
diff --git a/docs/README.md b/docs/README.md
index add4677..b2dc0fa 100644
--- a/docs/README.md
+++ b/docs/README.md
@@ -96,6 +96,7 @@ Here's a list of all the environment variables that can be used to configure the
 | PGSTREAM_POSTGRES_INITIAL_SNAPSHOT_TABLE_WORKERS    | 4                      | No       | Number of concurrent workers that will be used per table by the snapshotting process.                                                                                                                                                                                                                        |
 | PGSTREAM_POSTGRES_INITIAL_SNAPSHOT_BATCH_PAGE_SIZE  | 1000                   | No       | Size of the table page range which will be processed concurrently by the table workers from `PGSTREAM_POSTGRES_INITIAL_SNAPSHOT_TABLE_WORKERS`.                                                                                                                                                              |
 | PGSTREAM_POSTGRES_INITIAL_SNAPSHOT_WORKERS          | 1                      | No       | Number of schemas that will be processed in parallel by the snapshotting process.                                                                                                                                                                                                                            |
+| PGSTREAM_POSTGRES_INITIAL_SNAPSHOT_USE_SCHEMALOG    | False                  | No       | Forces the use of the `pgstream.schema_log` for the schema snapshot instead of using `pg_dump`/`pg_restore` for Postgres outputs.                                                                                                                                                                            |
 
 </details>
 
@@ -110,6 +111,7 @@ Here's a list of all the environment variables that can be used to configure the
 | PGSTREAM_POSTGRES_SNAPSHOT_TABLE_WORKERS   | 4       | No       | Number of concurrent workers that will be used per table by the snapshotting process.                                                                                                                                                                                                                        |
 | PGSTREAM_POSTGRES_SNAPSHOT_BATCH_PAGE_SIZE | 1000    | No       | Size of the table page range which will be processed concurrently by the table workers from `PGSTREAM_POSTGRES_SNAPSHOT_TABLE_WORKERS`.                                                                                                                                                                      |
 | PGSTREAM_POSTGRES_SNAPSHOT_WORKERS         | 1       | No       | Number of schemas that will be processed in parallel by the snapshotting process.                                                                                                                                                                                                                            |
+| PGSTREAM_POSTGRES_SNAPSHOT_USE_SCHEMALOG   | False   | No       | Forces the use of the `pgstream.schema_log` for the schema snapshot instead of using `pg_dump`/`pg_restore` for Postgres outputs.                                                                                                                                                                            |
 
 </details>
 
diff --git a/internal/postgres/mocks/mock_pg_querier.go b/internal/postgres/mocks/mock_pg_querier.go
index e7939a7..0dd0cca 100644
--- a/internal/postgres/mocks/mock_pg_querier.go
+++ b/internal/postgres/mocks/mock_pg_querier.go
@@ -43,5 +43,8 @@ func (m *Querier) ExecInTxWithOptions(ctx context.Context, fn func(tx postgres.T
 }
 
 func (m *Querier) Close(ctx context.Context) error {
-	return m.CloseFn(ctx)
+	if m.CloseFn != nil {
+		return m.CloseFn(ctx)
+	}
+	return nil
 }
diff --git a/internal/postgres/pg_conn_builder.go b/internal/postgres/pg_conn_builder.go
new file mode 100644
index 0000000..5cb27b7
--- /dev/null
+++ b/internal/postgres/pg_conn_builder.go
@@ -0,0 +1,15 @@
+// SPDX-License-Identifier: Apache-2.0
+
+package postgres
+
+import "context"
+
+type QuerierBuilder func(context.Context, string) (Querier, error)
+
+func ConnBuilder(ctx context.Context, url string) (Querier, error) {
+	return NewConn(ctx, url)
+}
+
+func ConnPoolBuilder(ctx context.Context, url string) (Querier, error) {
+	return NewConnPool(ctx, url)
+}
diff --git a/pkg/schemalog/mocks/store_mock.go b/pkg/schemalog/mocks/store_mock.go
index 37a6e09..1e5735c 100644
--- a/pkg/schemalog/mocks/store_mock.go
+++ b/pkg/schemalog/mocks/store_mock.go
@@ -43,7 +43,10 @@ func (m *Store) Ack(ctx context.Context, le *schemalog.LogEntry) error {
 }
 
 func (m *Store) Close() error {
-	return m.CloseFn()
+	if m.CloseFn != nil {
+		return m.CloseFn()
+	}
+	return nil
 }
 
 func (m *Store) GetInsertCalls() uint64 {
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
index 0d9c181..ca07bdf 100644
--- 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
@@ -8,18 +8,21 @@ import (
 
 	pglib "github.com/xataio/pgstream/internal/postgres"
 	loglib "github.com/xataio/pgstream/pkg/log"
+	"github.com/xataio/pgstream/pkg/schemalog"
+	schemalogpg "github.com/xataio/pgstream/pkg/schemalog/postgres"
 	"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
+	sourceURL      string
+	targetURL      string
+	pgDumpFn       pgdumpFn
+	pgRestoreFn    pgrestoreFn
+	schemalogStore schemalog.Store
+	connBuilder    pglib.QuerierBuilder
+	logger         loglib.Logger
 }
 
 type Config struct {
@@ -39,16 +42,12 @@ 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,
+		connBuilder: pglib.ConnBuilder,
 		logger:      loglib.NewNoopLogger(),
 	}
 
@@ -56,6 +55,10 @@ func NewSnapshotGenerator(ctx context.Context, c *Config, opts ...Option) (*Snap
 		opt(sg)
 	}
 
+	if err := sg.initialiseSchemaLogStore(ctx); err != nil {
+		return nil, err
+	}
+
 	return sg, nil
 }
 
@@ -84,8 +87,7 @@ func (s *SnapshotGenerator) CreateSnapshot(ctx context.Context, ss *snapshot.Sna
 	// 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 {
+		if err := s.createSchemaIfNotExists(ctx, ss.SchemaName); err != nil {
 			return err
 		}
 	}
@@ -95,11 +97,34 @@ func (s *SnapshotGenerator) CreateSnapshot(ctx context.Context, ss *snapshot.Sna
 		return err
 	}
 
+	// if we perform a schema snapshot using pg_dump/pg_restore, we need to make
+	// sure the schema_log table is updated accordingly with the schema view so
+	// that replication can work as expected if configured.
+	if s.schemalogStore != nil {
+		if _, err := s.schemalogStore.Insert(ctx, ss.SchemaName); err != nil {
+			return fmt.Errorf("inserting schemalog entry after schema snapshot: %w", err)
+		}
+	}
+
 	return nil
 }
 
 func (s *SnapshotGenerator) Close() error {
-	return s.targetConn.Close(context.Background())
+	if s.schemalogStore != nil {
+		return s.schemalogStore.Close()
+	}
+	return nil
+}
+
+func (s *SnapshotGenerator) createSchemaIfNotExists(ctx context.Context, schemaName string) error {
+	targetConn, err := s.connBuilder(ctx, s.targetURL)
+	if err != nil {
+		return err
+	}
+	defer targetConn.Close(context.Background())
+
+	_, err = targetConn.Exec(ctx, fmt.Sprintf("CREATE SCHEMA IF NOT EXISTS %s", schemaName))
+	return err
 }
 
 func (s *SnapshotGenerator) pgdumpOptions(ss *snapshot.Snapshot) pglib.PGDumpOptions {
@@ -123,3 +148,36 @@ func (s *SnapshotGenerator) pgrestoreOptions() pglib.PGRestoreOptions {
 		SchemaOnly:       true,
 	}
 }
+
+func (s *SnapshotGenerator) initialiseSchemaLogStore(ctx context.Context) error {
+	exists, err := s.schemalogExists(ctx)
+	if err != nil {
+		return err
+	}
+
+	if !exists {
+		return nil
+	}
+
+	s.schemalogStore, err = schemalogpg.NewStore(ctx, schemalogpg.Config{URL: s.sourceURL})
+	return err
+}
+
+func (s *SnapshotGenerator) schemalogExists(ctx context.Context) (bool, error) {
+	conn, err := s.connBuilder(ctx, s.sourceURL)
+	if err != nil {
+		return false, err
+	}
+	defer conn.Close(context.Background())
+
+	// check if the pgstream.schema_log table exists, if not, we can skip the initialisation
+	// of the schemalog store
+	existsQuery := "SELECT EXISTS (SELECT FROM information_schema.tables WHERE table_schema = $1 AND table_name = $2)"
+	var exists bool
+	err = conn.QueryRow(ctx, existsQuery, schemalog.SchemaName, schemalog.TableName).Scan(&exists)
+	if err != nil {
+		return false, fmt.Errorf("checking schemalog table existence: %w", err)
+	}
+
+	return exists, nil
+}


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
index 35b7b95..69ce3a2 100644
--- 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
@@ -5,12 +5,15 @@ package pgdumprestore
 import (
 	"context"
 	"errors"
+	"fmt"
 	"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/schemalog"
+	schemalogmocks "github.com/xataio/pgstream/pkg/schemalog/mocks"
 	"github.com/xataio/pgstream/pkg/snapshot"
 )
 
@@ -23,11 +26,13 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 	errTest := errors.New("oh noes")
 
 	tests := []struct {
-		name        string
-		snapshot    *snapshot.Snapshot
-		conn        pglib.Querier
-		pgdumpFn    pgdumpFn
-		pgrestoreFn pgrestoreFn
+		name           string
+		snapshot       *snapshot.Snapshot
+		conn           pglib.Querier
+		connBuilder    pglib.QuerierBuilder
+		pgdumpFn       pgdumpFn
+		pgrestoreFn    pgrestoreFn
+		schemalogStore schemalog.Store
 
 		wantErr error
 	}{
@@ -64,6 +69,26 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 
 			wantErr: nil,
 		},
+		{
+			name: "ok - no tables in public schema",
+			snapshot: &snapshot.Snapshot{
+				SchemaName: publicSchema,
+				TableNames: []string{},
+			},
+			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, errors.New("pgdumpFn: should not be called")
+			},
+			pgrestoreFn: func(po pglib.PGRestoreOptions, dump []byte) (string, error) {
+				return "", errors.New("pgrestoreFn: should not be called")
+			},
+
+			wantErr: nil,
+		},
 		{
 			name: "error - performing pgdump",
 			snapshot: &snapshot.Snapshot{
@@ -104,6 +129,24 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 
 			wantErr: errTest,
 		},
+		{
+			name: "error - getting target conn",
+			snapshot: &snapshot.Snapshot{
+				SchemaName: testSchema,
+				TableNames: []string{testTable},
+			},
+			connBuilder: func(ctx context.Context, s string) (pglib.Querier, error) {
+				return nil, 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,
+		},
 		{
 			name: "error - creating schema",
 			snapshot: &snapshot.Snapshot{
@@ -122,6 +165,33 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 				return "", errors.New("pgrestoreFn: should not be called")
 			},
 
+			wantErr: errTest,
+		},
+		{
+			name: "error - schemalog insertion fails",
+			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) {
+				return testDump, nil
+			},
+			pgrestoreFn: func(po pglib.PGRestoreOptions, dump []byte) (string, error) {
+				require.Equal(t, testDump, dump)
+				return "", nil
+			},
+			schemalogStore: &schemalogmocks.Store{
+				InsertFn: func(ctx context.Context, schemaName string) (*schemalog.LogEntry, error) {
+					return nil, errTest
+				},
+			},
+
 			wantErr: errTest,
 		},
 	}
@@ -131,16 +201,124 @@ func TestSnapshotGenerator_CreateSnapshot(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(),
+				sourceURL:      "source-url",
+				targetURL:      "target-url",
+				connBuilder:    func(ctx context.Context, s string) (pglib.Querier, error) { return tc.conn, nil },
+				pgDumpFn:       tc.pgdumpFn,
+				pgRestoreFn:    tc.pgrestoreFn,
+				schemalogStore: tc.schemalogStore,
+				logger:         log.NewNoopLogger(),
+			}
+
+			if tc.connBuilder != nil {
+				sg.connBuilder = tc.connBuilder
 			}
 
 			err := sg.CreateSnapshot(context.Background(), tc.snapshot)
 			require.ErrorIs(t, err, tc.wantErr)
+			sg.Close()
+		})
+	}
+}
+
+func TestSnapshotGenerator_schemalogExists(t *testing.T) {
+	t.Parallel()
+
+	errTest := errors.New("oh noes")
+
+	tests := []struct {
+		name        string
+		connBuilder pglib.QuerierBuilder
+
+		wantExists bool
+		wantErr    error
+	}{
+		{
+			name: "ok - schemalog exists",
+			connBuilder: func(ctx context.Context, s string) (pglib.Querier, error) {
+				return &mocks.Querier{
+					QueryRowFn: func(ctx context.Context, query string, args ...any) pglib.Row {
+						require.Equal(t, "SELECT EXISTS (SELECT FROM information_schema.tables WHERE table_schema = $1 AND table_name = $2)", query)
+						require.Equal(t, []any{schemalog.SchemaName, schemalog.TableName}, args)
+						return &mocks.Row{
+							ScanFn: func(args ...any) error {
+								require.Len(t, args, 1)
+								exists, ok := args[0].(*bool)
+								require.True(t, ok, fmt.Sprintf("exists, expected *bool, got %T", args[0]))
+								*exists = true
+								return nil
+							},
+						}
+					},
+				}, nil
+			},
+
+			wantExists: true,
+			wantErr:    nil,
+		},
+		{
+			name: "ok - schemalog does not exist",
+			connBuilder: func(ctx context.Context, s string) (pglib.Querier, error) {
+				return &mocks.Querier{
+					QueryRowFn: func(ctx context.Context, query string, args ...any) pglib.Row {
+						require.Equal(t, "SELECT EXISTS (SELECT FROM information_schema.tables WHERE table_schema = $1 AND table_name = $2)", query)
+						require.Equal(t, []any{schemalog.SchemaName, schemalog.TableName}, args)
+						return &mocks.Row{
+							ScanFn: func(args ...any) error {
+								require.Len(t, args, 1)
+								exists, ok := args[0].(*bool)
+								require.True(t, ok, fmt.Sprintf("exists, expected *bool, got %T", args[0]))
+								*exists = false
+								return nil
+							},
+						}
+					},
+				}, nil
+			},
+
+			wantExists: false,
+			wantErr:    nil,
+		},
+		{
+			name: "error - getting source connection",
+			connBuilder: func(ctx context.Context, s string) (pglib.Querier, error) {
+				return nil, errTest
+			},
+
+			wantExists: false,
+			wantErr:    errTest,
+		},
+		{
+			name: "error - scanning",
+			connBuilder: func(ctx context.Context, s string) (pglib.Querier, error) {
+				return &mocks.Querier{
+					QueryRowFn: func(ctx context.Context, query string, args ...any) pglib.Row {
+						require.Equal(t, "SELECT EXISTS (SELECT FROM information_schema.tables WHERE table_schema = $1 AND table_name = $2)", query)
+						require.Equal(t, []any{schemalog.SchemaName, schemalog.TableName}, args)
+						return &mocks.Row{
+							ScanFn: func(args ...any) error {
+								return errTest
+							},
+						}
+					},
+				}, nil
+			},
+
+			wantExists: false,
+			wantErr:    errTest,
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+			sg := SnapshotGenerator{
+				connBuilder: tc.connBuilder,
+			}
+
+			exists, err := sg.schemalogExists(context.Background())
+			require.ErrorIs(t, err, tc.wantErr)
+			require.Equal(t, tc.wantExists, exists)
 		})
 	}
 }