feat: Disable triggers options (#252)

Closes #228 * Implemented --disable-triggers - disables triggers on the restoring table * Implemented --use-session-replication-role-replica - kind of workaround option for AWS RDS * Implemented --superuser - uses superuser role provided to disable tigers or set session-replication-role to replica * Refactored TableRestorer and TableRestorerInsertFormat - implemented restoreBase struct that embeds with base functionality. * Implemented suite test helpers that provides base functionality to PG container in unit tests * Added simple storage mock * Covered Table restorers with tests
GreenmaskIO · Dec 7, 2024 · 33511ef · 33511ef
1 parent fc27a6a
commit 33511ef
Show file tree

Hide file tree

Showing 23 changed files with 1,791 additions and 265 deletions.
diff --git a/cmd/greenmask/cmd/dump/dump.go b/cmd/greenmask/cmd/dump/dump.go
@@ -66,7 +66,6 @@ var (
 // TODO: Check how does work mixed options - use-list + tables, etc.
 // TODO: Options currently are not implemented:
 //   - encoding
-//   - disable-triggers
 //   - lock-wait-timeout
 //   - no-sync
 //   - data-only
@@ -104,7 +103,6 @@ func init() {
 	)
 	Cmd.Flags().BoolP("no-owner", "O", false, "skip restoration of object ownership in plain-text format")
 	Cmd.Flags().BoolP("schema-only", "s", false, "dump only the schema, no data")
-	Cmd.Flags().StringP("superuser", "S", "", "superuser user name to use in plain-text format")
 	Cmd.Flags().StringSliceVarP(
 		&Config.Dump.PgDumpOptions.Table, "table", "t", []string{}, "dump the specified table(s) only",
 	)
@@ -113,7 +111,6 @@ func init() {
 	)
 	Cmd.Flags().BoolP("no-privileges", "X", false, "do not dump privileges (grant/revoke)")
 	Cmd.Flags().BoolP("disable-dollar-quoting", "", false, "disable dollar quoting, use SQL standard quoting")
-	Cmd.Flags().BoolP("disable-triggers", "", false, "disable triggers during data-only restore")
 	Cmd.Flags().BoolP(
 		"enable-row-security", "", false, "enable row security (dump only content user has access to)",
 	)
@@ -163,8 +160,8 @@ func init() {
 		"file", "jobs", "verbose", "compress", "dbname", "host", "username", "lock-wait-timeout", "no-sync",
 
 		"data-only", "blobs", "no-blobs", "clean", "create", "extension", "encoding", "schema", "exclude-schema",
-		"no-owner", "schema-only", "superuser", "table", "exclude-table", "no-privileges", "disable-dollar-quoting",
-		"disable-triggers", "enable-row-security", "exclude-table-data", "extra-float-digits", "if-exists",
+		"no-owner", "schema-only", "table", "exclude-table", "no-privileges", "disable-dollar-quoting",
+		"enable-row-security", "exclude-table-data", "extra-float-digits", "if-exists",
 		"include-foreign-data", "load-via-partition-root", "no-comments", "no-publications", "no-security-labels",
 		"no-subscriptions", "no-synchronized-snapshots", "no-tablespaces", "no-toast-compression",
 		"no-unlogged-table-data", "quote-all-identifiers", "section",

diff --git a/cmd/greenmask/cmd/restore/restore.go b/cmd/greenmask/cmd/restore/restore.go
@@ -150,12 +150,11 @@ func init() {
 	Cmd.Flags().BoolP("no-owner", "O", false, "skip restoration of object ownership")
 	Cmd.Flags().StringSliceVarP(&Config.Restore.PgRestoreOptions.Function, "function", "P", []string{}, "restore named function")
 	Cmd.Flags().BoolP("schema-only", "s", false, "restore only the schema, no data")
-	Cmd.Flags().StringP("superuser", "S", "", "superuser user name to use for disabling triggers")
 	Cmd.Flags().StringSliceVarP(&Config.Restore.PgRestoreOptions.Table, "table", "t", []string{}, "restore named relation (table, view, etc.)")
 	Cmd.Flags().StringSliceVarP(&Config.Restore.PgRestoreOptions.Trigger, "trigger", "T", []string{}, "restore named trigger")
 	Cmd.Flags().BoolP("no-privileges", "X", false, "skip restoration of access privileges (grant/revoke)")
 	Cmd.Flags().BoolP("single-transaction", "1", false, "restore as a single transaction")
-	Cmd.Flags().BoolP("disable-triggers", "", false, "disable triggers during data-only restore")
+	Cmd.Flags().BoolP("disable-triggers", "", false, "disable triggers during data section restore")
 	Cmd.Flags().BoolP("enable-row-security", "", false, "enable row security")
 	Cmd.Flags().BoolP("if-exists", "", false, "use IF EXISTS when dropping objects")
 	Cmd.Flags().BoolP("no-comments", "", false, "do not restore comments")
@@ -169,6 +168,12 @@ func init() {
 	Cmd.Flags().BoolP("strict-names", "", false, "restore named section (pre-data, data, or post-data) match at least one entity each")
 	Cmd.Flags().BoolP("use-set-session-authorization", "", false, "use SET SESSION AUTHORIZATION commands instead of ALTER OWNER commands to set ownership")
 	Cmd.Flags().BoolP("on-conflict-do-nothing", "", false, "add ON CONFLICT DO NOTHING to INSERT commands")
+	Cmd.Flags().StringP("superuser", "S", "", "superuser user name to use for disabling triggers")
+	Cmd.Flags().BoolP(
+		"use-session-replication-role-replica", "", false,
+		"use SET session_replication_role = 'replica' to disable triggers during data section restore"+
+			" (alternative for --disable-triggers)",
+	)
 	Cmd.Flags().BoolP("inserts", "", false, "restore data as INSERT commands, rather than COPY")
 	Cmd.Flags().BoolP("restore-in-order", "", false, "restore tables in topological order, ensuring that dependent tables are not restored until the tables they depend on have been restored")
 	Cmd.Flags().BoolP(
@@ -193,11 +198,11 @@ func init() {
 		"dbname", "file", "verbose",
 
 		"data-only", "clean", "create", "exit-on-error", "jobs", "list-format", "use-list", "schema", "exclude-schema",
-		"no-owner", "function", "schema-only", "superuser", "table", "trigger", "no-privileges", "single-transaction",
+		"no-owner", "function", "schema-only", "table", "trigger", "no-privileges", "single-transaction",
 		"disable-triggers", "enable-row-security", "if-exists", "no-comments", "no-data-for-failed-tables",
 		"no-security-labels", "no-subscriptions", "no-table-access-method", "no-tablespaces", "section",
 		"strict-names", "use-set-session-authorization", "inserts", "on-conflict-do-nothing", "restore-in-order",
-		"pgzip", "batch-size", "overriding-system-value",
+		"pgzip", "batch-size", "overriding-system-value", "superuser", "use-session-replication-role-replica",
 
 		"host", "port", "username",
 	} {

diff --git a/docs/commands/dump.md b/docs/commands/dump.md
@@ -20,7 +20,6 @@ Mostly it supports the same flags as the `pg_dump` utility, with some extra flag
   -a, --data-only                       dump only the data, not the schema
   -d, --dbname string                   database to dump (default "postgres")
       --disable-dollar-quoting          disable dollar quoting, use SQL standard quoting
-      --disable-triggers                disable triggers during data-only restore
       --enable-row-security             enable row security (dump only content user has access to)
   -E, --encoding string                 dump the data in encoding ENCODING
   -N, --exclude-schema strings          dump the specified schema(s) only
@@ -37,7 +36,7 @@ Mostly it supports the same flags as the `pg_dump` utility, with some extra flag
       --lock-wait-timeout int           fail after waiting TIMEOUT for a table lock (default -1)
   -B, --no-blobs                        exclude large objects in dump
       --no-comments                     do not dump comments
-  -O, --no-owner string                 skip restoration of object ownership in plain-text format
+  -O, --no-owner                        skip restoration of object ownership in plain-text format
   -X, --no-privileges                   do not dump privileges (grant/revoke)
       --no-publications                 do not dump publications
       --no-security-labels              do not dump security label assignments
@@ -51,12 +50,11 @@ Mostly it supports the same flags as the `pg_dump` utility, with some extra flag
   -p, --port int                        database server port number (default 5432)
       --quote-all-identifiers           quote all identifiers, even if not key words
   -n, --schema strings                  dump the specified schema(s) only
-  -s, --schema-only string              dump only the schema, no data
+  -s, --schema-only                     dump only the schema, no data
       --section string                  dump named section (pre-data, data, or post-data)
       --serializable-deferrable         wait until the dump can run without anomalies
       --snapshot string                 use given snapshot for the dump
       --strict-names                    require table and/or schema include patterns to match at least one entity each
-  -S, --superuser string                superuser user name to use in plain-text format
   -t, --table strings                   dump the specified table(s) only
       --test string                     connect as specified database user (default "postgres")
       --use-set-session-authorization   use SET SESSION AUTHORIZATION commands instead of ALTER OWNER commands to set ownership

diff --git a/docs/commands/restore.md b/docs/commands/restore.md
@@ -19,49 +19,50 @@ allowing you to configure the restoration process as needed.
 Mostly it supports the same flags as the `pg_restore` utility, with some extra flags for Greenmask-specific features.
 
 ```text title="Supported flags"
-      --batch-size int                  the number of rows to insert in a single batch during the COPY command (0 - all rows will be inserted in a single batch)
-  -c, --clean                           clean (drop) database objects before recreating
-  -C, --create                          create the target database
-  -a, --data-only                       restore only the data, no schema
-  -d, --dbname string                   connect to database name (default "postgres")
-      --disable-triggers                disable triggers during data-only restore
-      --enable-row-security             enable row security
-  -N, --exclude-schema strings          do not restore objects in this schema
-  -e, --exit-on-error                   exit on error, default is to continue
-  -f, --file string                     output file name (- for stdout)
-  -P, --function strings                restore named function
-  -h, --host string                     database server host or socket directory (default "/var/run/postgres")
-      --if-exists                       use IF EXISTS when dropping objects
-  -i, --index strings                   restore named index
-      --inserts                         restore data as INSERT commands, rather than COPY
-  -j, --jobs int                        use this many parallel jobs to restore (default 1)
-      --list-format string              use table of contents in format of text, json or yaml (default "text")
-      --no-comments                     do not restore comments
-      --no-data-for-failed-tables       do not restore data of tables that could not be created
-  -O, --no-owner string                 skip restoration of object ownership
-  -X, --no-privileges                   skip restoration of access privileges (grant/revoke)
-      --no-publications                 do not restore publications
-      --no-security-labels              do not restore security labels
-      --no-subscriptions                ddo not restore subscriptions
-      --no-table-access-method          do not restore table access methods
-      --no-tablespaces                  do not restore tablespace assignments
-      --on-conflict-do-nothing          add ON CONFLICT DO NOTHING to INSERT commands
-      --overriding-system-value         use OVERRIDING SYSTEM VALUE clause for INSERTs
-      --pgzip                           use pgzip decompression instead of gzip
-  -p, --port int                        database server port number (default 5432)
-      --restore-in-order                restore tables in topological order, ensuring that dependent tables are not restored until the tables they depend on have been restored
-  -n, --schema strings                  restore only objects in this schema
-  -s, --schema-only                     restore only the schema, no data
-      --section string                  restore named section (pre-data, data, or post-data)
-  -1, --single-transaction              restore as a single transaction
-      --strict-names                    restore named section (pre-data, data, or post-data) match at least one entity each
-  -S, --superuser string                superuser user name to use for disabling triggers
-  -t, --table strings                   restore named relation (table, view, etc.)
-  -T, --trigger strings                 restore named trigger
-  -L, --use-list string                 use table of contents from this file for selecting/ordering output
-      --use-set-session-authorization   use SET SESSION AUTHORIZATION commands instead of ALTER OWNER commands to set ownership
-  -U, --username string                 connect as specified database user (default "postgres")
-  -v, --verbose string                  verbose mode
+      --batch-size int                         the number of rows to insert in a single batch during the COPY command (0 - all rows will be inserted in a single batch)
+  -c, --clean                                  clean (drop) database objects before recreating
+  -C, --create                                 create the target database
+  -a, --data-only                              restore only the data, no schema
+  -d, --dbname string                          connect to database name (default "postgres")
+      --disable-triggers                       disable triggers during data section restore
+      --enable-row-security                    enable row security
+  -N, --exclude-schema strings                 do not restore objects in this schema
+  -e, --exit-on-error                          exit on error, default is to continue
+  -f, --file string                            output file name (- for stdout)
+  -P, --function strings                       restore named function
+  -h, --host string                            database server host or socket directory (default "/var/run/postgres")
+      --if-exists                              use IF EXISTS when dropping objects
+  -i, --index strings                          restore named index
+      --inserts                                restore data as INSERT commands, rather than COPY
+  -j, --jobs int                               use this many parallel jobs to restore (default 1)
+      --list-format string                     use table of contents in format of text, json or yaml (default "text")
+      --no-comments                            do not restore comments
+      --no-data-for-failed-tables              do not restore data of tables that could not be created
+  -O, --no-owner                               skip restoration of object ownership
+  -X, --no-privileges                          skip restoration of access privileges (grant/revoke)
+      --no-publications                        do not restore publications
+      --no-security-labels                     do not restore security labels
+      --no-subscriptions                       ddo not restore subscriptions
+      --no-table-access-method                 do not restore table access methods
+      --no-tablespaces                         do not restore tablespace assignments
+      --on-conflict-do-nothing                 add ON CONFLICT DO NOTHING to INSERT commands
+      --overriding-system-value                use OVERRIDING SYSTEM VALUE clause for INSERTs
+      --pgzip                                  use pgzip decompression instead of gzip
+  -p, --port int                               database server port number (default 5432)
+      --restore-in-order                       restore tables in topological order, ensuring that dependent tables are not restored until the tables they depend on have been restored
+  -n, --schema strings                         restore only objects in this schema
+  -s, --schema-only                            restore only the schema, no data
+      --section string                         restore named section (pre-data, data, or post-data)
+  -1, --single-transaction                     restore as a single transaction
+      --strict-names                           restore named section (pre-data, data, or post-data) match at least one entity each
+  -S, --superuser string                       superuser user name to use for disabling triggers
+  -t, --table strings                          restore named relation (table, view, etc.)
+  -T, --trigger strings                        restore named trigger
+  -L, --use-list string                        use table of contents from this file for selecting/ordering output
+      --use-session-replication-role-replica   use SET session_replication_role = 'replica' to disable triggers during data section restore (alternative for --disable-triggers)
+      --use-set-session-authorization          use SET SESSION AUTHORIZATION commands instead of ALTER OWNER commands to set ownership
+  -U, --username string                        connect as specified database user (default "postgres")
+  -v, --verbose string                         verbose mode
 ```
 
 ## Extra features

diff --git a/go.mod b/go.mod
@@ -89,6 +89,7 @@ require (
 	github.com/sourcegraph/conc v0.3.0 // indirect
 	github.com/spf13/afero v1.11.0 // indirect
 	github.com/spf13/pflag v1.0.5 // indirect
+	github.com/stretchr/objx v0.5.2 // indirect
 	github.com/subosito/gotenv v1.6.0 // indirect
 	github.com/tidwall/match v1.1.1 // indirect
 	github.com/tidwall/pretty v1.2.0 // indirect

diff --git a/internal/db/postgres/cmd/restore.go b/internal/db/postgres/cmd/restore.go
@@ -632,13 +632,10 @@ func (r *Restore) taskPusher(ctx context.Context, tasks chan restorers.RestoreTa
 							return fmt.Errorf("cannot get table definition from meta: %w", err)
 						}
 						task = restorers.NewTableRestorerInsertFormat(
-							entry, t, r.st, r.restoreOpt.ExitOnError, r.restoreOpt.OnConflictDoNothing,
-							r.cfg.ErrorExclusions, r.restoreOpt.Pgzip, r.restoreOpt.OverridingSystemValue,
+							entry, t, r.st, r.restoreOpt.ToDataSectionSettings(), r.cfg.ErrorExclusions,
 						)
 					} else {
-						task = restorers.NewTableRestorer(
-							entry, r.st, r.restoreOpt.ExitOnError, r.restoreOpt.Pgzip, r.restoreOpt.BatchSize,
-						)
+						task = restorers.NewTableRestorer(entry, r.st, r.restoreOpt.ToDataSectionSettings())
 					}
 
 				case toc.SequenceSetDesc: