Restore indexes in parallel during pg2pg schema snapshot

cmd/config/config_env.go

@@ -86,6 +86,7 @@ func bindEnvVars() {
 	viper.BindEnv("PGSTREAM_POSTGRES_SNAPSHOT_EXCLUDED_SECURITY_LABELS")
 	viper.BindEnv("PGSTREAM_POSTGRES_SNAPSHOT_REFRESH_MATERIALIZED_VIEWS")
 	viper.BindEnv("PGSTREAM_POSTGRES_SNAPSHOT_INDEX_CONSTRAINT_SESSION_SETTINGS")
+	viper.BindEnv("PGSTREAM_POSTGRES_SNAPSHOT_INDEX_RESTORE_WORKERS")
 	viper.BindEnv("PGSTREAM_POSTGRES_SNAPSHOT_INCLUDE_OBJECT_TYPES")
 	viper.BindEnv("PGSTREAM_POSTGRES_SNAPSHOT_EXCLUDE_OBJECT_TYPES")
 	viper.BindEnv("PGSTREAM_POSTGRES_SNAPSHOT_DISABLE_PROGRESS_TRACKING")
@@ -378,6 +379,7 @@ func parseSchemaSnapshotConfig(pgurl string) (*snapshotbuilder.SchemaSnapshotCon
 	if err != nil {
 		return nil, err
 	}
+
 	return &snapshotbuilder.SchemaSnapshotConfig{
 		DumpRestore: &pgdumprestore.Config{
 			SourcePGURL:                    pgurl,
@@ -393,6 +395,7 @@ func parseSchemaSnapshotConfig(pgurl string) (*snapshotbuilder.SchemaSnapshotCon
 			ExcludedSecurityLabels:         viper.GetStringSlice("PGSTREAM_POSTGRES_SNAPSHOT_EXCLUDED_SECURITY_LABELS"),
 			RefreshMaterializedViews:       viper.GetBool("PGSTREAM_POSTGRES_SNAPSHOT_REFRESH_MATERIALIZED_VIEWS"),
 			IndexConstraintSessionSettings: viper.GetStringSlice("PGSTREAM_POSTGRES_SNAPSHOT_INDEX_CONSTRAINT_SESSION_SETTINGS"),
+			IndexRestoreWorkers:            viper.GetUint("PGSTREAM_POSTGRES_SNAPSHOT_INDEX_RESTORE_WORKERS"),
 			IncludeObjectTypes:             viper.GetStringSlice("PGSTREAM_POSTGRES_SNAPSHOT_INCLUDE_OBJECT_TYPES"),
 			ExcludeObjectTypes:             viper.GetStringSlice("PGSTREAM_POSTGRES_SNAPSHOT_EXCLUDE_OBJECT_TYPES"),
 		},

cmd/config/config_env_test.go

@@ -6,6 +6,8 @@ import (
 	"os"
 	"testing"
 
+	"github.com/spf13/viper"
+
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
@@ -53,6 +55,7 @@ func Test_EnvVarsToStreamConfig(t *testing.T) {
 	os.Setenv("PGSTREAM_POSTGRES_SNAPSHOT_NO_OWNER", "true")
 	os.Setenv("PGSTREAM_POSTGRES_SNAPSHOT_MODE", "full")
 	os.Setenv("PGSTREAM_POSTGRES_SNAPSHOT_INDEX_CONSTRAINT_SESSION_SETTINGS", "maintenance_work_mem=4GB max_parallel_maintenance_workers=4 synchronous_commit=off statement_timeout=0 lock_timeout=0")
+	os.Setenv("PGSTREAM_POSTGRES_SNAPSHOT_INDEX_RESTORE_WORKERS", "4")
 	os.Setenv("PGSTREAM_POSTGRES_SNAPSHOT_DISABLE_PROGRESS_TRACKING", "true")
 
 	os.Setenv("PGSTREAM_KAFKA_READER_SERVERS", "localhost:9092")
@@ -130,3 +133,29 @@ func Test_EnvVarsToOtelConfig(t *testing.T) {
 
 	validateTestOtelConfig(t, otelConfig)
 }
+
+// Test_EnvVarsToSchemaSnapshotConfig_IndexRestoreWorkers resolves the setting
+// from the process environment on a viper with nothing else in it. The
+// surrounding tests share viper's global config map with the ones that load
+// test_config.env, so a variable that is never bound with BindEnv still
+// resolves there, out of the file, and a missing binding goes unnoticed.
+func Test_EnvVarsToSchemaSnapshotConfig_IndexRestoreWorkers(t *testing.T) {
+	reset := func() {
+		viper.Reset()
+		bindEnvVars()
+	}
+	reset()
+	t.Cleanup(reset)
+
+	t.Setenv("PGSTREAM_POSTGRES_SNAPSHOT_INDEX_RESTORE_WORKERS", "4")
+	cfg, err := parseSchemaSnapshotConfig("postgresql://user:password@localhost:5432/mydatabase")
+	require.NoError(t, err)
+	require.Equal(t, uint(4), cfg.DumpRestore.IndexRestoreWorkers)
+
+	// unset, the generator applies its own default
+	os.Unsetenv("PGSTREAM_POSTGRES_SNAPSHOT_INDEX_RESTORE_WORKERS")
+	reset()
+	cfg, err = parseSchemaSnapshotConfig("postgresql://user:password@localhost:5432/mydatabase")
+	require.NoError(t, err)
+	require.Equal(t, uint(0), cfg.DumpRestore.IndexRestoreWorkers)
+}

cmd/config/config_yaml.go

@@ -138,6 +138,7 @@ type PgDumpPgRestoreConfig struct {
 	ExcludedSecurityLabels         []string `mapstructure:"excluded_security_labels" yaml:"excluded_security_labels"`
 	RefreshMaterializedViews       bool     `mapstructure:"refresh_materialized_views" yaml:"refresh_materialized_views"`
 	IndexConstraintSessionSettings []string `mapstructure:"index_constraint_session_settings" yaml:"index_constraint_session_settings"`
+	IndexRestoreWorkers            uint     `mapstructure:"index_restore_workers" yaml:"index_restore_workers"`
 	IncludeObjectTypes             []string `mapstructure:"include_object_types" yaml:"include_object_types"`
 	ExcludeObjectTypes             []string `mapstructure:"exclude_object_types" yaml:"exclude_object_types"`
 }
@@ -657,6 +658,7 @@ func (c *YAMLConfig) parseSchemaSnapshotConfig() (*snapshotbuilder.SchemaSnapsho
 		streamSchemaCfg.DumpRestore.ExcludedSecurityLabels = schemaSnapshotCfg.PgDumpPgRestore.ExcludedSecurityLabels
 		streamSchemaCfg.DumpRestore.RefreshMaterializedViews = schemaSnapshotCfg.PgDumpPgRestore.RefreshMaterializedViews
 		streamSchemaCfg.DumpRestore.IndexConstraintSessionSettings = schemaSnapshotCfg.PgDumpPgRestore.IndexConstraintSessionSettings
+		streamSchemaCfg.DumpRestore.IndexRestoreWorkers = schemaSnapshotCfg.PgDumpPgRestore.IndexRestoreWorkers
 		streamSchemaCfg.DumpRestore.IncludeObjectTypes = schemaSnapshotCfg.PgDumpPgRestore.IncludeObjectTypes
 		streamSchemaCfg.DumpRestore.ExcludeObjectTypes = schemaSnapshotCfg.PgDumpPgRestore.ExcludeObjectTypes
 

cmd/config/helper_test.go

@@ -90,6 +90,7 @@ func validateTestStreamConfig(t *testing.T, streamConfig *stream.Config) {
 								"statement_timeout=0",
 								"lock_timeout=0",
 							},
+							IndexRestoreWorkers: 4,
 						},
 					},
 					Recorder: &builder.SnapshotRecorderConfig{

cmd/config/test/test_config.env

@@ -30,6 +30,7 @@ PGSTREAM_POSTGRES_SNAPSHOT_NO_PRIVILEGES=true
 PGSTREAM_POSTGRES_SNAPSHOT_EXCLUDED_SECURITY_LABELS="anon"
 PGSTREAM_POSTGRES_SNAPSHOT_REFRESH_MATERIALIZED_VIEWS=true
 PGSTREAM_POSTGRES_SNAPSHOT_INDEX_CONSTRAINT_SESSION_SETTINGS="maintenance_work_mem=4GB max_parallel_maintenance_workers=4 synchronous_commit=off statement_timeout=0 lock_timeout=0"
+PGSTREAM_POSTGRES_SNAPSHOT_INDEX_RESTORE_WORKERS=4
 PGSTREAM_POSTGRES_SNAPSHOT_DISABLE_PROGRESS_TRACKING=true
 
 # Kafka

cmd/config/test/test_config.yaml

@@ -34,6 +34,7 @@ source:
             - synchronous_commit=off
             - statement_timeout=0
             - lock_timeout=0
+          index_restore_workers: 4 # number of standalone indexes restored concurrently
       disable_progress_tracking: true # whether to disable progress tracking for the snapshot
     replication: # when mode is replication or snapshot_and_replication
       replication_slot: "pgstream_mydatabase_slot"

config_template.yaml

@@ -50,6 +50,11 @@ source:
             # - statement_timeout=0
             # - lock_timeout=0
             # - synchronous_commit=off # faster restore, but a target crash right after the restore can lose the final index/constraint commits
+          # Number of standalone CREATE INDEX/CREATE UNIQUE INDEX statements restored concurrently against the target. Statements that depend on an index
+          # (constraints, comments, partition attachments, REPLICA IDENTITY, CLUSTER) always wait for every index to be created first.
+          # Each worker holds one target connection and applies the index_constraint_session_settings above to it, so the target needs this many free
+          # connection slots and up to workers x maintenance_work_mem of memory. Defaults to 1 (sequential), maximum 32.
+          index_restore_workers: 1
           dump_file: pg_dump.sql # name of the file where the contents of the schema pg_dump command and output will be written for debugging purposes.
           # Granular object type filtering for schema snapshots. Only one of include_object_types or exclude_object_types can be set.
           # Available categories: tables, sequences, types, indexes, constraints, functions, views, materialized_views, triggers, event_triggers, policies, rules, comments, extensions, collations, text_search

docs/configuration.md

@@ -54,6 +54,11 @@ source:
             # - statement_timeout=0
             # - lock_timeout=0
             # - synchronous_commit=off # faster restore, but a target crash right after the restore can lose the final index/constraint commits
+          # Number of standalone CREATE INDEX/CREATE UNIQUE INDEX statements restored concurrently against the target. Statements that depend on an index
+          # (constraints, comments, partition attachments, REPLICA IDENTITY, CLUSTER) always wait for every index to be created first.
+          # Each worker holds one target connection and applies the index_constraint_session_settings above to it, so the target needs this many free
+          # connection slots and up to workers x maintenance_work_mem of memory. Defaults to 1 (sequential), maximum 32.
+          index_restore_workers: 1
           dump_file: pg_dump.sql # name of the file where the contents of the schema pg_dump command and output will be written for debugging purposes.
           # Granular object type filtering for schema snapshots. Only one of include_object_types or exclude_object_types can be set.
           # Available categories: tables, sequences, types, indexes, constraints, functions, views, materialized_views, triggers, event_triggers, policies, rules, comments, extensions, collations, text_search
@@ -248,6 +253,7 @@ Here's a list of all the environment variables that can be used to configure the
 | PGSTREAM_POSTGRES_SNAPSHOT_EXCLUDED_SECURITY_LABELS     | []                           | No       | When using `pg_dump`/`pg_restore` to snapshot schema for Postgres targets, list of providers whose security labels will be excluded.                                                                                                                                                                         |
 | PGSTREAM_POSTGRES_SNAPSHOT_REFRESH_MATERIALIZED_VIEWS   | False                        | No       | When using `pg_dump`/`pg_restore` to snapshot schema for Postgres targets, whether to refresh materialized views (REFRESH MATERIALIZED VIEW ... WITH DATA) after the table data has been restored.                                                                                                            |
 | PGSTREAM_POSTGRES_SNAPSHOT_INDEX_CONSTRAINT_SESSION_SETTINGS | []                       | No       | Space-separated PostgreSQL `name=value` session settings applied only while restoring indexes and constraints, for example `maintenance_work_mem=4GB max_parallel_maintenance_workers=4`. Each setting must be a whitespace-free `name=value` pair; invalid entries fail at startup. Unset or empty preserves existing behavior.                                           |
+| PGSTREAM_POSTGRES_SNAPSHOT_INDEX_RESTORE_WORKERS        | 1                            | No       | When using `pg_dump`/`pg_restore` to snapshot schema for Postgres targets, number of standalone `CREATE INDEX`/`CREATE UNIQUE INDEX` statements restored concurrently against the target. Statements that depend on an index (constraints added `USING INDEX`, comments, partition attachments, `REPLICA IDENTITY`, `CLUSTER`) always wait for every index to be created first. Each worker holds one target connection and applies `PGSTREAM_POSTGRES_SNAPSHOT_INDEX_CONSTRAINT_SESSION_SETTINGS` to it, so the target needs this many free connection slots and up to workers × `maintenance_work_mem` of memory. Defaults to 1 (sequential); values above 32 are rejected at startup.                                           |
 | PGSTREAM_POSTGRES_SNAPSHOT_INCLUDE_OBJECT_TYPES         | []                           | No       | When using `pg_dump`/`pg_restore` to snapshot schema for Postgres targets, list of object type categories to include in the schema snapshot. Everything else is excluded. Mutually exclusive with `PGSTREAM_POSTGRES_SNAPSHOT_EXCLUDE_OBJECT_TYPES`. See [object type filtering](#object-type-filtering).      |
 | PGSTREAM_POSTGRES_SNAPSHOT_EXCLUDE_OBJECT_TYPES         | []                           | No       | When using `pg_dump`/`pg_restore` to snapshot schema for Postgres targets, list of object type categories to exclude from the schema snapshot. Mutually exclusive with `PGSTREAM_POSTGRES_SNAPSHOT_INCLUDE_OBJECT_TYPES`. See [object type filtering](#object-type-filtering).                                 |
 | PGSTREAM_POSTGRES_SNAPSHOT_ROLE                         | ""                           | No       | When using `pg_dump`/`pg_restore` to snapshot schema for Postgres targets, role name to be used to create the dump.                                                                                                                                                                                          |

internal/postgres/pg_restore.go

@@ -426,20 +426,58 @@ func NewPGRestoreErrors(errs ...error) *PGRestoreErrors {
 	return pgrestoreErrs
 }
 
+// MergePGRestoreErrors combines the errors of several restores into a single
+// PGRestoreErrors that keeps each error in the bucket it was originally
+// classified into, so that HasCriticalErrors and IsRetryable describe the set
+// as a whole. It returns nil when none of the restores failed.
+//
+// Concurrent restores need this: classifying their combined failure by
+// whichever error happened to arrive first would make the retry decision
+// depend on scheduling, and a merged error built with NewPGRestoreErrors would
+// demote every nested restore error to critical, since a PGRestoreErrors
+// matches none of the classifications addError checks for.
+func MergePGRestoreErrors(errs ...error) error {
+	merged := &PGRestoreErrors{}
+	for _, err := range errs {
+		if err == nil {
+			continue
+		}
+		restoreErrs := &PGRestoreErrors{}
+		if errors.As(err, &restoreErrs) {
+			merged.ignoredErrs = append(merged.ignoredErrs, restoreErrs.ignoredErrs...)
+			merged.criticalErrs = append(merged.criticalErrs, restoreErrs.criticalErrs...)
+			merged.retryableErrs = append(merged.retryableErrs, restoreErrs.retryableErrs...)
+			continue
+		}
+		merged.addError(err)
+	}
+
+	if !merged.HasErrors() {
+		return nil
+	}
+	return merged
+}
+
 func (e PGRestoreErrors) Error() string {
 	if !e.HasErrors() {
 		return ""
 	}
+	return errors.Join(e.Unwrap()...).Error()
+}
+
+func (e PGRestoreErrors) HasErrors() bool {
+	return len(e.criticalErrs) > 0 || len(e.retryableErrs) > 0 || len(e.ignoredErrs) > 0
+}
 
+// Unwrap exposes the individual restore errors, so that a caller holding the
+// collection can still interrogate it with errors.Is and errors.As rather than
+// by matching on the joined message.
+func (e PGRestoreErrors) Unwrap() []error {
 	all := make([]error, 0, len(e.criticalErrs)+len(e.retryableErrs)+len(e.ignoredErrs))
 	all = append(all, e.criticalErrs...)
 	all = append(all, e.retryableErrs...)
 	all = append(all, e.ignoredErrs...)
-	return errors.Join(all...).Error()
-}
-
-func (e PGRestoreErrors) HasErrors() bool {
-	return len(e.criticalErrs) > 0 || len(e.retryableErrs) > 0 || len(e.ignoredErrs) > 0
+	return all
 }
 
 // HasCriticalErrors reports whether the restore hit an error that must not be

pkg/snapshot/generator/postgres/schema/pgdumprestore/exotic_dump_shapes_test.go

@@ -0,0 +1,97 @@
+// SPDX-License-Identifier: Apache-2.0
+
+package pgdumprestore
+
+import (
+	"os"
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+// TestParseDump_exoticShapes runs a real pg_dump --schema-only output through
+// the parser and the block partitioning that the concurrent index restore
+// depends on. The fixture is captured verbatim from postgres and covers the
+// shapes whose classification is not obvious from the code: an exclusion
+// constraint, a partial index with a multi-condition WHERE clause, a long
+// expression index, a constraint whose index is reused (UNIQUE), CLUSTER ON,
+// a table partition attachment and index partition attachments.
+//
+// Two properties matter for the parallel restore and neither is visible from
+// the parser alone: a statement must land in the group that can run it, and no
+// statement may be dropped or merged into another on the way there.
+func TestParseDump_exoticShapes(t *testing.T) {
+	t.Parallel()
+
+	dump, err := os.ReadFile("testdata/exotic_schema_dump.sql")
+	require.NoError(t, err)
+
+	sg := SnapshotGenerator{objectTypeFilter: &objectTypeFilter{}}
+	parsed := sg.parseDump(dump)
+	connectBlocks, indexBlocks, otherBlocks := partitionDumpBlocks(parsed.indicesAndConstraints, isIndexStatement)
+	require.Empty(t, connectBlocks)
+
+	// pg_dump emits every CREATE INDEX on a single line, however long the
+	// expression or WHERE clause, so each one becomes its own block and can be
+	// restored on its own connection
+	require.ElementsMatch(t, []string{
+		"CREATE INDEX parts_val_idx ON ONLY public.parts USING btree (val);",
+		"CREATE INDEX parts_0_val_idx ON public.parts_0 USING btree (val);",
+		"CREATE INDEX parts_1_val_idx ON public.parts_1 USING btree (val);",
+		"CREATE INDEX rooms_expr_idx ON public.rooms USING btree (lower(name), upper(email), ((val * (2)::numeric)), COALESCE(name, email, 'a-fairly-long-default-value'::text));",
+		"CREATE INDEX rooms_partial_idx ON public.rooms USING btree (name) WHERE ((val > (100)::numeric) AND (name IS NOT NULL) AND (email IS NOT NULL));",
+	}, indexBlocks)
+
+	// everything that builds its own index or references one has to wait for
+	// the index phase, and keeps its relative order
+	require.Equal(t, []string{
+		"ALTER TABLE ONLY public.rooms\n    ADD CONSTRAINT rooms_email_key UNIQUE (email);",
+		"ALTER TABLE public.rooms CLUSTER ON rooms_email_key;",
+		"ALTER TABLE ONLY public.rooms\n    ADD CONSTRAINT rooms_no_overlap EXCLUDE USING gist (room WITH =, during WITH &&);",
+		"COMMENT ON INDEX public.rooms_partial_idx IS 'partial';",
+		"ALTER INDEX public.parts_val_idx ATTACH PARTITION public.parts_0_val_idx;",
+		"ALTER INDEX public.parts_val_idx ATTACH PARTITION public.parts_1_val_idx;",
+	}, otherBlocks)
+
+	// a table partition attachment is not an index dependency and stays in the
+	// schema dump, which is restored before the data
+	require.Contains(t, string(parsed.filtered), "ALTER TABLE ONLY public.parts ATTACH PARTITION public.parts_0 FOR VALUES FROM (0) TO (10);")
+
+	// no index or constraint statement may be lost or glued to its neighbour
+	// on the way into a group. Statements outside that domain are excluded:
+	// the parser deliberately drops some of them (ownership of excluded roles,
+	// legacy PL/pgSQL handlers, filtered security labels).
+	restored := string(parsed.filtered) + string(parsed.indicesAndConstraints) + string(parsed.views)
+	for _, line := range strings.Split(string(dump), "\n") {
+		line = strings.TrimSpace(line)
+		if !isIndexStatement(line) && !strings.HasPrefix(line, "ADD CONSTRAINT") &&
+			!isAttachPartitionIndexStatement(line) && !strings.HasPrefix(line, "COMMENT ON INDEX") &&
+			!isClusterOnAlterTable(line) {
+			continue
+		}
+		require.Contains(t, restored, line, "statement dropped by the dump parsing")
+	}
+}
+
+// TestParseDump_singleLineAddConstraint covers the shape the block splitting
+// would be most sensitive to: an ALTER TABLE ... ADD CONSTRAINT on one line.
+// Real pg_dump wraps these onto two lines (see the fixture above), so this
+// exercises the defensive branch rather than an observed output, and pins that
+// consecutive constraints stay separate blocks instead of merging into one.
+func TestParseDump_singleLineAddConstraint(t *testing.T) {
+	t.Parallel()
+
+	sg := SnapshotGenerator{objectTypeFilter: &objectTypeFilter{}}
+	parsed := sg.parseDump([]byte(
+		"ALTER TABLE public.a ADD CONSTRAINT a_check CHECK (v > 0);\n" +
+			"ALTER TABLE public.b ADD CONSTRAINT b_check CHECK (v > 0);\n" +
+			"CREATE INDEX a_idx ON public.a USING btree (v);\n"))
+
+	_, indexBlocks, otherBlocks := partitionDumpBlocks(parsed.indicesAndConstraints, isIndexStatement)
+	require.Equal(t, []string{"CREATE INDEX a_idx ON public.a USING btree (v);"}, indexBlocks)
+	require.Equal(t, []string{
+		"ALTER TABLE public.a ADD CONSTRAINT a_check CHECK (v > 0);",
+		"ALTER TABLE public.b ADD CONSTRAINT b_check CHECK (v > 0);",
+	}, otherBlocks)
+}