STAC-23598: Adding check-and-finalize command to stackgraph

Author: viliakov

cmd/elasticsearch/restore-snapshot.go

@@ -65,7 +65,7 @@ func runRestore(appCtx *app.Context) error {
 	defer func() {
 		if len(scaledDeployments) > 0 {
 			appCtx.Logger.Println()
-			if err := scale.ScaleUp(appCtx.K8sClient, appCtx.Namespace, scaledDeployments, appCtx.Logger); err != nil {
+			if err := scale.ScaleUpFromAnnotations(appCtx.K8sClient, appCtx.Namespace, appCtx.Config.Elasticsearch.Restore.ScaleDownLabelSelector, appCtx.Logger); err != nil {
 				appCtx.Logger.Warningf("Failed to scale up deployments: %v", err)
 			}
 		}

cmd/root.go

@@ -17,7 +17,7 @@ var (
 // addBackupConfigFlags adds configuration flags needed for backup/restore operations
 // to commands that interact with data services (Elasticsearch, etc.)
 func addBackupConfigFlags(cmd *cobra.Command) {
-	cmd.PersistentFlags().StringVar(&flags.Namespace, "namespace", "", "Kubernetes namespace (required)")
+	cmd.PersistentFlags().StringVarP(&flags.Namespace, "namespace", "n", "", "Kubernetes namespace (required)")
 	cmd.PersistentFlags().StringVar(&flags.Kubeconfig, "kubeconfig", "", "Path to kubeconfig file (default: ~/.kube/config)")
 	cmd.PersistentFlags().BoolVar(&flags.Debug, "debug", false, "Enable debug output")
 	cmd.PersistentFlags().BoolVarP(&flags.Quiet, "quiet", "q", false, "Suppress operational messages (only show errors and data output)")

cmd/stackgraph/check_and_finalize.go

@@ -0,0 +1,144 @@
+package stackgraph
+
+import (
+	"fmt"
+	"os"
+
+	"github.com/spf13/cobra"
+	"github.com/stackvista/stackstate-backup-cli/internal/app"
+	"github.com/stackvista/stackstate-backup-cli/internal/foundation/config"
+	"github.com/stackvista/stackstate-backup-cli/internal/orchestration/scale"
+	batchv1 "k8s.io/api/batch/v1"
+)
+
+// Check and finalize command flags
+var (
+	checkJobName string
+	waitForJob   bool
+)
+
+func checkAndFinalizeCmd(globalFlags *config.CLIGlobalFlags) *cobra.Command {
+	cmd := &cobra.Command{
+		Use:   "check-and-finalize",
+		Short: "Check and finalize a Stackgraph restore job",
+		Long: `Check the status of a background Stackgraph restore job and clean up resources.
+
+This command is useful when a restore job was started with --background flag or was interrupted (Ctrl+C).
+It will check the job status, print logs if it failed, and clean up the job and PVC resources.
+
+Examples:
+  # Check job status without waiting
+  sts-backup stackgraph check-and-finalize --job stackgraph-restore-20250128t143000 -n my-namespace
+
+  # Wait for job completion and cleanup
+  sts-backup stackgraph check-and-finalize --job stackgraph-restore-20250128t143000 --wait -n my-namespace`,
+		Run: func(_ *cobra.Command, _ []string) {
+			appCtx, err := app.NewContext(globalFlags)
+			if err != nil {
+				_, _ = fmt.Fprintf(os.Stderr, "error: %v\n", err)
+				os.Exit(1)
+			}
+			if err := runCheckAndFinalize(appCtx); err != nil {
+				_, _ = fmt.Fprintf(os.Stderr, "error: %v\n", err)
+				os.Exit(1)
+			}
+		},
+	}
+
+	cmd.Flags().StringVarP(&checkJobName, "job", "j", "", "Stackgraph restore job name (required)")
+	cmd.Flags().BoolVarP(&waitForJob, "wait", "w", false, "Wait for job to complete before cleanup")
+	_ = cmd.MarkFlagRequired("job")
+
+	return cmd
+}
+
+func runCheckAndFinalize(appCtx *app.Context) error {
+	// Get job
+	appCtx.Logger.Infof("Checking status of job: %s", checkJobName)
+	job, err := appCtx.K8sClient.GetJob(appCtx.Namespace, checkJobName)
+	if err != nil {
+		return fmt.Errorf("failed to get job '%s': %w (job may not exist or has been deleted)", checkJobName, err)
+	}
+
+	// Check if job is already complete
+	completed, succeeded := isJobComplete(job)
+
+	if completed {
+		// Job already finished - print status and cleanup
+		return handleCompletedJob(appCtx, checkJobName, succeeded)
+	}
+
+	// Job still running
+	if waitForJob {
+		// Wait for completion, then cleanup
+		return waitAndFinalize(appCtx, checkJobName)
+	}
+
+	// Not waiting - just print status
+	printRunningJobStatus(appCtx.Logger, checkJobName, appCtx.Namespace, job.Status.Active)
+	return nil
+}
+
+// isJobComplete checks if job is in a terminal state
+func isJobComplete(job *batchv1.Job) (completed bool, succeeded bool) {
+	if job.Status.Succeeded > 0 {
+		return true, true
+	}
+	if job.Status.Failed > 0 {
+		return true, false
+	}
+	return false, false
+}
+
+// handleCompletedJob handles a job that's already complete
+func handleCompletedJob(appCtx *app.Context, jobName string, succeeded bool) error {
+	appCtx.Logger.Println()
+	if succeeded {
+		appCtx.Logger.Successf("Job completed successfully: %s", jobName)
+		appCtx.Logger.Println()
+
+		// Scale up deployments that were scaled down before restore
+		scaleDownLabelSelector := appCtx.Config.Stackgraph.Restore.ScaleDownLabelSelector
+		if err := scale.ScaleUpFromAnnotations(appCtx.K8sClient, appCtx.Namespace, scaleDownLabelSelector, appCtx.Logger); err != nil {
+			appCtx.Logger.Warningf("Failed to scale up deployments: %v", err)
+		}
+	} else {
+		appCtx.Logger.Errorf("Job failed: %s", jobName)
+		appCtx.Logger.Println()
+		appCtx.Logger.Infof("Fetching logs...")
+		appCtx.Logger.Println()
+		if err := printJobLogs(appCtx.K8sClient, appCtx.Namespace, jobName, appCtx.Logger); err != nil {
+			appCtx.Logger.Warningf("Failed to fetch logs: %v", err)
+		}
+	}
+
+	// Cleanup resources
+	appCtx.Logger.Println()
+	return cleanupRestoreResources(appCtx.K8sClient, appCtx.Namespace, jobName, appCtx.Logger)
+}
+
+// waitAndFinalize waits for job completion and then cleans up
+func waitAndFinalize(appCtx *app.Context, jobName string) error {
+	printWaitingMessage(appCtx.Logger, jobName, appCtx.Namespace)
+
+	if err := waitForJobCompletion(appCtx.K8sClient, appCtx.Namespace, jobName, appCtx.Logger); err != nil {
+		appCtx.Logger.Errorf("Job failed: %v", err)
+		// Still cleanup even if failed
+		appCtx.Logger.Println()
+		_ = cleanupRestoreResources(appCtx.K8sClient, appCtx.Namespace, jobName, appCtx.Logger)
+		return err
+	}
+
+	appCtx.Logger.Println()
+	appCtx.Logger.Successf("Job completed successfully: %s", jobName)
+	appCtx.Logger.Println()
+
+	// Scale up deployments that were scaled down before restore
+	scaleDownLabelSelector := appCtx.Config.Stackgraph.Restore.ScaleDownLabelSelector
+	if err := scale.ScaleUpFromAnnotations(appCtx.K8sClient, appCtx.Namespace, scaleDownLabelSelector, appCtx.Logger); err != nil {
+		appCtx.Logger.Warningf("Failed to scale up deployments: %v", err)
+	}
+
+	appCtx.Logger.Println()
+	return cleanupRestoreResources(appCtx.K8sClient, appCtx.Namespace, jobName, appCtx.Logger)
+}

cmd/stackgraph/restore.go

@@ -21,7 +21,6 @@ import (
 	"github.com/stackvista/stackstate-backup-cli/internal/orchestration/portforward"
 	"github.com/stackvista/stackstate-backup-cli/internal/orchestration/scale"
 	"github.com/stackvista/stackstate-backup-cli/internal/scripts"
-	batchv1 "k8s.io/api/batch/v1"
 	corev1 "k8s.io/api/core/v1"
 )
 
@@ -111,7 +110,7 @@ func runRestore(appCtx *app.Context) error {
 	defer func() {
 		if len(scaledDeployments) > 0 && !background {
 			appCtx.Logger.Println()
-			if err := scale.ScaleUp(appCtx.K8sClient, appCtx.Namespace, scaledDeployments, appCtx.Logger); err != nil {
+			if err := scale.ScaleUpFromAnnotations(appCtx.K8sClient, appCtx.Namespace, scaleDownLabelSelector, appCtx.Logger); err != nil {
 				appCtx.Logger.Warningf("Failed to scale up deployments: %v", err)
 			}
 		}
@@ -129,19 +128,18 @@ func runRestore(appCtx *app.Context) error {
 
 	jobName := fmt.Sprintf("%s-%s", jobNameTemplate, time.Now().Format("20060102t150405"))
 
-	job, pvc, err := createRestoreJob(appCtx.K8sClient, appCtx.Namespace, jobName, backupFile, appCtx.Config)
-	if err != nil {
+	if err = createRestoreJob(appCtx.K8sClient, appCtx.Namespace, jobName, backupFile, appCtx.Config); err != nil {
 		return fmt.Errorf("failed to create restore job: %w", err)
 	}
 
 	appCtx.Logger.Successf("Restore job created: %s", jobName)
 
 	if background {
-		printBackgroundModeInstructions(appCtx.Logger, jobName, appCtx.Namespace, scaleDownLabelSelector)
+		printRunningJobStatus(appCtx.Logger, jobName, appCtx.Namespace, 0)
 		return nil
 	}
 
-	return waitAndCleanupRestoreJob(appCtx.K8sClient, appCtx.Namespace, jobName, job, pvc, appCtx.Logger)
+	return waitAndCleanupRestoreJob(appCtx.K8sClient, appCtx.Namespace, jobName, appCtx.Logger)
 }
 
 // ensureRestoreResources ensures that required Kubernetes resources exist for the restore job
@@ -186,22 +184,56 @@ func ensureRestoreResources(k8sClient *k8s.Client, namespace string, config *con
 	return nil
 }
 
-// printBackgroundModeInstructions prints instructions for monitoring and cleanup in background mode
-func printBackgroundModeInstructions(log *logger.Logger, jobName, namespace, scaleDownSelector string) {
+// printWaitingMessage prints waiting message with instructions for interruption
+func printWaitingMessage(log *logger.Logger, jobName, namespace string) {
 	log.Println()
-	log.Infof("Job is running in background. Use the following commands to monitor and cleanup:")
-	log.Infof("  Monitor: kubectl logs --follow job/%s -n %s", jobName, namespace)
-	log.Infof("  Cleanup: kubectl delete job,pvc %s -n %s", jobName, namespace)
-	log.Infof("  Scale up: kubectl scale --replicas=1 deployments --selector=%s -n %s", scaleDownSelector, namespace)
+	log.Infof("Waiting for restore job to complete (this may take several minutes)...")
 	log.Println()
-	log.Warningf("IMPORTANT: The restore job may take significant time to complete.")
-	log.Warningf("Remember to scale up deployments after the job completes successfully!")
+	log.Infof("You can safely interrupt this command with Ctrl+C.")
+	log.Infof("To check status, scale up the required deployments and cleanup later, run:")
+	log.Infof("  sts-backup stackgraph check-and-finalize --job %s --wait -n %s", jobName, namespace)
 }
 
-// waitAndCleanupRestoreJob waits for job completion and cleans up resources
-func waitAndCleanupRestoreJob(k8sClient *k8s.Client, namespace, jobName string, job *batchv1.Job, pvc *corev1.PersistentVolumeClaim, log *logger.Logger) error {
+// printRunningJobStatus prints status and instructions for a running job
+func printRunningJobStatus(log *logger.Logger, jobName, namespace string, activePods int32) {
 	log.Println()
-	log.Infof("Waiting for restore job to complete (this may take several minutes)...")
+	log.Infof("Job is running in background: %s", jobName)
+	if activePods > 0 {
+		log.Infof("  Active pods: %d", activePods)
+	}
+	log.Println()
+	log.Infof("Monitoring commands:")
+	log.Infof("  kubectl logs --follow job/%s -n %s", jobName, namespace)
+	log.Infof("  kubectl get job %s -n %s", jobName, namespace)
+	log.Println()
+	log.Infof("To wait for completion, scaling up the necessary deployments and cleanup, run:")
+	log.Infof("  sts-backup stackgraph check-and-finalize --job %s --wait -n %s", jobName, namespace)
+}
+
+// cleanupRestoreResources cleans up job and PVC resources
+func cleanupRestoreResources(k8sClient *k8s.Client, namespace, jobName string, log *logger.Logger) error {
+	log.Infof("Cleaning up job and PVC...")
+
+	// Delete job
+	if err := k8sClient.DeleteJob(namespace, jobName); err != nil {
+		log.Warningf("Failed to delete job: %v", err)
+	} else {
+		log.Successf("Job deleted: %s", jobName)
+	}
+
+	// Delete PVC (same name as job)
+	if err := k8sClient.DeletePVC(namespace, jobName); err != nil {
+		log.Warningf("Failed to delete PVC: %v", err)
+	} else {
+		log.Successf("PVC deleted: %s", jobName)
+	}
+
+	return nil
+}
+
+// waitAndCleanupRestoreJob waits for job completion and cleans up resources
+func waitAndCleanupRestoreJob(k8sClient *k8s.Client, namespace, jobName string, log *logger.Logger) error {
+	printWaitingMessage(log, jobName, namespace)
 
 	if err := waitForJobCompletion(k8sClient, namespace, jobName, log); err != nil {
 		log.Errorf("Job failed: %v", err)
@@ -214,16 +246,9 @@ func waitAndCleanupRestoreJob(k8sClient *k8s.Client, namespace, jobName string,
 	log.Println()
 	log.Successf("Restore completed successfully")
 
-	// Cleanup job and PVC
-	log.Infof("Cleaning up job and PVC...")
-	if err := k8sClient.DeleteJob(namespace, job.Name); err != nil {
-		log.Warningf("Failed to delete job: %v", err)
-	}
-	if err := k8sClient.DeletePVC(namespace, pvc.Name); err != nil {
-		log.Warningf("Failed to delete PVC: %v", err)
-	}
-
-	return nil
+	// Cleanup job and PVC using shared function
+	log.Println()
+	return cleanupRestoreResources(k8sClient, namespace, jobName, log)
 }
 
 // getLatestBackup retrieves the most recent backup from S3
@@ -305,7 +330,7 @@ func buildPVCSpec(name string, config *config.Config, labels map[string]string)
 }
 
 // createRestoreJob creates a Kubernetes Job and PVC for restoring from backup
-func createRestoreJob(k8sClient *k8s.Client, namespace, jobName, backupFile string, config *config.Config) (*batchv1.Job, *corev1.PersistentVolumeClaim, error) {
+func createRestoreJob(k8sClient *k8s.Client, namespace, jobName, backupFile string, config *config.Config) error {
 	defaultMode := int32(configMapDefaultFileMode)
 
 	// Merge common labels with resource-specific labels
@@ -316,7 +341,7 @@ func createRestoreJob(k8sClient *k8s.Client, namespace, jobName, backupFile stri
 	pvcSpec := buildPVCSpec(jobName, config, pvcLabels)
 	pvc, err := k8sClient.CreatePVC(namespace, pvcSpec)
 	if err != nil {
-		return nil, nil, fmt.Errorf("failed to create PVC: %w", err)
+		return fmt.Errorf("failed to create PVC: %w", err)
 	}
 
 	// Build job spec using configuration
@@ -339,14 +364,14 @@ func createRestoreJob(k8sClient *k8s.Client, namespace, jobName, backupFile stri
 	}
 
 	// Create job
-	job, err := k8sClient.CreateBackupJob(namespace, spec)
+	_, err = k8sClient.CreateBackupJob(namespace, spec)
 	if err != nil {
 		// Cleanup PVC if job creation fails
 		_ = k8sClient.DeletePVC(namespace, pvc.Name)
-		return nil, nil, fmt.Errorf("failed to create job: %w", err)
+		return fmt.Errorf("failed to create job: %w", err)
 	}
 
-	return job, pvc, nil
+	return nil
 }
 
 // buildRestoreEnvVars constructs environment variables for the restore job

cmd/stackgraph/stackgraph.go

@@ -13,6 +13,7 @@ func Cmd(globalFlags *config.CLIGlobalFlags) *cobra.Command {
 
 	cmd.AddCommand(listCmd(globalFlags))
 	cmd.AddCommand(restoreCmd(globalFlags))
+	cmd.AddCommand(checkAndFinalizeCmd(globalFlags))
 
 	return cmd
 }