Restarting jobs after revert of an upgrade operation (#1502)

Author: jcchr

platform/services/installer/app/commands/upgrade.py

@@ -354,7 +354,7 @@ def execute_upgrade(config: UpgradeConfig) -> None:  # noqa: PLR0915
         perform_k3s_upgrade(config)
 
     reset_custom_signal_handler(handler_state)
-
+    revert = False
     try:
         prepare_upgrade(config, current_platform_version, data_folder, location)
         perform_upgrade(config)
@@ -364,13 +364,15 @@ def execute_upgrade(config: UpgradeConfig) -> None:  # noqa: PLR0915
         click.secho("\n" + UpgradeCmdTexts.upgrade_failed, fg="red")
         click.secho("\nReverting to the previous state...", fg="yellow")
         rollback_platform(config, location, data_folder)
+        revert = True
         click.secho("\n" + UpgradeCmdTexts.revert_succeeded, fg="green")
         sys.exit(1)
     except PrepareUpgradeException:
         logger.error("Platform preparation failed.")
         click.secho("\n" + UpgradeCmdTexts.preparation_failed, fg="red")
         click.secho("\nReverting to the previous state...", fg="yellow")
         rollback_platform(config, location, data_folder)
+        revert = True
         click.secho("\n" + UpgradeCmdTexts.revert_succeeded, fg="green")
         sys.exit(1)
     finally:
@@ -383,7 +385,7 @@ def execute_upgrade(config: UpgradeConfig) -> None:  # noqa: PLR0915
             shutil.rmtree(OFFLINE_TOOLS_DIR)
         try:
             with click_spinner.spinner():
-                restore_platform(config=config)
+                restore_platform(config=config, revert=revert)
         except Exception:
             logger.error("Finalization tasks failed.")
             click.echo(UpgradeCmdTexts.finalization_failed)

platform/services/installer/app/platform_utils/management/management.py

@@ -4,6 +4,7 @@
 """Platform management functions"""
 
 import logging
+from time import sleep
 
 import kubernetes
 from kubernetes.client.exceptions import ApiException
@@ -166,7 +167,7 @@ def _replace_statefulset_with_fetch_retry(
         raise
 
 
-def _patch_platform(config: UpgradeConfig, replicas: int) -> None:  # noqa: C901, PLR0912, PLR0915
+def _patch_platform(config: UpgradeConfig, replicas: int, revert: bool = False) -> None:  # noqa: C901, PLR0912, PLR0915
     """Patch Platform by stopping or starting all Platform's pods."""
     KubernetesConfigHandler(kube_config=config.kube_config.value)
 
@@ -257,6 +258,9 @@ def _patch_platform(config: UpgradeConfig, replicas: int) -> None:  # noqa: C901
         for deployment in deployments:
             _ensure_desired_replicas(core_api=core_api, obj=deployment, replicas=replicas)
 
+        if revert:
+            restart_jobs_after_revert()
+
         for stateful_set in stateful_sets:
             _ensure_desired_replicas(core_api=core_api, obj=stateful_set, replicas=replicas)
 
@@ -285,9 +289,9 @@ def stop_platform(config: UpgradeConfig) -> None:
     _patch_platform(config=config, replicas=0)
 
 
-def restore_platform(config: UpgradeConfig) -> None:
+def restore_platform(config: UpgradeConfig, revert: bool = False) -> None:
     """Restore Platform by starting all Platform's pods."""
-    _patch_platform(config=config, replicas=1)
+    _patch_platform(config=config, replicas=1, revert=revert)
 
 
 def get_ordered_deployments(apps_api: kubernetes.client.AppsV1Api) -> list[kubernetes.client.V1Deployment]:
@@ -479,3 +483,55 @@ def are_pods_present():
             except PodStillPresent as err:
                 logger.exception(err)
                 raise
+
+
+def restart_jobs_after_revert() -> None:
+    """Restarts specific jobs after revert operation."""
+    for job_name in (
+        "impt-etcd-auth",
+        "impt-kafka-provisioning",
+    ):
+        _restart_job(job_name=job_name, namespace=PLATFORM_NAMESPACE)
+
+
+def _restart_job(job_name: str, namespace: str) -> None:
+    """Restart a Kubernetes Job by deleting it and making its copy."""
+    logger.info(f"Restarting job: {job_name} in namespace: {namespace}")
+    with kubernetes.client.ApiClient() as client:
+        batch_api = kubernetes.client.BatchV1Api(client)
+        try:
+            # Get the existing Job manifest
+            job = batch_api.read_namespaced_job(name=job_name, namespace=namespace)
+            # Delete the existing Job
+            batch_api.delete_namespaced_job(
+                name=job_name,
+                namespace=namespace,
+                body=kubernetes.client.V1DeleteOptions(propagation_policy="Foreground"),
+            )
+            # Wait for the Job to be fully deleted
+            for _ in range(10):
+                try:
+                    batch_api.read_namespaced_job(name=job_name, namespace=namespace)
+                except ApiException as e:
+                    if e.status == 404:
+                        break
+                sleep(1)
+            # Remove fields that should not be set on creation
+            job.metadata.creation_timestamp = None
+            job.metadata.resource_version = None
+            job.metadata.uid = None
+            job.status = None
+            if job.metadata.labels.get("controller-uid"):
+                del job.metadata.labels["controller-uid"]
+            if job.metadata.labels.get("batch.kubernetes.io/controller-uid"):
+                del job.metadata.labels["batch.kubernetes.io/controller-uid"]
+            if job.spec.template.metadata.labels.get("controller-uid"):
+                del job.spec.template.metadata.labels["controller-uid"]
+            if job.spec.template.metadata.labels.get("batch.kubernetes.io/controller-uid"):
+                del job.spec.template.metadata.labels["batch.kubernetes.io/controller-uid"]
+            job.spec.selector = None
+            batch_api.create_namespaced_job(namespace=namespace, body=job)
+        except ApiException as e:
+            logger.error(f"Exception when restarting job {job_name}: {e}")
+        except Exception as e:
+            logger.error(f"Unexpected error when restarting job {job_name}: {e}")