Repository: xataio/pgstream
Pull request #284 (2025-04-15 15:49:59)

[Snapshots] Fix repeatable snapshots
This PR updates the pgdump/pgrestore schema snapshot generator with the following changes:
- The snapshot tables are not passed as pg_dump filter, but instead excluded_tables are retrieved for the snapshot schema and used, so that any non table object is dumped along with the selected tables.
- The pg_dump command now uses plain text format and the clean option instead of using `pg_restore` with clean, since the command fails when trying to delete the public schema. With `psql` we can parse the output and ignore the error, since the rest of types are recreated appropriately.

These two changes should improve the experience of repeatable snapshots for postgres to postgres.

github-actions[bot]: ### Merging this branch changes the coverage (1 decrease, 1 increase)

| Impacted Packages | Coverage Δ | :robot: |
|-------------------|------------|---------|
| github.com/xataio/pgstream/internal/postgres | 1.01% (**-0.06%**) | :thumbsdown: |
| github.com/xataio/pgstream/pkg/snapshot/generator/postgres/schema/pgdumprestore | 82.61% (**+6.05%**) | :thumbsup: |

---

<details>

<summary>Coverage by file</summary>

### Changed files (no unit tests)

| Changed File | Coverage Δ | Total | Covered | Missed | :robot: |
|--------------|------------|-------|---------|--------|---------|
| github.com/xataio/pgstream/internal/postgres/pg_dump.go | 0.00% (ø) | 26 (+3) | 0 | 26 (+3) |  |
| github.com/xataio/pgstream/internal/postgres/pg_restore.go | 0.00% (ø) | 52 (+8) | 0 | 52 (+8) |  |
| github.com/xataio/pgstream/pkg/snapshot/generator/postgres/schema/pgdumprestore/snapshot_pg_dump_restore_generator.go | 82.61% (**+6.05%**) | 92 (+28) | 76 (+27) | 16 (+1) | :thumbsup: |

_Please note that the "Total", "Covered", and "Missed" counts above refer to ***code statements*** instead of lines of code. The value in brackets refers to the test coverage of that file in the old version of the code._

### Changed unit test files

- github.com/xataio/pgstream/internal/postgres/pg_dump_pg_restore_integration_test.go
- github.com/xataio/pgstream/pkg/snapshot/generator/postgres/schema/pgdumprestore/snapshot_pg_dump_restore_generator_test.go

</details>


Patch:
diff --git a/internal/postgres/pg_dump.go b/internal/postgres/pg_dump.go
index 29e326e..e9c55a1 100644
--- a/internal/postgres/pg_dump.go
+++ b/internal/postgres/pg_dump.go
@@ -25,6 +25,8 @@ type PGDumpOptions struct {
 	SchemaOnly bool
 	// do not dump privileges (grant/revoke)
 	NoPrivileges bool
+	// Clean all the objects that will be dumped
+	Clean bool
 	// Options to pass to pg_dump
 	Options []string
 }
@@ -58,6 +60,11 @@ func (opts *PGDumpOptions) ToArgs() []string {
 		options = append(options, fmt.Sprintf("--exclude-table=%v", table))
 	}
 
+	if opts.Clean {
+		options = append(options, "--clean")
+		options = append(options, "--if-exists")
+	}
+
 	options = append(options, opts.Options...)
 	return options
 }
diff --git a/internal/postgres/pg_restore.go b/internal/postgres/pg_restore.go
index 43a1dd1..5a0b179 100644
--- a/internal/postgres/pg_restore.go
+++ b/internal/postgres/pg_restore.go
@@ -14,6 +14,7 @@ import (
 
 const (
 	pgRestoreCmd = "pg_restore"
+	psqlCmd      = "psql"
 )
 
 type PGRestoreOptions struct {
@@ -23,6 +24,8 @@ type PGRestoreOptions struct {
 	SchemaOnly bool
 	// Clean all the objects that will be restored
 	Clean bool
+	// Format (c custom, d directory, t tar, p plain text)
+	Format string
 	// Options to pass to pg_restore
 	Options []string
 }
@@ -45,10 +48,20 @@ func (opts PGRestoreOptions) toArgs() []string {
 	return options
 }
 
+func (opts PGRestoreOptions) toPSQLArgs() []string {
+	return []string{opts.ConnectionString}
+}
+
 // Func RunPGRestore runs pg_restore command with the given options and returns
 // the result.
 func RunPGRestore(opts PGRestoreOptions, dump []byte) (string, error) {
-	cmd := exec.Command(pgRestoreCmd, opts.toArgs()...) //nolint:gosec
+	var cmd *exec.Cmd
+	switch opts.Format {
+	case "c":
+		cmd = exec.Command(pgRestoreCmd, opts.toArgs()...) //nolint:gosec
+	default:
+		cmd = exec.Command(psqlCmd, opts.toPSQLArgs()...) //nolint:gosec
+	}
 
 	stdin, err := cmd.StdinPipe()
 	if err != nil {
@@ -74,15 +87,20 @@ func parsePgRestoreOutputErrs(out []byte) error {
 	scanner := bufio.NewScanner(bytes.NewReader(out))
 	for scanner.Scan() {
 		line := scanner.Text()
-		if strings.Contains(line, "pg_restore: error:") {
+		if strings.Contains(line, "pg_restore: error:") || strings.Contains(line, "ERROR") {
 			switch {
 			case strings.Contains(line, "already exists"),
 				strings.Contains(line, "multiple primary keys for table"):
 				errs.addError(&ErrRelationAlreadyExists{Details: line})
+			case strings.Contains(line, "cannot drop schema public because other objects depend on it"):
+				errs.addError(&ErrConstraintViolation{Details: line})
 			default:
 				errs.addError(errors.New(line))
 			}
 		}
+		if strings.Contains(line, "psql: error:") {
+			errs.addError(errors.New(line))
+		}
 	}
 	return errs
 }
@@ -122,8 +140,10 @@ func (e *PGRestoreErrors) addError(err error) {
 	}
 
 	var errAlreadyExists *ErrRelationAlreadyExists
+	var errConstraintViolation *ErrConstraintViolation
 	switch {
-	case errors.As(err, &errAlreadyExists):
+	case errors.As(err, &errAlreadyExists),
+		errors.As(err, &errConstraintViolation):
 		e.ignoredErrs = append(e.ignoredErrs, err)
 	default:
 		e.criticalErrs = append(e.criticalErrs, err)
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 35062de..a8acd75 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
@@ -6,6 +6,8 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"slices"
+	"strings"
 
 	pglib "github.com/xataio/pgstream/internal/postgres"
 	loglib "github.com/xataio/pgstream/pkg/log"
@@ -40,7 +42,10 @@ type (
 
 type Option func(s *SnapshotGenerator)
 
-const publicSchema = "public"
+const (
+	publicSchema = "public"
+	wildcard     = "*"
+)
 
 // NewSnapshotGenerator will return a postgres schema snapshot generator that
 // uses pg_dump and pg_restore to sync the schema of two postgres databases
@@ -82,7 +87,13 @@ func (s *SnapshotGenerator) CreateSnapshot(ctx context.Context, ss *snapshot.Sna
 	if ss.SchemaName == publicSchema && len(ss.TableNames) == 0 {
 		return nil
 	}
-	dump, err := s.pgDumpFn(s.pgdumpOptions(ss))
+
+	pgdumpOpts, err := s.pgdumpOptions(ctx, ss)
+	if err != nil {
+		return fmt.Errorf("preparing pg_dump options: %w", err)
+	}
+
+	dump, err := s.pgDumpFn(*pgdumpOpts)
 	if err != nil {
 		return err
 	}
@@ -104,7 +115,8 @@ func (s *SnapshotGenerator) CreateSnapshot(ctx context.Context, ss *snapshot.Sna
 			if pgrestoreErr.HasCriticalErrors() {
 				return err
 			}
-			s.logger.Warn(nil, "pg_restore errors ignored", loglib.Fields{"errors_ignored": len(pgrestoreErr.GetIgnoredErrors())})
+			ignoredErrors := pgrestoreErr.GetIgnoredErrors()
+			s.logger.Warn(nil, fmt.Sprintf("pg_restore: %d errors ignored", len(ignoredErrors)), loglib.Fields{"errors_ignored": ignoredErrors})
 		default:
 			return err
 		}
@@ -140,34 +152,80 @@ func (s *SnapshotGenerator) createSchemaIfNotExists(ctx context.Context, schemaN
 	return err
 }
 
-func (s *SnapshotGenerator) pgdumpOptions(ss *snapshot.Snapshot) pglib.PGDumpOptions {
-	opts := pglib.PGDumpOptions{
+func (s *SnapshotGenerator) pgrestoreOptions() pglib.PGRestoreOptions {
+	return pglib.PGRestoreOptions{
+		ConnectionString: s.targetURL,
+		SchemaOnly:       true,
+		Clean:            s.cleanTargetDB,
+		Format:           "p",
+	}
+}
+
+func (s *SnapshotGenerator) pgdumpOptions(ctx context.Context, ss *snapshot.Snapshot) (*pglib.PGDumpOptions, error) {
+	opts := &pglib.PGDumpOptions{
 		ConnectionString: s.sourceURL,
-		Format:           "c",
+		Format:           "p",
 		SchemaOnly:       true,
 		Schemas:          []string{pglib.QuoteIdentifier(ss.SchemaName)},
+		Clean:            s.cleanTargetDB,
 	}
 
-	const wildcard = "*"
-	for _, table := range ss.TableNames {
-		if table == wildcard {
-			// wildcard means all tables in the schema, so no table filter
-			// required
-			opts.Tables = nil
-			break
+	// wildcard means all tables in the schema, so no table filter required
+	if slices.Contains(ss.TableNames, wildcard) {
+		return opts, nil
+	}
+
+	// we use the excluded tables flag to make sure we still dump non table
+	// objects for the schema in question. If we use the tables filter, only
+	// those tables are dumped, and any related non table objects will not be
+	// dumped, causing the restore to fail due to missing related objects.
+	if len(ss.TableNames) > 0 {
+		var err error
+		opts.ExcludeTables, err = s.pgdumpExcludedTables(ctx, ss.SchemaName, ss.TableNames)
+		if err != nil {
+			return nil, err
 		}
-		opts.Tables = append(opts.Tables, pglib.QuoteQualifiedIdentifier(ss.SchemaName, table))
 	}
 
-	return opts
+	return opts, nil
 }
 
-func (s *SnapshotGenerator) pgrestoreOptions() pglib.PGRestoreOptions {
-	return pglib.PGRestoreOptions{
-		ConnectionString: s.targetURL,
-		SchemaOnly:       true,
-		Clean:            s.cleanTargetDB,
+func (s *SnapshotGenerator) pgdumpExcludedTables(ctx context.Context, schemaName string, includeTables []string) ([]string, error) {
+	conn, err := s.connBuilder(ctx, s.sourceURL)
+	if err != nil {
+		return nil, err
+	}
+	defer conn.Close(ctx)
+
+	paramRefs := make([]string, 0, len(includeTables))
+	tableParams := make([]any, 0, len(includeTables))
+	for i, table := range includeTables {
+		tableParams = append(tableParams, table)
+		paramRefs = append(paramRefs, fmt.Sprintf("$%d", i+1))
+	}
+
+	// get all tables in the schema that are not in the include list
+	query := fmt.Sprintf("SELECT tablename FROM pg_tables WHERE schemaname = '%s' AND tablename NOT IN (%s)", schemaName, strings.Join(paramRefs, ","))
+	rows, err := conn.Query(ctx, query, tableParams...)
+	if err != nil {
+		return nil, fmt.Errorf("retrieving tables from schema: %w", err)
+	}
+	defer rows.Close()
+
+	excludeTables := []string{}
+	for rows.Next() {
+		tableName := ""
+		if err := rows.Scan(&tableName); err != nil {
+			return nil, fmt.Errorf("scanning table name: %w", err)
+		}
+		excludeTables = append(excludeTables, pglib.QuoteIdentifier(tableName))
 	}
+
+	if err := rows.Err(); err != nil {
+		return nil, err
+	}
+
+	return excludeTables, nil
 }
 
 func (s *SnapshotGenerator) initialiseSchemaLogStore(ctx context.Context) error {


Test patch:
diff --git a/internal/postgres/pg_dump_pg_restore_integration_test.go b/internal/postgres/pg_dump_pg_restore_integration_test.go
index edc83ee..2cbd1e8 100644
--- a/internal/postgres/pg_dump_pg_restore_integration_test.go
+++ b/internal/postgres/pg_dump_pg_restore_integration_test.go
@@ -19,61 +19,73 @@ func Test_pgdump_pgrestore(t *testing.T) {
 
 	ctx := context.Background()
 
-	var sourcePGURL, targetPGURL string
-	cleanup, err := testcontainers.SetupPostgresContainer(ctx, &sourcePGURL, testcontainers.Postgres17)
-	require.NoError(t, err)
-	defer cleanup()
-
-	cleanup2, err := testcontainers.SetupPostgresContainer(ctx, &targetPGURL, testcontainers.Postgres17)
-	require.NoError(t, err)
-	defer cleanup2()
-
-	sourceConn, err := NewConn(ctx, sourcePGURL)
-	require.NoError(t, err)
-	targetConn, err := NewConn(ctx, targetPGURL)
-	require.NoError(t, err)
-
-	testSchema := "test_schema"
-	testTable1 := "test_table_1"
-	testTable2 := "test_table_2"
-
-	// create schema + tables
-	_, err = sourceConn.Exec(ctx, fmt.Sprintf("create schema %s", testSchema))
-	require.NoError(t, err)
-	_, err = sourceConn.Exec(ctx, fmt.Sprintf("create table %s.%s(id serial primary key, name text)", testSchema, testTable1))
-	require.NoError(t, err)
-	_, err = sourceConn.Exec(ctx, fmt.Sprintf("create table %s.%s(id serial primary key, name text)", testSchema, testTable2))
-	require.NoError(t, err)
-	// insert data
-	_, err = sourceConn.Exec(ctx, fmt.Sprintf("insert into %s.%s(name) values('a'),('b'),('c')", testSchema, testTable1))
-	require.NoError(t, err)
-
-	pgdumpOpts := PGDumpOptions{
-		ConnectionString: sourcePGURL,
-		Format:           "c",
-		SchemaOnly:       true,
-		Schemas:          []string{testSchema},
-		ExcludeTables:    []string{testSchema + "." + testTable2},
-	}
+	run := func(t *testing.T, format string) {
+		var sourcePGURL, targetPGURL string
+		cleanup, err := testcontainers.SetupPostgresContainer(ctx, &sourcePGURL, testcontainers.Postgres17)
+		require.NoError(t, err)
+		defer cleanup()
+
+		cleanup2, err := testcontainers.SetupPostgresContainer(ctx, &targetPGURL, testcontainers.Postgres17)
+		require.NoError(t, err)
+		defer cleanup2()
+
+		sourceConn, err := NewConn(ctx, sourcePGURL)
+		require.NoError(t, err)
+		targetConn, err := NewConn(ctx, targetPGURL)
+		require.NoError(t, err)
+
+		testSchema := "test_schema"
+		testTable1 := "test_table_1"
+		testTable2 := "test_table_2"
+
+		// create schema + tables
+		_, err = sourceConn.Exec(ctx, fmt.Sprintf("create schema %s", testSchema))
+		require.NoError(t, err)
+		_, err = sourceConn.Exec(ctx, fmt.Sprintf("create table %s.%s(id serial primary key, name text)", testSchema, testTable1))
+		require.NoError(t, err)
+		_, err = sourceConn.Exec(ctx, fmt.Sprintf("create table %s.%s(id serial primary key, name text)", testSchema, testTable2))
+		require.NoError(t, err)
+		// insert data
+		_, err = sourceConn.Exec(ctx, fmt.Sprintf("insert into %s.%s(name) values('a'),('b'),('c')", testSchema, testTable1))
+		require.NoError(t, err)
+
+		pgdumpOpts := PGDumpOptions{
+			ConnectionString: sourcePGURL,
+			Format:           format,
+			SchemaOnly:       true,
+			Clean:            true,
+			Schemas:          []string{testSchema},
+			ExcludeTables:    []string{testSchema + "." + testTable2},
+		}
+
+		dump, err := RunPGDump(pgdumpOpts)
+		require.NoError(t, err)
+
+		pgrestoreOpts := PGRestoreOptions{
+			ConnectionString: targetPGURL,
+			SchemaOnly:       true,
+			Clean:            true,
+			Format:           format,
+		}
 
-	dump, err := RunPGDump(pgdumpOpts)
-	require.NoError(t, err)
+		_, err = RunPGRestore(pgrestoreOpts, dump)
+		require.NoError(t, err)
 
-	pgrestoreOpts := PGRestoreOptions{
-		ConnectionString: targetPGURL,
-		SchemaOnly:       true,
+		// schema only pgdump, no data should be available but the schema and
+		// selected table should exist.
+		var count int
+		err = targetConn.QueryRow(ctx, fmt.Sprintf("select count(*) from %s.%s", testSchema, testTable1)).Scan(&count)
+		require.NoError(t, err)
+		require.Equal(t, 0, count)
+		// test table 2 should not exist
+		err = targetConn.QueryRow(ctx, fmt.Sprintf("select count(*) from %s.%s", testSchema, testTable2)).Scan(&count)
+		require.Error(t, err)
 	}
 
-	_, err = RunPGRestore(pgrestoreOpts, dump)
-	require.NoError(t, err)
-
-	// schema only pgdump, no data should be available but the schema and
-	// selected table should exist.
-	var count int
-	err = targetConn.QueryRow(ctx, fmt.Sprintf("select count(*) from %s.%s", testSchema, testTable1)).Scan(&count)
-	require.NoError(t, err)
-	require.Equal(t, 0, count)
-	// test table 2 should not exist
-	err = targetConn.QueryRow(ctx, fmt.Sprintf("select count(*) from %s.%s", testSchema, testTable2)).Scan(&count)
-	require.Error(t, err)
+	t.Run("custom format - pg_restore", func(t *testing.T) {
+		run(t, "c")
+	})
+	t.Run("plain format - psql", func(t *testing.T) {
+		run(t, "p")
+	})
 }
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 a28232a..6bc8ee4 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
@@ -23,6 +23,7 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 	testDump := []byte("test dump")
 	testSchema := "test_schema"
 	testTable := "test_table"
+	excludedTable := "excluded_test_table"
 	errTest := errors.New("oh noes")
 
 	tests := []struct {
@@ -47,14 +48,31 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 					require.Equal(t, "CREATE SCHEMA IF NOT EXISTS "+testSchema, query)
 					return pglib.CommandTag{}, nil
 				},
+				QueryFn: func(ctx context.Context, query string, args ...any) (pglib.Rows, error) {
+					require.Equal(t, "SELECT tablename FROM pg_tables WHERE schemaname = 'test_schema' AND tablename NOT IN ($1)", query)
+					require.Equal(t, []any{testTable}, args)
+					return &mocks.Rows{
+						CloseFn: func() {},
+						NextFn:  func(i uint) bool { return i == 1 },
+						ScanFn: func(dest ...any) error {
+							require.Len(t, dest, 1)
+							tableName, ok := dest[0].(*string)
+							require.True(t, ok)
+							*tableName = excludedTable
+							return nil
+						},
+						ErrFn: func() error { return nil },
+					}, nil
+				},
 			},
 			pgdumpFn: func(po pglib.PGDumpOptions) ([]byte, error) {
 				require.Equal(t, pglib.PGDumpOptions{
 					ConnectionString: "source-url",
-					Format:           "c",
+					Format:           "p",
+					Clean:            false,
 					SchemaOnly:       true,
 					Schemas:          []string{pglib.QuoteIdentifier(testSchema)},
-					Tables:           []string{pglib.QuoteQualifiedIdentifier(testSchema, testTable)},
+					ExcludeTables:    []string{pglib.QuoteIdentifier(excludedTable)},
 				}, po)
 				return testDump, nil
 			},
@@ -62,6 +80,7 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 				require.Equal(t, pglib.PGRestoreOptions{
 					ConnectionString: "target-url",
 					SchemaOnly:       true,
+					Format:           "p",
 				}, po)
 				require.Equal(t, testDump, dump)
 				return "", nil
@@ -84,7 +103,7 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 			pgdumpFn: func(po pglib.PGDumpOptions) ([]byte, error) {
 				require.Equal(t, pglib.PGDumpOptions{
 					ConnectionString: "source-url",
-					Format:           "c",
+					Format:           "p",
 					SchemaOnly:       true,
 					Schemas:          []string{pglib.QuoteIdentifier(testSchema)},
 				}, po)
@@ -94,6 +113,7 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 				require.Equal(t, pglib.PGRestoreOptions{
 					ConnectionString: "target-url",
 					SchemaOnly:       true,
+					Format:           "p",
 				}, po)
 				require.Equal(t, testDump, dump)
 				return "", nil
@@ -121,6 +141,90 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 
 			wantErr: nil,
 		},
+		{
+			name: "error - querying excluded tables",
+			snapshot: &snapshot.Snapshot{
+				SchemaName: testSchema,
+				TableNames: []string{testTable},
+			},
+			conn: &mocks.Querier{
+				QueryFn: func(ctx context.Context, query string, args ...any) (pglib.Rows, error) {
+					return nil, errTest
+				},
+			},
+			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: errTest,
+		},
+		{
+			name: "error - scanning excluded tables",
+			snapshot: &snapshot.Snapshot{
+				SchemaName: testSchema,
+				TableNames: []string{testTable},
+			},
+			conn: &mocks.Querier{
+				QueryFn: func(ctx context.Context, query string, args ...any) (pglib.Rows, error) {
+					require.Equal(t, "SELECT tablename FROM pg_tables WHERE schemaname = 'test_schema' AND tablename NOT IN ($1)", query)
+					return &mocks.Rows{
+						CloseFn: func() {},
+						NextFn:  func(i uint) bool { return i == 1 },
+						ScanFn: func(dest ...any) error {
+							return errTest
+						},
+						ErrFn: func() error { return nil },
+					}, nil
+				},
+			},
+			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: errTest,
+		},
+		{
+			name: "error - excluded tables rows",
+			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
+				},
+				QueryFn: func(ctx context.Context, query string, args ...any) (pglib.Rows, error) {
+					require.Equal(t, "SELECT tablename FROM pg_tables WHERE schemaname = 'test_schema' AND tablename NOT IN ($1)", query)
+					return &mocks.Rows{
+						CloseFn: func() {},
+						NextFn:  func(i uint) bool { return i == 1 },
+						ScanFn: func(dest ...any) error {
+							require.Len(t, dest, 1)
+							tableName, ok := dest[0].(*string)
+							require.True(t, ok)
+							*tableName = excludedTable
+							return nil
+						},
+						ErrFn: func() error { return errTest },
+					}, nil
+				},
+			},
+			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: errTest,
+		},
 		{
 			name: "error - performing pgdump",
 			snapshot: &snapshot.Snapshot{
@@ -131,6 +235,21 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 				ExecFn: func(ctx context.Context, i uint, query string, args ...any) (pglib.CommandTag, error) {
 					return pglib.CommandTag{}, errors.New("ExecFn: should not be called")
 				},
+				QueryFn: func(ctx context.Context, query string, args ...any) (pglib.Rows, error) {
+					require.Equal(t, "SELECT tablename FROM pg_tables WHERE schemaname = 'test_schema' AND tablename NOT IN ($1)", query)
+					return &mocks.Rows{
+						CloseFn: func() {},
+						NextFn:  func(i uint) bool { return i == 1 },
+						ScanFn: func(dest ...any) error {
+							require.Len(t, dest, 1)
+							tableName, ok := dest[0].(*string)
+							require.True(t, ok)
+							*tableName = excludedTable
+							return nil
+						},
+						ErrFn: func() error { return nil },
+					}, nil
+				},
 			},
 			pgdumpFn: func(po pglib.PGDumpOptions) ([]byte, error) {
 				return nil, errTest
@@ -151,6 +270,21 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 				ExecFn: func(ctx context.Context, i uint, query string, args ...any) (pglib.CommandTag, error) {
 					return pglib.CommandTag{}, errors.New("ExecFn: should not be called")
 				},
+				QueryFn: func(ctx context.Context, query string, args ...any) (pglib.Rows, error) {
+					require.Equal(t, "SELECT tablename FROM pg_tables WHERE schemaname = 'public' AND tablename NOT IN ($1)", query)
+					return &mocks.Rows{
+						CloseFn: func() {},
+						NextFn:  func(i uint) bool { return i == 1 },
+						ScanFn: func(dest ...any) error {
+							require.Len(t, dest, 1)
+							tableName, ok := dest[0].(*string)
+							require.True(t, ok)
+							*tableName = excludedTable
+							return nil
+						},
+						ErrFn: func() error { return nil },
+					}, nil
+				},
 			},
 			pgdumpFn: func(po pglib.PGDumpOptions) ([]byte, error) {
 				return testDump, nil
@@ -171,6 +305,21 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 				ExecFn: func(ctx context.Context, i uint, query string, args ...any) (pglib.CommandTag, error) {
 					return pglib.CommandTag{}, errors.New("ExecFn: should not be called")
 				},
+				QueryFn: func(ctx context.Context, query string, args ...any) (pglib.Rows, error) {
+					require.Equal(t, "SELECT tablename FROM pg_tables WHERE schemaname = 'public' AND tablename NOT IN ($1)", query)
+					return &mocks.Rows{
+						CloseFn: func() {},
+						NextFn:  func(i uint) bool { return i == 1 },
+						ScanFn: func(dest ...any) error {
+							require.Len(t, dest, 1)
+							tableName, ok := dest[0].(*string)
+							require.True(t, ok)
+							*tableName = excludedTable
+							return nil
+						},
+						ErrFn: func() error { return nil },
+					}, nil
+				},
 			},
 			pgdumpFn: func(po pglib.PGDumpOptions) ([]byte, error) {
 				return testDump, nil
@@ -191,6 +340,21 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 				ExecFn: func(ctx context.Context, i uint, query string, args ...any) (pglib.CommandTag, error) {
 					return pglib.CommandTag{}, errors.New("ExecFn: should not be called")
 				},
+				QueryFn: func(ctx context.Context, query string, args ...any) (pglib.Rows, error) {
+					require.Equal(t, "SELECT tablename FROM pg_tables WHERE schemaname = 'public' AND tablename NOT IN ($1)", query)
+					return &mocks.Rows{
+						CloseFn: func() {},
+						NextFn:  func(i uint) bool { return i == 1 },
+						ScanFn: func(dest ...any) error {
+							require.Len(t, dest, 1)
+							tableName, ok := dest[0].(*string)
+							require.True(t, ok)
+							*tableName = excludedTable
+							return nil
+						},
+						ErrFn: func() error { return nil },
+					}, nil
+				},
 			},
 			pgdumpFn: func(po pglib.PGDumpOptions) ([]byte, error) {
 				return testDump, nil
@@ -208,7 +372,26 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 				TableNames: []string{testTable},
 			},
 			connBuilder: func(ctx context.Context, s string) (pglib.Querier, error) {
-				return nil, errTest
+				if s == "target-url" {
+					return nil, errTest
+				}
+				return &mocks.Querier{
+					QueryFn: func(ctx context.Context, query string, args ...any) (pglib.Rows, error) {
+						require.Equal(t, "SELECT tablename FROM pg_tables WHERE schemaname = 'test_schema' AND tablename NOT IN ($1)", query)
+						return &mocks.Rows{
+							CloseFn: func() {},
+							NextFn:  func(i uint) bool { return i == 1 },
+							ScanFn: func(dest ...any) error {
+								require.Len(t, dest, 1)
+								tableName, ok := dest[0].(*string)
+								require.True(t, ok)
+								*tableName = excludedTable
+								return nil
+							},
+							ErrFn: func() error { return nil },
+						}, nil
+					},
+				}, nil
 			},
 			pgdumpFn: func(po pglib.PGDumpOptions) ([]byte, error) {
 				return testDump, nil
@@ -229,6 +412,21 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 				ExecFn: func(ctx context.Context, i uint, query string, args ...any) (pglib.CommandTag, error) {
 					return pglib.CommandTag{}, errTest
 				},
+				QueryFn: func(ctx context.Context, query string, args ...any) (pglib.Rows, error) {
+					require.Equal(t, "SELECT tablename FROM pg_tables WHERE schemaname = 'test_schema' AND tablename NOT IN ($1)", query)
+					return &mocks.Rows{
+						CloseFn: func() {},
+						NextFn:  func(i uint) bool { return i == 1 },
+						ScanFn: func(dest ...any) error {
+							require.Len(t, dest, 1)
+							tableName, ok := dest[0].(*string)
+							require.True(t, ok)
+							*tableName = excludedTable
+							return nil
+						},
+						ErrFn: func() error { return nil },
+					}, nil
+				},
 			},
 			pgdumpFn: func(po pglib.PGDumpOptions) ([]byte, error) {
 				return testDump, nil
@@ -250,6 +448,21 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 					require.Equal(t, "CREATE SCHEMA IF NOT EXISTS "+testSchema, query)
 					return pglib.CommandTag{}, nil
 				},
+				QueryFn: func(ctx context.Context, query string, args ...any) (pglib.Rows, error) {
+					require.Equal(t, "SELECT tablename FROM pg_tables WHERE schemaname = 'test_schema' AND tablename NOT IN ($1)", query)
+					return &mocks.Rows{
+						CloseFn: func() {},
+						NextFn:  func(i uint) bool { return i == 1 },
+						ScanFn: func(dest ...any) error {
+							require.Len(t, dest, 1)
+							tableName, ok := dest[0].(*string)
+							require.True(t, ok)
+							*tableName = excludedTable
+							return nil
+						},
+						ErrFn: func() error { return nil },
+					}, nil
+				},
 			},
 			pgdumpFn: func(po pglib.PGDumpOptions) ([]byte, error) {
 				return testDump, nil
@@ -287,7 +500,9 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 			}
 
 			err := sg.CreateSnapshot(context.Background(), tc.snapshot)
-			require.Equal(t, err, tc.wantErr)
+			if !errors.Is(err, tc.wantErr) {
+				require.Equal(t, err, tc.wantErr)
+			}
 			sg.Close()
 		})
 	}