Skip to content

Commit 0ff34b0

Browse files
Sarumyanagneum
authored andcommitted
Add built-in databaseRename option for snapshot jobs
1 parent 0fe7abe commit 0ff34b0

12 files changed

Lines changed: 462 additions & 5 deletions

README.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -106,6 +106,7 @@ Read more:
106106
- Data provisioning & retrieval
107107
- Physical (pg_basebackup, WAL-G, pgBackRest) and logical (dump/restore) provisioning
108108
- Partial data retrieval in logical mode (specific databases/tables)
109+
- Database renaming during snapshot creation (`databaseRename` option)
109110
- Continuous update in physical mode
110111
- Periodic full refresh in logical mode without downtime
111112
- Recovery & management

engine/configs/config.example.logical_generic.yml

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -135,7 +135,9 @@ retrieval: # Data retrieval: initial sync and ongoing updates. Two methods:
135135
options:
136136
<<: *db_configs # Adjust PostgreSQL configuration
137137
preprocessingScript: "" # Pre-processing script for data scrubbing/masking; e.g., "/tmp/scripts/custom.sh"
138-
138+
databaseRename: # Rename databases before finalizing snapshot; runs after preprocessingScript; default: empty (disabled)
139+
# mydb_prod: mydb_dblab # Rename "mydb_prod" to "mydb_dblab"
140+
139141
dataPatching: # Pre-processing SQL queries for data patching
140142
<<: *db_container
141143
queryPreprocessing:

engine/configs/config.example.logical_rds_iam.yml

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -135,7 +135,9 @@ retrieval: # Data retrieval: initial sync and ongoing updates. Two methods:
135135
options:
136136
<<: *db_configs # Adjust PostgreSQL configuration
137137
preprocessingScript: "" # Pre-processing script for data scrubbing/masking; e.g., "/tmp/scripts/custom.sh"
138-
138+
databaseRename: # Rename databases before finalizing snapshot; runs after preprocessingScript; default: empty (disabled)
139+
# mydb_prod: mydb_dblab # Rename "mydb_prod" to "mydb_dblab"
140+
139141
dataPatching: # Pre-processing SQL queries for data patching
140142
<<: *db_container
141143
queryPreprocessing:

engine/configs/config.example.physical_generic.yml

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -103,6 +103,9 @@ retrieval: # Data retrieval: initial sync and ongoing updates. Two methods:
103103
# recovery_target_action: 'promote'
104104
# recovery_target_timeline: 'latest'
105105
preprocessingScript: "" # Shell script path to execute before finalizing snapshot; example: "/tmp/scripts/custom.sh"; default: "" (disabled)
106+
databaseRename: # Rename databases before finalizing snapshot; runs after preprocessingScript; default: empty (disabled)
107+
# example_production: example_dblab # Rename "example_production" to "example_dblab"
108+
# analytics_prod: analytics_dblab
106109
scheduler: # Snapshot scheduling and retention policy configuration
107110
snapshot: # Snapshot creation scheduling
108111
timetable: "0 */6 * * *" # Cron expression defining snapshot schedule: https://en.wikipedia.org/wiki/Cron#Overview

engine/configs/config.example.physical_pgbackrest.yml

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -118,6 +118,9 @@ retrieval: # Data retrieval: initial sync and ongoing updates. Two methods:
118118
# recovery_target_timeline: 'latest'
119119

120120
preprocessingScript: "" # Shell script path to execute before finalizing snapshot; example: "/tmp/scripts/custom.sh"; default: "" (disabled)
121+
databaseRename: # Rename databases before finalizing snapshot; runs after preprocessingScript; default: empty (disabled)
122+
# example_production: example_dblab # Rename "example_production" to "example_dblab"
123+
# analytics_prod: analytics_dblab
121124
scheduler: # Snapshot scheduling and retention policy configuration
122125
snapshot: # Snapshot creation scheduling
123126
timetable: "0 */6 * * *" # Cron expression defining snapshot schedule: https://en.wikipedia.org/wiki/Cron#Overview

