Add built-in databaseRename option for snapshot jobs

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:

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

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")
 	}

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")

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
+}

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)
+		})
+	}
+}

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"