Merge pull request #33 from GreenmaskIO/restore_command_fix

restore and show-dump command fixes

internal/db/postgres/cmd/restore.go

@@ -23,6 +23,7 @@ import (
 	"fmt"
 	"io"
 	"os"
+	"os/exec"
 	"path"
 	"regexp"
 	"slices"
@@ -74,6 +75,9 @@ type Restore struct {
 	dumpIdList []int32
 	tocObj     *toc.Toc
 	tmpDir     string
+
+	preDataClenUpToc  string
+	postDataClenUpToc string
 }
 
 func NewRestore(
@@ -222,9 +226,19 @@ func (r *Restore) sortAndFilterEntriesByRestoreList() error {
 }
 
 func (r *Restore) preDataRestore(ctx context.Context, conn *pgx.Conn) error {
-	// Do not restore this section if implicitly provided
+	// pg_dump has a limitation:
+	//	If we want to use --cleanup command then this command must be performed for whole schema (--schema-only)
+	//  without --section parameter. For avoiding cascade dropping we need to run pg_restore with --schema-only --clean
+	//  and then remove all post-data objects manually. If we call pg_restore with --section=pre-data --clean then it
+	//  causes errors because we need to drop post-data dependencies before dropping pre-data
+	//
+	//	In current implementation Greenmask modifies toc.dat file by removing create statement in post-data section
+	//  and applies this toc.dat in the pre-data section restoration. The post-data restoration uses original
+	//  (non modified) toc.dat file.
+
+	// Do not restore this section if implicitly provided another section
 	if r.restoreOpt.DataOnly ||
-		r.restoreOpt.Section != "" && r.restoreOpt.Section != preDataSection {
+		(r.restoreOpt.Section != "" && r.restoreOpt.Section != preDataSection) {
 		return nil
 	}
 
@@ -233,12 +247,30 @@ func (r *Restore) preDataRestore(ctx context.Context, conn *pgx.Conn) error {
 		return err
 	}
 
-	// Execute pre-data section restore using pg_restore
 	options := *r.restoreOpt
-	options.Section = "pre-data"
-	options.DirPath = r.tmpDir
+
+	if r.restoreOpt.Clean && r.restoreOpt.Section == "" {
+		// Handling parameters for --clean
+		options.SchemaOnly = true
+
+		// Build clean up toc for dropping dependant objects in post-data stage without restoration them
+		// right now
+		var err error
+		r.preDataClenUpToc, r.postDataClenUpToc, err = r.prepareCleanupToc()
+		if err != nil {
+			return fmt.Errorf("cannot prepare clean up toc: %w", err)
+		}
+		options.DirPath = r.preDataClenUpToc
+	} else {
+		options.DirPath = r.tmpDir
+		options.Section = "pre-data"
+	}
+
 	if err := r.pgRestore.Run(ctx, &options); err != nil {
-		return fmt.Errorf("cannot restore pre-data section using pg_restore: %w", err)
+		var exitErr *exec.ExitError
+		if r.restoreOpt.ExitOnError || (errors.As(err, &exitErr) && exitErr.ExitCode() != 1) {
+			return fmt.Errorf("cannot restore pre-data section using pg_restore: %w", err)
+		}
 	}
 
 	// Execute PreData After scripts
@@ -249,6 +281,63 @@ func (r *Restore) preDataRestore(ctx context.Context, conn *pgx.Conn) error {
 	return nil
 }
 
+// prepareCleanupToc - replaces create statements in post-data section with SELECT 1 and stores in tmp directory
+func (r *Restore) prepareCleanupToc() (string, string, error) {
+	preDataCleanUpToc := r.tocObj.Copy()
+	postDataCleanUpToc := r.tocObj.Copy()
+
+	statementReplacements := ";"
+
+	for idx := range preDataCleanUpToc.Entries {
+		log.Debug().Int("a", idx)
+		preEntry := preDataCleanUpToc.Entries[idx]
+		if preEntry.Section == toc.SectionPostData && preEntry.Defn != nil {
+			preEntry.Defn = &statementReplacements
+		}
+
+		postEntry := postDataCleanUpToc.Entries[idx]
+		if postEntry.Section == toc.SectionPostData && postEntry.DropStmt != nil {
+			postEntry.DropStmt = &statementReplacements
+		}
+	}
+
+	preDatadirName := path.Join(r.tmpDir, "pre_data_clean_up_toc")
+	postDatadirName := path.Join(r.tmpDir, "post_data_clean_up_toc")
+
+	// Pre-data section
+
+	if err := os.Mkdir(preDatadirName, 0700); err != nil {
+		return "", "", fmt.Errorf("cannot create pre-data clean up toc directory: %w", err)
+	}
+
+	f1, err := os.Create(path.Join(preDatadirName, "toc.dat"))
+	if err != nil {
+		return "", "", fmt.Errorf("cannot create clean up toc file: %w", err)
+	}
+	defer f1.Close()
+
+	if err = toc.NewWriter(f1).Write(preDataCleanUpToc); err != nil {
+		return "", "", fmt.Errorf("cannot write clean up toc: %w", err)
+	}
+
+	// post-data section
+	if err = os.Mkdir(postDatadirName, 0700); err != nil {
+		return "", "", fmt.Errorf("cannot create post-data clean up toc directory: %w", err)
+	}
+
+	f2, err := os.Create(path.Join(postDatadirName, "toc.dat"))
+	if err != nil {
+		return "", "", fmt.Errorf("cannot create clean up toc file: %w", err)
+	}
+	defer f2.Close()
+
+	if err = toc.NewWriter(f2).Write(postDataCleanUpToc); err != nil {
+		return "", "", fmt.Errorf("cannot write clean up toc: %w", err)
+	}
+
+	return preDatadirName, postDatadirName, nil
+}
+
 func (r *Restore) dataRestore(ctx context.Context, conn *pgx.Conn) error {
 	// Execute Data Before scripts
 
@@ -366,6 +455,11 @@ func (r *Restore) postDataRestore(ctx context.Context, conn *pgx.Conn) error {
 	options := *r.restoreOpt
 	options.Section = "post-data"
 	options.DirPath = r.tmpDir
+
+	if r.postDataClenUpToc != "" {
+		options.DirPath = r.postDataClenUpToc
+	}
+
 	if err := r.pgRestore.Run(ctx, &options); err != nil {
 		return fmt.Errorf("cannot restore post-data section using pg_restore: %w", err)
 	}

internal/db/postgres/cmd/show_dump.go

@@ -38,7 +38,7 @@ const (
 )
 
 var templateString = `;
-; Archive created at {{ .Header.CreationDate.TableFormat "2006-01-02 15:04:05 UTC" }}
+; Archive created at {{ .Header.CreationDate.Format "2006-01-02 15:04:05 UTC" }}
 ;     dbname: {{ .Header.DbName }}
 ;     TOC Entries: {{ .Header.TocEntriesCount }}
 ;     Compression: {{ .Header.Compression }}

internal/db/postgres/toc/entry.go

@@ -61,3 +61,42 @@ type Entry struct {
 	//OriginalSize   int64
 	//CompressedSize int64
 }
+
+func (e *Entry) Copy() *Entry {
+	res := NewObj(*e)
+	if e.Tag != nil {
+		res.Tag = NewObj(*e.Tag)
+	}
+	if e.Namespace != nil {
+		res.Namespace = NewObj(*e.Namespace)
+	}
+	if e.Tablespace != nil {
+		res.Tablespace = NewObj(*e.Tablespace)
+	}
+	if e.Tableam != nil {
+		res.Tableam = NewObj(*e.Tableam)
+	}
+	if e.Owner != nil {
+		res.Owner = NewObj(*e.Owner)
+	}
+	if e.Desc != nil {
+		res.Desc = NewObj(*e.Desc)
+	}
+	if e.Defn != nil {
+		res.Defn = NewObj(*e.Defn)
+	}
+	if e.DropStmt != nil {
+		res.DropStmt = NewObj(*e.DropStmt)
+	}
+	if e.CopyStmt != nil {
+		res.CopyStmt = NewObj(*e.CopyStmt)
+	}
+	if e.FileName != nil {
+		res.FileName = NewObj(*e.FileName)
+	}
+	return res
+}
+
+func NewObj[T string | Toc | Header | Entry](v T) *T {
+	return &v
+}

internal/db/postgres/toc/header.go

@@ -54,3 +54,17 @@ type Header struct {
 	TocCount             int32
 	MaxDumpId            int32
 }
+
+func (h *Header) Copy() *Header {
+	res := NewObj(*h)
+	if h.ArchDbName != nil {
+		res.ArchDbName = NewObj(*h.ArchDbName)
+	}
+	if h.ArchiveRemoteVersion != nil {
+		res.ArchiveRemoteVersion = NewObj(*h.ArchiveRemoteVersion)
+	}
+	if h.ArchiveDumpVersion != nil {
+		res.ArchiveDumpVersion = NewObj(*h.ArchiveDumpVersion)
+	}
+	return res
+}

internal/db/postgres/toc/toc.go

@@ -18,3 +18,17 @@ type Toc struct {
 	Header  *Header
 	Entries []*Entry
 }
+
+func (t *Toc) Copy() *Toc {
+
+	entries := make([]*Entry, len(t.Entries))
+
+	for i, entry := range t.Entries {
+		entries[i] = entry.Copy()
+	}
+
+	return &Toc{
+		Header:  t.Header.Copy(),
+		Entries: entries,
+	}
+}

internal/utils/cmd_runner/cmd_runner.go

@@ -88,7 +88,7 @@ func Run(ctx context.Context, logger *zerolog.Logger, name string, args ...strin
 		defer outWriter.Close()
 		defer errWriter.Close()
 		if err := cmd.Wait(); err != nil {
-			return fmt.Errorf("external command runtime error: %w", err)
+			return err
 		}
 		return nil
 	})