fix: logical dump and restore issues (#152)

Author: agneum

pkg/retrieval/engine/postgres/initialize/logical/dump.go

@@ -65,7 +65,7 @@ type DumpJob struct {
 
 // DumpOptions defines a logical dump options.
 type DumpOptions struct {
-	DumpFile     string         `yaml:"dumpLocation"`
+	DumpLocation string         `yaml:"dumpLocation"`
 	DockerImage  string         `yaml:"dockerImage"`
 	Connection   Connection     `yaml:"connection"`
 	Source       Source         `yaml:"source"`
@@ -273,36 +273,19 @@ func (d *DumpJob) Run(ctx context.Context) (err error) {
 		log.Msg("Partial dump will be run. Tables for dumping: ", strings.Join(d.Partial.Tables, ", "))
 	}
 
-	var output io.Writer = os.Stdout
-
-	if d.DumpOptions.DumpFile != "" {
-		dumpFile, err := os.Create(d.getDumpContainerPath())
-		if err != nil {
-			return errors.Wrap(err, "failed to create file")
-		}
-
-		defer func() {
-			if err := dumpFile.Close(); err != nil {
-				log.Err("failed to close dump file", err)
-			}
-		}()
-
-		output = dumpFile
-	}
-
-	if err := d.performDumpCommand(ctx, output, cont.ID, execCommand.ID); err != nil {
+	if err := d.performDumpCommand(ctx, os.Stdout, cont.ID, execCommand.ID); err != nil {
 		return errors.Wrap(err, "failed to dump a database")
 	}
 
-	if err := d.markDatabaseData(); err != nil {
-		return errors.Wrap(err, "failed to mark the created dump")
-	}
-
 	if d.DumpOptions.Restore != nil {
+		if err := d.markDatabaseData(); err != nil {
+			return errors.Wrap(err, "failed to mark the created dump")
+		}
+
 		if err := recalculateStats(ctx, d.dockerClient, cont.ID, buildAnalyzeCommand(Connection{
 			DBName:   d.config.db.DBName,
 			Username: defaults.Username,
-		})); err != nil {
+		}, d.DumpOptions.ParallelJobs)); err != nil {
 			return errors.Wrap(err, "failed to recalculate statistics after restore")
 		}
 	}
@@ -345,16 +328,16 @@ func (d *DumpJob) performDumpCommand(ctx context.Context, cmdOutput io.Writer, c
 	return nil
 }
 
-func (d *DumpJob) getDumpContainerPath() string {
-	return d.DumpFile
-}
-
 func (d *DumpJob) getEnvironmentVariables() []string {
 	envs := []string{
-		"PGDATA=" + d.globalCfg.DataDir,
 		"POSTGRES_HOST_AUTH_METHOD=trust",
 	}
 
+	// Avoid initialization of PostgreSQL directory in case of preparing of a dump.
+	if d.DumpOptions.Restore != nil {
+		envs = append(envs, "PGDATA="+d.globalCfg.DataDir)
+	}
+
 	if d.DumpOptions.Source.Type == sourceTypeLocal && d.DumpOptions.Source.Connection.Port == defaults.Port {
 		log.Msg(fmt.Sprintf("The default PostgreSQL port is busy, trying to use an alternative one: %d", reservePort))
 		envs = append(envs, "PGPORT="+strconv.Itoa(reservePort))
@@ -365,6 +348,7 @@ func (d *DumpJob) getEnvironmentVariables() []string {
 
 func (d *DumpJob) buildContainerConfig() *container.Config {
 	return &container.Config{
+		Labels:      map[string]string{"label": tools.DBLabControlLabel},
 		Env:         d.getEnvironmentVariables(),
 		Image:       d.DockerImage,
 		Healthcheck: health.GetConfig(),
@@ -388,11 +372,11 @@ func (d *DumpJob) buildHostConfig() (*container.HostConfig, error) {
 func (d *DumpJob) getMountVolumes() []mount.Mount {
 	mounts := d.dumper.GetMounts()
 
-	if d.DumpOptions.DumpFile != "" {
+	if d.DumpOptions.DumpLocation != "" {
 		mounts = append(mounts, mount.Mount{
 			Type:   mount.TypeBind,
-			Source: filepath.Dir(d.DumpOptions.DumpFile),
-			Target: filepath.Dir(d.DumpOptions.DumpFile),
+			Source: filepath.Dir(d.DumpOptions.DumpLocation),
+			Target: filepath.Dir(d.DumpOptions.DumpLocation),
 		})
 	}
 
@@ -426,21 +410,27 @@ func (d *DumpJob) getExecEnvironmentVariables() []string {
 }
 
 func (d *DumpJob) buildLogicalDumpCommand() []string {
-	dumpCmd := []string{"pg_dump", "-C", "-Fc"}
+	format := "custom"
+
+	if d.DumpOptions.ParallelJobs > defaultParallelJobs {
+		format = "directory"
+	}
 
 	optionalArgs := map[string]string{
-		"-h": d.config.db.Host,
-		"-p": strconv.Itoa(d.config.db.Port),
-		"-U": d.config.db.Username,
-		"-d": d.config.db.DBName,
-		"-j": strconv.Itoa(d.DumpOptions.ParallelJobs),
+		"--host":     d.config.db.Host,
+		"--port":     strconv.Itoa(d.config.db.Port),
+		"--username": d.config.db.Username,
+		"--dbname":   d.config.db.DBName,
+		"--jobs":     strconv.Itoa(d.DumpOptions.ParallelJobs),
 	}
-	dumpCmd = append(dumpCmd, prepareCmdOptions(optionalArgs)...)
+
+	dumpCmd := append([]string{"pg_dump", "--create", "--format", format}, prepareCmdOptions(optionalArgs)...)
 
 	for _, table := range d.Partial.Tables {
-		dumpCmd = append(dumpCmd, "-t", table)
+		dumpCmd = append(dumpCmd, "--table", table)
 	}
 
+	// Define if restore directly or export to dump location.
 	if d.DumpOptions.Restore != nil {
 		dumpCmd = append(dumpCmd, d.buildLogicalRestoreCommand()...)
 		cmd := strings.Join(dumpCmd, " ")
@@ -450,11 +440,13 @@ func (d *DumpJob) buildLogicalDumpCommand() []string {
 		return []string{"sh", "-c", cmd}
 	}
 
+	dumpCmd = append(dumpCmd, "--file", d.DumpOptions.DumpLocation)
+
 	return dumpCmd
 }
 
 func (d *DumpJob) buildLogicalRestoreCommand() []string {
-	restoreCmd := []string{"|", "pg_restore", "-U", defaults.Username, "-C", "-d", defaults.DBName, "--no-privileges"}
+	restoreCmd := []string{"|", "pg_restore", "--username", defaults.Username, "--create", "--dbname", defaults.DBName, "--no-privileges"}
 
 	if d.Restore.ForceInit {
 		restoreCmd = append(restoreCmd, "--clean", "--if-exists")

pkg/retrieval/engine/postgres/initialize/logical/logical.go

@@ -7,6 +7,7 @@ package logical
 
 import (
 	"context"
+	"strconv"
 
 	"github.com/docker/docker/api/types"
 	"github.com/docker/docker/client"
@@ -40,12 +41,13 @@ func recalculateStats(ctx context.Context, dockerClient *client.Client, contID s
 	return nil
 }
 
-func buildAnalyzeCommand(conn Connection) []string {
+func buildAnalyzeCommand(conn Connection, parallelJobs int) []string {
 	analyzeCmd := []string{
-		"psql",
-		"-U", conn.Username,
-		"-d", conn.DBName,
-		"-c", "vacuum freeze analyze;",
+		"vacuumdb",
+		"--analyze",
+		"--jobs", strconv.Itoa(parallelJobs),
+		"--username", conn.Username,
+		"--dbname", conn.DBName,
 	}
 
 	return analyzeCmd

pkg/retrieval/engine/postgres/initialize/logical/restore.go

@@ -148,10 +148,6 @@ func (r *RestoreJob) Run(ctx context.Context) (err error) {
 
 	log.Msg(fmt.Sprintf("Running container: %s. ID: %v", r.restoreContainerName(), cont.ID))
 
-	if err := r.markDatabase(ctx, cont.ID); err != nil {
-		return errors.Wrap(err, "failed to mark the database")
-	}
-
 	if err := tools.CheckContainerReadiness(ctx, r.dockerClient, cont.ID); err != nil {
 		return errors.Wrap(err, "failed to readiness check")
 	}
@@ -182,10 +178,14 @@ func (r *RestoreJob) Run(ctx context.Context) (err error) {
 		return errors.Wrap(err, "failed to exec restore command")
 	}
 
+	if err := r.markDatabase(ctx, cont.ID); err != nil {
+		return errors.Wrap(err, "failed to mark the database")
+	}
+
 	if err := recalculateStats(ctx, r.dockerClient, cont.ID, buildAnalyzeCommand(Connection{
-		DBName:   r.RestoreOptions.DBName,
 		Username: defaults.Username,
-	})); err != nil {
+		DBName:   r.RestoreOptions.DBName,
+	}, r.RestoreOptions.ParallelJobs)); err != nil {
 		return errors.Wrap(err, "failed to recalculate statistics after restore")
 	}
 
@@ -196,8 +196,10 @@ func (r *RestoreJob) Run(ctx context.Context) (err error) {
 
 func (r *RestoreJob) buildContainerConfig() *container.Config {
 	return &container.Config{
+		Labels: map[string]string{"label": tools.DBLabControlLabel},
 		Env: []string{
 			"PGDATA=" + r.globalCfg.DataDir,
+			"POSTGRES_HOST_AUTH_METHOD=trust",
 		},
 		Image:       r.RestoreOptions.DockerImage,
 		Healthcheck: health.GetConfig(),
@@ -244,7 +246,7 @@ func (r *RestoreJob) markDatabase(ctx context.Context, contID string) error {
 }
 
 func (r *RestoreJob) retrieveDataStateAt(ctx context.Context, contID string) (string, error) {
-	restoreMetaCmd := []string{"sh", "-c", "pg_restore -l " + r.RestoreOptions.DumpFile + " | head -n 10"}
+	restoreMetaCmd := []string{"sh", "-c", "pg_restore --list " + r.RestoreOptions.DumpFile + " | head -n 10"}
 
 	log.Dbg("Running a restore metadata command: ", restoreMetaCmd)
 
@@ -273,18 +275,19 @@ func (r *RestoreJob) retrieveDataStateAt(ctx context.Context, contID string) (st
 }
 
 func (r *RestoreJob) buildLogicalRestoreCommand() []string {
-	restoreCmd := []string{"pg_restore", "-U", defaults.Username, "-C"}
+	restoreCmd := []string{"pg_restore", "--username", defaults.Username, "--dbname", defaults.DBName, "--create", "--no-privileges"}
 
 	if r.ForceInit {
-		restoreCmd = append(restoreCmd, "-d", defaults.DBName, "--clean", "--if-exists")
-	} else {
-		restoreCmd = append(restoreCmd, "-d", r.RestoreOptions.DBName)
+		restoreCmd = append(restoreCmd, "--clean", "--if-exists")
 	}
+	//else {
+	//	restoreCmd = append(restoreCmd)
+	//}
 
-	restoreCmd = append(restoreCmd, "-j", strconv.Itoa(r.ParallelJobs))
+	restoreCmd = append(restoreCmd, "--jobs", strconv.Itoa(r.ParallelJobs))
 
 	for _, table := range r.Partial.Tables {
-		restoreCmd = append(restoreCmd, "-t", table)
+		restoreCmd = append(restoreCmd, "--table", table)
 	}
 
 	restoreCmd = append(restoreCmd, r.RestoreOptions.DumpFile)

pkg/retrieval/engine/postgres/initialize/logical/restore_test.go

@@ -24,14 +24,14 @@ func TestRestoreCommandBuilding(t *testing.T) {
 				ForceInit:    false,
 				DumpFile:     "/tmp/db.dump",
 			},
-			Command: []string{"pg_restore", "-U", "postgres", "-C", "-d", "testDB", "-j", "1", "/tmp/db.dump"},
+			Command: []string{"pg_restore", "--username", "postgres", "--dbname", "postgres", "--create", "--no-privileges", "--jobs", "1", "/tmp/db.dump"},
 		},
 		{
 			CopyOptions: RestoreOptions{
 				ParallelJobs: 4,
 				ForceInit:    true,
 			},
-			Command: []string{"pg_restore", "-U", "postgres", "-C", "-d", "postgres", "--clean", "--if-exists", "-j", "4", ""},
+			Command: []string{"pg_restore", "--username", "postgres", "--dbname", "postgres", "--create", "--no-privileges", "--clean", "--if-exists", "--jobs", "4", ""},
 		},
 		{
 			CopyOptions: RestoreOptions{
@@ -40,7 +40,7 @@ func TestRestoreCommandBuilding(t *testing.T) {
 				Partial:      Partial{Tables: []string{"test", "users"}},
 				DumpFile:     "/tmp/db.dump",
 			},
-			Command: []string{"pg_restore", "-U", "postgres", "-C", "-d", "testDB", "-j", "1", "-t", "test", "-t", "users", "/tmp/db.dump"},
+			Command: []string{"pg_restore", "--username", "postgres", "--dbname", "postgres", "--create", "--no-privileges", "--jobs", "1", "--table", "test", "--table", "users", "/tmp/db.dump"},
 		},
 	}
 

pkg/retrieval/engine/postgres/initialize/physical/physical.go

@@ -360,8 +360,9 @@ func (r *RestoreJob) getEnvironmentVariables() []string {
 
 func (r *RestoreJob) buildContainerConfig() *container.Config {
 	return &container.Config{
-		Env:   r.getEnvironmentVariables(),
-		Image: r.DockerImage,
+		Labels: map[string]string{"label": tools.DBLabControlLabel},
+		Env:    r.getEnvironmentVariables(),
+		Image:  r.DockerImage,
 	}
 }
 

pkg/retrieval/engine/postgres/initialize/snapshot/physical.go

@@ -465,7 +465,7 @@ func (p *PhysicalInitial) adjustRecoveryConfiguration(pgVersion, clonePGDataDir
 
 func (p *PhysicalInitial) buildContainerConfig(clonePath, promoteImage string) *container.Config {
 	return &container.Config{
-		Labels: map[string]string{"label": "dblab_control"},
+		Labels: map[string]string{"label": tools.DBLabControlLabel},
 		Env: []string{
 			"PGDATA=" + clonePath,
 			"POSTGRES_HOST_AUTH_METHOD=trust",

pkg/retrieval/engine/postgres/initialize/tools/tools.go

@@ -38,6 +38,9 @@ const (
 
 	// SyncInstanceContainerPrefix defines a sync container name.
 	SyncInstanceContainerPrefix = "dblab_sync_"
+
+	// DBLabControlLabel defines a label to mark service containers.
+	DBLabControlLabel = "dblab_control"
 )
 
 // IsEmptyDirectory checks whether a directory is empty.