Repository: xataio/pgstream
Pull request #356 (2025-05-26 15:22:32)

[Snapshots] Use exclude schemas on schema `pg_dump`
This PR updates the logic behind the `pg_dump` to use the `exclude-schema` option instead of the schema filter, so that any objects that do not belong to any particular schema (such as extensions) are snapshotted correctly. Otherwise, if the schema filter is provided, the `pg_dump` will only contain schema linked objects, providing an incomplete view of the source database, and potentially leading to issues on the target due to missing objects/dependencies.

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

| Impacted Packages | Coverage Δ | :robot: |
|-------------------|------------|---------|
| github.com/xataio/pgstream/internal/postgres | 5.33% (**-0.05%**) | :thumbsdown: |
| github.com/xataio/pgstream/pkg/snapshot/generator/postgres/schema/pgdumprestore | 76.42% (**+1.18%**) | :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% (ø) | 28 (+2) | 0 | 28 (+2) |  |
| github.com/xataio/pgstream/pkg/snapshot/generator/postgres/schema/pgdumprestore/snapshot_pg_dump_restore_generator.go | 76.42% (**+1.18%**) | 123 (+22) | 94 (+18) | 29 (+4) | :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/pkg/snapshot/generator/postgres/schema/pgdumprestore/snapshot_pg_dump_restore_generator_test.go

</details>


Patch:
diff --git a/cmd/config/config_env.go b/cmd/config/config_env.go
index cbb0229..99151f9 100644
--- a/cmd/config/config_env.go
+++ b/cmd/config/config_env.go
@@ -61,6 +61,7 @@ func init() {
 	viper.BindEnv("PGSTREAM_POSTGRES_WRITER_DISABLE_TRIGGERS")
 	viper.BindEnv("PGSTREAM_POSTGRES_WRITER_ON_CONFLICT_ACTION")
 	viper.BindEnv("PGSTREAM_POSTGRES_WRITER_BULK_INGEST_ENABLED")
+	viper.BindEnv("PGSTREAM_POSTGRES_SNAPSHOT_INCLUDE_GLOBAL_DB_OBJECTS")
 
 	viper.BindEnv("PGSTREAM_KAFKA_READER_SERVERS")
 	viper.BindEnv("PGSTREAM_KAFKA_WRITER_SERVERS")
@@ -221,9 +222,10 @@ func parseSchemaSnapshotConfig(pgurl string) snapshotbuilder.SchemaSnapshotConfi
 	if pgTargetURL != "" && !useSchemaLog {
 		return snapshotbuilder.SchemaSnapshotConfig{
 			DumpRestore: &pgdumprestore.Config{
-				SourcePGURL:   pgurl,
-				TargetPGURL:   pgTargetURL,
-				CleanTargetDB: viper.GetBool("PGSTREAM_POSTGRES_SNAPSHOT_CLEAN_TARGET_DB"),
+				SourcePGURL:            pgurl,
+				TargetPGURL:            pgTargetURL,
+				CleanTargetDB:          viper.GetBool("PGSTREAM_POSTGRES_SNAPSHOT_CLEAN_TARGET_DB"),
+				IncludeGlobalDBObjects: viper.GetBool("PGSTREAM_POSTGRES_SNAPSHOT_INCLUDE_GLOBAL_DB_OBJECTS"),
 			},
 		}
 	}
diff --git a/cmd/config/config_yaml.go b/cmd/config/config_yaml.go
index 2192368..b476543 100644
--- a/cmd/config/config_yaml.go
+++ b/cmd/config/config_yaml.go
@@ -97,7 +97,8 @@ type SnapshotSchemaConfig struct {
 }
 
 type PgDumpPgRestoreConfig struct {
-	CleanTargetDB bool `mapstructure:"clean_target_db" yaml:"clean_target_db"`
+	CleanTargetDB          bool `mapstructure:"clean_target_db" yaml:"clean_target_db"`
+	IncludeGlobalDBObjects bool `mapstructure:"include_global_db_objects" yaml:"include_global_db_objects"`
 }
 
 type ReplicationConfig struct {
@@ -497,6 +498,7 @@ func (c *YAMLConfig) parseSchemaSnapshotConfig() (snapshotbuilder.SchemaSnapshot
 
 		if schemaSnapshotCfg.PgDumpPgRestore != nil {
 			streamSchemaCfg.DumpRestore.CleanTargetDB = schemaSnapshotCfg.PgDumpPgRestore.CleanTargetDB
+			streamSchemaCfg.DumpRestore.IncludeGlobalDBObjects = schemaSnapshotCfg.PgDumpPgRestore.IncludeGlobalDBObjects
 		}
 
 		return streamSchemaCfg, nil
diff --git a/docs/README.md b/docs/README.md
index 0e24429..4418589 100644
--- a/docs/README.md
+++ b/docs/README.md
@@ -111,6 +111,7 @@ source:
         mode: pgdump_pgrestore # options are pgdump_pgrestore or schemalog
         pgdump_pgrestore:
           clean_target_db: true # whether to clean the target database before restoring
+		  include_global_db_objects: true # whether to include database global objects, such as extensions or triggers, on the schema snapshot
     replication: # when mode is replication or snapshot_and_replication
       replication_slot: "pgstream_mydatabase_slot"
   kafka:
@@ -246,16 +247,17 @@ Here's a list of all the environment variables that can be used to configure the
 <details>
   ##### <summary>Postgres Snapshoter</summary>
 
-| Environment Variable                       | Default | Required | Description                                                                                                                                                                                                                                                                                                  |
-| ------------------------------------------ | ------- | -------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ |
-| PGSTREAM_POSTGRES_SNAPSHOT_LISTENER_URL    | N/A     | Yes      | URL of the Postgres database to snapshot.                                                                                                                                                                                                                                                                    |
-| PGSTREAM_POSTGRES_SNAPSHOT_TABLES          | ""      | No       | Tables for which there will be an initial snapshot generated. The syntax supports wildcards. Tables without a schema defined will be applied the public schema. Example: for `public.test_table` and all tables in the `test_schema` schema, the value would be the following: `"test_table test_schema.\*"` |
-| PGSTREAM_POSTGRES_SNAPSHOT_SCHEMA_WORKERS  | 4       | No       | Number of tables per schema that will be processed in parallel by the snapshotting process.                                                                                                                                                                                                                  |
-| 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 targets.                                                                                                                                                                            |
-| PGSTREAM_POSTGRES_SNAPSHOT_CLEAN_TARGET_DB | False   | No       | When using `pg_dump`/`pg_restore` to snapshot schema for Postgres targets, option to issue commands to DROP all the objects that will be restored.                                                                                                                                                           |
+| Environment Variable                                 | Default | Required | Description                                                                                                                                                                                                                                                                                                  |
+| ---------------------------------------------------- | ------- | -------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ |
+| PGSTREAM_POSTGRES_SNAPSHOT_LISTENER_URL              | N/A     | Yes      | URL of the Postgres database to snapshot.                                                                                                                                                                                                                                                                    |
+| PGSTREAM_POSTGRES_SNAPSHOT_TABLES                    | ""      | No       | Tables for which there will be an initial snapshot generated. The syntax supports wildcards. Tables without a schema defined will be applied the public schema. Example: for `public.test_table` and all tables in the `test_schema` schema, the value would be the following: `"test_table test_schema.\*"` |
+| PGSTREAM_POSTGRES_SNAPSHOT_SCHEMA_WORKERS            | 4       | No       | Number of tables per schema that will be processed in parallel by the snapshotting process.                                                                                                                                                                                                                  |
+| 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 targets.                                                                                                                                                                            |
+| PGSTREAM_POSTGRES_SNAPSHOT_CLEAN_TARGET_DB           | False   | No       | When using `pg_dump`/`pg_restore` to snapshot schema for Postgres targets, option to issue commands to DROP all the objects that will be restored.                                                                                                                                                           |
+| PGSTREAM_POSTGRES_SNAPSHOT_INCLUDE_GLOBAL_DB_OBJECTS | False   | No       | When using `pg_dump`/`pg_restore` to snapshot schema for Postgres targets, option to snapshot all global database objects outside of the selected schema (such as extensions, triggers, etc).                                                                                                                |
 
 </details>
 
diff --git a/internal/postgres/pg_dump.go b/internal/postgres/pg_dump.go
index 72c74e1..e3dc2f2 100644
--- a/internal/postgres/pg_dump.go
+++ b/internal/postgres/pg_dump.go
@@ -18,6 +18,8 @@ type PGDumpOptions struct {
 	Format string
 	// Schemas to export
 	Schemas []string
+	// Schemas to be excluded from the dump, regex patterns supported
+	ExcludeSchemas []string
 	// Tables to export
 	Tables []string
 	// Tables to be excluded from the dump, regex patterns supported
@@ -53,6 +55,10 @@ func (opts *PGDumpOptions) ToArgs() []string {
 		options = append(options, fmt.Sprintf("--schema=%v", schema))
 	}
 
+	for _, schema := range opts.ExcludeSchemas {
+		options = append(options, fmt.Sprintf("--exclude-schema=%v", schema))
+	}
+
 	for _, table := range opts.Tables {
 		options = append(options, fmt.Sprintf("--table=%v", table))
 	}
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 51aff65..7f40e44 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
@@ -25,21 +25,25 @@ import (
 // SnapshotGenerator generates postgres schema snapshots using pg_dump and
 // pg_restore
 type SnapshotGenerator struct {
-	sourceURL      string
-	targetURL      string
-	pgDumpFn       pglib.PGDumpFn
-	pgRestoreFn    pglib.PGRestoreFn
-	schemalogStore schemalog.Store
-	connBuilder    pglib.QuerierBuilder
-	cleanTargetDB  bool
-	logger         loglib.Logger
-	generator      generator.SnapshotGenerator
+	sourceURL              string
+	targetURL              string
+	pgDumpFn               pglib.PGDumpFn
+	pgRestoreFn            pglib.PGRestoreFn
+	schemalogStore         schemalog.Store
+	connBuilder            pglib.QuerierBuilder
+	cleanTargetDB          bool
+	includeGlobalDBObjects bool
+	logger                 loglib.Logger
+	generator              generator.SnapshotGenerator
 }
 
 type Config struct {
 	SourcePGURL   string
 	TargetPGURL   string
 	CleanTargetDB bool
+	// if set to true the snapshot will include all database objects, not tied
+	// to any particular schema, such as extensions or triggers.
+	IncludeGlobalDBObjects bool
 }
 
 type Option func(s *SnapshotGenerator)
@@ -59,13 +63,14 @@ const (
 // 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) {
 	sg := &SnapshotGenerator{
-		sourceURL:     c.SourcePGURL,
-		targetURL:     c.TargetPGURL,
-		pgDumpFn:      pglib.RunPGDump,
-		pgRestoreFn:   pglib.RunPGRestore,
-		connBuilder:   pglib.ConnBuilder,
-		cleanTargetDB: c.CleanTargetDB,
-		logger:        loglib.NewNoopLogger(),
+		sourceURL:              c.SourcePGURL,
+		targetURL:              c.TargetPGURL,
+		pgDumpFn:               pglib.RunPGDump,
+		pgRestoreFn:            pglib.RunPGRestore,
+		connBuilder:            pglib.ConnBuilder,
+		cleanTargetDB:          c.CleanTargetDB,
+		includeGlobalDBObjects: c.IncludeGlobalDBObjects,
+		logger:                 loglib.NewNoopLogger(),
 	}
 
 	if err := sg.initialiseSchemaLogStore(ctx); err != nil {
@@ -240,10 +245,23 @@ func (s *SnapshotGenerator) pgdumpOptions(ctx context.Context, ss *snapshot.Snap
 		ConnectionString: s.sourceURL,
 		Format:           "p",
 		SchemaOnly:       true,
-		Schemas:          []string{pglib.QuoteIdentifier(ss.SchemaName)},
+		Schemas:          []string{ss.SchemaName},
 		Clean:            s.cleanTargetDB,
 	}
 
+	if s.includeGlobalDBObjects {
+		// instead of using the schema filter, we use the exclude schemas filter
+		// to make sure extensions and other database global objects are
+		// created. pg_dump will not include them when using the schema filter,
+		// since they do not belong to the schema.
+		var err error
+		opts.ExcludeSchemas, err = s.pgdumpExcludedSchemas(ctx, ss.SchemaName)
+		if err != nil {
+			return nil, err
+		}
+		opts.Schemas = nil
+	}
+
 	// wildcard means all tables in the schema, so no table filter required
 	if slices.Contains(ss.TableNames, wildcard) {
 		return opts, nil
@@ -302,6 +320,37 @@ func (s *SnapshotGenerator) pgdumpExcludedTables(ctx context.Context, schemaName
 	return excludeTables, nil
 }
 
+func (s *SnapshotGenerator) pgdumpExcludedSchemas(ctx context.Context, schemaName string) ([]string, error) {
+	conn, err := s.connBuilder(ctx, s.sourceURL)
+	if err != nil {
+		return nil, err
+	}
+	defer conn.Close(ctx)
+
+	// get all schemas in the database that are not in the snapshot request
+	const query = "SELECT schema_name FROM information_schema.schemata WHERE schema_name != $1"
+	rows, err := conn.Query(ctx, query, schemaName)
+	if err != nil {
+		return nil, fmt.Errorf("retrieving tables from schema: %w", err)
+	}
+	defer rows.Close()
+
+	excludeSchemas := []string{}
+	for rows.Next() {
+		schema := ""
+		if err := rows.Scan(&schema); err != nil {
+			return nil, fmt.Errorf("scanning schema name: %w", err)
+		}
+		excludeSchemas = append(excludeSchemas, pglib.QuoteIdentifier(schema))
+	}
+
+	if err := rows.Err(); err != nil {
+		return nil, err
+	}
+
+	return excludeSchemas, nil
+}
+
 func (s *SnapshotGenerator) initialiseSchemaLogStore(ctx context.Context) error {
 	exists, err := s.schemalogExists(ctx)
 	if err != nil {


Test patch:
diff --git a/cmd/config/config_env_test.go b/cmd/config/config_env_test.go
index 53af481..a9c4a0f 100644
--- a/cmd/config/config_env_test.go
+++ b/cmd/config/config_env_test.go
@@ -41,6 +41,7 @@ func Test_EnvVarsToStreamConfig(t *testing.T) {
 	os.Setenv("PGSTREAM_POSTGRES_SNAPSHOT_TABLE_WORKERS", "4")
 	os.Setenv("PGSTREAM_POSTGRES_SNAPSHOT_BATCH_PAGE_SIZE", "1000")
 	os.Setenv("PGSTREAM_POSTGRES_SNAPSHOT_CLEAN_TARGET_DB", "true")
+	os.Setenv("PGSTREAM_POSTGRES_SNAPSHOT_INCLUDE_GLOBAL_DB_OBJECTS", "true")
 
 	os.Setenv("PGSTREAM_KAFKA_READER_SERVERS", "localhost:9092")
 	os.Setenv("PGSTREAM_KAFKA_TOPIC_NAME", "mytopic")
diff --git a/cmd/config/helper_test.go b/cmd/config/helper_test.go
index bd174be..2d00449 100644
--- a/cmd/config/helper_test.go
+++ b/cmd/config/helper_test.go
@@ -26,6 +26,7 @@ func validateTestStreamConfig(t *testing.T, streamConfig *stream.Config) {
 	assert.Equal(t, uint(1000), streamConfig.Listener.Postgres.Snapshot.Generator.BatchPageSize)
 	assert.NotNil(t, streamConfig.Listener.Postgres.Snapshot.Schema.DumpRestore)
 	assert.True(t, streamConfig.Listener.Postgres.Snapshot.Schema.DumpRestore.CleanTargetDB)
+	assert.True(t, streamConfig.Listener.Postgres.Snapshot.Schema.DumpRestore.IncludeGlobalDBObjects)
 	assert.NotNil(t, streamConfig.Listener.Postgres.Snapshot.Recorder)
 	assert.Equal(t, "postgresql://user:password@localhost:5432/mytargetdatabase", streamConfig.Listener.Postgres.Snapshot.Recorder.SnapshotStoreURL)
 	assert.True(t, streamConfig.Listener.Postgres.Snapshot.Recorder.RepeatableSnapshots)
diff --git a/cmd/config/test/test_config.env b/cmd/config/test/test_config.env
index e52dac5..99f471d 100644
--- a/cmd/config/test/test_config.env
+++ b/cmd/config/test/test_config.env
@@ -11,6 +11,7 @@ PGSTREAM_POSTGRES_SNAPSHOT_SCHEMA_WORKERS=4
 PGSTREAM_POSTGRES_SNAPSHOT_TABLE_WORKERS=4
 PGSTREAM_POSTGRES_SNAPSHOT_BATCH_PAGE_SIZE=1000
 PGSTREAM_POSTGRES_SNAPSHOT_CLEAN_TARGET_DB=true
+PGSTREAM_POSTGRES_SNAPSHOT_INCLUDE_GLOBAL_DB_OBJECTS=true
 
 # Kafka
 PGSTREAM_KAFKA_READER_SERVERS="localhost:9092"
diff --git a/cmd/config/test/test_config.yaml b/cmd/config/test/test_config.yaml
index 5c6b70c..9102b1d 100644
--- a/cmd/config/test/test_config.yaml
+++ b/cmd/config/test/test_config.yaml
@@ -17,6 +17,7 @@ source:
         mode: pgdump_pgrestore # options are pgdump_pgrestore or schemalog
         pgdump_pgrestore:
           clean_target_db: true # whether to clean the target database before restoring
+          include_global_db_objects: true # whether to include database global objects, such as extensions or triggers, on the schema snapshot
     replication: # when mode is replication or snapshot_and_replication
       replication_slot: "pgstream_mydatabase_slot"
   kafka:
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 484f829..8f3930b 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
@@ -26,8 +26,52 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 	testSchema := "test_schema"
 	testTable := "test_table"
 	excludedTable := "excluded_test_table"
+	excludedSchema := "excluded_test_schema"
 	errTest := errors.New("oh noes")
 
+	validQuerier := func() *mocks.Querier {
+		return &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) {
+				switch query {
+				case "SELECT schema_name FROM information_schema.schemata WHERE schema_name != $1":
+					require.Equal(t, []any{testSchema}, args)
+					return &mocks.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 = excludedSchema
+							return nil
+						},
+						ErrFn: func() error { return nil },
+					}, nil
+				case "SELECT tablename FROM pg_tables WHERE schemaname = 'test_schema' AND tablename NOT IN ($1)":
+					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
+				default:
+					return nil, fmt.Errorf("unexpected query: %s", query)
+				}
+			},
+		}
+	}
+
 	tests := []struct {
 		name           string
 		snapshot       *snapshot.Snapshot
@@ -45,35 +89,14 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 				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)
-					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
-				},
-			},
+			conn: validQuerier(),
 			pgdumpFn: func(_ context.Context, po pglib.PGDumpOptions) ([]byte, error) {
 				require.Equal(t, pglib.PGDumpOptions{
 					ConnectionString: "source-url",
 					Format:           "p",
 					Clean:            false,
 					SchemaOnly:       true,
-					Schemas:          []string{pglib.QuoteIdentifier(testSchema)},
+					ExcludeSchemas:   []string{pglib.QuoteIdentifier(excludedSchema)},
 					ExcludeTables:    []string{pglib.QuoteIdentifier(excludedTable)},
 				}, po)
 				return testDump, nil
@@ -96,18 +119,13 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 				SchemaName: testSchema,
 				TableNames: []string{"*"},
 			},
-			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
-				},
-			},
+			conn: validQuerier(),
 			pgdumpFn: func(_ context.Context, po pglib.PGDumpOptions) ([]byte, error) {
 				require.Equal(t, pglib.PGDumpOptions{
 					ConnectionString: "source-url",
 					Format:           "p",
 					SchemaOnly:       true,
-					Schemas:          []string{pglib.QuoteIdentifier(testSchema)},
+					ExcludeSchemas:   []string{pglib.QuoteIdentifier(excludedSchema)},
 				}, po)
 				return testDump, nil
 			},
@@ -144,7 +162,7 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 			wantErr: nil,
 		},
 		{
-			name: "error - querying excluded tables",
+			name: "error - querying excluded schemas",
 			snapshot: &snapshot.Snapshot{
 				SchemaName: testSchema,
 				TableNames: []string{testTable},
@@ -171,15 +189,34 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 			},
 			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
+					switch query {
+					case "SELECT schema_name FROM information_schema.schemata WHERE schema_name != $1":
+						require.Equal(t, []any{testSchema}, args)
+						return &mocks.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 = excludedSchema
+								return nil
+							},
+							ErrFn: func() error { return nil },
+						}, nil
+					case "SELECT tablename FROM pg_tables WHERE schemaname = 'test_schema' AND tablename NOT IN ($1)":
+						require.Equal(t, []any{testTable}, args)
+						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
+					default:
+						return nil, fmt.Errorf("unexpected query: %s", query)
+					}
 				},
 			},
 			pgdumpFn: func(_ context.Context, po pglib.PGDumpOptions) ([]byte, error) {
@@ -203,19 +240,38 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 					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
+					switch query {
+					case "SELECT schema_name FROM information_schema.schemata WHERE schema_name != $1":
+						require.Equal(t, []any{testSchema}, args)
+						return &mocks.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 = excludedSchema
+								return nil
+							},
+							ErrFn: func() error { return nil },
+						}, nil
+					case "SELECT tablename FROM pg_tables WHERE schemaname = 'test_schema' AND tablename NOT IN ($1)":
+						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 errTest },
+						}, nil
+					default:
+						return nil, fmt.Errorf("unexpected query: %s", query)
+					}
 				},
 			},
 			pgdumpFn: func(_ context.Context, po pglib.PGDumpOptions) ([]byte, error) {
@@ -233,26 +289,7 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 				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")
-				},
-				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
-				},
-			},
+			conn: validQuerier(),
 			pgdumpFn: func(_ context.Context, po pglib.PGDumpOptions) ([]byte, error) {
 				return nil, errTest
 			},
