Repository: xataio/pgstream
Pull request #362 (2025-05-27 11:58:54)

[Snapshots] Create indices/constraints after data dump
This PR updates the snapshot generator to restore the indices/constraints after the data snapshot, improving the performance of the bulk inserts, since they don't have to be checked/maintained on every insert. 

The snapshot aggregator is replaced by a schema generator decorator instead. The schema generator is configurable so that it can be created with or without data snapshotting. Currently both schema and data need to be snapshotted, but this builds the groundwork for the support for optional schema/data snapshots (https://github.com/xataio/pgstream/issues/231), which will be tackled in a separate PR.



Patch:
diff --git a/pkg/snapshot/generator/mocks/mock_snapshot_generator.go b/pkg/snapshot/generator/mocks/mock_snapshot_generator.go
index a012d80..fa78a5b 100644
--- a/pkg/snapshot/generator/mocks/mock_snapshot_generator.go
+++ b/pkg/snapshot/generator/mocks/mock_snapshot_generator.go
@@ -9,11 +9,13 @@ import (
 )
 
 type Generator struct {
-	CreateSnapshotFn func(ctx context.Context, snapshot *snapshot.Snapshot) error
-	CloseFn          func() error
+	CreateSnapshotFn    func(ctx context.Context, snapshot *snapshot.Snapshot) error
+	CloseFn             func() error
+	createSnapshotCalls uint
 }
 
 func (m *Generator) CreateSnapshot(ctx context.Context, snapshot *snapshot.Snapshot) error {
+	m.createSnapshotCalls++
 	return m.CreateSnapshotFn(ctx, snapshot)
 }
 
@@ -23,3 +25,7 @@ func (m *Generator) Close() error {
 	}
 	return nil
 }
