From 6aac8d00477d6b9fefb50399c444b58f5be73419 Mon Sep 17 00:00:00 2001 From: Denis Date: Fri, 13 Feb 2026 18:13:01 +0200 Subject: [PATCH] feat: add built-in databaseRename option for snapshot jobs Add native databaseRename configuration to physicalSnapshot and logicalSnapshot jobs, eliminating the need for custom preprocessing scripts for this common operation. The feature spins up a temporary container and runs ALTER DATABASE RENAME for each configured mapping. Co-Authored-By: Claude Opus 4.6 --- .../config.example.logical_generic.yml | 4 +- .../configs/config.example.physical_walg.yml | 3 + .../engine/postgres/snapshot/logical.go | 7 + .../engine/postgres/snapshot/physical.go | 7 + .../engine/postgres/snapshot/rename.go | 163 ++++++++++++++++++ .../engine/postgres/snapshot/rename_test.go | 62 +++++++ .../engine/postgres/tools/cont/container.go | 2 + 7 files changed, 247 insertions(+), 1 deletion(-) create mode 100644 engine/internal/retrieval/engine/postgres/snapshot/rename.go create mode 100644 engine/internal/retrieval/engine/postgres/snapshot/rename_test.go diff --git a/engine/configs/config.example.logical_generic.yml b/engine/configs/config.example.logical_generic.yml index 8618e472..bff93df2 100644 --- a/engine/configs/config.example.logical_generic.yml +++ b/engine/configs/config.example.logical_generic.yml @@ -135,7 +135,9 @@ retrieval: # Data retrieval: initial sync and ongoing updates. Two methods: options: <<: *db_configs # Adjust PostgreSQL configuration preprocessingScript: "" # Pre-processing script for data scrubbing/masking; e.g., "/tmp/scripts/custom.sh" - + databaseRename: # Rename databases before finalizing snapshot; runs after preprocessingScript; default: empty (disabled) + # mydb_prod: mydb_dblab # Rename "mydb_prod" to "mydb_dblab" + dataPatching: # Pre-processing SQL queries for data patching <<: *db_container queryPreprocessing: diff --git a/engine/configs/config.example.physical_walg.yml b/engine/configs/config.example.physical_walg.yml index 297c7609..a8fe25de 100644 --- a/engine/configs/config.example.physical_walg.yml +++ b/engine/configs/config.example.physical_walg.yml @@ -104,6 +104,9 @@ retrieval: # Data retrieval: initial sync and ongoing updates. Two methods: # recovery_target_timeline: 'latest' preprocessingScript: "" # Shell script path to execute before finalizing snapshot; example: "/tmp/scripts/custom.sh"; default: "" (disabled) + databaseRename: # Rename databases before finalizing snapshot; runs after preprocessingScript; default: empty (disabled) + # example_production: example_dblab # Rename "example_production" to "example_dblab" + # analytics_prod: analytics_dblab scheduler: # Snapshot scheduling and retention policy configuration snapshot: # Snapshot creation scheduling timetable: "0 */6 * * *" # Cron expression defining snapshot schedule: https://en.wikipedia.org/wiki/Cron#Overview diff --git a/engine/internal/retrieval/engine/postgres/snapshot/logical.go b/engine/internal/retrieval/engine/postgres/snapshot/logical.go index 042d28f9..1efea6f6 100644 --- a/engine/internal/retrieval/engine/postgres/snapshot/logical.go +++ b/engine/internal/retrieval/engine/postgres/snapshot/logical.go @@ -63,6 +63,7 @@ type LogicalInitial struct { type LogicalOptions struct { DataPatching DataPatching `yaml:"dataPatching"` PreprocessingScript string `yaml:"preprocessingScript"` + DatabaseRename map[string]string `yaml:"databaseRename"` Configs map[string]string `yaml:"configs"` Schedule Scheduler `yaml:"schedule"` } @@ -127,6 +128,12 @@ func (s *LogicalInitial) Run(ctx context.Context) error { } } + if len(s.options.DatabaseRename) > 0 { + if err := runDatabaseRename(ctx, s.dockerClient, s.engineProps, s.globalCfg, s.fsPool.DataDir(), s.options.DatabaseRename); err != nil { + return errors.Wrap(err, "failed to rename databases") + } + } + if err := s.touchConfigFiles(); err != nil { return errors.Wrap(err, "failed to create PostgreSQL configuration files") } diff --git a/engine/internal/retrieval/engine/postgres/snapshot/physical.go b/engine/internal/retrieval/engine/postgres/snapshot/physical.go index 2d95e699..8c95ab40 100644 --- a/engine/internal/retrieval/engine/postgres/snapshot/physical.go +++ b/engine/internal/retrieval/engine/postgres/snapshot/physical.go @@ -114,6 +114,7 @@ type PhysicalOptions struct { SkipStartSnapshot bool `yaml:"skipStartSnapshot"` Promotion Promotion `yaml:"promotion"` PreprocessingScript string `yaml:"preprocessingScript"` + DatabaseRename map[string]string `yaml:"databaseRename"` Configs map[string]string `yaml:"configs"` Sysctls map[string]string `yaml:"sysctls"` Envs map[string]string `yaml:"envs"` @@ -388,6 +389,12 @@ func (p *PhysicalInitial) run(ctx context.Context) (err error) { } } + if len(p.options.DatabaseRename) > 0 { + if err := runDatabaseRename(ctx, p.dockerClient, p.engineProps, p.globalCfg, cloneDataDir, p.options.DatabaseRename); err != nil { + return errors.Wrap(err, "failed to rename databases") + } + } + // Mark database data. if err := p.markDatabaseData(); err != nil { return errors.Wrap(err, "failed to mark the prepared data") diff --git a/engine/internal/retrieval/engine/postgres/snapshot/rename.go b/engine/internal/retrieval/engine/postgres/snapshot/rename.go new file mode 100644 index 00000000..ae833aed --- /dev/null +++ b/engine/internal/retrieval/engine/postgres/snapshot/rename.go @@ -0,0 +1,163 @@ +/* +2024 © Postgres.ai +*/ + +package snapshot + +import ( + "context" + "fmt" + + "github.com/docker/docker/api/types/container" + "github.com/docker/docker/api/types/filters" + "github.com/docker/docker/client" + "github.com/pkg/errors" + + "gitlab.com/postgres-ai/database-lab/v3/internal/diagnostic" + "gitlab.com/postgres-ai/database-lab/v3/internal/retrieval/engine/postgres/tools" + "gitlab.com/postgres-ai/database-lab/v3/internal/retrieval/engine/postgres/tools/cont" + "gitlab.com/postgres-ai/database-lab/v3/internal/retrieval/engine/postgres/tools/health" + "gitlab.com/postgres-ai/database-lab/v3/pkg/config/global" + "gitlab.com/postgres-ai/database-lab/v3/pkg/log" +) + +const renameContainerPrefix = "dblab_rename_" + +// runDatabaseRename renames databases using ALTER DATABASE in a temporary container. +func runDatabaseRename( + ctx context.Context, + dockerClient *client.Client, + engineProps *global.EngineProps, + globalCfg *global.Config, + dataDir string, + renames map[string]string, +) error { + if len(renames) == 0 { + return nil + } + + connDB := globalCfg.Database.Name() + + if err := validateDatabaseRenames(renames, connDB); err != nil { + return err + } + + pgVersion, err := tools.DetectPGVersion(dataDir) + if err != nil { + return errors.Wrap(err, "failed to detect postgres version") + } + + image := fmt.Sprintf("postgresai/extended-postgres:%g", pgVersion) + + if err := tools.PullImage(ctx, dockerClient, image); err != nil { + return errors.Wrap(err, "failed to pull image for database rename") + } + + pwd, err := tools.GeneratePassword() + if err != nil { + return errors.Wrap(err, "failed to generate password") + } + + hostConfig, err := cont.BuildHostConfig(ctx, dockerClient, dataDir, nil) + if err != nil { + return errors.Wrap(err, "failed to build host config") + } + + containerName := renameContainerPrefix + engineProps.InstanceID + + containerID, err := tools.CreateContainerIfMissing(ctx, dockerClient, containerName, + &container.Config{ + Labels: map[string]string{ + cont.DBLabControlLabel: cont.DBLabRenameLabel, + cont.DBLabInstanceIDLabel: engineProps.InstanceID, + cont.DBLabEngineNameLabel: engineProps.ContainerName, + }, + Env: []string{ + "PGDATA=" + dataDir, + "POSTGRES_PASSWORD=" + pwd, + }, + Image: image, + Healthcheck: health.GetConfig( + globalCfg.Database.User(), + connDB, + ), + }, + hostConfig, + ) + if err != nil { + return fmt.Errorf("failed to create rename container: %w", err) + } + + defer tools.RemoveContainer(ctx, dockerClient, containerID, cont.StopPhysicalTimeout) + + defer func() { + if err != nil { + tools.PrintContainerLogs(ctx, dockerClient, containerName) + tools.PrintLastPostgresLogs(ctx, dockerClient, containerName, dataDir) + + filterArgs := filters.NewArgs( + filters.KeyValuePair{Key: "label", + Value: fmt.Sprintf("%s=%s", cont.DBLabControlLabel, cont.DBLabRenameLabel)}) + + if diagErr := diagnostic.CollectDiagnostics(ctx, dockerClient, filterArgs, containerName, dataDir); diagErr != nil { + log.Err("failed to collect rename container diagnostics", diagErr) + } + } + }() + + log.Msg(fmt.Sprintf("Running rename container: %s. ID: %v", containerName, containerID)) + + if err = dockerClient.ContainerStart(ctx, containerID, container.StartOptions{}); err != nil { + return errors.Wrap(err, "failed to start rename container") + } + + log.Msg("Waiting for rename container readiness") + log.Msg(fmt.Sprintf("View logs using the command: %s %s", tools.ViewLogsCmd, containerName)) + + if err = tools.CheckContainerReadiness(ctx, dockerClient, containerID); err != nil { + return errors.Wrap(err, "rename container readiness check failed") + } + + for oldName, newName := range renames { + log.Msg(fmt.Sprintf("Renaming database %q to %q", oldName, newName)) + + cmd := buildRenameCommand(globalCfg.Database.User(), connDB, oldName, newName) + + output, execErr := tools.ExecCommandWithOutput(ctx, dockerClient, containerID, container.ExecOptions{Cmd: cmd}) + if execErr != nil { + err = errors.Wrapf(execErr, "failed to rename database %q to %q", oldName, newName) + return err + } + + log.Msg("Rename result: ", output) + } + + if err = tools.RunCheckpoint(ctx, dockerClient, containerID, globalCfg.Database.User(), connDB); err != nil { + return errors.Wrap(err, "failed to run checkpoint after rename") + } + + if err = tools.StopPostgres(ctx, dockerClient, containerID, dataDir, tools.DefaultStopTimeout); err != nil { + return errors.Wrap(err, "failed to stop postgres after rename") + } + + return nil +} + +func buildRenameCommand(username, connDB, oldName, newName string) []string { + return []string{ + "psql", + "-U", username, + "-d", connDB, + "-XAtc", fmt.Sprintf(`ALTER DATABASE "%s" RENAME TO "%s"`, oldName, newName), + } +} + +func validateDatabaseRenames(renames map[string]string, connDB string) error { + for oldName := range renames { + if oldName == connDB { + return fmt.Errorf("cannot rename database %q: it is used as the connection database", oldName) + } + } + + return nil +} diff --git a/engine/internal/retrieval/engine/postgres/snapshot/rename_test.go b/engine/internal/retrieval/engine/postgres/snapshot/rename_test.go new file mode 100644 index 00000000..b762188e --- /dev/null +++ b/engine/internal/retrieval/engine/postgres/snapshot/rename_test.go @@ -0,0 +1,62 @@ +package snapshot + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestValidateDatabaseRenames(t *testing.T) { + tests := []struct { + name string + renames map[string]string + connDB string + wantErr bool + }{ + {name: "empty map", renames: map[string]string{}, connDB: "postgres", wantErr: false}, + {name: "valid renames", renames: map[string]string{"prod_db": "dblab_db"}, connDB: "postgres", wantErr: false}, + {name: "multiple valid renames", renames: map[string]string{"db1": "db1_new", "db2": "db2_new"}, connDB: "postgres", wantErr: false}, + {name: "rename matches connDB", renames: map[string]string{"postgres": "pg_renamed"}, connDB: "postgres", wantErr: true}, + {name: "one of multiple matches connDB", renames: map[string]string{"safe_db": "new_safe", "postgres": "renamed"}, connDB: "postgres", wantErr: true}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := validateDatabaseRenames(tt.renames, tt.connDB) + if tt.wantErr { + require.Error(t, err) + assert.Contains(t, err.Error(), "connection database") + } else { + require.NoError(t, err) + } + }) + } +} + +func TestBuildRenameCommand(t *testing.T) { + tests := []struct { + name string + username string + connDB string + oldName string + newName string + expected []string + }{ + { + name: "simple rename", username: "postgres", connDB: "postgres", oldName: "prod_db", newName: "dblab_db", + expected: []string{"psql", "-U", "postgres", "-d", "postgres", "-XAtc", `ALTER DATABASE "prod_db" RENAME TO "dblab_db"`}, + }, + { + name: "special characters in name", username: "admin", connDB: "management", oldName: "my-db", newName: "my_db", + expected: []string{"psql", "-U", "admin", "-d", "management", "-XAtc", `ALTER DATABASE "my-db" RENAME TO "my_db"`}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := buildRenameCommand(tt.username, tt.connDB, tt.oldName, tt.newName) + assert.Equal(t, tt.expected, result) + }) + } +} diff --git a/engine/internal/retrieval/engine/postgres/tools/cont/container.go b/engine/internal/retrieval/engine/postgres/tools/cont/container.go index 0c39db56..df250b11 100644 --- a/engine/internal/retrieval/engine/postgres/tools/cont/container.go +++ b/engine/internal/retrieval/engine/postgres/tools/cont/container.go @@ -58,6 +58,8 @@ const ( DBLabEmbeddedUILabel = "dblab_embedded_ui" // DBLabFoundationLabel defines a label value to mark foundation containers. DBLabFoundationLabel = "dblab_foundation" + // DBLabRenameLabel defines a label value for database rename containers. + DBLabRenameLabel = "dblab_rename" // DBLabRunner defines a label to mark runner containers. DBLabRunner = "dblab_runner"