Merge branch 'feat/physical-restorer-custom-options' into 'master'

feat(retrieval): pass custom flags to physical restore tools

Closes #760

See merge request postgres-ai/database-lab!1182

Author: Sarumyan

engine/configs/config.example.physical_pgbackrest.yml

@@ -91,6 +91,8 @@ retrieval:  # Data retrieval: initial sync and ongoing updates. Two methods:
             # standby_mode: on
             # recovery_target_timeline: 'latest'
 
+        # These envs reach the restore and sync containers only; the promotion container reads the
+        # physicalSnapshot envs below, so anything WAL replay needs must be set in both places.
         envs: # Environment variables for pgBackRest; see https://pgbackrest.org/user-guide.html
           PGBACKREST_LOG_LEVEL_CONSOLE: detail # Log level; options: off, error, warn, info, detail, debug, trace
           PGBACKREST_PROCESS_MAX: 2 # Maximum number of processes to use for compression/decompression
@@ -111,6 +113,21 @@ retrieval:  # Data retrieval: initial sync and ongoing updates. Two methods:
         pgbackrest: # pgBackRest specific configuration
           stanza: stanzaName # Stanza name (must match the stanza configured in your pgBackRest setup)
           delta: false # Use delta restore; set to true for incremental restores from last backup
+          # Uncomment to append extra flags to the "pgbackrest restore" command.
+          # The list is joined into a shell command line, so entries follow shell quoting rules: quote a value
+          # containing whitespace or glob characters.
+          # Do not repeat --type, --stanza, --pg1-path or --delta: pgBackRest rejects a single-valued option set
+          # twice ("cannot be set multiple times"). List options such as --db-include may be repeated freely.
+          # --recovery-option may be repeated too, but where it ends up depends on the Postgres version, and
+          # neither outcome is what a reader expects. On 12 and newer pgBackRest writes it to
+          # postgresql.auto.conf, read after the recovery config DBLab generates, so it overrides DBLab's
+          # restore_command; on 11 and older it goes to recovery.conf, which DBLab then truncates and rewrites,
+          # so it is dropped without a word. Never set recovery_target* through it either way - promotion relies
+          # on recovery_target=immediate with recovery_target_action=promote.
+          # Choose a repository with the PGBACKREST_REPO env rather than --repo, since restore and archive-get
+          # both read it.
+          # customOptions:
+          #   - "--db-include=mydb"
 
     physicalSnapshot:
       options:

engine/configs/config.example.physical_walg.yml

@@ -91,12 +91,24 @@ retrieval:  # Data retrieval: initial sync and ongoing updates. Two methods:
             # standby_mode: on
             # recovery_target_timeline: 'latest'
 
+        # These envs reach the restore and sync containers only; the promotion container reads the
+        # physicalSnapshot envs below, so anything WAL replay needs must be set in both places.
         envs: # Environment variables for WAL-G; see https://github.com/wal-g/wal-g/blob/master/docs/README.md
           WALG_GS_PREFIX: "gs://{BUCKET}/{SCOPE}" # Google Storage prefix for WAL-G backups
           GOOGLE_APPLICATION_CREDENTIALS: "/tmp/sa.json" # Path to Google service account credentials
 
         walg: # WAL-G specific configuration
           backupName: LATEST # Which backup to restore; use "LATEST" for most recent backup
+          # Uncomment to append extra flags to the "wal-g backup-fetch" command, for options wal-g exposes
+          # only as flags (--mask, --restore-spec, --restore-only); anything with a WALG_* variable belongs in envs.
+          # The list is joined into a shell command line, so entries follow shell quoting rules: quote a value
+          # containing whitespace or glob characters, e.g. - '--restore-only=my_db,"another db"'.
+          # The flags reach backup-fetch alone - not the backup-list that resolves LATEST, and not the wal-fetch
+          # in restore_command - so choose a storage with WALG_TARGET_STORAGE, which all three read, not --target-storage.
+          # --target-user-data cannot be combined with backupName, in either form: wal-g rejects the flag and
+          # the WALG_FETCH_TARGET_USER_DATA env alike when a backup name is also given.
+          # customOptions:
+          #   - "--restore-only=mydb"
 
     physicalSnapshot:
       options:

engine/internal/retrieval/engine/postgres/physical/pgbackrest.go

@@ -7,6 +7,7 @@ package physical
 import (
 	"context"
 	"fmt"
+	"strings"
 
 	"gitlab.com/postgres-ai/database-lab/v3/internal/retrieval/engine/postgres/tools/defaults"
 )
@@ -21,8 +22,9 @@ type pgbackrest struct {
 }
 
 type pgbackrestOptions struct {
-	Stanza string `yaml:"stanza"`
-	Delta  bool   `yaml:"delta"`
+	Stanza        string   `yaml:"stanza"`
+	Delta         bool     `yaml:"delta"`
+	CustomOptions []string `yaml:"customOptions"`
 }
 
 func newPgBackRest(options pgbackrestOptions) *pgbackrest {
@@ -40,6 +42,10 @@ func (p *pgbackrest) GetRestoreCommand() string {
 		restoreCmd += " --delta"
 	}
 
+	if len(p.options.CustomOptions) > 0 {
+		restoreCmd += " " + strings.Join(p.options.CustomOptions, " ")
+	}
+
 	return restoreCmd
 }
 

engine/internal/retrieval/engine/postgres/physical/pgbackrest_test.go

@@ -66,6 +66,31 @@ func TestPgBackRestRestoreCommand_DifferentStanzas(t *testing.T) {
 	}
 }
 
