Skip to content

Commit b8c5e93

Browse files
agneumagneum
committed
Merge branch '361-custom-options' into 'master'
feat(engine): add customOptions for pg_dump and pg_restore commands (#361) Closes #361 See merge request postgres-ai/database-lab!660
2 parents 1787ddd + 9800db8 commit b8c5e93

File tree

10 files changed

+203
-18
lines changed

10 files changed

+203
-18
lines changed

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

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -255,6 +255,16 @@ retrieval:
255255
# # Option to adjust PostgreSQL configuration for a logical dump job.
256256
# # It's useful if a dumped database contains non-standard extensions.
257257
# <<: *db_configs
258+
# # Custom options for pg_restore command.
259+
# customOptions:
260+
# - "--no-privileges"
261+
# - "--no-owner"
262+
# - "--exit-on-error"
263+
264+
# Custom options for pg_dump command.
265+
customOptions:
266+
# - --no-publications
267+
# - --no-subscriptions
258268

259269
# Restores PostgreSQL database from the provided dump. If you use this block, do not use
260270
# "restore" option in the "logicalDump" job.
@@ -305,6 +315,13 @@ retrieval:
305315
# Inline SQL. Queries run after scripts placed in 'queryPath'.
306316
inline: ""
307317

318+
# Custom options for pg_restore command.
319+
customOptions:
320+
- "--no-tablespaces"
321+
- "--no-privileges"
322+
- "--no-owner"
323+
- "--exit-on-error"
324+
308325
logicalSnapshot:
309326
options:
310327
# Adjust PostgreSQL configuration

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

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -252,6 +252,15 @@ retrieval:
252252
# # Option to adjust PostgreSQL configuration for a logical dump job.
253253
# # It's useful if a dumped database contains non-standard extensions.
254254
# <<: *db_configs
255+
# # Custom options for pg_restore command.
256+
# customOptions:
257+
# - "--no-privileges"
258+
# - "--no-owner"
259+
# - "--exit-on-error"
260+
261+
# Custom options for pg_dump command.
262+
customOptions:
263+
- "--exclude-schema=rdsdms"
255264

256265
# Restores PostgreSQL database from the provided dump. If you use this block, do not use
257266
# "restore" option in the "logicalDump" job.
@@ -306,6 +315,13 @@ retrieval:
306315
# Inline SQL. Queries run after scripts placed in 'queryPath'.
307316
inline: ""
308317

318+
# Custom options for pg_restore command.
319+
customOptions:
320+
- "--no-tablespaces"
321+
- "--no-privileges"
322+
- "--no-owner"
323+
- "--exit-on-error"
324+
309325
logicalSnapshot:
310326
options:
311327
# Adjust PostgreSQL configuration

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

Lines changed: 10 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -88,6 +88,7 @@ type DumpOptions struct {
8888
Databases map[string]DumpDefinition `yaml:"databases"`
8989
ParallelJobs int `yaml:"parallelJobs"`
9090
Restore ImmediateRestore `yaml:"immediateRestore"`
91+
CustomOptions []string `yaml:"customOptions"`
9192
}
9293

9394
// Source describes source of data to dump.
@@ -133,9 +134,10 @@ type Connection struct {
133134

134135
// ImmediateRestore contains options for direct data restore without saving the dump file on disk.
135136
type ImmediateRestore struct {
136-
Enabled bool `yaml:"enabled"`
137-
ForceInit bool `yaml:"forceInit"`
138-
Configs map[string]string `yaml:"configs"`
137+
Enabled bool `yaml:"enabled"`
138+
ForceInit bool `yaml:"forceInit"`
139+
Configs map[string]string `yaml:"configs"`
140+
CustomOptions []string `yaml:"customOptions"`
139141
}
140142

141143
// NewDumpJob creates a new DumpJob.
@@ -694,6 +696,8 @@ func (d *DumpJob) buildLogicalDumpCommand(dbName string, dump DumpDefinition) []
694696
dumpCmd = append(dumpCmd, "--exclude-table", table)
695697
}
696698

699+
dumpCmd = append(dumpCmd, d.DumpOptions.CustomOptions...)
700+
697701
// Define if restore directly or export to dump location.
698702
if d.DumpOptions.Restore.Enabled {
699703
dumpCmd = append(dumpCmd, "--format", customFormat)
@@ -711,8 +715,7 @@ func (d *DumpJob) buildLogicalDumpCommand(dbName string, dump DumpDefinition) []
711715
}
712716

713717
func (d *DumpJob) buildLogicalRestoreCommand(dbName string) []string {
714-
restoreCmd := []string{"|", "pg_restore", "--username", d.globalCfg.Database.User(), "--dbname", defaults.DBName,
715-
"--no-privileges", "--no-owner", "--exit-on-error"}
718+
restoreCmd := []string{"|", "pg_restore", "--username", d.globalCfg.Database.User(), "--dbname", defaults.DBName}
716719

717720
if dbName != defaults.DBName {
718721
// To avoid recreating of the default database.
@@ -723,6 +726,8 @@ func (d *DumpJob) buildLogicalRestoreCommand(dbName string) []string {
723726
restoreCmd = append(restoreCmd, "--clean", "--if-exists")
724727
}
725728

729+
restoreCmd = append(restoreCmd, d.DumpOptions.Restore.CustomOptions...)
730+
726731
return restoreCmd
727732
}
728733

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

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -104,6 +104,7 @@ type RestoreOptions struct {
104104
ParallelJobs int `yaml:"parallelJobs"`
105105
Configs map[string]string `yaml:"configs"`
106106
QueryPreprocessing query.PreprocessorCfg `yaml:"queryPreprocessing"`
107+
CustomOptions []string `yaml:"customOptions"`
107108
}
108109

109110
// Partial defines tables and rules for a partial logical restore.
@@ -736,8 +737,7 @@ func (r *RestoreJob) buildPlainTextCommand(dumpName string, definition DumpDefin
736737
}
737738

738739
func (r *RestoreJob) buildPGRestoreCommand(dumpName string, definition DumpDefinition) []string {
739-
restoreCmd := []string{"pg_restore", "--username", r.globalCfg.Database.User(), "--dbname", defaults.DBName,
740-
"--no-privileges", "--no-owner", "--exit-on-error"}
740+
restoreCmd := []string{"pg_restore", "--username", r.globalCfg.Database.User(), "--dbname", defaults.DBName}
741741

742742
if definition.dbName != defaults.DBName {
743743
// To avoid recreating of the default database.
@@ -760,6 +760,8 @@ func (r *RestoreJob) buildPGRestoreCommand(dumpName string, definition DumpDefin
760760

761761
restoreCmd = append(restoreCmd, r.getDumpLocation(definition.Format, dumpName))
762762

763+
restoreCmd = append(restoreCmd, r.RestoreOptions.CustomOptions...)
764+
763765
return restoreCmd
764766
}
765767

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

Lines changed: 42 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -40,25 +40,27 @@ func TestRestoreCommandBuilding(t *testing.T) {
4040
Format: customFormat,
4141
},
4242
},
43-
DumpLocation: "/tmp/db.dump",
43+
DumpLocation: "/tmp/db.dump",
44+
CustomOptions: []string{"--no-privileges", "--no-owner", "--exit-on-error"},
4445
},
45-
command: []string{"pg_restore", "--username", "john", "--dbname", "postgres", "--no-privileges", "--no-owner", "--exit-on-error", "--create", "--jobs", "1", "/tmp/db.dump"},
46+
command: []string{"pg_restore", "--username", "john", "--dbname", "postgres", "--create", "--jobs", "1", "/tmp/db.dump", "--no-privileges", "--no-owner", "--exit-on-error"},
4647
},
4748
{
4849
copyOptions: RestoreOptions{
4950
ParallelJobs: 4,
5051
ForceInit: true,
5152
},
52-
command: []string{"pg_restore", "--username", "john", "--dbname", "postgres", "--no-privileges", "--no-owner", "--exit-on-error", "--create", "--clean", "--if-exists", "--jobs", "4", ""},
53+
command: []string{"pg_restore", "--username", "john", "--dbname", "postgres", "--create", "--clean", "--if-exists", "--jobs", "4"},
5354
},
5455
{
5556
copyOptions: RestoreOptions{
56-
ParallelJobs: 2,
57-
ForceInit: false,
58-
Databases: map[string]DumpDefinition{"testDB": {}},
59-
DumpLocation: "/tmp/db.dump",
57+
ParallelJobs: 2,
58+
ForceInit: false,
59+
Databases: map[string]DumpDefinition{"testDB": {}},
60+
DumpLocation: "/tmp/db.dump",
61+
CustomOptions: []string{"--no-privileges", "--no-owner", "--exit-on-error"},
6062
},
61-
command: []string{"pg_restore", "--username", "john", "--dbname", "postgres", "--no-privileges", "--no-owner", "--exit-on-error", "--create", "--jobs", "2", "/tmp/db.dump/testDB"},
63+
command: []string{"pg_restore", "--username", "john", "--dbname", "postgres", "--create", "--jobs", "2", "/tmp/db.dump/testDB", "--no-privileges", "--no-owner", "--exit-on-error"},
6264
},
6365
{
6466
copyOptions: RestoreOptions{
@@ -69,9 +71,10 @@ func TestRestoreCommandBuilding(t *testing.T) {
6971
Format: directoryFormat,
7072
},
7173
},
72-
DumpLocation: "/tmp/db.dump",
74+
DumpLocation: "/tmp/db.dump",
75+
CustomOptions: []string{"--no-privileges", "--no-owner", "--exit-on-error"},
7376
},
74-
command: []string{"pg_restore", "--username", "john", "--dbname", "postgres", "--no-privileges", "--no-owner", "--exit-on-error", "--create", "--jobs", "1", "--table", "test", "--table", "users", "/tmp/db.dump/testDB"},
77+
command: []string{"pg_restore", "--username", "john", "--dbname", "postgres", "--create", "--jobs", "1", "--table", "test", "--table", "users", "/tmp/db.dump/testDB", "--no-privileges", "--no-owner", "--exit-on-error"},
7578
},
7679
{
7780
copyOptions: RestoreOptions{
@@ -133,6 +136,11 @@ func TestDumpCommandBuilding(t *testing.T) {
133136
Password: "secret",
134137
},
135138
},
139+
globalCfg: &global.Config{
140+
Database: global.Database{
141+
Username: "postgres",
142+
},
143+
},
136144
}
137145

