feat(restore): Halt on error by default

clevyr · Feb 20, 2024 · c6467c4 · c6467c4
1 parent 21d2887
commit c6467c4
Show file tree

Hide file tree

Showing 8 changed files with 25 additions and 0 deletions.
diff --git a/cmd/restore/restore.go b/cmd/restore/restore.go
@@ -50,6 +50,7 @@ Supported Input Filetypes:
 	flags.Quiet(cmd, &action.Quiet)
 	flags.RemoteGzip(cmd)
 	flags.Analyze(cmd)
+	flags.HaltOnError(cmd)
 	flags.Spinner(cmd, &action.Spinner)
 	cmd.Flags().BoolVarP(&action.Force, consts.ForceFlag, "f", false, "Do not prompt before restore")
 
@@ -85,6 +86,8 @@ func preRun(cmd *cobra.Command, args []string) (err error) {
 	flags.BindJobPodLabels(cmd)
 	flags.BindNoJob(cmd)
 	flags.BindSpinner(cmd)
+	flags.BindHaltOnError(cmd)
+	action.HaltOnError = viper.GetBool(consts.HaltOnErrorKey)
 	action.Spinner = viper.GetString(consts.SpinnerKey)
 
 	if len(args) > 0 {

diff --git a/docs/kubedb_restore.md b/docs/kubedb_restore.md
@@ -23,6 +23,7 @@ kubedb restore filename [flags]
   -d, --dbname string                   Database name to use (default discovered)
   -f, --force                           Do not prompt before restore
   -F, --format string                   Output file format One of (gzip|custom|plain) (default "gzip")
+      --halt-on-error                   Halt on error (Postgres only) (default true)
   -h, --help                            help for restore
       --job-pod-labels stringToString   Pod labels to add to the job (default [])
       --no-job                          Database commands will be run in the database pod instead of a dedicated job

diff --git a/internal/config/flags/database.go b/internal/config/flags/database.go
@@ -119,6 +119,16 @@ func BindAnalyze(cmd *cobra.Command) {
 	}
 }
 
+func HaltOnError(cmd *cobra.Command) {
+	cmd.Flags().Bool(consts.HaltOnErrorFlag, true, "Halt on error (Postgres only)")
+}
+
+func BindHaltOnError(cmd *cobra.Command) {
+	if err := viper.BindPFlag(consts.HaltOnErrorKey, cmd.Flags().Lookup(consts.HaltOnErrorFlag)); err != nil {
+		panic(err)
+	}
+}
+
 func listTables(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
 	conf := config.Exec{DisableHeaders: true}
 

diff --git a/internal/config/restore.go b/internal/config/restore.go
@@ -8,4 +8,5 @@ type Restore struct {
 	NoOwner           bool
 	Force             bool
 	Spinner           string
+	HaltOnError       bool
 }
diff --git a/internal/consts/flags.go b/internal/consts/flags.go
@@ -15,6 +15,7 @@ const (
 	ExcludeTableFlag      = "exclude-table"
 	ExcludeTableDataFlag  = "exclude-table-data"
 	AnalyzeFlag           = "analyze"
+	HaltOnErrorFlag       = "halt-on-error"
 
 	DirectoryFlag = "directory"
 

diff --git a/internal/consts/viper.go b/internal/consts/viper.go
@@ -2,6 +2,7 @@ package consts
 
 const (
 	AnalyzeKey         = "restore.analyze"
+	HaltOnErrorKey     = "restore.halt-on-error"
 	SpinnerKey         = "spinner.name"
 	KubeconfigKey      = "kubernetes.kubeconfig"
 	JobPodLabelsKey    = "kubernetes.job-pod-labels"

diff --git a/internal/database/postgres/postgres.go b/internal/database/postgres/postgres.go
@@ -263,6 +263,9 @@ func (Postgres) RestoreCommand(conf config.Restore, inputFormat sqlformat.Format
 		if conf.Quiet {
 			cmd.Push("--quiet", "--output=/dev/null")
 		}
+		if conf.HaltOnError {
+			cmd.Push("--set=ON_ERROR_STOP=1")
+		}
 	case sqlformat.Custom:
 		cmd.Push("pg_restore", "--format=custom", "--clean", "--exit-on-error")
 		if conf.NoOwner {

diff --git a/internal/database/postgres/postgres_test.go b/internal/database/postgres/postgres_test.go
@@ -260,6 +260,11 @@ func TestPostgres_RestoreCommand(t *testing.T) {
 			args{config.Restore{Global: config.Global{Port: 1234}}, sqlformat.Plain},
 			command.NewBuilder(pgpassword, "psql", "--host=", "--username=", "--dbname=", "--port=1234"),
 		},
+		{
+			"halt_on_error",
+			args{config.Restore{HaltOnError: true}, sqlformat.Plain},
+			command.NewBuilder(pgpassword, "psql", "--set=ON_ERROR_STOP=1", "--host=", "--username=", "--dbname="),
+		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {