Fix create database in workspace installation (#2224)

## Changes
Database creation should be done asynchronously. Mypy probably missed
this as the method did not have any type hints. Introduced in #1920

---------

Co-authored-by: Serge Smertin <[email protected]>

Author: JCZuurmond

src/databricks/labs/ucx/install.py

@@ -490,12 +490,17 @@ def config(self):
     def folder(self):
         return self._installation.install_folder()
 
-    def run(self):
+    def run(self) -> bool:
+        """Run workflow installation.
+
+        Returns
+            bool :
+                True, installation finished. False installation did not finish.
+        """
         logger.info(f"Installing UCX v{self._product_info.version()}")
-        install_tasks = [self._create_database()]  # Need the database before creating the dashboards
-        install_tasks.extend(self._create_dashboards())
+        install_tasks = [self._create_database_and_dashboards, self._workflows_installer.create_jobs]
         Threads.strict("installing components", install_tasks)
-        readme_url = self._workflows_installer.create_jobs()
+        readme_url = self._installation.workspace_link("README")
         if not self._is_account_install and self._prompts.confirm(f"Open job overview in your browser? {readme_url}"):
             webbrowser.open(readme_url)
         logger.info(f"Installation completed successfully! Please refer to the {readme_url} for the next steps.")
@@ -504,6 +509,10 @@ def run(self):
             self._trigger_workflow("assessment")
         return True
 
+    def _create_database_and_dashboards(self) -> None:
+        self._create_database()  # Need the database before creating the dashboards
+        Threads.strict("installing dashboards", list(self._get_create_dashboard_tasks()))
+
     def _create_database(self):
         try:
             deploy_schema(self._sql_backend, self._config.inventory_database)
@@ -519,8 +528,8 @@ def _create_database(self):
                 raise BadRequest(msg) from err
             raise err
 
-    def _create_dashboards(self) -> Iterable[Callable[[], None]]:
-        """Create the lakeview dashboards from the SQL queries in the queries subfolders"""
+    def _get_create_dashboard_tasks(self) -> Iterable[Callable[[], None]]:
+        """Get the tasks to create Lakeview dashboards from the SQL queries in the queries subfolders"""
         logger.info("Creating dashboards...")
         dashboard_folder_remote = f"{self._installation.install_folder()}/dashboards"
         try:
@@ -544,7 +553,7 @@ def _create_dashboards(self) -> Iterable[Callable[[], None]]:
 
     # TODO: Confirm the assumption below is correct
     # An InternalError may occur when the dashboard is being published and the database does not exists
-    @retried(on=[InternalError], timeout=timedelta(minutes=2))
+    @retried(on=[InternalError], timeout=timedelta(minutes=4))
     def _create_dashboard(self, folder: Path, *, parent_path: str | None = None) -> None:
         """Create a lakeview dashboard from the SQL queries in the folder"""
         logger.info(f"Creating dashboard in {folder}...")

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

@@ -417,11 +417,11 @@ def __init__(
         self._this_file = Path(__file__)
         super().__init__(config, installation, ws)
 
-    def create_jobs(self):
+    def create_jobs(self) -> None:
         remote_wheels = self._upload_wheel()
         desired_workflows = {t.workflow for t in self._tasks if t.cloud_compatible(self._ws.config)}
-        wheel_runner = None
 
+        wheel_runner = ""
         if self._config.override_clusters:
             wheel_runner = self._upload_wheel_runner(remote_wheels)
         for workflow_name in desired_workflows:
@@ -446,7 +446,7 @@ def create_jobs(self):
 
         self._install_state.save()
         self._create_debug(remote_wheels)
-        return self._create_readme()
+        self._create_readme()
 
     @property
     def _config_file(self):
@@ -464,7 +464,7 @@ def _get_test_purge_time(self) -> str:
         timeout = TEST_NIGHTLY_CI_JOBS_PURGE_TIMEOUT if self._is_nightly() else TEST_JOBS_PURGE_TIMEOUT
         return (datetime.utcnow() + timeout).strftime("%Y%m%d%H")
 
-    def _create_readme(self) -> str:
+    def _create_readme(self) -> None:
         debug_notebook_link = self._installation.workspace_markdown_link('debug notebook', 'DEBUG.py')
         markdown = [
             "# UCX - The Unity Catalog Migration Assistant",
@@ -496,8 +496,6 @@ def _create_readme(self) -> str:
         preamble = ["# Databricks notebook source", "# MAGIC %md"]
         intro = "\n".join(preamble + [f"# MAGIC {line}" for line in markdown])
         self._installation.upload('README.py', intro.encode('utf8'))
-        readme_url = self._installation.workspace_link('README')
-        return readme_url
 
     def _create_dashboard_links(self, step_name, dashboard_link):
         dashboards_per_step = [d for d in self._install_state.dashboards.keys() if d.startswith(step_name)]
@@ -570,7 +568,7 @@ def _upload_wheel(self):
             wheel_paths = [f"/Workspace{wheel}" for wheel in wheel_paths]
             return wheel_paths
 
-    def _upload_wheel_runner(self, remote_wheels: list[str]):
+    def _upload_wheel_runner(self, remote_wheels: list[str]) -> str:
         # TODO: we have to be doing this workaround until ES-897453 is solved in the platform
         remote_wheels_str = " ".join(remote_wheels)
         code = TEST_RUNNER_NOTEBOOK.format(remote_wheel=remote_wheels_str, config_file=self._config_file).encode("utf8")

tests/unit/test_install.py

@@ -228,7 +228,7 @@ def test_create_database(ws, caplog, mock_installation, any_prompt):
             raise e.errs[0]
 
     assert "Kindly uninstall and reinstall UCX" in str(failure.value)
-    wheels.upload_to_wsfs.assert_not_called()
+    wheels.upload_to_wsfs.assert_called_once()
 
 
 def test_install_cluster_override_jobs(ws, mock_installation):