Add integration test, fetch parameter from current job run

Author: pritishpai

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

@@ -1,4 +1,5 @@
 import logging
+from functools import cached_property
 
 from databricks.sdk.service.jobs import JobParameterDefinition
 
@@ -11,8 +12,21 @@
 
 class Assessment(Workflow):  # pylint: disable=too-many-public-methods
     def __init__(self):
-        self._force_refresh_param = JobParameterDefinition(name="force_refresh", default=False)
-        super().__init__('assessment', [self._force_refresh_param])
+        super().__init__('assessment', [JobParameterDefinition(name="force_refresh", default=False)])
+
+    @staticmethod
+    def _get_force_refresh(ctx: RuntimeContext) -> bool:
+        """Extracts the force_refresh parameter from the job run parameters."""
+        force_refresh = False
+        job_id = ctx.install_state.jobs["assessment"]
+        job_runs = [job_run for job_run in ctx.workspace_client.jobs.list_runs(active_only=True, job_id=job_id)]
+        job_parameters = job_runs[0].job_parameters if job_runs else []
+        for job_parameter in job_parameters:
+            if job_parameter.name == "force_refresh" and job_parameter.value is not None:
+                force_refresh = job_parameter.value.lower() in {"true", "1"}
+                logger.info("Found force_refresh parameter in job run: %s", force_refresh)
+                break
+        return force_refresh
 
     @job_task
     def crawl_tables(self, ctx: RuntimeContext):
@@ -21,8 +35,7 @@ def crawl_tables(self, ctx: RuntimeContext):
         `$inventory_database.tables`. Note that the `inventory_database` is set in the configuration file. The metadata
         stored is then used in the subsequent tasks and workflows to, for example, find all Hive Metastore tables that
         cannot easily be migrated to Unity Catalog."""
-        force_refresh = ctx.named_parameters.get("force_refresh", "False").lower() in {"true", "1"}
-        ctx.tables_crawler.snapshot(force_refresh=force_refresh)
+        ctx.tables_crawler.snapshot(force_refresh=self._get_force_refresh(ctx))
 
     @job_task
     def crawl_udfs(self, ctx: RuntimeContext):

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

@@ -249,13 +249,14 @@ def __init__(self, ws: WorkspaceClient, install_state: InstallState):
         self._ws = ws
         self._install_state = install_state
 
-    def run_workflow(self, step: str, skip_job_wait: bool = False, max_wait: timedelta = timedelta(minutes=20)) -> int:
+    def run_workflow(self, step: str, skip_job_wait: bool = False, max_wait: timedelta = timedelta(minutes=20),
+                     named_parameters: dict[str, str] | None = None ) -> int:
         # this dunder variable is hiding this method from tracebacks, making it cleaner
         # for the user to see the actual error without too much noise.
         __tracebackhide__ = True  # pylint: disable=unused-variable
         job_id = int(self._install_state.jobs[step])
         logger.debug(f"starting {step} job: {self._ws.config.host}#job/{job_id}")
-        job_initial_run = self._ws.jobs.run_now(job_id)
+        job_initial_run = self._ws.jobs.run_now(job_id, job_parameters=named_parameters)
         run_id = job_initial_run.run_id
         if not run_id:
             raise NotFound(f"job run not found for {step}")

tests/integration/assessment/test_workflows.py

@@ -49,3 +49,12 @@ def test_running_real_assessment_job(
     query = f"SELECT * FROM {installation_ctx.inventory_database}.workflow_problems"
     workflow_problems_without_path = [problem for problem in sql_backend.fetch(query) if problem["path"] == "UNKNOWN"]
     assert not workflow_problems_without_path
+
+    post_assessment_table = installation_ctx.make_table(schema_name=source_schema.name, ctas="SELECT 1 AS one")
+    installation_ctx.deployed_workflows.run_workflow(workflow, skip_job_wait=False, named_parameters={"force_refresh": "true"})
+    assert installation_ctx.deployed_workflows.validate_step(workflow), f"Workflow failed: {workflow}"
+    expected_tables = {managed_table.name, external_table.name, tmp_table.name, view.name, non_delta.name, post_assessment_table.name}
+
+    query = f"SELECT * FROM {installation_ctx.inventory_database}.tables"
+    tables_after_assessment_rerun = [table.name for table in sql_backend.fetch(query)]
+    assert tables_after_assessment_rerun == expected_tables