Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly
Sign up
Appearance settings

Commit 0fed6f9

Browse files
committed
Merge branch '520-skip-policies' into 'master'
feat: add option to skip policies during restore (#520) Closes #520 See merge request postgres-ai/database-lab!769
2 parents 7629636 + b765f40 commit 0fed6f9

File tree

4 files changed

+72
-6
lines changed

4 files changed

+72
-6
lines changed

‎engine/configs/config.example.logical_generic.yml‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -315,6 +315,9 @@ retrieval:
315315
- "--no-owner"
316316
- "--exit-on-error"
317317

318+
# Option to skip policies during restore.
319+
skipPolicies: true
320+
318321
logicalSnapshot:
319322
options:
320323
# Adjust PostgreSQL configuration

‎engine/configs/config.example.logical_rds_iam.yml‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -315,6 +315,9 @@ retrieval:
315315
- "--no-owner"
316316
- "--exit-on-error"
317317

318+
# Option to skip policies during restore.
319+
skipPolicies: true
320+
318321
logicalSnapshot:
319322
options:
320323
# Adjust PostgreSQL configuration

‎engine/internal/retrieval/engine/postgres/logical/restore.go‎

Lines changed: 65 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -105,6 +105,7 @@ type RestoreOptions struct {
105105
Configs map[string]string `yaml:"configs"`
106106
QueryPreprocessing query.PreprocessorCfg `yaml:"queryPreprocessing"`
107107
CustomOptions []string `yaml:"customOptions"`
108+
SkipPolicies bool `yaml:"skipPolicies"`
108109
}
109110

110111
// Partial defines tables and rules for a partial logical restore.
@@ -501,7 +502,48 @@ func (r *RestoreJob) restoreDB(ctx context.Context, contID, dbName string, dbDef
501502
}
502503
}
503504

504-
restoreCommand := r.buildLogicalRestoreCommand(dbName, dbDefinition)
505+
var (
506+
tmpListFile *os.File
507+
err error
508+
)
509+
510+
if r.RestoreOptions.SkipPolicies {
511+
tmpListFile, err = os.CreateTemp(r.getDumpLocation(dbDefinition.Format, dbDefinition.dbName), dbName+"-list-file*")
512+
if err != nil {
513+
return fmt.Errorf("failed to create temporary list file: %w", err)
514+
}
515+
516+
defer func() {
517+
if err := os.Remove(tmpListFile.Name()); err != nil {
518+
log.Dbg("Cannot remove temporary file:", err)
519+
}
520+
}()
521+
defer func() { _ = tmpListFile.Close() }()
522+
523+
dumpLocation := r.getDumpLocation(dbDefinition.Format, dbDefinition.dbName)
524+
525+
if dbDefinition.Format != directoryFormat {
526+
dumpLocation = r.RestoreOptions.DumpLocation
527+
}
528+
529+
preCmd := []string{"bash", "-c", `pg_restore -l ` + dumpLocation + " | grep -v POLICY > " + tmpListFile.Name()}
530+
531+
log.Msg("Running preparatory command to create list file for "+dbName, preCmd)
532+
533+
output, err := tools.ExecCommandWithOutput(ctx, r.dockerClient, contID, types.ExecConfig{
534+
Tty: true,
535+
Cmd: preCmd,
536+
Env: []string{"PGAPPNAME=" + dleRetrieval},
537+
})
538+
539+
if err != nil {
540+
log.Dbg(output)
541+
542+
return fmt.Errorf("failed to perform preparatory command: %w", err)
543+
}
544+
}
545+
546+
restoreCommand := r.buildLogicalRestoreCommand(dbName, dbDefinition, tmpListFile)
505547
log.Msg("Running restore command for "+dbName, restoreCommand)
506548

507549
output, err := tools.ExecCommandWithOutput(ctx, r.dockerClient, contID, types.ExecConfig{
@@ -699,12 +741,12 @@ func (r *RestoreJob) updateDataStateAt() {
699741
r.fsPool.SetDSA(dsaTime)
700742
}
701743

702-
func (r *RestoreJob) buildLogicalRestoreCommand(dumpName string, definition DumpDefinition) []string {
744+
func (r *RestoreJob) buildLogicalRestoreCommand(dumpName string, definition DumpDefinition, listFile*os.File) []string {
703745
if definition.Format == plainFormat {
704746
return r.buildPlainTextCommand(dumpName, definition)
705747
}
706748

707-
return r.buildPGRestoreCommand(dumpName, definition)
749+
return r.buildPGRestoreCommand(dumpName, definition, listFile)
708750
}
709751

710752
func (r *RestoreJob) buildPlainTextCommand(dumpName string, definition DumpDefinition) []string {
@@ -729,7 +771,7 @@ func (r *RestoreJob) buildPlainTextCommand(dumpName string, definition DumpDefin
729771
}
730772
}
731773

732-
func (r *RestoreJob) buildPGRestoreCommand(dumpName string, definition DumpDefinition) []string {
774+
func (r *RestoreJob) buildPGRestoreCommand(dumpName string, definition DumpDefinition, listFile*os.File) []string {
733775
// Using the default database name because the database for connection must exist.
734776
restoreCmd := []string{"pg_restore", "--username", r.globalCfg.Database.User(), "--dbname", defaults.DBName}
735777

@@ -750,7 +792,25 @@ func (r *RestoreJob) buildPGRestoreCommand(dumpName string, definition DumpDefin
750792

751793
restoreCmd = append(restoreCmd, r.getDumpLocation(definition.Format, dumpName))
752794

753-
restoreCmd = append(restoreCmd, r.RestoreOptions.CustomOptions...)
795+
customOptions := r.RestoreOptions.CustomOptions
796+
797+
// Skip policies: https://gitlab.com/postgres-ai/database-lab/-/merge_requests/769
798+
if listFile != nil {
799+
restoreCmd = append(restoreCmd, fmt.Sprintf("--use-list=%s", listFile.Name()))
800+
801+
customOptions = []string{}
802+
803+
for _, customOption := range r.RestoreOptions.CustomOptions {
804+
// Exclude -L (--use-list) in customOptions to avoid conflicts if skipping policies is enabled.
805+
if strings.HasPrefix(customOption, "-L") || strings.HasPrefix(customOption, "--use-list") {
806+
continue
807+
}
808+
809+
customOptions = append(customOptions, customOption)
810+
}
811+
}
812+
813+
restoreCmd = append(restoreCmd, customOptions...)
754814

755815
return restoreCmd
756816
}

‎engine/internal/retrieval/engine/postgres/logical/restore_test.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -116,7 +116,7 @@ func TestRestoreCommandBuilding(t *testing.T) {
116116
logicalJob.RestoreOptions = tc.copyOptions
117117
logicalJob.isDumpLocationDir = tc.isDumpLocationDir
118118
for dbName, definition := range tc.copyOptions.Databases {
119-
restoreCommand := logicalJob.buildLogicalRestoreCommand(dbName, definition)
119+
restoreCommand := logicalJob.buildLogicalRestoreCommand(dbName, definition, nil)
120120
assert.Equal(t, restoreCommand, tc.command)
121121
}
122122
}

0 commit comments

Comments
(0)

AltStyle によって変換されたページ (->オリジナル) /