138146
testCases := []struct {
@@ -152,8 +160,31 @@ func TestDumpCommandBuilding(t *testing.T) {
152160
},
153161
},
154162
},
163+
CustomOptions: []string{"--exclude-scheme=test-scheme"},
164+
},
165+
command: []string{"pg_dump", "--create", "--host", "localhost", "--port", "5432", "--username", "john", "--dbname", "testDB", "--jobs", "1", "--table", "test", "--table", "users", "--exclude-table", "test2", "--exclude-table", "users2", "--exclude-scheme=test-scheme", "--format", "directory", "--file", "/tmp/db.dump/testDB"},
166+
},
167+
{
168+
copyOptions: DumpOptions{
169+
170+
ParallelJobs: 1,
171+
DumpLocation: "/tmp/db.dump",
172+
Databases: map[string]DumpDefinition{
173+
"testDB": {
174+
Tables: []string{"test", "users"},
175+
ExcludeTables: []string{
176+
"test2",
177+
"users2",
178+
},
179+
},
180+
},
181+
Restore: ImmediateRestore{
182+
Enabled: true,
183+
CustomOptions: []string{"--no-privileges", "--no-owner", "--exit-on-error"},
184+
},
185+
CustomOptions: []string{"--exclude-scheme=test-scheme"},
155186
},
156-
command: []string{"pg_dump", "--create", "--host", "localhost", "--port", "5432", "--username", "john", "--dbname", "testDB", "--jobs", "1", "--table", "test", "--table", "users", "--exclude-table", "test2", "--exclude-table", "users2", "--format", "directory", "--file", "/tmp/db.dump/testDB"},
187+
command: []string{"sh", "-c", "pg_dump --create --host localhost --port 5432 --username john --dbname testDB --jobs 1 --table test --table users --exclude-table test2 --exclude-table users2 --exclude-scheme=test-scheme --format custom | pg_restore --username postgres --dbname postgres --create --no-privileges --no-owner --exit-on-error"},
157188
},
158189
}
159190

