Fixed handling of potentially corrupt state.json of UCX workflows (#2673)

Fix #2667

Author: nfx

src/databricks/labs/ucx/install.py

@@ -29,7 +29,7 @@
 from databricks.labs.lsql.dashboards import DashboardMetadata, Dashboards
 from databricks.labs.lsql.deployment import SchemaDeployer
 from databricks.sdk import WorkspaceClient, AccountClient
-from databricks.sdk.core import with_user_agent_extra
+from databricks.sdk.useragent import with_extra
 from databricks.sdk.errors import (
     AlreadyExists,
     BadRequest,
@@ -87,7 +87,6 @@
 NUM_USER_ATTEMPTS = 10  # number of attempts user gets at answering a question
 
 logger = logging.getLogger(__name__)
-with_user_agent_extra("cmd", "install")
 
 
 def deploy_schema(sql_backend: SqlBackend, inventory_schema: str):
@@ -679,17 +678,7 @@ def _remove_secret_scope(self):
             logger.error("Secret scope already deleted")
 
     def _remove_jobs(self):
-        logger.info("Deleting jobs")
-        if not self._install_state.jobs:
-            logger.error("No jobs present or jobs already deleted")
-            return
-        for step_name, job_id in self._install_state.jobs.items():
-            try:
-                logger.info(f"Deleting {step_name} job_id={job_id}.")
-                self._ws.jobs.delete(job_id)
-            except InvalidParameterValue:
-                logger.error(f"Already deleted: {step_name} job_id={job_id}.")
-                continue
+        self._workflows_installer.remove_jobs()
 
     def _remove_warehouse(self):
         try:
@@ -894,6 +883,7 @@ def _get_collection_workspace(
 
 
 if __name__ == "__main__":
+    with_extra("cmd", "install")
     logger = get_logger(__file__)
     if is_in_debug():
         logging.getLogger('databricks').setLevel(logging.DEBUG)

src/databricks/labs/ucx/installer/workflows.py

@@ -478,19 +478,46 @@ def create_jobs(self) -> None:
                 )
             self._deploy_workflow(workflow_name, settings)
 
-        for workflow_name, job_id in self._install_state.jobs.items():
-            if workflow_name not in desired_workflows:
-                try:
-                    logger.info(f"Removing job_id={job_id}, as it is no longer needed")
-                    self._ws.jobs.delete(job_id)
-                except InvalidParameterValue:
-                    logger.warning(f"step={workflow_name} does not exist anymore for some reason")
-                    continue
-
+        self.remove_jobs(keep=desired_workflows)
         self._install_state.save()
         self._create_debug(remote_wheels)
         self._create_readme()
 
+    def remove_jobs(self, *, keep: set[str] | None = None) -> None:
+        for workflow_name, job_id in self._install_state.jobs.items():
+            if keep and workflow_name in keep:
+                continue
+            try:
+                if not self._is_managed_job_failsafe(int(job_id)):
+                    logger.warning(f"Corrupt installation state. Skipping job_id={job_id} as it is not managed by UCX")
+                    continue
+                logger.info(f"Removing job_id={job_id}, as it is no longer needed")
+                self._ws.jobs.delete(job_id)
+            except InvalidParameterValue:
+                logger.warning(f"step={workflow_name} does not exist anymore for some reason")
+                continue
+
+    # see https://github.com/databrickslabs/ucx/issues/2667
+    def _is_managed_job_failsafe(self, job_id: int) -> bool:
+        install_folder = self._installation.install_folder()
+        try:
+            return self._is_managed_job(job_id, install_folder)
+        except ResourceDoesNotExist:
+            return False
+        except InvalidParameterValue:
+            return False
+
+    def _is_managed_job(self, job_id: int, install_folder: str) -> bool:
+        job = self._ws.jobs.get(job_id)
+        if not job.settings or not job.settings.tasks:
+            return False
+        for task in job.settings.tasks:
+            if task.notebook_task and task.notebook_task.notebook_path.startswith(install_folder):
+                return True
+            if task.python_wheel_task and task.python_wheel_task.package_name == "databricks_labs_ucx":
+                return True
+        return False
+
     @property
     def _config_file(self):
         return f"{self._installation.install_folder()}/config.yml"

src/databricks/labs/ucx/source_code/known.json

@@ -1824,86 +1824,7 @@
     "databricks.labs.pylint": []
   },
   "databricks-labs-pytester": {
-    "databricks": [],
-    "databricks.labs": [],
     "databricks.labs.pytester": [],
-    "databricks.labs.pytester.fixtures": [],
-    "databricks.labs.pytester.fixtures.baseline": [
-      {
-        "code": "direct-filesystem-access",
-        "message": "The use of direct filesystem references is deprecated: /"
-      }
-    ],
-    "databricks.labs.pytester.fixtures.catalog": [
-      {
-        "code": "direct-filesystem-access",
-        "message": "The use of direct filesystem references is deprecated: /"
-      },
-      {
-        "code": "direct-filesystem-access",
-        "message": "The use of direct filesystem references is deprecated: dbfs:/tmp/ucx_test_"
-      },
-      {
-        "code": "direct-filesystem-access",
-        "message": "The use of direct filesystem references is deprecated: dbfs:/user/hive/warehouse/"
-      },
-      {
-        "code": "direct-filesystem-access",
-        "message": "The use of direct filesystem references is deprecated: s3a://databricks-datasets-oregon/delta-sharing/share/open-datasets.share"
-      }
-    ],
-    "databricks.labs.pytester.fixtures.compute": [],
-    "databricks.labs.pytester.fixtures.connect": [],
-    "databricks.labs.pytester.fixtures.environment": [],
-    "databricks.labs.pytester.fixtures.iam": [
-      {
-        "code": "direct-filesystem-access",
-        "message": "The use of direct filesystem references is deprecated: /members"
-      },
-      {
-        "code": "direct-filesystem-access",
-        "message": "The use of direct filesystem references is deprecated: /users/groups/"
-      }
-    ],
-    "databricks.labs.pytester.fixtures.ml": [
-      {
-        "code": "direct-filesystem-access",
-        "message": "The use of direct filesystem references is deprecated: /"
-      },
-      {
-        "code": "direct-filesystem-access",
-        "message": "The use of direct filesystem references is deprecated: /api/2.0/feature-store/feature-tables/create"
-      },
-      {
-        "code": "direct-filesystem-access",
-        "message": "The use of direct filesystem references is deprecated: /api/2.0/feature-store/feature-tables/delete"
-      }
-    ],
-    "databricks.labs.pytester.fixtures.notebooks": [
-      {
-        "code": "direct-filesystem-access",
-        "message": "The use of direct filesystem references is deprecated: /Repos/"
-      },
-      {
-        "code": "direct-filesystem-access",
-        "message": "The use of direct filesystem references is deprecated: /Users/"
-      },
-      {
-        "code": "direct-filesystem-access",
-        "message": "The use of direct filesystem references is deprecated: /dummy-"
-      },
-      {
-        "code": "direct-filesystem-access",
-        "message": "The use of direct filesystem references is deprecated: /sdk-"
-      }
-    ],
-    "databricks.labs.pytester.fixtures.permissions": [],
-    "databricks.labs.pytester.fixtures.plugin": [],
-    "databricks.labs.pytester.fixtures.redash": [],
-    "databricks.labs.pytester.fixtures.secrets": [],
-    "databricks.labs.pytester.fixtures.sql": [],
-    "databricks.labs.pytester.fixtures.unwrap": [],
-    "databricks.labs.pytester.fixtures.watchdog": [],
     "tests.resources.hatchling-whl.src.hatchling_whl": []
   },
   "databricks-labs-ucx": {
@@ -30000,4 +29921,4 @@
     "zipp.compat.py310": [],
     "zipp.glob": []
   }
-}
+}