+
+func (m *Generator) CreateSnapshotCalls() uint {
+	return m.createSnapshotCalls
+}
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 1a30080..51aff65 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
@@ -3,6 +3,8 @@
 package pgdumprestore
 
 import (
+	"bufio"
+	"bytes"
 	"context"
 	"errors"
 	"fmt"
@@ -17,6 +19,7 @@ import (
 	schemaloginstrumentation "github.com/xataio/pgstream/pkg/schemalog/instrumentation"
 	schemalogpg "github.com/xataio/pgstream/pkg/schemalog/postgres"
 	"github.com/xataio/pgstream/pkg/snapshot"
+	"github.com/xataio/pgstream/pkg/snapshot/generator"
 )
 
 // SnapshotGenerator generates postgres schema snapshots using pg_dump and
@@ -30,6 +33,7 @@ type SnapshotGenerator struct {
 	connBuilder    pglib.QuerierBuilder
 	cleanTargetDB  bool
 	logger         loglib.Logger
+	generator      generator.SnapshotGenerator
 }
 
 type Config struct {
@@ -40,6 +44,12 @@ type Config struct {
 
 type Option func(s *SnapshotGenerator)
 
+type dump struct {
+	full                  []byte
+	filtered              []byte
+	indicesAndConstraints []byte
+}
+
 const (
 	publicSchema = "public"
 	wildcard     = "*"
@@ -77,6 +87,12 @@ func WithLogger(logger loglib.Logger) Option {
 	}
 }
 
+func WithSnapshotGenerator(g generator.SnapshotGenerator) Option {
+	return func(sg *SnapshotGenerator) {
+		sg.generator = g
+	}
+}
+
 func WithInstrumentation(i *otel.Instrumentation) Option {
 	return func(sg *SnapshotGenerator) {
 		var err error
@@ -103,16 +119,66 @@ func (s *SnapshotGenerator) CreateSnapshot(ctx context.Context, ss *snapshot.Sna
 		return nil
 	}
 
+	dump, err := s.pgdump(ctx, ss)
+	if err != nil {
+		return err
+	}
+
+	// if there's no further snapshotting happening, we can apply the full dump,
+	// no need to apply the constraints/indices separately.
+	if s.generator == nil {
+		return s.pgrestore(ctx, ss, dump.full)
+	}
+
+	// otherwise, we need to apply the filtered schema dump first, then call the
+	// wrapped snapshot generator, and apply the indices and constraints last.
+	// This will make the data snapshot faster, since there will be no
+	// constraints to be updated/checked on each insert.
+	if err := s.pgrestore(ctx, ss, dump.filtered); err != nil {
+		return err
+	}
+
+	if err := s.generator.CreateSnapshot(ctx, ss); err != nil {
+		return err
+	}
+
+	s.logger.Info("restoring schema indices and constraints", loglib.Fields{"schema": ss.SchemaName, "tables": ss.TableNames})
+	// apply the indices and constraints when the wrapped generator has finished
+	return s.pgrestore(ctx, ss, dump.indicesAndConstraints)
+}
+
+func (s *SnapshotGenerator) Close() error {
+	if s.schemalogStore != nil {
+		return s.schemalogStore.Close()
+	}
+
+	if s.generator != nil {
+		return s.generator.Close()
+	}
+
+	return nil
+}
+
+func (s *SnapshotGenerator) pgdump(ctx context.Context, ss *snapshot.Snapshot) (*dump, error) {
 	pgdumpOpts, err := s.pgdumpOptions(ctx, ss)
 	if err != nil {
-		return fmt.Errorf("preparing pg_dump options: %w", err)
+		return nil, fmt.Errorf("preparing pg_dump options: %w", err)
 	}
 
-	dump, err := s.pgDumpFn(ctx, *pgdumpOpts)
+	d, err := s.pgDumpFn(ctx, *pgdumpOpts)
 	if err != nil {
-		return err
+		return nil, err
 	}
 
+	filtered, constraints := s.extractIndicesAndConstraints(d)
+	return &dump{
+		full:                  d,
+		filtered:              filtered,
+		indicesAndConstraints: constraints,
+	}, nil
+}
+
+func (s *SnapshotGenerator) pgrestore(ctx context.Context, ss *snapshot.Snapshot, dump []byte) error {
 	// 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)
@@ -122,7 +188,7 @@ func (s *SnapshotGenerator) CreateSnapshot(ctx context.Context, ss *snapshot.Sna
 		}
 	}
 
-	_, err = s.pgRestoreFn(ctx, s.pgrestoreOptions(), dump)
+	_, err := s.pgRestoreFn(ctx, s.pgrestoreOptions(), dump)
 	pgrestoreErr := &pglib.PGRestoreErrors{}
 	if err != nil {
 		switch {
@@ -149,13 +215,6 @@ func (s *SnapshotGenerator) CreateSnapshot(ctx context.Context, ss *snapshot.Sna
 	return nil
 }
 
-func (s *SnapshotGenerator) Close() error {
-	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 {
@@ -275,3 +334,46 @@ func (s *SnapshotGenerator) schemalogExists(ctx context.Context) (bool, error) {
 
 	return exists, nil
 }
+
+func (s *SnapshotGenerator) extractIndicesAndConstraints(d []byte) ([]byte, []byte) {
+	scanner := bufio.NewScanner(bytes.NewReader(d))
+	scanner.Split(bufio.ScanLines)
+	indicesAndConstraints := strings.Builder{}
+	filteredDump := strings.Builder{}
+	alterTable := ""
+	for scanner.Scan() {
+		line := scanner.Text()
+		switch {
+		case alterTable != "":
+			// check if the previous alter table line is split in two lines and matches a constraint
+			if strings.Contains(line, "ADD CONSTRAINT") {
+				indicesAndConstraints.WriteString(alterTable)
+				indicesAndConstraints.WriteString("\n")
+				indicesAndConstraints.WriteString(line)
+				indicesAndConstraints.WriteString("\n\n")
+			} else {
+				filteredDump.WriteString(alterTable)
+				filteredDump.WriteString("\n")
+				filteredDump.WriteString(line)
+				filteredDump.WriteString("\n")
+			}
+			alterTable = ""
+		case strings.HasPrefix(line, "CREATE INDEX"),
+			strings.HasPrefix(line, "CREATE UNIQUE INDEX"),
+			strings.HasPrefix(line, "CREATE CONSTRAINT"),
+			strings.HasPrefix(line, "CREATE TRIGGER"):
+			indicesAndConstraints.WriteString(line)
+			indicesAndConstraints.WriteString("\n\n")
+		case strings.HasPrefix(line, "ALTER TABLE") && strings.Contains(line, "ADD CONSTRAINT"):
+			indicesAndConstraints.WriteString(line)
+		case strings.HasPrefix(line, "ALTER TABLE") && !strings.HasSuffix(line, ";"):
+			// keep it in case the alter table is provided in two lines (pg_dump format)
+			alterTable = line
+		default:
+			filteredDump.WriteString(line)
+			filteredDump.WriteString("\n")
+		}
+	}
+
+	return []byte(filteredDump.String()), []byte(indicesAndConstraints.String())
+}
diff --git a/pkg/snapshot/generator/postgres/schema/schemalog/snapshot_schemalog_generator.go b/pkg/snapshot/generator/postgres/schema/schemalog/snapshot_schemalog_generator.go
index 4a31f47..da9f247 100644
--- a/pkg/snapshot/generator/postgres/schema/schemalog/snapshot_schemalog_generator.go
+++ b/pkg/snapshot/generator/postgres/schema/schemalog/snapshot_schemalog_generator.go
@@ -9,20 +9,36 @@ import (
 	"github.com/xataio/pgstream/internal/json"
 	"github.com/xataio/pgstream/pkg/schemalog"
 	"github.com/xataio/pgstream/pkg/snapshot"
+	"github.com/xataio/pgstream/pkg/snapshot/generator"
 )
 
 type SnapshotGenerator struct {
 	schemalogStore schemalog.Store
 	marshaler      func(any) ([]byte, error)
 	processRow     snapshot.RowProcessor
+	generator      generator.SnapshotGenerator
 }
 
-func NewSnapshotGenerator(schemalogStore schemalog.Store, processRow snapshot.RowProcessor) *SnapshotGenerator {
-	return &SnapshotGenerator{
+type Option func(*SnapshotGenerator)
+
+func NewSnapshotGenerator(schemalogStore schemalog.Store, processRow snapshot.RowProcessor, opts ...Option) *SnapshotGenerator {
+	sg := &SnapshotGenerator{
 		schemalogStore: schemalogStore,
 		processRow:     processRow,
 		marshaler:      json.Marshal,
 	}
+
+	for _, opt := range opts {
+		opt(sg)
+	}
+
+	return sg
+}
+
+func WithSnapshotGenerator(g generator.SnapshotGenerator) Option {
+	return func(sg *SnapshotGenerator) {
+		sg.generator = g
+	}
 }
 
 func (s *SnapshotGenerator) CreateSnapshot(ctx context.Context, ss *snapshot.Snapshot) error {
@@ -43,10 +59,17 @@ func (s *SnapshotGenerator) CreateSnapshot(ctx context.Context, ss *snapshot.Sna
 		return snapshot.NewErrors(err)
 	}
 
+	if s.generator != nil {
+		return s.generator.CreateSnapshot(ctx, ss)
+	}
+
 	return nil
 }
 
 func (s *SnapshotGenerator) Close() error {
+	if s.generator != nil {
+		return s.generator.Close()
+	}
 	return nil
 }
 
diff --git a/pkg/snapshot/generator/snapshot_generator_aggregator.go b/pkg/snapshot/generator/snapshot_generator_aggregator.go
deleted file mode 100644
index c0b55f9..0000000
--- a/pkg/snapshot/generator/snapshot_generator_aggregator.go
+++ /dev/null
@@ -1,50 +0,0 @@
-// SPDX-License-Identifier: Apache-2.0
-
-package generator
-
-import (
-	"context"
-	"errors"
-
-	"github.com/xataio/pgstream/pkg/snapshot"
-)
-
-// GeneratorAggregator aggregates snapshot generators
-type GeneratorAggregator struct {
-	generators []SnapshotGenerator
-}
-
-// NewAggregator will aggregate the generators on input.
-func NewAggregator(generators []SnapshotGenerator) *GeneratorAggregator {
-	return &GeneratorAggregator{
-		generators: generators,
-	}
-}
-
-// CreateSnapshot will call the CreateSnapshot method for each of the aggregated
-// snapshot generators sequentially, in the order in which they were provided.
-// If one of them fails, an error will be returned without continuing with the
-// rest of generator calls.
-func (a *GeneratorAggregator) CreateSnapshot(ctx context.Context, ss *snapshot.Snapshot) error {
-	for _, gen := range a.generators {
-		if err := gen.CreateSnapshot(ctx, ss); err != nil {
-			return err
-		}
-	}
-	return nil
-}
-
-// Close will aggregate the errors of the generator close calls.
-func (a *GeneratorAggregator) Close() error {
-	var closeErrs error
-	for _, gen := range a.generators {
-		if err := gen.Close(); err != nil {
-			if closeErrs != nil {
-				closeErrs = errors.Join(closeErrs, err)
-			} else {
-				closeErrs = err
-			}
-		}
-	}
-	return closeErrs
-}
diff --git a/pkg/wal/listener/snapshot/builder/wal_listener_snapshot_generator_builder.go b/pkg/wal/listener/snapshot/builder/wal_listener_snapshot_generator_builder.go
index 0575ad1..835cc4f 100644
--- a/pkg/wal/listener/snapshot/builder/wal_listener_snapshot_generator_builder.go
+++ b/pkg/wal/listener/snapshot/builder/wal_listener_snapshot_generator_builder.go
@@ -36,35 +36,32 @@ var errSchemaSnapshotNotConfigured = errors.New("no schema snapshot has been con
 //
 // ┌───────────────────────────────────────────────────┐
 // │               Snapshot Generator                  │
-// │ ┌─────────┐ ┌────────┐ ┌────────────────────────┐ │
-// │ │         │ │        │ │       Aggregator       │ │
-// │ │         │ │        │ │  ┌──────────────────┐  │ │
-// │ │         │ │        │ │  │ Schema Snapshot  │  │ │
-// │ │Snapshot │ │ Table  │ │  └─────────┬────────┘  │ │
-// │ │Activity │─▶ Finder ├▶│            │           │ │
-// │ │Recorder │ │        │ │  ┌─────────▼────────┐  │ │
-// │ │         │ │        │ │  │  Data Snapshot   │  │ │
-// │ │         │ │        │ │  └──────────────────┘  │ │
-// │ └─────────┘ └────────┘ └────────────────────────┘ │
+// │ ┌─────────┐  ┌─────────┐  ┌────────┐  ┌─────────┐ │
+// │ │         │  │         │  │        │  │         │ │
+// │ │         │  │         │  │        │  │         │ │
+// │ │         │  │         │  │        │  │         │ │
+// │ │Snapshot │  │ Schema  │  │ Table  │  │  Data   │ │
+// │ │Activity │─▶│Snapshot ├─▶│ Finder ├─▶│Snapshot │ │
+// │ │Recorder │  │         │  │        │  │         │ │
+// │ │         │  │         │  │        │  │         │ │
+// │ │         │  │         │  │        │  │         │ │
+// │ └─────────┘  └─────────┘  └────────┘  └─────────┘ │
 // └───────────────────────────────────────────────────┘
 
 func NewSnapshotGenerator(ctx context.Context, cfg *SnapshotListenerConfig, processEvent listener.ProcessWalEvent, logger loglib.Logger, instrumentation *otel.Instrumentation) (listenersnapshot.Generator, error) {
 	processEventAdapter := adapter.NewProcessEventAdapter(processEvent)
-	// postgres schema snapshot generator
-	schemaSnapshotGenerator, err := newSchemaSnapshotGenerator(ctx, &cfg.Schema, processEventAdapter.ProcessRow, logger, instrumentation)
-	if err != nil {
-		return nil, err
-	}
 
-	// postgres data snapshot generator
-	var dataSnapshotGenerator generator.SnapshotGenerator
+	var g generator.SnapshotGenerator
+	var err error
+
+	// postgres data snapshot generator layer
 	opts := []pgsnapshotgenerator.Option{
 		pgsnapshotgenerator.WithLogger(logger),
 	}
 	if instrumentation.IsEnabled() {
 		opts = append(opts, pgsnapshotgenerator.WithInstrumentation(instrumentation))
 	}
-	dataSnapshotGenerator, err = pgsnapshotgenerator.NewSnapshotGenerator(ctx, &cfg.Generator, processEventAdapter.ProcessRow, opts...)
+	g, err = pgsnapshotgenerator.NewSnapshotGenerator(ctx, &cfg.Generator, processEventAdapter.ProcessRow, opts...)
 	if err != nil {
 		return nil, err
 	}
@@ -73,14 +70,16 @@ func NewSnapshotGenerator(ctx context.Context, cfg *SnapshotListenerConfig, proc
 	if instrumentation.IsEnabled() {
 		finderOpts = append(finderOpts, pgtablefinder.WithInstrumentation(instrumentation))
 	}
-	dataSnapshotGenerator, err = pgtablefinder.NewSnapshotTableFinder(ctx, cfg.Generator.URL, dataSnapshotGenerator, finderOpts...)
+	g, err = pgtablefinder.NewSnapshotTableFinder(ctx, cfg.Generator.URL, g, finderOpts...)
 	if err != nil {
 		return nil, err
 	}
 
-	var g generator.SnapshotGenerator
-	// snapshot generator aggregator
-	g = generator.NewAggregator([]generator.SnapshotGenerator{schemaSnapshotGenerator, dataSnapshotGenerator})
+	// postgres schema snapshot generator layer
+	g, err = newSchemaSnapshotGenerator(ctx, &cfg.Schema, g, processEventAdapter.ProcessRow, logger, instrumentation)
+	if err != nil {
+		return nil, err
+	}
 
 	if cfg.Recorder != nil && cfg.Recorder.SnapshotStoreURL != "" {
 		// snapshot activity recorder layer
@@ -105,7 +104,7 @@ func NewSnapshotGenerator(ctx context.Context, cfg *SnapshotListenerConfig, proc
 	return adapter.NewSnapshotGeneratorAdapter(&cfg.Adapter, g, adapter.WithLogger(logger)), nil
 }
 
-func newSchemaSnapshotGenerator(ctx context.Context, cfg *SchemaSnapshotConfig, processRow snapshot.RowProcessor, logger loglib.Logger, instrumentation *otel.Instrumentation) (generator.SnapshotGenerator, error) {
+func newSchemaSnapshotGenerator(ctx context.Context, cfg *SchemaSnapshotConfig, g generator.SnapshotGenerator, processRow snapshot.RowProcessor, logger loglib.Logger, instrumentation *otel.Instrumentation) (generator.SnapshotGenerator, error) {
 	switch {
 	case cfg.SchemaLogStore != nil:
 		// postgres schemalog schema snapshot generator
@@ -119,11 +118,15 @@ func newSchemaSnapshotGenerator(ctx context.Context, cfg *SchemaSnapshotConfig,
 		if instrumentation.IsEnabled() {
 			schemaLogStore = schemaloginstrumentation.NewStore(schemaLogStore, instrumentation)
 		}
-		return schemalogsnapshotgenerator.NewSnapshotGenerator(schemaLogStore, processRow), nil
+		return schemalogsnapshotgenerator.NewSnapshotGenerator(
+			schemaLogStore,
+			processRow,
+			schemalogsnapshotgenerator.WithSnapshotGenerator(g)), nil
 	case cfg.DumpRestore != nil:
 		// postgres pgdump/pgrestore schema snapshot generator
 		opts := []pgdumprestoregenerator.Option{
 			pgdumprestoregenerator.WithLogger(logger),
+			pgdumprestoregenerator.WithSnapshotGenerator(g),
 		}
 		if instrumentation.IsEnabled() {
 			opts = append(opts, pgdumprestoregenerator.WithInstrumentation(instrumentation))


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 7513047..484f829 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
@@ -6,6 +6,8 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"os"
+	"strings"
 	"testing"
 
 	"github.com/stretchr/testify/require"
@@ -609,3 +611,28 @@ func TestSnapshotGenerator_schemalogExists(t *testing.T) {
 		})
 	}
 }
+
+func TestSnapshotGenerator_extractIndicesAndConstraints(t *testing.T) {
+	t.Parallel()
+
+	// Read the file as []byte
+	dumpBytes, err := os.ReadFile("test/test_dump.sql")
+	require.NoError(t, err)
+
+	wantFilteredDumpBytes, err := os.ReadFile("test/test_dump_filtered.sql")
+	require.NoError(t, err)
+
+	wantConstraintsBytes, err := os.ReadFile("test/test_dump_constraints.sql")
+	require.NoError(t, err)
+
+	sg := &SnapshotGenerator{}
+	filtered, constraints := sg.extractIndicesAndConstraints(dumpBytes)
+
+	filteredStr := strings.Trim(string(filtered), "\n")
+	wantFilteredStr := strings.Trim(string(wantFilteredDumpBytes), "\n")
+	constraintsStr := strings.Trim(string(constraints), "\n")
+	wantConstraintsStr := strings.Trim(string(wantConstraintsBytes), "\n")
+
+	require.Equal(t, wantFilteredStr, filteredStr)
+	require.Equal(t, wantConstraintsStr, constraintsStr)
+}
diff --git a/pkg/snapshot/generator/postgres/schema/pgdumprestore/test/test_dump.sql b/pkg/snapshot/generator/postgres/schema/pgdumprestore/test/test_dump.sql
new file mode 100644
index 0000000..f6430a4
--- /dev/null
+++ b/pkg/snapshot/generator/postgres/schema/pgdumprestore/test/test_dump.sql
@@ -0,0 +1,306 @@
+--
+-- PostgreSQL database dump
+--
+
+-- Dumped from database version 17.2
+-- Dumped by pg_dump version 17.5
+
+SET statement_timeout = 0;
+SET lock_timeout = 0;
+SET idle_in_transaction_session_timeout = 0;
+SET transaction_timeout = 0;
+SET client_encoding = 'UTF8';
+SET standard_conforming_strings = on;
+SELECT pg_catalog.set_config('search_path', '', false);
+SET check_function_bodies = false;
+SET xmloption = content;
+SET client_min_messages = warning;
+SET row_security = off;
+
+ALTER TABLE IF EXISTS ONLY musicbrainz.alternative_medium DROP CONSTRAINT IF EXISTS alternative_medium_fk_medium;
+ALTER TABLE IF EXISTS ONLY musicbrainz.alternative_medium DROP CONSTRAINT IF EXISTS alternative_medium_fk_alternative_release;
+DROP TRIGGER IF EXISTS a_del_alternative_release ON musicbrainz.alternative_release;
+DROP TRIGGER IF EXISTS a_del_alternative_medium_track ON musicbrainz.alternative_medium_track;
+DROP INDEX IF EXISTS musicbrainz.alternative_release_idx_artist_credit;
+DROP INDEX IF EXISTS musicbrainz.alternative_medium_idx_alternative_release;
+ALTER TABLE IF EXISTS ONLY musicbrainz.alternative_medium_track DROP CONSTRAINT IF EXISTS alternative_medium_track_pkey;
+ALTER TABLE IF EXISTS ONLY musicbrainz.alternative_medium DROP CONSTRAINT IF EXISTS alternative_medium_pkey;
+ALTER TABLE IF EXISTS musicbrainz.alternative_release ALTER COLUMN id DROP DEFAULT;
+ALTER TABLE IF EXISTS musicbrainz.alternative_medium ALTER COLUMN id DROP DEFAULT;
+DROP SEQUENCE IF EXISTS musicbrainz.alternative_medium_id_seq;
+DROP TABLE IF EXISTS musicbrainz.alternative_medium;
+DROP TABLE IF EXISTS musicbrainz.edit;
+DROP TEXT SEARCH CONFIGURATION IF EXISTS musicbrainz.mb_simple;
+DROP AGGREGATE IF EXISTS musicbrainz.median(integer);
+DROP FUNCTION IF EXISTS musicbrainz.a_del_alternative_medium_track();
+DROP FUNCTION IF EXISTS musicbrainz._median(integer[]);
+DROP TYPE IF EXISTS musicbrainz.edit_note_status;
+DROP TYPE IF EXISTS musicbrainz.cover_art_presence;
+DROP EXTENSION IF EXISTS cube;
+DROP COLLATION IF EXISTS musicbrainz.musicbrainz;
+DROP SCHEMA IF EXISTS musicbrainz;
+--
+-- Name: musicbrainz; Type: SCHEMA; Schema: -; Owner: postgres
+--
+
+CREATE SCHEMA musicbrainz;
+
+
+ALTER SCHEMA musicbrainz OWNER TO postgres;
+
+--
+-- Name: musicbrainz; Type: COLLATION; Schema: musicbrainz; Owner: postgres
+--
+
+CREATE COLLATION musicbrainz.musicbrainz (provider = icu, locale = 'und-u-kf-lower-kn-true');
+
+
+ALTER COLLATION musicbrainz.musicbrainz OWNER TO postgres;
+
+--
+-- Name: cube; Type: EXTENSION; Schema: -; Owner: -
+--
+
+CREATE EXTENSION IF NOT EXISTS cube WITH SCHEMA public;
+
+
+--
+-- Name: EXTENSION cube; Type: COMMENT; Schema: -; Owner:
+--
+
+COMMENT ON EXTENSION cube IS 'data type for multidimensional cubes';
+
+
+--
+-- Name: cover_art_presence; Type: TYPE; Schema: musicbrainz; Owner: postgres
+--
+
+CREATE TYPE musicbrainz.cover_art_presence AS ENUM (
+    'absent',
+    'present',
+    'darkened'
+);
+
+
+ALTER TYPE musicbrainz.cover_art_presence OWNER TO postgres;
+
+--
+-- Name: edit_note_status; Type: TYPE; Schema: musicbrainz; Owner: postgres
+--
+
+CREATE TYPE musicbrainz.edit_note_status AS ENUM (
+    'deleted',
+    'edited'
+);
+
+
+ALTER TYPE musicbrainz.edit_note_status OWNER TO postgres;
+
+--
+-- Name: _median(integer[]); Type: FUNCTION; Schema: musicbrainz; Owner: postgres
+--
+
+CREATE FUNCTION musicbrainz._median(integer[]) RETURNS integer
+    LANGUAGE sql IMMUTABLE
+    AS $_$
+  WITH q AS (
+      SELECT val
+      FROM unnest($1) val
+      WHERE VAL IS NOT NULL
+      ORDER BY val
+  )
+  SELECT val
+  FROM q
+  LIMIT 1
+  -- Subtracting (n + 1) % 2 creates a left bias
+  OFFSET greatest(0, floor((select count(*) FROM q) / 2.0) - ((select count(*) + 1 FROM q) % 2));
+$_$;
+
+
+ALTER FUNCTION musicbrainz._median(integer[]) OWNER TO postgres;
+
+--
+-- Name: a_del_alternative_medium_track(); Type: FUNCTION; Schema: musicbrainz; Owner: postgres
+--
+
+CREATE FUNCTION musicbrainz.a_del_alternative_medium_track() RETURNS trigger
+    LANGUAGE plpgsql
+    AS $$
+BEGIN
+    PERFORM dec_ref_count('alternative_track', OLD.alternative_track, 1);
+    RETURN NULL;
+END;
+$$;
+
+
+ALTER FUNCTION musicbrainz.a_del_alternative_medium_track() OWNER TO postgres;
+
+--
+-- Name: median(integer); Type: AGGREGATE; Schema: musicbrainz; Owner: postgres
+--
+
+CREATE AGGREGATE musicbrainz.median(integer) (
+    SFUNC = array_append,
+    STYPE = integer[],
+    INITCOND = '{}',
+    FINALFUNC = musicbrainz._median
+);
+
+
+ALTER AGGREGATE musicbrainz.median(integer) OWNER TO postgres;
+
+--
+-- Name: mb_simple; Type: TEXT SEARCH CONFIGURATION; Schema: musicbrainz; Owner: postgres
+--
+
+CREATE TEXT SEARCH CONFIGURATION musicbrainz.mb_simple (
+    PARSER = pg_catalog."default" );
+
+ALTER TEXT SEARCH CONFIGURATION musicbrainz.mb_simple
+    ADD MAPPING FOR asciiword WITH simple;
+
+
+ALTER TEXT SEARCH CONFIGURATION musicbrainz.mb_simple OWNER TO postgres;
+
+--
+-- Name: edit; Type: TABLE; Schema: musicbrainz; Owner: postgres
+--
+
+CREATE TABLE musicbrainz.edit (
+    id integer NOT NULL,
+    editor integer NOT NULL,
+    type smallint NOT NULL,
+    status smallint NOT NULL,
+    autoedit smallint DEFAULT 0 NOT NULL,
+    open_time timestamp with time zone DEFAULT now(),
+    close_time timestamp with time zone,
+    expire_time timestamp with time zone NOT NULL,
+    language integer,
+    quality smallint DEFAULT 1 NOT NULL
+);
+
+
+ALTER TABLE musicbrainz.edit OWNER TO postgres;
+
+--
+-- Name: alternative_medium; Type: TABLE; Schema: musicbrainz; Owner: postgres
+--
+
+CREATE TABLE musicbrainz.alternative_medium (
+    id integer NOT NULL,
+    medium integer NOT NULL,
+    alternative_release integer NOT NULL,
+    name character varying,
+    CONSTRAINT alternative_medium_name_check CHECK (((name)::text <> ''::text))
+);
+
+
+ALTER TABLE musicbrainz.alternative_medium OWNER TO postgres;
+
+--
+-- Name: alternative_medium_id_seq; Type: SEQUENCE; Schema: musicbrainz; Owner: postgres
+--
+
+CREATE SEQUENCE musicbrainz.alternative_medium_id_seq
+    AS integer
+    START WITH 1
+    INCREMENT BY 1
+    NO MINVALUE
+    NO MAXVALUE
+    CACHE 1;
+
+
+ALTER SEQUENCE musicbrainz.alternative_medium_id_seq OWNER TO postgres;
+
+--
+-- Name: alternative_medium id; Type: DEFAULT; Schema: musicbrainz; Owner: postgres
+--
+
+ALTER TABLE ONLY musicbrainz.alternative_medium ALTER COLUMN id SET DEFAULT nextval('musicbrainz.alternative_medium_id_seq'::regclass);
+
+
+--
+-- Name: alternative_release id; Type: DEFAULT; Schema: musicbrainz; Owner: postgres
+--
+
+ALTER TABLE ONLY musicbrainz.alternative_release ALTER COLUMN id SET DEFAULT nextval('musicbrainz.alternative_release_id_seq'::regclass);
+
+
+--
+-- Name: alternative_medium alternative_medium_pkey; Type: CONSTRAINT; Schema: musicbrainz; Owner: postgres
+--
+
+ALTER TABLE ONLY musicbrainz.alternative_medium
+    ADD CONSTRAINT alternative_medium_pkey PRIMARY KEY (id);
+
+
+--
+-- Name: alternative_medium_track alternative_medium_track_pkey; Type: CONSTRAINT; Schema: musicbrainz; Owner: postgres
+--
+
+ALTER TABLE ONLY musicbrainz.alternative_medium_track
+    ADD CONSTRAINT alternative_medium_track_pkey PRIMARY KEY (alternative_medium, track);
+
+
+
+ALTER TABLE ONLY musicbrainz.alternative_medium_track
+    ADD COLUMN alternative_track integer NOT NULL;
+
+--
+-- Name: area_alias_idx_txt; Type: INDEX; Schema: musicbrainz; Owner: postgres
+--
+
+CREATE INDEX area_alias_idx_txt ON musicbrainz.area_alias USING gin (musicbrainz.mb_simple_tsvector((name)::text));
+
+
+--
+-- Name: area_alias_type_idx_gid; Type: INDEX; Schema: musicbrainz; Owner: postgres
+--
+
+CREATE UNIQUE INDEX area_alias_type_idx_gid ON musicbrainz.area_alias_type USING btree (gid);
+
+
+--
+-- Name: alternative_medium_track a_del_alternative_medium_track; Type: TRIGGER; Schema: musicbrainz; Owner: postgres
+--
+
+CREATE TRIGGER a_del_alternative_medium_track AFTER DELETE ON musicbrainz.alternative_medium_track FOR EACH ROW EXECUTE FUNCTION musicbrainz.a_del_alternative_medium_track();
+
+
+--
+-- Name: alternative_release a_del_alternative_release; Type: TRIGGER; Schema: musicbrainz; Owner: postgres
+--
+
+CREATE TRIGGER a_del_alternative_release AFTER DELETE ON musicbrainz.alternative_release FOR EACH ROW EXECUTE FUNCTION musicbrainz.a_del_alternative_release_or_track();
+
+--
+-- Name: release apply_artist_release_group_pending_updates; Type: TRIGGER; Schema: musicbrainz; Owner: postgres
+--
+
+CREATE CONSTRAINT TRIGGER apply_artist_release_group_pending_updates AFTER INSERT OR DELETE OR UPDATE ON musicbrainz.release DEFERRABLE INITIALLY DEFERRED FOR EACH ROW EXECUTE FUNCTION musicbrainz.apply_artist_release_group_pending_updates();
+
+
+--
+-- Name: release_group apply_artist_release_group_pending_updates; Type: TRIGGER; Schema: musicbrainz; Owner: postgres
+--
+
+CREATE CONSTRAINT TRIGGER apply_artist_release_group_pending_updates AFTER INSERT OR DELETE OR UPDATE ON musicbrainz.release_group DEFERRABLE INITIALLY DEFERRED FOR EACH ROW EXECUTE FUNCTION musicbrainz.apply_artist_release_group_pending_updates();
+
+--
+-- Name: alternative_medium alternative_medium_fk_alternative_release; Type: FK CONSTRAINT; Schema: musicbrainz; Owner: postgres
+--
+
+ALTER TABLE ONLY musicbrainz.alternative_medium
+    ADD CONSTRAINT alternative_medium_fk_alternative_release FOREIGN KEY (alternative_release) REFERENCES musicbrainz.alternative_release(id);
+
+
+--
+-- Name: alternative_medium alternative_medium_fk_medium; Type: FK CONSTRAINT; Schema: musicbrainz; Owner: postgres
+--
+
+ALTER TABLE ONLY musicbrainz.alternative_medium
+    ADD CONSTRAINT alternative_medium_fk_medium FOREIGN KEY (medium) REFERENCES musicbrainz.medium(id);
+
+--
+-- PostgreSQL database dump complete
+--
diff --git a/pkg/snapshot/generator/postgres/schema/pgdumprestore/test/test_dump_constraints.sql b/pkg/snapshot/generator/postgres/schema/pgdumprestore/test/test_dump_constraints.sql
new file mode 100644
index 0000000..5bdd504
--- /dev/null
+++ b/pkg/snapshot/generator/postgres/schema/pgdumprestore/test/test_dump_constraints.sql
@@ -0,0 +1,23 @@
+ALTER TABLE ONLY musicbrainz.alternative_medium
+    ADD CONSTRAINT alternative_medium_pkey PRIMARY KEY (id);
+
+ALTER TABLE ONLY musicbrainz.alternative_medium_track
+    ADD CONSTRAINT alternative_medium_track_pkey PRIMARY KEY (alternative_medium, track);
+
+CREATE INDEX area_alias_idx_txt ON musicbrainz.area_alias USING gin (musicbrainz.mb_simple_tsvector((name)::text));
+
+CREATE UNIQUE INDEX area_alias_type_idx_gid ON musicbrainz.area_alias_type USING btree (gid);
+
+CREATE TRIGGER a_del_alternative_medium_track AFTER DELETE ON musicbrainz.alternative_medium_track FOR EACH ROW EXECUTE FUNCTION musicbrainz.a_del_alternative_medium_track();
+
+CREATE TRIGGER a_del_alternative_release AFTER DELETE ON musicbrainz.alternative_release FOR EACH ROW EXECUTE FUNCTION musicbrainz.a_del_alternative_release_or_track();
+
+CREATE CONSTRAINT TRIGGER apply_artist_release_group_pending_updates AFTER INSERT OR DELETE OR UPDATE ON musicbrainz.release DEFERRABLE INITIALLY DEFERRED FOR EACH ROW EXECUTE FUNCTION musicbrainz.apply_artist_release_group_pending_updates();
+
+CREATE CONSTRAINT TRIGGER apply_artist_release_group_pending_updates AFTER INSERT OR DELETE OR UPDATE ON musicbrainz.release_group DEFERRABLE INITIALLY DEFERRED FOR EACH ROW EXECUTE FUNCTION musicbrainz.apply_artist_release_group_pending_updates();
+
+ALTER TABLE ONLY musicbrainz.alternative_medium
+    ADD CONSTRAINT alternative_medium_fk_alternative_release FOREIGN KEY (alternative_release) REFERENCES musicbrainz.alternative_release(id);
+
+ALTER TABLE ONLY musicbrainz.alternative_medium
+    ADD CONSTRAINT alternative_medium_fk_medium FOREIGN KEY (medium) REFERENCES musicbrainz.medium(id);
diff --git a/pkg/snapshot/generator/postgres/schema/pgdumprestore/test/test_dump_filtered.sql b/pkg/snapshot/generator/postgres/schema/pgdumprestore/test/test_dump_filtered.sql
new file mode 100644
index 0000000..fd06bc7
--- /dev/null
+++ b/pkg/snapshot/generator/postgres/schema/pgdumprestore/test/test_dump_filtered.sql
@@ -0,0 +1,292 @@
+--
+-- PostgreSQL database dump
+--
+
+-- Dumped from database version 17.2
+-- Dumped by pg_dump version 17.5
+
+SET statement_timeout = 0;
+SET lock_timeout = 0;
+SET idle_in_transaction_session_timeout = 0;
+SET transaction_timeout = 0;
+SET client_encoding = 'UTF8';
+SET standard_conforming_strings = on;
+SELECT pg_catalog.set_config('search_path', '', false);
+SET check_function_bodies = false;
+SET xmloption = content;
+SET client_min_messages = warning;
+SET row_security = off;
+
+ALTER TABLE IF EXISTS ONLY musicbrainz.alternative_medium DROP CONSTRAINT IF EXISTS alternative_medium_fk_medium;
+ALTER TABLE IF EXISTS ONLY musicbrainz.alternative_medium DROP CONSTRAINT IF EXISTS alternative_medium_fk_alternative_release;
+DROP TRIGGER IF EXISTS a_del_alternative_release ON musicbrainz.alternative_release;
+DROP TRIGGER IF EXISTS a_del_alternative_medium_track ON musicbrainz.alternative_medium_track;
+DROP INDEX IF EXISTS musicbrainz.alternative_release_idx_artist_credit;
+DROP INDEX IF EXISTS musicbrainz.alternative_medium_idx_alternative_release;
+ALTER TABLE IF EXISTS ONLY musicbrainz.alternative_medium_track DROP CONSTRAINT IF EXISTS alternative_medium_track_pkey;
+ALTER TABLE IF EXISTS ONLY musicbrainz.alternative_medium DROP CONSTRAINT IF EXISTS alternative_medium_pkey;
+ALTER TABLE IF EXISTS musicbrainz.alternative_release ALTER COLUMN id DROP DEFAULT;
+ALTER TABLE IF EXISTS musicbrainz.alternative_medium ALTER COLUMN id DROP DEFAULT;
+DROP SEQUENCE IF EXISTS musicbrainz.alternative_medium_id_seq;
+DROP TABLE IF EXISTS musicbrainz.alternative_medium;
+DROP TABLE IF EXISTS musicbrainz.edit;
+DROP TEXT SEARCH CONFIGURATION IF EXISTS musicbrainz.mb_simple;
+DROP AGGREGATE IF EXISTS musicbrainz.median(integer);
+DROP FUNCTION IF EXISTS musicbrainz.a_del_alternative_medium_track();
+DROP FUNCTION IF EXISTS musicbrainz._median(integer[]);
+DROP TYPE IF EXISTS musicbrainz.edit_note_status;
+DROP TYPE IF EXISTS musicbrainz.cover_art_presence;
+DROP EXTENSION IF EXISTS cube;
+DROP COLLATION IF EXISTS musicbrainz.musicbrainz;
+DROP SCHEMA IF EXISTS musicbrainz;
+--
+-- Name: musicbrainz; Type: SCHEMA; Schema: -; Owner: postgres
+--
+
+CREATE SCHEMA musicbrainz;
+
+
+ALTER SCHEMA musicbrainz OWNER TO postgres;
+
+--
+-- Name: musicbrainz; Type: COLLATION; Schema: musicbrainz; Owner: postgres
+--
+
+CREATE COLLATION musicbrainz.musicbrainz (provider = icu, locale = 'und-u-kf-lower-kn-true');
+
+
+ALTER COLLATION musicbrainz.musicbrainz OWNER TO postgres;
+
+--
+-- Name: cube; Type: EXTENSION; Schema: -; Owner: -
+--
+
+CREATE EXTENSION IF NOT EXISTS cube WITH SCHEMA public;
+
+
+--
+-- Name: EXTENSION cube; Type: COMMENT; Schema: -; Owner:
+--
+
+COMMENT ON EXTENSION cube IS 'data type for multidimensional cubes';
+
+
+--
+-- Name: cover_art_presence; Type: TYPE; Schema: musicbrainz; Owner: postgres
+--
+
+CREATE TYPE musicbrainz.cover_art_presence AS ENUM (
+    'absent',
+    'present',
+    'darkened'
+);
+
+
+ALTER TYPE musicbrainz.cover_art_presence OWNER TO postgres;
+
+--
+-- Name: edit_note_status; Type: TYPE; Schema: musicbrainz; Owner: postgres
+--
+
+CREATE TYPE musicbrainz.edit_note_status AS ENUM (
+    'deleted',
+    'edited'
+);
+
+
+ALTER TYPE musicbrainz.edit_note_status OWNER TO postgres;
+
+--
+-- Name: _median(integer[]); Type: FUNCTION; Schema: musicbrainz; Owner: postgres
+--
+
+CREATE FUNCTION musicbrainz._median(integer[]) RETURNS integer
+    LANGUAGE sql IMMUTABLE
+    AS $_$
+  WITH q AS (
+      SELECT val
+      FROM unnest($1) val
+      WHERE VAL IS NOT NULL
+      ORDER BY val
+  )
+  SELECT val
+  FROM q
+  LIMIT 1
+  -- Subtracting (n + 1) % 2 creates a left bias
+  OFFSET greatest(0, floor((select count(*) FROM q) / 2.0) - ((select count(*) + 1 FROM q) % 2));
+$_$;
+
+
+ALTER FUNCTION musicbrainz._median(integer[]) OWNER TO postgres;
+
+--
+-- Name: a_del_alternative_medium_track(); Type: FUNCTION; Schema: musicbrainz; Owner: postgres
+--
+
+CREATE FUNCTION musicbrainz.a_del_alternative_medium_track() RETURNS trigger
+    LANGUAGE plpgsql
+    AS $$
+BEGIN
+    PERFORM dec_ref_count('alternative_track', OLD.alternative_track, 1);
+    RETURN NULL;
+END;
+$$;
+
+
+ALTER FUNCTION musicbrainz.a_del_alternative_medium_track() OWNER TO postgres;
+
+--
+-- Name: median(integer); Type: AGGREGATE; Schema: musicbrainz; Owner: postgres
+--
+
+CREATE AGGREGATE musicbrainz.median(integer) (
+    SFUNC = array_append,
+    STYPE = integer[],
+    INITCOND = '{}',
+    FINALFUNC = musicbrainz._median
+);
+
+
+ALTER AGGREGATE musicbrainz.median(integer) OWNER TO postgres;
+
+--
+-- Name: mb_simple; Type: TEXT SEARCH CONFIGURATION; Schema: musicbrainz; Owner: postgres
+--
+
+CREATE TEXT SEARCH CONFIGURATION musicbrainz.mb_simple (
+    PARSER = pg_catalog."default" );
+
+ALTER TEXT SEARCH CONFIGURATION musicbrainz.mb_simple
+    ADD MAPPING FOR asciiword WITH simple;
+
+
+ALTER TEXT SEARCH CONFIGURATION musicbrainz.mb_simple OWNER TO postgres;
+
+--
+-- Name: edit; Type: TABLE; Schema: musicbrainz; Owner: postgres
+--
+
+CREATE TABLE musicbrainz.edit (
+    id integer NOT NULL,
+    editor integer NOT NULL,
+    type smallint NOT NULL,
+    status smallint NOT NULL,
+    autoedit smallint DEFAULT 0 NOT NULL,
+    open_time timestamp with time zone DEFAULT now(),
+    close_time timestamp with time zone,
+    expire_time timestamp with time zone NOT NULL,
+    language integer,
+    quality smallint DEFAULT 1 NOT NULL
+);
+
+
+ALTER TABLE musicbrainz.edit OWNER TO postgres;
+
+--
+-- Name: alternative_medium; Type: TABLE; Schema: musicbrainz; Owner: postgres
+--
+
+CREATE TABLE musicbrainz.alternative_medium (
+    id integer NOT NULL,
+    medium integer NOT NULL,
+    alternative_release integer NOT NULL,
+    name character varying,
+    CONSTRAINT alternative_medium_name_check CHECK (((name)::text <> ''::text))
+);
+
+
+ALTER TABLE musicbrainz.alternative_medium OWNER TO postgres;
+
+--
+-- Name: alternative_medium_id_seq; Type: SEQUENCE; Schema: musicbrainz; Owner: postgres
+--
+
+CREATE SEQUENCE musicbrainz.alternative_medium_id_seq
+    AS integer
+    START WITH 1
+    INCREMENT BY 1
+    NO MINVALUE
+    NO MAXVALUE
+    CACHE 1;
+
+
+ALTER SEQUENCE musicbrainz.alternative_medium_id_seq OWNER TO postgres;
+
+--
+-- Name: alternative_medium id; Type: DEFAULT; Schema: musicbrainz; Owner: postgres
+--
+
+ALTER TABLE ONLY musicbrainz.alternative_medium ALTER COLUMN id SET DEFAULT nextval('musicbrainz.alternative_medium_id_seq'::regclass);
+
+
+--
+-- Name: alternative_release id; Type: DEFAULT; Schema: musicbrainz; Owner: postgres
+--
+
+ALTER TABLE ONLY musicbrainz.alternative_release ALTER COLUMN id SET DEFAULT nextval('musicbrainz.alternative_release_id_seq'::regclass);
+
+
+--
+-- Name: alternative_medium alternative_medium_pkey; Type: CONSTRAINT; Schema: musicbrainz; Owner: postgres
+--
+
+
+
+--
+-- Name: alternative_medium_track alternative_medium_track_pkey; Type: CONSTRAINT; Schema: musicbrainz; Owner: postgres
+--
+
+
+
+
+ALTER TABLE ONLY musicbrainz.alternative_medium_track
+    ADD COLUMN alternative_track integer NOT NULL;
+
+--
+-- Name: area_alias_idx_txt; Type: INDEX; Schema: musicbrainz; Owner: postgres
+--
+
+
+
+--
+-- Name: area_alias_type_idx_gid; Type: INDEX; Schema: musicbrainz; Owner: postgres
+--
+
+
+
+--
+-- Name: alternative_medium_track a_del_alternative_medium_track; Type: TRIGGER; Schema: musicbrainz; Owner: postgres
+--
+
+
+
+--
+-- Name: alternative_release a_del_alternative_release; Type: TRIGGER; Schema: musicbrainz; Owner: postgres
+--
+
+
+--
+-- Name: release apply_artist_release_group_pending_updates; Type: TRIGGER; Schema: musicbrainz; Owner: postgres
+--
+
+
+
+--
+-- Name: release_group apply_artist_release_group_pending_updates; Type: TRIGGER; Schema: musicbrainz; Owner: postgres
+--
+
+
+--
+-- Name: alternative_medium alternative_medium_fk_alternative_release; Type: FK CONSTRAINT; Schema: musicbrainz; Owner: postgres
+--
+
+
+
+--
+-- Name: alternative_medium alternative_medium_fk_medium; Type: FK CONSTRAINT; Schema: musicbrainz; Owner: postgres
+--
+
+
+--
+-- PostgreSQL database dump complete
+--
diff --git a/pkg/snapshot/generator/postgres/schema/schemalog/snapshot_schemalog_generator_test.go b/pkg/snapshot/generator/postgres/schema/schemalog/snapshot_schemalog_generator_test.go
index 9244e40..0c0ec7b 100644
--- a/pkg/snapshot/generator/postgres/schema/schemalog/snapshot_schemalog_generator_test.go
+++ b/pkg/snapshot/generator/postgres/schema/schemalog/snapshot_schemalog_generator_test.go
@@ -15,6 +15,7 @@ import (
 	"github.com/xataio/pgstream/pkg/schemalog"
 	schemalogmocks "github.com/xataio/pgstream/pkg/schemalog/mocks"
 	"github.com/xataio/pgstream/pkg/snapshot"
+	generatormocks "github.com/xataio/pgstream/pkg/snapshot/generator/mocks"
 )
 
 func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
@@ -70,8 +71,10 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 		schemaStore *schemalogmocks.Store
 		processRow  snapshot.RowProcessor
 		marshaler   func(any) ([]byte, error)
+		generator   *generatormocks.Generator
 
-		wantErr error
+		wantGeneratorCalls uint
+		wantErr            error
 	}{
 		{
 			name: "ok",
@@ -87,6 +90,28 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 
 			wantErr: nil,
 		},
+		{
+			name: "ok - with generator",
+			schemaStore: &schemalogmocks.Store{
+				InsertFn: func(ctx context.Context, schemaName string) (*schemalog.LogEntry, error) {
+					return testSchemaLog, nil
+				},
+			},
+			processRow: func(ctx context.Context, r *snapshot.Row) error {
+				require.Equal(t, testRow, r)
+				return nil
+			},
+			generator: &generatormocks.Generator{
+				CreateSnapshotFn: func(ctx context.Context, ss *snapshot.Snapshot) error {
+					require.Equal(t, testSnapshot, ss)
+					return nil
+				},
+				CloseFn: func() error { return nil },
+			},
+
+			wantGeneratorCalls: 1,
+			wantErr:            nil,
+		},
 		{
 			name: "error - inserting schema log",
 			schemaStore: &schemalogmocks.Store{
@@ -134,7 +159,12 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 		t.Run(tc.name, func(t *testing.T) {
 			t.Parallel()
 
-			g := NewSnapshotGenerator(tc.schemaStore, tc.processRow)
+			opts := []Option{}
+			if tc.generator != nil {
+				opts = append(opts, WithSnapshotGenerator(tc.generator))
+			}
+
+			g := NewSnapshotGenerator(tc.schemaStore, tc.processRow, opts...)
 			defer g.Close()
 
 			if tc.marshaler != nil {
@@ -143,6 +173,10 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 
 			err := g.CreateSnapshot(context.Background(), testSnapshot)
 			require.Equal(t, err, tc.wantErr)
+
+			if tc.generator != nil {
+				require.Equal(t, tc.wantGeneratorCalls, tc.generator.CreateSnapshotCalls())
+			}
 		})
 	}
 }
diff --git a/pkg/snapshot/generator/snapshot_generator_aggregator_test.go b/pkg/snapshot/generator/snapshot_generator_aggregator_test.go
deleted file mode 100644
index 69b6420..0000000
--- a/pkg/snapshot/generator/snapshot_generator_aggregator_test.go
+++ /dev/null
@@ -1,105 +0,0 @@
-// SPDX-License-Identifier: Apache-2.0
-
-package generator
-
-import (
-	"context"
-	"errors"
-	"testing"
-
-	"github.com/stretchr/testify/require"
-	"github.com/xataio/pgstream/pkg/snapshot"
-)
-
-func TestAggregator_CreateSnapshot(t *testing.T) {
-	t.Parallel()
-
-	testSnapshot := &snapshot.Snapshot{
-		SchemaName: "test-schema",
-		TableNames: []string{"table1", "table2"},
-	}
-
-	newGenerator := func(err error) *mockGenerator {
-		return &mockGenerator{
-			createSnapshotFn: func(ctx context.Context, ss *snapshot.Snapshot) error {
-				require.Equal(t, testSnapshot, ss)
-				return err
-			},
-		}
-	}
-
-	errTest := errors.New("oh noes")
-
-	tests := []struct {
-		name       string
-		generators []SnapshotGenerator
-
-		wantErr error
-	}{
-		{
-			name:       "ok",
-			generators: []SnapshotGenerator{newGenerator(nil), newGenerator(nil)},
-			wantErr:    nil,
-		},
-		{
-			name:       "error on first generator",
-			generators: []SnapshotGenerator{newGenerator(errTest), newGenerator(errors.New("should not be called"))},
-			wantErr:    errTest,
-		},
-	}
-
-	for _, tc := range tests {
-		t.Run(tc.name, func(t *testing.T) {
-			t.Parallel()
-
-			agg := NewAggregator(tc.generators)
-
-			err := agg.CreateSnapshot(context.Background(), testSnapshot)
-			require.Equal(t, err, tc.wantErr)
-		})
-	}
-}
-
-func TestAggregator_Close(t *testing.T) {
-	t.Parallel()
-
-	newGenerator := func(err error) *mockGenerator {
-		return &mockGenerator{
-			closeFn: func() error {
-				return err
-			},
-		}
-	}
-
-	errTest1 := errors.New("oh noes")
-	errTest2 := errors.New("oh nose")
-
-	tests := []struct {
-		name       string
-		generators []SnapshotGenerator
-
-		wantErr error
-	}{
-		{
-			name:       "ok",
-			generators: []SnapshotGenerator{newGenerator(nil), newGenerator(nil)},
-			wantErr:    nil,
-		},
-		{
-			name:       "error on first generator",
-			generators: []SnapshotGenerator{newGenerator(errTest1), newGenerator(errTest2)},
-			wantErr:    errors.Join(errTest1, errTest2),
-		},
-	}
-
-	for _, tc := range tests {
-		t.Run(tc.name, func(t *testing.T) {
-			t.Parallel()
-
-			agg := NewAggregator(tc.generators)
-
-			err := agg.Close()
-			require.Equal(t, err, tc.wantErr)
-		})
-	}
-}