Merge branch '224-parallel-job-configs' into 'master'

agneum · agneum · commit 5b95c2dc6ae5 · 2021-01-21T04:34:22.000Z
fix: incompatible options for parallel jobs in logical retrieval steps (#224) Closes #224 See merge request postgres-ai/database-lab!238
diff --git a/pkg/retrieval/engine/postgres/logical/dump.go b/pkg/retrieval/engine/postgres/logical/dump.go
@@ -48,6 +48,10 @@ const (
 	// Container network modes.
 	networkModeDefault = container.NetworkMode("default")
 	networkModeHost    = container.NetworkMode("host")
+
+	// PostgreSQL pg_dump formats.
+	customFormat    = "custom"
+	directoryFormat = "directory"
 )
 
 // DumpJob declares a job for logical dumping.
@@ -397,12 +401,6 @@ func (d *DumpJob) getExecEnvironmentVariables() []string {
 }
 
 func (d *DumpJob) buildLogicalDumpCommand() []string {
-	format := "custom"
-
-	if d.DumpOptions.ParallelJobs > defaultParallelJobs {
-		format = "directory"
-	}
-
 	optionalArgs := map[string]string{
 		"--host":     d.config.db.Host,
 		"--port":     strconv.Itoa(d.config.db.Port),
@@ -411,14 +409,15 @@ func (d *DumpJob) buildLogicalDumpCommand() []string {
 		"--jobs":     strconv.Itoa(d.DumpOptions.ParallelJobs),
 	}
 
-	dumpCmd := append([]string{"pg_dump", "--create", "--format", format}, prepareCmdOptions(optionalArgs)...)
+	dumpCmd := append([]string{"pg_dump", "--create"}, prepareCmdOptions(optionalArgs)...)
 
 	for _, table := range d.Partial.Tables {
 		dumpCmd = append(dumpCmd, "--table", table)
 	}
 
 	// Define if restore directly or export to dump location.
 	if d.DumpOptions.Restore != nil {
+		dumpCmd = append(dumpCmd, "--format", customFormat)
 		dumpCmd = append(dumpCmd, d.buildLogicalRestoreCommand()...)
 		cmd := strings.Join(dumpCmd, " ")
 
@@ -427,7 +426,7 @@ func (d *DumpJob) buildLogicalDumpCommand() []string {
 		return []string{"sh", "-c", cmd}
 	}
 
-	dumpCmd = append(dumpCmd, "--file", d.DumpOptions.DumpLocation)
+	dumpCmd = append(dumpCmd, "--format", directoryFormat, "--file", d.DumpOptions.DumpLocation)
 
 	return dumpCmd
 }
diff --git a/pkg/retrieval/engine/postgres/logical/restore.go b/pkg/retrieval/engine/postgres/logical/restore.go
@@ -37,7 +37,9 @@ const (
 
 	// const defines restore options.
 	restoreContainerPrefix = "dblab_lr_"
-	defaultParallelJobs    = 1
+
+	// defaultParallelJobs declares a default number of parallel jobs for logical dump and restore.
+	defaultParallelJobs = 1
 )
 
 // RestoreJob defines a logical restore job.
diff --git a/pkg/retrieval/retrieval.go b/pkg/retrieval/retrieval.go
@@ -40,6 +40,7 @@ type Retrieval struct {
 	scheduler     *cron.Cron
 	retrieveMutex sync.Mutex
 	ctxCancel     context.CancelFunc
+	jobSpecs      map[string]config.JobSpec
 }
 
 // New creates a new data retrieval.
@@ -50,6 +51,7 @@ func New(cfg *dblabCfg.Config, docker *client.Client, pm *pool.Manager, runner r
 		docker:      docker,
 		poolManager: pm,
 		runner:      runner,
+		jobSpecs:    make(map[string]config.JobSpec, len(cfg.Retrieval.Jobs)),
 	}
 }
 
@@ -147,6 +149,7 @@ func (r *Retrieval) parseJobs(fsm pool.FSManager) error {
 		}
 
 		jobSpec.Name = jobName
+		r.jobSpecs[jobName] = jobSpec
 
 		jobCfg := config.JobConfig{
 			Spec:   jobSpec,
@@ -172,16 +175,10 @@ func (r *Retrieval) addJob(job components.JobRunner) {
 }
 
 func (r *Retrieval) validate() error {
-	jobsList := make(map[string]struct{}, len(r.jobs))
+	_, hasLogicalRestore := r.jobSpecs[logical.RestoreJobType]
+	_, hasPhysicalRestore := r.jobSpecs[physical.RestoreJobType]
 
-	for _, job := range r.jobs {
-		jobsList[job.Name()] = struct{}{}
-	}
-
-	_, hasLogical := jobsList[logical.RestoreJobType]
-	_, hasPhysical := jobsList[physical.RestoreJobType]
-
-	if hasLogical && hasPhysical {
+	if hasLogicalRestore && hasPhysicalRestore {
 		return errors.New("must not contain physical and logical restore jobs simultaneously")
 	}
 
diff --git a/pkg/retrieval/retrieval_test.go b/pkg/retrieval/retrieval_test.go
@@ -0,0 +1,55 @@
+package retrieval
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+
+	"gitlab.com/postgres-ai/database-lab/pkg/retrieval/config"
+)
+
+func TestParallelJobSpecs(t *testing.T) {
+	testCases := []map[string]config.JobSpec{
+		{
+			"logicalRestore": {},
+		},
+		{
+			"physicalRestore": {},
+		},
+		{
+			"logicalDump": {},
+		},
+		{
+			"logicalDump":    {},
+			"logicalRestore": {},
+		},
+	}
+
+	for _, tc := range testCases {
+		r := Retrieval{
+			jobSpecs: tc,
+		}
+
+		err := r.validate()
+		assert.Nil(t, err)
+	}
+
+}
+
+func TestInvalidParallelJobSpecs(t *testing.T) {
+	testCases := []map[string]config.JobSpec{
+		{
+			"physicalRestore": {},
+			"logicalRestore":  {},
+		},
+	}
+
+	for _, tc := range testCases {
+		r := Retrieval{
+			jobSpecs: tc,
+		}
+
+		err := r.validate()
+		assert.Error(t, err)
+	}
+}

Original file line number	Diff line number	Diff line change
`@@ -37,7 +37,9 @@ const (`
`37`	`37`
`38`	`38`	`// const defines restore options.`
`39`	`39`	`restoreContainerPrefix = "dblab_lr_"`
`40`		`- defaultParallelJobs = 1`
	`40`	`+`
	`41`	`+ // defaultParallelJobs declares a default number of parallel jobs for logical dump and restore.`
	`42`	`+ defaultParallelJobs = 1`
`41`	`43`	`)`
`42`	`44`
`43`	`45`	`// RestoreJob defines a logical restore job.`
Original file line number	Diff line number	Diff line change
`@@ -40,6 +40,7 @@ type Retrieval struct {`
`40`	`40`	`scheduler *cron.Cron`
`41`	`41`	`retrieveMutex sync.Mutex`
`42`	`42`	`ctxCancel context.CancelFunc`
	`43`	`+ jobSpecs map[string]config.JobSpec`
`43`	`44`	`}`
`44`	`45`
`45`	`46`	`// New creates a new data retrieval.`
`@@ -50,6 +51,7 @@ func New(cfg dblabCfg.Config, docker client.Client, pm *pool.Manager, runner r`
`50`	`51`	`docker: docker,`
`51`	`52`	`poolManager: pm,`
`52`	`53`	`runner: runner,`
	`54`	`+ jobSpecs: make(map[string]config.JobSpec, len(cfg.Retrieval.Jobs)),`
`53`	`55`	`}`
`54`	`56`	`}`
`55`	`57`
`@@ -147,6 +149,7 @@ func (r *Retrieval) parseJobs(fsm pool.FSManager) error {`
`147`	`149`	`}`
`148`	`150`
`149`	`151`	`jobSpec.Name = jobName`
	`152`	`+ r.jobSpecs[jobName] = jobSpec`
`150`	`153`
`151`	`154`	`jobCfg := config.JobConfig{`
`152`	`155`	`Spec: jobSpec,`
`@@ -172,16 +175,10 @@ func (r *Retrieval) addJob(job components.JobRunner) {`
`172`	`175`	`}`
`173`	`176`
`174`	`177`	`func (r *Retrieval) validate() error {`
`175`		`- jobsList := make(map[string]struct{}, len(r.jobs))`
	`178`	`+ _, hasLogicalRestore := r.jobSpecs[logical.RestoreJobType]`
	`179`	`+ _, hasPhysicalRestore := r.jobSpecs[physical.RestoreJobType]`
`176`	`180`
`177`		`- for _, job := range r.jobs {`
`178`		`- jobsList[job.Name()] = struct{}{}`
`179`		`- }`
`180`		`-`
`181`		`- _, hasLogical := jobsList[logical.RestoreJobType]`
`182`		`- _, hasPhysical := jobsList[physical.RestoreJobType]`
`183`		`-`
`184`		`- if hasLogical && hasPhysical {`
	`181`	`+ if hasLogicalRestore && hasPhysicalRestore {`
`185`	`182`	`return errors.New("must not contain physical and logical restore jobs simultaneously")`
`186`	`183`	`}`
`187`	`184`