Add job timeout CLI arg

DavideRutigliano · DavideRutigliano · commit 6080bc6f6405 · 2026-03-02T15:48:52.000+01:00
diff --git a/cmd/clickhouse/check_and_finalize.go b/cmd/clickhouse/check_and_finalize.go
@@ -117,7 +117,7 @@ func waitAndFinalize(appCtx *app.Context, chClient clickhouse.Interface, operati
 		}
 	}
 
-	if err := restore.WaitForAPIRestore(checkStatusFn, defaultPollInterval, defaultRestoreTimeout, appCtx.Logger); err != nil {
+	if err := restore.WaitForAPIRestore(checkStatusFn, defaultPollInterval, appCtx.JobTimeout, appCtx.Logger); err != nil {
 		return err
 	}
 
diff --git a/cmd/root.go b/cmd/root.go
@@ -27,6 +27,7 @@ func addBackupConfigFlags(cmd *cobra.Command) {
 	cmd.PersistentFlags().StringVar(&flags.ConfigMapName, "configmap", "suse-observability-backup-config", "ConfigMap name containing backup configuration")
 	cmd.PersistentFlags().StringVar(&flags.SecretName, "secret", "suse-observability-backup-config", "Secret name containing backup configuration")
 	cmd.PersistentFlags().StringVarP(&flags.OutputFormat, "output", "o", "table", "Output format (table, json)")
+	cmd.PersistentFlags().IntVar(&flags.JobTimeout, "job-timeout", 30, "Timeout in minutes for waiting on job/restore completion")
 	_ = cmd.MarkPersistentFlagRequired("namespace")
 }
 
diff --git a/cmd/settings/check_and_finalize.go b/cmd/settings/check_and_finalize.go
@@ -54,5 +54,6 @@ func runCheckAndFinalize(appCtx *app.Context) error {
 		CleanupPVC:    false,
 		WaitForJob:    waitForJob,
 		Log:           appCtx.Logger,
+		Timeout:       appCtx.JobTimeout,
 	})
 }
diff --git a/cmd/settings/list.go b/cmd/settings/list.go
@@ -185,7 +185,7 @@ func getBackupListFromPVC(appCtx *app.Context) ([]BackupFileInfo, error) {
 		}
 	}()
 
-	if err := restore.WaitForJobCompletion(appCtx.K8sClient, appCtx.Namespace, jobName, appCtx.Logger); err != nil {
+	if err := restore.WaitForJobCompletion(appCtx.K8sClient, appCtx.Namespace, jobName, appCtx.Logger, appCtx.JobTimeout); err != nil {
 		return nil, err
 	}
 
diff --git a/cmd/settings/restore.go b/cmd/settings/restore.go
@@ -126,13 +126,13 @@ func runRestore(appCtx *app.Context) error {
 		return nil
 	}
 
-	return waitAndCleanupRestoreJob(appCtx.K8sClient, appCtx.Namespace, jobName, appCtx.Logger)
+	return waitAndCleanupRestoreJob(appCtx.K8sClient, appCtx.Namespace, jobName, appCtx.Logger, appCtx.JobTimeout)
 }
 
 // waitAndCleanupRestoreJob waits for job completion and cleans up resources
-func waitAndCleanupRestoreJob(k8sClient *k8s.Client, namespace, jobName string, log *logger.Logger) error {
+func waitAndCleanupRestoreJob(k8sClient *k8s.Client, namespace, jobName string, log *logger.Logger, jobTimeout time.Duration) error {
 	restore.PrintWaitingMessage(log, "settings", jobName, namespace)
-	return restore.WaitAndCleanup(k8sClient, namespace, jobName, log, false)
+	return restore.WaitAndCleanup(k8sClient, namespace, jobName, log, false, jobTimeout)
 }
 
 // getLatestBackup retrieves the most recent backup from all sources (S3 and PVC)
diff --git a/cmd/stackgraph/check_and_finalize.go b/cmd/stackgraph/check_and_finalize.go
@@ -54,5 +54,6 @@ func runCheckAndFinalize(appCtx *app.Context) error {
 		CleanupPVC:    true,
 		WaitForJob:    waitForJob,
 		Log:           appCtx.Logger,
+		Timeout:       appCtx.JobTimeout,
 	})
 }
diff --git a/cmd/stackgraph/restore.go b/cmd/stackgraph/restore.go
@@ -133,13 +133,13 @@ func runRestore(appCtx *app.Context) error {
 		return nil
 	}
 
-	return waitAndCleanupRestoreJob(appCtx.K8sClient, appCtx.Namespace, jobName, appCtx.Logger)
+	return waitAndCleanupRestoreJob(appCtx.K8sClient, appCtx.Namespace, jobName, appCtx.Logger, appCtx.JobTimeout)
 }
 
 // waitAndCleanupRestoreJob waits for job completion and cleans up resources
-func waitAndCleanupRestoreJob(k8sClient *k8s.Client, namespace, jobName string, log *logger.Logger) error {
+func waitAndCleanupRestoreJob(k8sClient *k8s.Client, namespace, jobName string, log *logger.Logger, jobTimeout time.Duration) error {
 	restore.PrintWaitingMessage(log, "stackgraph", jobName, namespace)
-	return restore.WaitAndCleanup(k8sClient, namespace, jobName, log, true)
+	return restore.WaitAndCleanup(k8sClient, namespace, jobName, log, true, jobTimeout)
 }
 
 // getLatestBackup retrieves the most recent backup from S3
diff --git a/cmd/victoriametrics/check_and_finalize.go b/cmd/victoriametrics/check_and_finalize.go
@@ -54,5 +54,6 @@ func runCheckAndFinalize(appCtx *app.Context) error {
 		CleanupPVC:    false,
 		WaitForJob:    waitForJob,
 		Log:           appCtx.Logger,
+		Timeout:       appCtx.JobTimeout,
 	})
 }
diff --git a/cmd/victoriametrics/restore.go b/cmd/victoriametrics/restore.go
@@ -130,13 +130,13 @@ func runRestore(appCtx *app.Context) error {
 		return nil
 	}
 
-	return waitAndCleanupRestoreJob(appCtx.K8sClient, appCtx.Namespace, jobName, appCtx.Logger)
+	return waitAndCleanupRestoreJob(appCtx.K8sClient, appCtx.Namespace, jobName, appCtx.Logger, appCtx.JobTimeout)
 }
 
 // waitAndCleanupRestoreJob waits for job completion and cleans up resources
-func waitAndCleanupRestoreJob(k8sClient *k8s.Client, namespace, jobName string, log *logger.Logger) error {
+func waitAndCleanupRestoreJob(k8sClient *k8s.Client, namespace, jobName string, log *logger.Logger, jobTimeout time.Duration) error {
 	restore.PrintWaitingMessage(log, "victoria-metrics", jobName, namespace)
-	return restore.WaitAndCleanup(k8sClient, namespace, jobName, log, false)
+	return restore.WaitAndCleanup(k8sClient, namespace, jobName, log, false, jobTimeout)
 }
 
 // getLatestBackup retrieves the most recent backup from S3
diff --git a/internal/app/app.go b/internal/app/app.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"os"
+	"time"
 
 	"github.com/stackvista/stackstate-backup-cli/internal/clients/clickhouse"
 	"github.com/stackvista/stackstate-backup-cli/internal/clients/elasticsearch"
@@ -16,15 +17,16 @@ import (
 
 // Context holds all dependencies for cli commands
 type Context struct {
-	K8sClient *k8s.Client
-	Namespace string
-	S3Client  s3.Interface
-	ESClient  elasticsearch.Interface
-	CHClient  clickhouse.Interface
-	Config    *config.Config
-	Logger    *logger.Logger
-	Formatter *output.Formatter
-	Context   context.Context
+	K8sClient  *k8s.Client
+	Namespace  string
+	S3Client   s3.Interface
+	ESClient   elasticsearch.Interface
+	CHClient   clickhouse.Interface
+	Config     *config.Config
+	Logger     *logger.Logger
+	Formatter  *output.Formatter
+	Context    context.Context
+	JobTimeout time.Duration
 }
 
 // NewContext creates production dependencies
@@ -68,14 +70,15 @@ func NewContext(flags *config.CLIGlobalFlags) (*Context, error) {
 	formatter := output.NewFormatter(os.Stdout, flags.OutputFormat)
 
 	return &Context{
-		K8sClient: k8sClient,
-		Namespace: flags.Namespace,
-		Config:    cfg,
-		S3Client:  s3Client,
-		ESClient:  esClient,
-		CHClient:  chClient,
-		Logger:    logger.New(flags.Quiet, flags.Debug),
-		Formatter: formatter,
-		Context:   context.Background(),
+		K8sClient:  k8sClient,
+		Namespace:  flags.Namespace,
+		Config:     cfg,
+		S3Client:   s3Client,
+		ESClient:   esClient,
+		CHClient:   chClient,
+		Logger:     logger.New(flags.Quiet, flags.Debug),
+		Formatter:  formatter,
+		Context:    context.Background(),
+		JobTimeout: time.Duration(flags.JobTimeout) * time.Minute,
 	}, nil
 }
diff --git a/internal/foundation/config/config.go b/internal/foundation/config/config.go
@@ -349,6 +349,7 @@ type CLIGlobalFlags struct {
 	ConfigMapName string
 	SecretName    string
 	OutputFormat  string // table, json
+	JobTimeout    int    // job completion timeout in minutes
 }
 
 func NewCLIGlobalFlags() *CLIGlobalFlags {
diff --git a/internal/orchestration/restore/finalize.go b/internal/orchestration/restore/finalize.go
@@ -2,6 +2,7 @@ package restore
 
 import (
 	"fmt"
+	"time"
 
 	"github.com/stackvista/stackstate-backup-cli/internal/clients/k8s"
 	"github.com/stackvista/stackstate-backup-cli/internal/foundation/logger"
@@ -74,13 +75,14 @@ type WaitAndFinalizeParams struct {
 	ScaleSelector string
 	CleanupPVC    bool
 	Log           *logger.Logger
+	Timeout       time.Duration
 }
 
 // WaitAndFinalize waits for job completion and then cleans up
 func WaitAndFinalize(params WaitAndFinalizeParams) error {
 	PrintWaitingMessage(params.Log, params.ServiceName, params.JobName, params.Namespace)
 
-	if err := WaitForJobCompletion(params.K8sClient, params.Namespace, params.JobName, params.Log); err != nil {
+	if err := WaitForJobCompletion(params.K8sClient, params.Namespace, params.JobName, params.Log, params.Timeout); err != nil {
 		params.Log.Errorf("Job failed: %v", err)
 		// Still cleanup even if failed
 		params.Log.Println()
@@ -113,6 +115,7 @@ type CheckAndFinalizeParams struct {
 	CleanupPVC    bool
 	WaitForJob    bool
 	Log           *logger.Logger
+	Timeout       time.Duration
 }
 
 // CheckAndFinalize checks the status of a background restore job and cleans up resources
@@ -157,6 +160,7 @@ func CheckAndFinalize(params CheckAndFinalizeParams) error {
 			ScaleSelector: params.ScaleSelector,
 			CleanupPVC:    params.CleanupPVC,
 			Log:           params.Log,
+			Timeout:       params.Timeout,
 		})
 	}
 
diff --git a/internal/orchestration/restore/job.go b/internal/orchestration/restore/job.go
@@ -14,8 +14,45 @@ const (
 )
 
 // WaitForJobCompletion waits for a Kubernetes job to complete
-func WaitForJobCompletion(k8sClient *k8s.Client, namespace, jobName string, log *logger.Logger) error {
-	timeout := time.After(defaultJobCompletionTimeout)
+func WaitForJobCompletion(k8sClient *k8s.Client, namespace, jobName string, log *logger.Logger, jobTimeout time.Duration) error {
+	if jobTimeout <= 0 {
+		jobTimeout = defaultJobCompletionTimeout
+	}
+
+	// Initial check before starting the ticker
+	checkJob := func() (bool, error) {
+		job, err := k8sClient.GetJob(namespace, jobName)
+		if err != nil {
+			return false, fmt.Errorf("failed to get job status: %w", err)
+		}
+
+		if job.Status.Succeeded > 0 {
+			return true, nil
+		}
+
+		if job.Status.Failed > 0 {
+			// Get and print logs from failed job
+			log.Println()
+			log.Errorf("Job failed. Fetching logs...")
+			log.Println()
+			if err := PrintJobLogs(k8sClient, namespace, jobName, log); err != nil {
+				log.Warningf("Failed to fetch job logs: %v", err)
+			}
+			return false, fmt.Errorf("job failed")
+		}
+
+		log.Debugf("Job status: Active=%d, Succeeded=%d, Failed=%d",
+			job.Status.Active, job.Status.Succeeded, job.Status.Failed)
+		return false, nil
+	}
+
+	// Perform initial check
+	done, err := checkJob()
+	if err != nil || done {
+		return err
+	}
+
+	timeout := time.After(jobTimeout)
 	ticker := time.NewTicker(defaultJobStatusCheckInterval)
 	defer ticker.Stop()
 
@@ -24,28 +61,10 @@ func WaitForJobCompletion(k8sClient *k8s.Client, namespace, jobName string, log
 		case <-timeout:
 			return fmt.Errorf("timeout waiting for job to complete")
 		case <-ticker.C:
-			job, err := k8sClient.GetJob(namespace, jobName)
-			if err != nil {
-				return fmt.Errorf("failed to get job status: %w", err)
+			done, err := checkJob()
+			if err != nil || done {
+				return err
 			}
-
-			if job.Status.Succeeded > 0 {
-				return nil
-			}
-
-			if job.Status.Failed > 0 {
-				// Get and print logs from failed job
-				log.Println()
-				log.Errorf("Job failed. Fetching logs...")
-				log.Println()
-				if err := PrintJobLogs(k8sClient, namespace, jobName, log); err != nil {
-					log.Warningf("Failed to fetch job logs: %v", err)
-				}
-				return fmt.Errorf("job failed")
-			}
-
-			log.Debugf("Job status: Active=%d, Succeeded=%d, Failed=%d",
-				job.Status.Active, job.Status.Succeeded, job.Status.Failed)
 		}
 	}
 }
@@ -112,8 +131,8 @@ func PrintRunningJobStatus(log *logger.Logger, serviceName, jobName, namespace s
 }
 
 // WaitAndCleanup waits for job completion and cleans up resources
-func WaitAndCleanup(k8sClient *k8s.Client, namespace, jobName string, log *logger.Logger, cleanupPVC bool) error {
-	if err := WaitForJobCompletion(k8sClient, namespace, jobName, log); err != nil {
+func WaitAndCleanup(k8sClient *k8s.Client, namespace, jobName string, log *logger.Logger, cleanupPVC bool, jobTimeout time.Duration) error {
+	if err := WaitForJobCompletion(k8sClient, namespace, jobName, log, jobTimeout); err != nil {
 		log.Errorf("Job failed: %v", err)
 		log.Println()
 		log.Infof("Cleanup commands:")
diff --git a/internal/orchestration/restore/job_test.go b/internal/orchestration/restore/job_test.go

Original file line number	Diff line number	Diff line change
`@@ -117,7 +117,7 @@ func waitAndFinalize(appCtx *app.Context, chClient clickhouse.Interface, operati`
`117`	`117`	`}`
`118`	`118`	`}`
`119`	`119`
`120`		`- if err := restore.WaitForAPIRestore(checkStatusFn, defaultPollInterval, defaultRestoreTimeout, appCtx.Logger); err != nil {`
	`120`	`+ if err := restore.WaitForAPIRestore(checkStatusFn, defaultPollInterval, appCtx.JobTimeout, appCtx.Logger); err != nil {`
`121`	`121`	`return err`
`122`	`122`	`}`
`123`	`123`
Original file line number	Diff line number	Diff line change
`@@ -27,6 +27,7 @@ func addBackupConfigFlags(cmd *cobra.Command) {`
`27`	`27`	`cmd.PersistentFlags().StringVar(&flags.ConfigMapName, "configmap", "suse-observability-backup-config", "ConfigMap name containing backup configuration")`
`28`	`28`	`cmd.PersistentFlags().StringVar(&flags.SecretName, "secret", "suse-observability-backup-config", "Secret name containing backup configuration")`
`29`	`29`	`cmd.PersistentFlags().StringVarP(&flags.OutputFormat, "output", "o", "table", "Output format (table, json)")`
	`30`	`+ cmd.PersistentFlags().IntVar(&flags.JobTimeout, "job-timeout", 30, "Timeout in minutes for waiting on job/restore completion")`
`30`	`31`	`_ = cmd.MarkPersistentFlagRequired("namespace")`
`31`	`32`	`}`
`32`	`33`
Original file line number	Diff line number	Diff line change
`@@ -54,5 +54,6 @@ func runCheckAndFinalize(appCtx *app.Context) error {`
`54`	`54`	`CleanupPVC: false,`
`55`	`55`	`WaitForJob: waitForJob,`
`56`	`56`	`Log: appCtx.Logger,`
	`57`	`+ Timeout: appCtx.JobTimeout,`
`57`	`58`	`})`
`58`	`59`	`}`
Original file line number	Diff line number	Diff line change
`@@ -185,7 +185,7 @@ func getBackupListFromPVC(appCtx *app.Context) ([]BackupFileInfo, error) {`
`185`	`185`	`}`
`186`	`186`	`}()`
`187`	`187`
`188`		`- if err := restore.WaitForJobCompletion(appCtx.K8sClient, appCtx.Namespace, jobName, appCtx.Logger); err != nil {`
	`188`	`+ if err := restore.WaitForJobCompletion(appCtx.K8sClient, appCtx.Namespace, jobName, appCtx.Logger, appCtx.JobTimeout); err != nil {`
`189`	`189`	`return nil, err`
`190`	`190`	`}`
`191`	`191`
Original file line number	Diff line number	Diff line change
`@@ -126,13 +126,13 @@ func runRestore(appCtx *app.Context) error {`
`126`	`126`	`return nil`
`127`	`127`	`}`
`128`	`128`
`129`		`- return waitAndCleanupRestoreJob(appCtx.K8sClient, appCtx.Namespace, jobName, appCtx.Logger)`
	`129`	`+ return waitAndCleanupRestoreJob(appCtx.K8sClient, appCtx.Namespace, jobName, appCtx.Logger, appCtx.JobTimeout)`
`130`	`130`	`}`
`131`	`131`
`132`	`132`	`// waitAndCleanupRestoreJob waits for job completion and cleans up resources`
`133`		`-func waitAndCleanupRestoreJob(k8sClient k8s.Client, namespace, jobName string, log logger.Logger) error {`
	`133`	`+func waitAndCleanupRestoreJob(k8sClient k8s.Client, namespace, jobName string, log logger.Logger, jobTimeout time.Duration) error {`
`134`	`134`	`restore.PrintWaitingMessage(log, "settings", jobName, namespace)`
`135`		`- return restore.WaitAndCleanup(k8sClient, namespace, jobName, log, false)`
	`135`	`+ return restore.WaitAndCleanup(k8sClient, namespace, jobName, log, false, jobTimeout)`
`136`	`136`	`}`
`137`	`137`
`138`	`138`	`// getLatestBackup retrieves the most recent backup from all sources (S3 and PVC)`
Original file line number	Diff line number	Diff line change
`@@ -133,13 +133,13 @@ func runRestore(appCtx *app.Context) error {`
`133`	`133`	`return nil`
`134`	`134`	`}`
`135`	`135`
`136`		`- return waitAndCleanupRestoreJob(appCtx.K8sClient, appCtx.Namespace, jobName, appCtx.Logger)`
	`136`	`+ return waitAndCleanupRestoreJob(appCtx.K8sClient, appCtx.Namespace, jobName, appCtx.Logger, appCtx.JobTimeout)`
`137`	`137`	`}`
`138`	`138`
`139`	`139`	`// waitAndCleanupRestoreJob waits for job completion and cleans up resources`
`140`		`-func waitAndCleanupRestoreJob(k8sClient k8s.Client, namespace, jobName string, log logger.Logger) error {`
	`140`	`+func waitAndCleanupRestoreJob(k8sClient k8s.Client, namespace, jobName string, log logger.Logger, jobTimeout time.Duration) error {`
`141`	`141`	`restore.PrintWaitingMessage(log, "stackgraph", jobName, namespace)`
`142`		`- return restore.WaitAndCleanup(k8sClient, namespace, jobName, log, true)`
	`142`	`+ return restore.WaitAndCleanup(k8sClient, namespace, jobName, log, true, jobTimeout)`
`143`	`143`	`}`
`144`	`144`
`145`	`145`	`// getLatestBackup retrieves the most recent backup from S3`
Original file line number	Diff line number	Diff line change
`@@ -130,13 +130,13 @@ func runRestore(appCtx *app.Context) error {`
`130`	`130`	`return nil`
`131`	`131`	`}`
`132`	`132`
`133`		`- return waitAndCleanupRestoreJob(appCtx.K8sClient, appCtx.Namespace, jobName, appCtx.Logger)`
	`133`	`+ return waitAndCleanupRestoreJob(appCtx.K8sClient, appCtx.Namespace, jobName, appCtx.Logger, appCtx.JobTimeout)`
`134`	`134`	`}`
`135`	`135`
`136`	`136`	`// waitAndCleanupRestoreJob waits for job completion and cleans up resources`
`137`		`-func waitAndCleanupRestoreJob(k8sClient k8s.Client, namespace, jobName string, log logger.Logger) error {`
	`137`	`+func waitAndCleanupRestoreJob(k8sClient k8s.Client, namespace, jobName string, log logger.Logger, jobTimeout time.Duration) error {`
`138`	`138`	`restore.PrintWaitingMessage(log, "victoria-metrics", jobName, namespace)`
`139`		`- return restore.WaitAndCleanup(k8sClient, namespace, jobName, log, false)`
	`139`	`+ return restore.WaitAndCleanup(k8sClient, namespace, jobName, log, false, jobTimeout)`
`140`	`140`	`}`
`141`	`141`
`142`	`142`	`// getLatestBackup retrieves the most recent backup from S3`