feat: include/restore statistics in snapshots (PG18+)

Author: radim

cmd/snapshot.go

@@ -389,19 +389,25 @@ func runSnapshotRestore() error {
 		return fmt.Errorf("database connection failed: %w", err)
 	}
 
+	// Connect to check database state and version
+	db, err := regresql.OpenDB(cfg.PgUri)
+	if err != nil {
+		return fmt.Errorf("failed to connect: %w", err)
+	}
+	defer db.Close()
+
 	// Check if database has existing tables and warn if --clean not specified
 	if !snapshotClean {
-		db, err := regresql.OpenDB(cfg.PgUri)
-		if err == nil {
-			defer db.Close()
-			var tableCount int
-			db.QueryRow("SELECT COUNT(*) FROM information_schema.tables WHERE table_schema = 'public' AND table_type = 'BASE TABLE'").Scan(&tableCount)
-			if tableCount > 0 {
-				return fmt.Errorf("database has %d existing table(s). Use --clean to drop them before restore, or manually clear the database", tableCount)
-			}
+		var tableCount int
+		db.QueryRow("SELECT COUNT(*) FROM information_schema.tables WHERE table_schema = 'public' AND table_type = 'BASE TABLE'").Scan(&tableCount)
+		if tableCount > 0 {
+			return fmt.Errorf("database has %d existing table(s). Use --clean to drop them before restore, or manually clear the database", tableCount)
 		}
 	}
 
+	// Detect PostgreSQL version for statistics support
+	serverCtx, _ := regresql.CaptureServerContext(db)
+
 	inputPath := snapshotInput
 	if inputPath == "" {
 		inputPath = regresql.GetSnapshotPath(cfg.Snapshot, snapshotCwd)
@@ -422,10 +428,12 @@ func runSnapshotRestore() error {
 		return err
 	}
 
+	withStats := serverCtx != nil && serverCtx.MajorVersion() >= 18
 	opts := regresql.RestoreOptions{
-		InputPath: inputPath,
-		Format:    format,
-		Clean:     snapshotClean,
+		InputPath:      inputPath,
+		Format:         format,
+		Clean:          snapshotClean,
+		WithStatistics: withStats,
 	}
 
 	fmt.Printf("Restoring database snapshot...\n")
@@ -434,6 +442,9 @@ func runSnapshotRestore() error {
 	if snapshotClean {
 		fmt.Printf("  Mode:     clean (drop existing objects)\n")
 	}
+	if withStats {
+		fmt.Printf("  Stats:    restoring optimizer statistics (PG18+)\n")
+	}
 	fmt.Println()
 
 	if err := regresql.RestoreSnapshot(cfg.PgUri, opts); err != nil {

regresql/snapshot.go

@@ -75,10 +75,11 @@ type (
 	SnapshotFormat string
 
 	SnapshotOptions struct {
-		OutputPath string
-		Format     SnapshotFormat
-		SchemaOnly bool
-		Section    string
+		OutputPath     string
+		Format         SnapshotFormat
+		SchemaOnly     bool
+		Section        string
+		WithStatistics bool // PostgreSQL 18+: include optimizer statistics
 	}
 
 	SectionsOptions struct {
@@ -102,6 +103,7 @@ type (
 		Format         SnapshotFormat
 		Clean          bool   // drop existing objects before restore
 		TargetDatabase string // override database name from connection string
+		WithStatistics bool   // PostgreSQL 18+: restore optimizer statistics
 	}
 )
 
@@ -319,6 +321,11 @@ func buildPgDumpArgs(pguri string, opts SnapshotOptions) []string {
 		args = append(args, "--section", opts.Section)
 	}
 
+	// PostgreSQL 18+: include optimizer statistics in dump (requires pg_dump 18+)
+	if opts.WithStatistics && parseToolMajorVersion("pg_dump") >= 18 {
+		args = append(args, "--statistics")
+	}
+
 	return args
 }
 
@@ -588,6 +595,11 @@ func restoreWithPgRestore(pguri string, opts RestoreOptions, format SnapshotForm
 		args = append(args, "--clean", "--if-exists")
 	}
 
+	// PostgreSQL 18+: restore optimizer statistics (requires pg_restore 18+)
+	if opts.WithStatistics && parseToolMajorVersion("pg_restore") >= 18 {
+		args = append(args, "--statistics")
+	}
+
 	switch format {
 	case FormatCustom:
 		args = append(args, "--format=custom")

regresql/snapshot_build.go

@@ -147,13 +147,22 @@ func BuildSnapshot(basePgUri string, root string, opts SnapshotBuildOptions) (*s
 		return nil, fmt.Errorf("failed to capture server context: %w", err)
 	}
 
+	// Run ANALYZE to ensure statistics are up to date before pg_dump
+	if opts.Verbose {
+		fmt.Printf("Running ANALYZE...\n")
+	}
+	if _, err := db.Exec("ANALYZE"); err != nil {
+		return nil, fmt.Errorf("failed to analyze database: %w", err)
+	}
+
 	if opts.Verbose {
 		fmt.Printf("Capturing snapshot with pg_dump...\n")
 	}
 
 	info, err := CaptureSnapshot(tempDB.PgUri, SnapshotOptions{
-		OutputPath: opts.OutputPath,
-		Format:     opts.Format,
+		OutputPath:     opts.OutputPath,
+		Format:         opts.Format,
+		WithStatistics: serverCtx.MajorVersion() >= 18,
 	})
 	if err != nil {
 		return nil, fmt.Errorf("failed to capture snapshot: %w", err)