feat: add restore_database config option for snapshot restore

Author: radim

regresql/config.go

@@ -35,6 +35,7 @@ type (
 		MigrationCommand string   `yaml:"migration_command,omitempty"` // external command to run migrations (e.g., goose, migrate)
 		Fixtures         []string `yaml:"fixtures,omitempty"`          // SQL/YAML fixture files for snapshot build
 		AutoRestore      *bool    `yaml:"auto_restore,omitempty"`      // restore snapshot before test (default: true if path is set)
+		RestoreDatabase  string   `yaml:"restore_database,omitempty"`  // override target database for restore
 	}
 )
 

regresql/regresql.go

@@ -146,8 +146,9 @@ func autoRestore(cfg config, root string, noRestore bool) {
 	}
 	fmt.Printf("Restoring snapshot: %s\n", snapshotPath)
 	opts := RestoreOptions{
-		InputPath: snapshotPath,
-		Clean:     true,
+		InputPath:      snapshotPath,
+		Clean:          true,
+		TargetDatabase: cfg.Snapshot.RestoreDatabase,
 	}
 	if err := RestoreSnapshot(cfg.PgUri, opts); err != nil {
 		fmt.Printf("Error: failed to restore snapshot: %s\n", err)

regresql/snapshot.go

@@ -5,6 +5,7 @@ import (
 	"encoding/hex"
 	"fmt"
 	"io"
+	"net/url"
 	"os"
 	"os/exec"
 	"path/filepath"
@@ -61,9 +62,10 @@ type (
 	}
 
 	RestoreOptions struct {
-		InputPath string
-		Format    SnapshotFormat
-		Clean     bool // drop existing objects before restore
+		InputPath      string
+		Format         SnapshotFormat
+		Clean          bool   // drop existing objects before restore
+		TargetDatabase string // override database name from connection string
 	}
 )
 
@@ -428,21 +430,41 @@ func ShouldAutoRestore(cfg *SnapshotConfig) bool {
 	return *cfg.AutoRestore
 }
 
+// replaceDatabase returns a new connection string with a different database
+func replaceDatabase(pguri, newDB string) (string, error) {
+	u, err := url.Parse(pguri)
+	if err != nil {
+		return "", err
+	}
+	u.Path = "/" + newDB
+	return u.String(), nil
+}
+
 // RestoreSnapshot restores a database snapshot using pg_restore or psql
 func RestoreSnapshot(pguri string, opts RestoreOptions) error {
 	if _, err := os.Stat(opts.InputPath); os.IsNotExist(err) {
 		return fmt.Errorf("snapshot file not found: %s", opts.InputPath)
 	}
 
+	// Override target database if specified
+	targetURI := pguri
+	if opts.TargetDatabase != "" {
+		var err error
+		targetURI, err = replaceDatabase(pguri, opts.TargetDatabase)
+		if err != nil {
+			return fmt.Errorf("failed to set target database: %w", err)
+		}
+	}
+
 	format := opts.Format
 	if format == "" {
 		format = DetectSnapshotFormat(opts.InputPath)
 	}
 
 	if format == FormatPlain {
-		return restoreWithPsql(pguri, opts)
+		return restoreWithPsql(targetURI, opts)
 	}
-	return restoreWithPgRestore(pguri, opts, format)
+	return restoreWithPgRestore(targetURI, opts, format)
 }
 
 func restoreWithPgRestore(pguri string, opts RestoreOptions, format SnapshotFormat) error {