engine/configs/config.example.physical_walg.yml

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -104,6 +104,9 @@ retrieval: # Data retrieval: initial sync and ongoing updates. Two methods:
104104
# recovery_target_timeline: 'latest'
105105

106106
preprocessingScript: "" # Shell script path to execute before finalizing snapshot; example: "/tmp/scripts/custom.sh"; default: "" (disabled)
107+
databaseRename: # Rename databases before finalizing snapshot; runs after preprocessingScript; default: empty (disabled)
108+
# example_production: example_dblab # Rename "example_production" to "example_dblab"
109+
# analytics_prod: analytics_dblab
107110
scheduler: # Snapshot scheduling and retention policy configuration
108111
snapshot: # Snapshot creation scheduling
109112
timetable: "0 */6 * * *" # Cron expression defining snapshot schedule: https://en.wikipedia.org/wiki/Cron#Overview

engine/internal/retrieval/engine/postgres/snapshot/logical.go

Lines changed: 22 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -63,6 +63,7 @@ type LogicalInitial struct {
6363
type LogicalOptions struct {
6464
DataPatching DataPatching `yaml:"dataPatching"`
6565
PreprocessingScript string `yaml:"preprocessingScript"`
66+
DatabaseRename map[string]string `yaml:"databaseRename"`
6667
Configs map[string]string `yaml:"configs"`
6768
Schedule Scheduler `yaml:"schedule"`
6869
}
@@ -121,6 +122,12 @@ func (s *LogicalInitial) ReportActivity(_ context.Context) (*activity.Activity,
121122

122123
// Run starts the job.
123124
func (s *LogicalInitial) Run(ctx context.Context) error {
125+
if len(s.options.DatabaseRename) > 0 {
126+
if err := validateDatabaseRenames(s.options.DatabaseRename, s.globalCfg.Database.Name()); err != nil {
127+
return fmt.Errorf("invalid database rename configuration: %w", err)
128+
}
129+
}
130+
124131
if s.options.PreprocessingScript != "" {
125132
if err := runPreprocessingScript(s.options.PreprocessingScript); err != nil {
126133
return err
@@ -144,7 +151,7 @@ func (s *LogicalInitial) Run(ctx context.Context) error {
144151
return errors.Wrap(err, "failed to store PostgreSQL configs for the snapshot")
145152
}
146153

147-
if s.queryProcessor != nil {
154+
if s.queryProcessor != nil || len(s.options.DatabaseRename) > 0 {
148155
if err := s.runPreprocessingQueries(ctx, dataDir); err != nil {
149156
return errors.Wrap(err, "failed to run preprocessing queries")
150157
}
@@ -265,8 +272,20 @@ func (s *LogicalInitial) runPreprocessingQueries(ctx context.Context, dataDir st
265272
return errors.Wrap(err, "failed to readiness check")
266273
}
267274

268-
if err := s.queryProcessor.ApplyPreprocessingQueries(ctx, containerID); err != nil {
269-
return errors.Wrap(err, "failed to run preprocessing queries")
275+
if s.queryProcessor != nil {
276+
if err := s.queryProcessor.ApplyPreprocessingQueries(ctx, containerID); err != nil {
277+
return errors.Wrap(err, "failed to run preprocessing queries")
278+
}
279+
}
280+
281+
if len(s.options.DatabaseRename) > 0 {
282+
if err := executeDatabaseRenames(
283+
ctx, s.dockerClient, containerID,
284+
s.globalCfg.Database.User(), s.globalCfg.Database.Name(),
285+
s.options.DatabaseRename,
286+
); err != nil {
287+
return fmt.Errorf("failed to rename databases: %w", err)
288+
}
270289
}
271290

272291
if err := tools.RunCheckpoint(ctx, s.dockerClient, containerID, s.globalCfg.Database.User(), s.globalCfg.Database.Name()); err != nil {

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

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -114,6 +114,7 @@ type PhysicalOptions struct {
114114
SkipStartSnapshot bool `yaml:"skipStartSnapshot"`
115115
Promotion Promotion `yaml:"promotion"`
116116
PreprocessingScript string `yaml:"preprocessingScript"`
117+
DatabaseRename map[string]string `yaml:"databaseRename"`
117118
Configs map[string]string `yaml:"configs"`
118119
Sysctls map[string]string `yaml:"sysctls"`
119120
Envs map[string]string `yaml:"envs"`
@@ -318,6 +319,12 @@ func (p *PhysicalInitial) run(ctx context.Context) (err error) {
318319
default:
319320
}
320321

322+
if len(p.options.DatabaseRename) > 0 {
323+
if err := validateDatabaseRenames(p.options.DatabaseRename, p.globalCfg.Database.Name()); err != nil {
324+
return fmt.Errorf("invalid database rename configuration: %w", err)
325+
}
326+
}
327+
321328
p.dbMark.DataStateAt = extractDataStateAt(p.dbMarker)
322329

323330
// Snapshot data.
@@ -388,6 +395,18 @@ func (p *PhysicalInitial) run(ctx context.Context) (err error) {
388395
}
389396
}
390397

398+
if !p.options.Promotion.Enabled && len(p.options.DatabaseRename) > 0 {
399+
if err := runDatabaseRename(ctx, renameParams{
400+
dockerClient: p.dockerClient,
401+
engineProps: p.engineProps,
402+
globalCfg: p.globalCfg,
403+
dataDir: cloneDataDir,
404+
renames: p.options.DatabaseRename,
405+
}); err != nil {
406+
return errors.Wrap(err, "failed to rename databases")
407+
}
408+
}
409+
391410
// Mark database data.
392411
if err := p.markDatabaseData(); err != nil {
393412
return errors.Wrap(err, "failed to mark the prepared data")
@@ -696,6 +715,16 @@ func (p *PhysicalInitial) promoteInstance(ctx context.Context, clonePath string,
696715
}
697716
}
698717

718+
if len(p.options.DatabaseRename) > 0 {
719+
if err := executeDatabaseRenames(
720+
ctx, p.dockerClient, containerID,
721+
p.globalCfg.Database.User(), p.globalCfg.Database.Name(),
722+
p.options.DatabaseRename,
723+
); err != nil {
724+
return fmt.Errorf("failed to rename databases: %w", err)
725+
}
726+
}
727+
699728
if err := tools.RunCheckpoint(ctx, p.dockerClient, containerID, p.globalCfg.Database.User(), p.globalCfg.Database.Name()); err != nil {
700729
return errors.Wrap(err, "failed to run checkpoint")
701730
}
Lines changed: 220 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,220 @@
1+
/*
2+
2026 © Postgres.ai
3+
*/
4+
5+
package snapshot
6+
7+
import (
8+
"context"
9+
"fmt"
10+
"regexp"
11+
12+
"github.com/docker/docker/api/types/container"
13+
"github.com/docker/docker/api/types/filters"
14+
"github.com/docker/docker/client"
15+
16+
"gitlab.com/postgres-ai/database-lab/v3/internal/diagnostic"
17+
"gitlab.com/postgres-ai/database-lab/v3/internal/retrieval/engine/postgres/tools"
18+
"gitlab.com/postgres-ai/database-lab/v3/internal/retrieval/engine/postgres/tools/cont"
19+
"gitlab.com/postgres-ai/database-lab/v3/internal/retrieval/engine/postgres/tools/health"
20+
"gitlab.com/postgres-ai/database-lab/v3/pkg/config/global"
21+
"gitlab.com/postgres-ai/database-lab/v3/pkg/log"
22+
)
23+
24+
const (
25+
renameContainerPrefix = "dblab_rename_"
26+
dbNameRules = "must start with a letter or underscore" +
27+
" and contain only letters, digits, underscores, and hyphens"
28+
)
29+
30+
var validDBNameRegex = regexp.MustCompile(`^[a-zA-Z_][a-zA-Z0-9_-]*$`)
31+
32+
// renameParams holds parameters for running database renames in a standalone container.
33+
type renameParams struct {
34+
dockerClient *client.Client
35+
engineProps *global.EngineProps
36+
globalCfg *global.Config
37+
dataDir string
38+
renames map[string]string
39+
}
40+
41+
// executeDatabaseRenames runs ALTER DATABASE RENAME statements inside an already-running container.
42+
func executeDatabaseRenames(
43+
ctx context.Context,
44+
dockerClient *client.Client,
45+
containerID, user, connDB string,
46+
renames map[string]string,
47+
) error {
48+
for oldName, newName := range renames {
49+
log.Msg(fmt.Sprintf("Renaming database %q to %q", oldName, newName))
50+
51+
cmd := buildRenameCommand(user, connDB, oldName, newName)
52+
53+
output, err := tools.ExecCommandWithOutput(ctx, dockerClient, containerID, container.ExecOptions{Cmd: cmd})
54+
if err != nil {
55+
return fmt.Errorf("failed to rename database %q to %q: %w", oldName, newName, err)
56+
}
57+
58+
log.Msg("Rename result: ", output)
59+
}
60+
61+
return nil
62+
}
63+
64+
// runDatabaseRename renames databases using ALTER DATABASE in a temporary container.
65+
func runDatabaseRename(ctx context.Context, params renameParams) error {
66+
if len(params.renames) == 0 {
67+
return nil
68+
}
69+
70+
connDB := params.globalCfg.Database.Name()
71+
72+
if err := validateDatabaseRenames(params.renames, connDB); err != nil {
73+
return err
74+
}
75+
76+
pgVersion, err := tools.DetectPGVersion(params.dataDir)
77+
if err != nil {
78+
return fmt.Errorf("failed to detect postgres version: %w", err)
79+
}
80+
81+
image := fmt.Sprintf("postgresai/extended-postgres:%g", pgVersion)
82+
83+
if err := tools.PullImage(ctx, params.dockerClient, image); err != nil {
84+
return fmt.Errorf("failed to pull image for database rename: %w", err)
85+
}
86+
87+
pwd, err := tools.GeneratePassword()
88+
if err != nil {
89+
return fmt.Errorf("failed to generate password: %w", err)
90+
}
91+
92+
hostConfig, err := cont.BuildHostConfig(ctx, params.dockerClient, params.dataDir, nil)
93+
if err != nil {
94+
return fmt.Errorf("failed to build host config: %w", err)
95+
}
96+
97+
containerName := renameContainerPrefix + params.engineProps.InstanceID
98+
99+
containerID, err := tools.CreateContainerIfMissing(ctx, params.dockerClient, containerName,
100+
&container.Config{
101+
Labels: map[string]string{
102+
cont.DBLabControlLabel: cont.DBLabRenameLabel,
103+
cont.DBLabInstanceIDLabel: params.engineProps.InstanceID,
104+
cont.DBLabEngineNameLabel: params.engineProps.ContainerName,
105+
},
106+
Env: []string{
107+
"PGDATA=" + params.dataDir,
108+
"POSTGRES_PASSWORD=" + pwd,
109+
},
110+
Image: image,
111+
Healthcheck: health.GetConfig(
112+
params.globalCfg.Database.User(),
113+
connDB,
114+
),
115+
},
116+
hostConfig,
117+
)
118+
if err != nil {
119+
return fmt.Errorf("failed to create rename container: %w", err)
120+
}
121+
122+
defer tools.RemoveContainer(ctx, params.dockerClient, containerID, cont.StopPhysicalTimeout)
123+
124+
defer func() {
125+
if err != nil {
126+
tools.PrintContainerLogs(ctx, params.dockerClient, containerName)
127+
tools.PrintLastPostgresLogs(ctx, params.dockerClient, containerName, params.dataDir)
128+
129+
filterArgs := filters.NewArgs(
130+
filters.KeyValuePair{Key: "label",
131+
Value: fmt.Sprintf("%s=%s", cont.DBLabControlLabel, cont.DBLabRenameLabel)})
132+
133+
if diagErr := diagnostic.CollectDiagnostics(ctx, params.dockerClient, filterArgs, containerName, params.dataDir); diagErr != nil {
134+
log.Err("failed to collect rename container diagnostics", diagErr)
135+
}
136+
}
137+
}()
138+
139+
log.Msg(fmt.Sprintf("Running rename container: %s. ID: %v", containerName, containerID))
140+
141+
if err = params.dockerClient.ContainerStart(ctx, containerID, container.StartOptions{}); err != nil {
142+
return fmt.Errorf("failed to start rename container: %w", err)
143+
}
144+
145+
log.Msg("Waiting for rename container readiness")
146+
log.Msg(fmt.Sprintf("View logs using the command: %s %s", tools.ViewLogsCmd, containerName))
147+
148+
if err = tools.CheckContainerReadiness(ctx, params.dockerClient, containerID); err != nil {
149+
return fmt.Errorf("rename container readiness check failed: %w", err)
150+
}
151+
152+
user := params.globalCfg.Database.User()
153+
154+
if err = executeDatabaseRenames(ctx, params.dockerClient, containerID, user, connDB, params.renames); err != nil {
155+
return err
156+
}
157+
158+
if err = tools.RunCheckpoint(ctx, params.dockerClient, containerID, user, connDB); err != nil {
159+
return fmt.Errorf("failed to run checkpoint after rename: %w", err)
160+
}
161+
162+
if err = tools.StopPostgres(ctx, params.dockerClient, containerID, params.dataDir, tools.DefaultStopTimeout); err != nil {
163+
log.Msg("failed to stop postgres after rename", err)
164+
}
165+
166+
return nil
167+
}
168+
169+
func buildRenameCommand(username, connDB, oldName, newName string) []string {
170+
return []string{
171+
"psql",
172+
"-U", username,
173+
"-d", connDB,
174+
"-XAtc", fmt.Sprintf(`ALTER DATABASE "%s" RENAME TO "%s"`, oldName, newName),
175+
}
176+
}
177+
178+
func validateDatabaseRenames(renames map[string]string, connDB string) error {
179+
targets := make(map[string]string, len(renames))
180+
181+
for oldName, newName := range renames {
182+
if oldName == "" || newName == "" {
183+
return fmt.Errorf("database rename names must not be empty")
184+
}
185+
186+
if !validDBNameRegex.MatchString(oldName) {
187+
return fmt.Errorf("invalid database name %q: %s", oldName, dbNameRules)
188+
}
189+
190+
if !validDBNameRegex.MatchString(newName) {
191+
return fmt.Errorf("invalid database name %q: %s", newName, dbNameRules)
192+
}
193+
194+
if oldName == connDB {
195+
return fmt.Errorf("cannot rename database %q: it is used as the connection database", oldName)
196+
}
197+
198+
if newName == connDB {
199+
return fmt.Errorf("cannot rename database %q to %q: target name is the connection database", oldName, newName)
200+
}
201+
202+
if oldName == newName {
203+
return fmt.Errorf("cannot rename database %q to itself", oldName)
204+
}
205+
206+
if prev, ok := targets[newName]; ok {
207+
return fmt.Errorf("duplicate rename target %q: both %q and %q rename to the same database", newName, prev, oldName)
208+
}
209+
210+
targets[newName] = oldName
211+
}
212+
213+
for oldName, newName := range renames {
214+
if _, ok := renames[newName]; ok {
215+
return fmt.Errorf("chained rename conflict: %q renames to %q, which is also a rename source", oldName, newName)
216+
}
217+
}
218+
219+
return nil
220+
}

0 commit comments

Comments
 (0)