‎engine/internal/srv/config.go

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ import (
55
"context"
66
"fmt"
77
"net/http"
8+
"regexp"
89
"time"
910

1011
"github.com/docker/docker/api/types"
@@ -341,6 +342,14 @@ func (s *Server) validateConfig(
341342
return err
342343
}
343344

345+
if err := validateCustomOptions(proj.DumpCustomOptions); err != nil {
346+
return fmt.Errorf("invalid custom dump options: %w", err)
347+
}
348+
349+
if err := validateCustomOptions(proj.RestoreCustomOptions); err != nil {
350+
return fmt.Errorf("invalid custom restore options: %w", err)
351+
}
352+
344353
if proj.DockerImage != nil {
345354
stream, err := s.docker.ImagePull(ctx, *proj.DockerImage, types.ImagePullOptions{})
346355
if err != nil {
@@ -355,3 +364,25 @@ func (s *Server) validateConfig(
355364

356365
return nil
357366
}
367+
368+
var (
369+
isValidCustomOption = regexp.MustCompile("^[A-Za-z0-9-_=\"]+$").MatchString
370+
errInvalidOption = fmt.Errorf("due to security reasons, current implementation of custom options supports only " +
371+
"letters, numbers, hyphen, underscore, equal sign, and double quotes")
372+
errInvalidOptionType = fmt.Errorf("invalid type of custom option")
373+
)
374+
375+
func validateCustomOptions(customOptions []interface{}) error {
376+
for _, opt := range customOptions {
377+
castedValue, ok := opt.(string)
378+
if !ok {
379+
return fmt.Errorf("%w: %q", errInvalidOptionType, opt)
380+
}
381+
382+
if !isValidCustomOption(castedValue) {
383+
return fmt.Errorf("invalid option %q: %w", castedValue, errInvalidOption)
384+
}
385+
}
386+
387+
return nil
388+
}

‎engine/internal/srv/config_test.go

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,45 @@
1+
package srv
2+
3+
import (
4+
"testing"
5+
6+
"github.com/stretchr/testify/require"
7+
)
8+
9+
func TestCustomOptions(t *testing.T) {
10+
testCases := []struct {
11+
customOptions []interface{}
12+
expectedResult error
13+
}{
14+
{
15+
customOptions: []interface{}{"--verbose"},
16+
expectedResult: nil,
17+
},
18+
{
19+
customOptions: []interface{}{"--exclude-scheme=test_scheme"},
20+
expectedResult: nil,
21+
},
22+
{
23+
customOptions: []interface{}{`--exclude-scheme="test_scheme"`},
24+
expectedResult: nil,
25+
},
26+
{
27+
customOptions: []interface{}{"--table=$(echo 'test')"},
28+
expectedResult: errInvalidOption,
29+
},
30+
{
31+
customOptions: []interface{}{"--table=test&table"},
32+
expectedResult: errInvalidOption,
33+
},
34+
{
35+
customOptions: []interface{}{5},
36+
expectedResult: errInvalidOptionType,
37+
},
38+
}
39+
40+
for _, tc := range testCases {
41+
validationResult := validateCustomOptions(tc.customOptions)
42+
43+
require.ErrorIs(t, validationResult, tc.expectedResult)
44+
}
45+
}

‎engine/pkg/models/configuration.go

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,4 +25,6 @@ type ConfigProjection struct {
2525
DBList map[string]interface{} `proj:"retrieval.spec.logicalDump.options.databases,createKey"`
2626
DumpParallelJobs *int64 `proj:"retrieval.spec.logicalDump.options.parallelJobs"`
2727
RestoreParallelJobs *int64 `proj:"retrieval.spec.logicalRestore.options.parallelJobs"`
28+
DumpCustomOptions []interface{} `proj:"retrieval.spec.logicalDump.options.customOptions"`
29+
RestoreCustomOptions []interface{} `proj:"retrieval.spec.logicalRestore.options.customOptions"`
2830
}

‎engine/pkg/util/projection/yaml.go

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -75,6 +75,14 @@ func (y *yamlSoft) Set(set FieldSet) error {
7575
return nil
7676
}
7777

78+
if seq, ok := set.Value.([]interface{}); ok {
79+
if err := node.Encode(seq); err != nil {
80+
return fmt.Errorf("cannot encode slice: %w", err)
81+
}
82+
83+
return nil
84+
}
85+
7886
conv, err := ptypes.Convert(set.Value, ptypes.String)
7987
if err != nil {
8088
return err
@@ -109,6 +117,10 @@ func (y *yamlSoft) Get(get FieldGet) (interface{}, error) {
109117
return convertMap(node)
110118
}
111119

120+
if node.Tag == "!!seq" {
121+
return convertSlice(node)
122+
}
123+
112124
typed, err := ptypes.Convert(node.Value, get.Type)
113125
if err != nil {
114126
return nil, err
@@ -191,6 +203,8 @@ func ptypeToNodeTag(t ptypes.Type) string {
191203
return "!!bool"
192204
case ptypes.Map:
193205
return "!!map"
206+
case ptypes.Slice:
207+
return "!!seq"
194208
default:
195209
return ""
196210
}
@@ -208,6 +222,8 @@ func nodeTagToPType(nodeTag string) ptypes.Type {
208222
return ptypes.Bool
209223
case "!!map":
210224
return ptypes.Map
225+
case "!!seq":
226+
return ptypes.Slice
211227
default:
212228
return ptypes.Invalid
213229
}

0 commit comments

Comments
 (0)