diff --git a/src/databricks/labs/ucx/install.py b/src/databricks/labs/ucx/install.py index 069e675c52..e2346b5c5c 100644 --- a/src/databricks/labs/ucx/install.py +++ b/src/databricks/labs/ucx/install.py @@ -113,6 +113,10 @@ def _run_configured(self): def run_for_config( ws: WorkspaceClient, config: WorkspaceConfig, *, prefix="ucx", override_clusters: dict[str, str] | None = None ) -> "WorkspaceInstaller": + def run_for_config( + ws: WorkspaceClient, config: WorkspaceConfig, *, prefix="ucx", override_clusters: dict[str, str] | None = None + ) -> "WorkspaceInstaller": + logger.info(f"Installing UCX v{__version__} on {ws.config.host}") workspace_installer = WorkspaceInstaller(ws, prefix=prefix, promtps=False) logger.info(f"Installing UCX v{workspace_installer._version} on {ws.config.host}") workspace_installer._config = config @@ -542,6 +546,17 @@ def _job_settings(self, step_name: str, dbfs_path: str): "tasks": [self._job_task(task, dbfs_path) for task in tasks], } + @staticmethod + def _apply_cluster_overrides(settings: dict[str, any], overrides: dict[str, str]) -> dict: + settings["job_clusters"] = [_ for _ in settings["job_clusters"] if _.job_cluster_key not in overrides] + for job_task in settings["tasks"]: + if job_task.job_cluster_key is None: + continue + if job_task.job_cluster_key in overrides: + job_task.existing_cluster_id = overrides[job_task.job_cluster_key] + job_task.job_cluster_key = None + return settings + def _upload_wheel_runner(self, remote_wheel: str): # TODO: we have to be doing this workaround until ES-897453 is solved in the platform path = f"{self._install_folder}/wheels/wheel-test-runner-{self._version}.py" diff --git a/tests/integration/test_installation.py b/tests/integration/test_installation.py index 9cef4ae49f..82e92205a4 100644 --- a/tests/integration/test_installation.py +++ b/tests/integration/test_installation.py @@ -77,7 +77,8 @@ def test_logs_are_available(ws, sql_backend, env_or_skip, make_random): def test_jobs_with_no_inventory_database( ws, - sql_backend, + sql_exec, + sql_fetch_all, make_cluster_policy, make_cluster_policy_permissions, make_ucx_group,