src/databricks/labs/ucx/uninstall.py

@@ -1,13 +1,15 @@
 import logging
 
 from databricks.sdk import WorkspaceClient
+from databricks.sdk.useragent import with_extra
 
 from databricks.labs.ucx.__about__ import __version__
 from databricks.labs.ucx.install import WorkspaceInstallation
 
 logger = logging.getLogger("databricks.labs.ucx.install")
 
 if __name__ == "__main__":
+    with_extra("cmd", "uninstall")
     logger.setLevel("INFO")
     ws = WorkspaceClient(product="ucx", product_version=__version__)
     installer = WorkspaceInstallation.current(ws)

tests/unit/install/test_install.py

@@ -34,6 +34,7 @@
 import databricks.labs.ucx.installer.mixins
 import databricks.labs.ucx.uninstall  # noqa
 from databricks.labs.ucx.config import WorkspaceConfig
+from databricks.labs.ucx.framework.tasks import Task
 from databricks.labs.ucx.install import AccountInstaller, WorkspaceInstallation, WorkspaceInstaller, extract_major_minor
 from databricks.labs.ucx.installer.workflows import DeployedWorkflows, WorkflowsDeployment
 from databricks.labs.ucx.runtime import Workflows
@@ -86,7 +87,7 @@ def mock_installation_extra_jobs():
         {
             'state.json': {
                 'resources': {
-                    'jobs': {"assessment": "123", "extra_job": "123"},
+                    'jobs': {"assessment": "123", "extra_job": "124", "other_job": "125"},
                     'dashboards': {'assessment_main': 'abc', 'assessment_estimates': 'def'},
                 }
             }
@@ -615,9 +616,10 @@ def test_remove_jobs_with_state_missing_job(ws, caplog, mock_installation_with_j
         config, mock_installation_with_jobs, install_state, sql_backend, ws, workflows_installer, prompts, PRODUCT_INFO
     )
 
-    with caplog.at_level('ERROR'):
+    with caplog.at_level('WARNING'):
         workspace_installation.uninstall()
-        assert 'Already deleted: assessment job_id=123.' in caplog.messages
+        failure = 'Corrupt installation state. Skipping job_id=123 as it is not managed by UCX'
+        assert failure in caplog.messages
 
     mock_installation_with_jobs.assert_removed()
     wheels.upload_to_wsfs.assert_not_called()
@@ -1282,14 +1284,18 @@ def test_remove_jobs(ws, caplog, mock_installation_extra_jobs, any_prompt):
     sql_backend = MockBackend()
     install_state = InstallState.from_installation(mock_installation_extra_jobs)
     wheels = create_autospec(WheelsV2)
+
+    def dummy_task(*_):
+        pass
+
     workflows_installation = WorkflowsDeployment(
         WorkspaceConfig(inventory_database="...", policy_id='123'),
         mock_installation_extra_jobs,
         install_state,
         ws,
         wheels,
         PRODUCT_INFO,
-        [],
+        [Task('assessment', 'some', '...', dummy_task)],
     )
 
     workspace_installation = WorkspaceInstallation(
@@ -1303,9 +1309,28 @@ def test_remove_jobs(ws, caplog, mock_installation_extra_jobs, any_prompt):
         PRODUCT_INFO,
     )
 
-    workspace_installation.run()
-    ws.jobs.delete.assert_called_with("123")
+    def job_side_effect(job_id):
+        tasks = {
+            123: [jobs.Task('x', notebook_task=jobs.NotebookTask(notebook_path='~/mock/assessment'))],
+            124: [jobs.Task('y', python_wheel_task=jobs.PythonWheelTask('databricks_labs_ucx', 'runtime'))],
+            125: [jobs.Task('z', notebook_task=jobs.NotebookTask(notebook_path='outside-of-ucx'))],
+        }
+        return jobs.Job(
+            settings=jobs.JobSettings(
+                tasks=tasks[job_id],
+            ),
+        )
+
+    ws.jobs.get.side_effect = job_side_effect
+
+    with caplog.at_level('WARNING'):
+        workspace_installation.run()
+
+    job_deletes = {_.args[0] for _ in ws.jobs.delete.mock_calls}
+    assert len(job_deletes) == 1
+    assert '124' in job_deletes
     wheels.upload_to_wsfs.assert_called()
+    assert 'Corrupt installation state. Skipping job_id=125 as it is not managed by UCX' in caplog.messages
 
 
 def test_remove_jobs_already_deleted(ws, caplog, mock_installation_extra_jobs, any_prompt):