+func TestPgBackRestRestoreCommand_CustomOptions(t *testing.T) {
+	baseCmd := "sudo -Eu postgres pgbackrest --type=standby --pg1-path=${PGDATA} --stanza=stanzaName restore " +
+		"--recovery-option=restore_command='pgbackrest --pg1-path=${PGDATA} --stanza=stanzaName archive-get %f %p'"
+
+	testCases := []struct {
+		name          string
+		delta         bool
+		customOptions []string
+		expected      string
+	}{
+		{name: "no custom options", customOptions: nil, expected: baseCmd},
+		{name: "empty custom options", customOptions: []string{}, expected: baseCmd},
+		{name: "single option", customOptions: []string{"--db-include=mydb"}, expected: baseCmd + " --db-include=mydb"},
+		{name: "multiple options", customOptions: []string{"--db-include=mydb", "--process-max=4"}, expected: baseCmd + " --db-include=mydb --process-max=4"},
+		{name: "appended after delta", delta: true, customOptions: []string{"--db-include=mydb"}, expected: baseCmd + " --delta --db-include=mydb"},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			p := newPgBackRest(pgbackrestOptions{Stanza: "stanzaName", Delta: tc.delta, CustomOptions: tc.customOptions})
+			assert.Equal(t, tc.expected, p.GetRestoreCommand())
+		})
+	}
+}
+
 func TestPgBackRestRecoveryConfig_VersionBoundary(t *testing.T) {
 	p := newPgBackRest(pgbackrestOptions{Stanza: "mydb"})
 

engine/internal/retrieval/engine/postgres/physical/physical_test.go

@@ -11,6 +11,8 @@ import (
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+
+	"gitlab.com/postgres-ai/database-lab/v3/internal/provision/resources"
 )
 
 func TestInitParamsExtraction(t *testing.T) {
@@ -83,3 +85,33 @@ some_other_param setting:             value
 	require.NoError(t, err)
 	assert.Empty(t, settings)
 }
+
+func TestRestoreJobReload_CustomOptions(t *testing.T) {
+	r := &RestoreJob{fsPool: &resources.Pool{MountDir: "/", PoolDirName: "pgdata"}}
+
+	err := r.Reload(map[string]interface{}{
+		"tool": walgTool,
+		"walg": map[string]interface{}{
+			"backupName":    "LATEST",
+			"customOptions": []interface{}{"--mask", "'base/*'"},
+		},
+		"pgbackrest": map[string]interface{}{
+			"stanza":        "main",
+			"customOptions": []interface{}{"--db-include=mydb"},
+		},
+	})
+	require.NoError(t, err)
+
+	assert.Equal(t, []string{"--mask", "'base/*'"}, r.WALG.CustomOptions)
+	assert.Equal(t, []string{"--db-include=mydb"}, r.PgBackRest.CustomOptions)
+
+	walgRestorer, err := r.getRestorer(nil, walgTool)
+	require.NoError(t, err)
+	assert.Equal(t, "wal-g backup-fetch /pgdata LATEST --mask 'base/*'", walgRestorer.GetRestoreCommand())
+
+	pgBackRestRestorer, err := r.getRestorer(nil, pgbackrestTool)
+	require.NoError(t, err)
+	assert.Equal(t, "sudo -Eu postgres pgbackrest --type=standby --pg1-path=${PGDATA} --stanza=main restore "+
+		"--recovery-option=restore_command='pgbackrest --pg1-path=${PGDATA} --stanza=main archive-get %f %p' --db-include=mydb",
+		pgBackRestRestorer.GetRestoreCommand())
+}

engine/internal/retrieval/engine/postgres/physical/wal_g.go

@@ -33,7 +33,8 @@ type walg struct {
 }
 
 type walgOptions struct {
-	BackupName string `yaml:"backupName"`
+	BackupName    string   `yaml:"backupName"`
+	CustomOptions []string `yaml:"customOptions"`
 }
 
 func newWALG(dockerClient *client.Client, pgDataDir string, options walgOptions) *walg {
@@ -49,7 +50,13 @@ func newWALG(dockerClient *client.Client, pgDataDir string, options walgOptions)
 
 // GetRestoreCommand returns a command to restore data.
 func (w *walg) GetRestoreCommand() string {
-	return fmt.Sprintf("wal-g backup-fetch %s %s", w.pgDataDir, w.parsedBackupName)
+	restoreCmd := fmt.Sprintf("wal-g backup-fetch %s %s", w.pgDataDir, w.parsedBackupName)
+
+	if len(w.options.CustomOptions) > 0 {
+		restoreCmd += " " + strings.Join(w.options.CustomOptions, " ")
+	}
+
+	return restoreCmd
 }
 
 // GetRecoveryConfig returns a recovery config to restore data.

engine/internal/retrieval/engine/postgres/physical/wal_g_test.go

@@ -51,6 +51,28 @@ func TestWALGGetRestoreCommand(t *testing.T) {
 	}
 }
 
+func TestWALGGetRestoreCommand_CustomOptions(t *testing.T) {
+	testCases := []struct {
+		name          string
+		customOptions []string
+		expected      string
+	}{
+		{name: "no custom options", customOptions: nil, expected: "wal-g backup-fetch /pgdata LATEST"},
+		{name: "empty custom options", customOptions: []string{}, expected: "wal-g backup-fetch /pgdata LATEST"},
+		{name: "single option", customOptions: []string{"--reverse-unpack"}, expected: "wal-g backup-fetch /pgdata LATEST --reverse-unpack"},
+		{name: "option with value", customOptions: []string{"--mask", "'base/*'"}, expected: "wal-g backup-fetch /pgdata LATEST --mask 'base/*'"},
+		{name: "multiple options", customOptions: []string{"--reverse-unpack", "--skip-redundant-tars"}, expected: "wal-g backup-fetch /pgdata LATEST --reverse-unpack --skip-redundant-tars"},
+		{name: "quotes are passed through unmodified", customOptions: []string{`--restore-only=my_db,"another db"`}, expected: `wal-g backup-fetch /pgdata LATEST --restore-only=my_db,"another db"`},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			w := newWALG(nil, "/pgdata", walgOptions{BackupName: "LATEST", CustomOptions: tc.customOptions})
+			assert.Equal(t, tc.expected, w.GetRestoreCommand())
+		})
+	}
+}
+
 func TestWALGVersionParse_TableDriven(t *testing.T) {
 	testCases := []struct {
 		name     string