@@ -265,29 +302,10 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 		{
 			name: "error - performing pgrestore",
 			snapshot: &snapshot.Snapshot{
-				SchemaName: publicSchema,
+				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")
-				},
-				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
-				},
-			},
+			conn: validQuerier(),
 			pgdumpFn: func(_ context.Context, po pglib.PGDumpOptions) ([]byte, error) {
 				return testDump, nil
 			},
@@ -300,29 +318,10 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 		{
 			name: "error - performing pgrestore output critical error",
 			snapshot: &snapshot.Snapshot{
-				SchemaName: publicSchema,
+				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")
-				},
-				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
-				},
-			},
+			conn: validQuerier(),
 			pgdumpFn: func(_ context.Context, po pglib.PGDumpOptions) ([]byte, error) {
 				return testDump, nil
 			},
@@ -335,29 +334,10 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 		{
 			name: "error - performing pgrestore output ignored errors",
 			snapshot: &snapshot.Snapshot{
-				SchemaName: publicSchema,
+				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")
-				},
-				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
-				},
-			},
+			conn: validQuerier(),
 			pgdumpFn: func(_ context.Context, po pglib.PGDumpOptions) ([]byte, error) {
 				return testDump, nil
 			},
@@ -377,23 +357,7 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 				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
+				return validQuerier(), nil
 			},
 			pgdumpFn: func(_ context.Context, po pglib.PGDumpOptions) ([]byte, error) {
 				return testDump, nil
@@ -410,26 +374,13 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 				SchemaName: testSchema,
 				TableNames: []string{testTable},
 			},
-			conn: &mocks.Querier{
-				ExecFn: func(ctx context.Context, i uint, query string, args ...any) (pglib.CommandTag, error) {
+			conn: func() *mocks.Querier {
+				c := validQuerier()
+				c.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
-				},
-			},
+				}
+				return c
+			}(),
 			pgdumpFn: func(_ context.Context, po pglib.PGDumpOptions) ([]byte, error) {
 				return testDump, nil
 			},
@@ -445,27 +396,7 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 				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 nil },
-					}, nil
-				},
-			},
+			conn: validQuerier(),
 			pgdumpFn: func(_ context.Context, po pglib.PGDumpOptions) ([]byte, error) {
 				return testDump, nil
 			},
@@ -488,13 +419,14 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 			t.Parallel()
 
 			sg := SnapshotGenerator{
-				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(),
+				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(),
+				includeGlobalDBObjects: true,
 			}
 
 			if tc.connBuilder != nil {
@@ -503,7 +435,7 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 
 			err := sg.CreateSnapshot(context.Background(), tc.snapshot)
 			if !errors.Is(err, tc.wantErr) {
-				require.Equal(t, err, tc.wantErr)
+				require.Equal(t, tc.wantErr, err)
 			}
 			sg.Close()
 		})