Merge branch 'main' into feature/##1938

Author: aminmovahed-db

Makefile

@@ -23,7 +23,7 @@ integration:
 	hatch run integration
 
 coverage:
-	hatch run coverage && open htmlcov/index.html
+	hatch run coverage; status=$$?; [ -e "htmlcov/index.html" ] && open htmlcov/index.html; exit $$status
 
 known:
 	hatch run python src/databricks/labs/ucx/source_code/known.py

README.md

@@ -92,6 +92,7 @@ See [contributing instructions](CONTRIBUTING.md) to help improve this project.
 * [Metastore related commands](#metastore-related-commands)
   * [`show-all-metastores` command](#show-all-metastores-command)
   * [`assign-metastore` command](#assign-metastore-command)
+  * [`create-ucx-catalog` command](#create-ucx-catalog-command)
 * [Table migration commands](#table-migration-commands)
   * [`principal-prefix-access` command](#principal-prefix-access-command)
     * [Access for AWS S3 Buckets](#access-for-aws-s3-buckets)
@@ -1169,9 +1170,23 @@ a region, and you want to see which ones are available for assignment.
 databricks labs ucx assign-metastore --workspace-id <workspace-id> [--metastore-id <metastore-id>]
 ```
 
-This command assigns a metastore to a workspace with `workspace-id`. If there is only a single metastore in the workspace
-region, it will be automatically assigned to the workspace. If there are multiple metastores available, you need to specify
-the metastore id of the metastore you want to assign to the workspace.
+This command assigns a metastore to a workspace with `--workspace-id`. If there is only a single metastore in the
+workspace region, the command automatically assigns that metastore to the workspace. If there are multiple metastores
+available, the command prompts for specification of the metastore (id) you want to assign to the workspace.
+
+[[back to top](#databricks-labs-ucx)]
+
+## `create-ucx-catalog` command
+
+```commandline
+databricks labs ucx create-ucx-catalog
+16:12:59  INFO [d.l.u.hive_metastore.catalog_schema] Validating UC catalog: ucx
+Please provide storage location url for catalog: ucx (default: metastore): ...
+16:13:01  INFO [d.l.u.hive_metastore.catalog_schema] Creating UC catalog: ucx
+```
+
+Create and setup UCX artifact catalog. Amongst other things, the artifacts are used for tracking the migration progress
+across workspaces.
 
 # Table migration commands
 

labs.yml

@@ -261,15 +261,18 @@ commands:
 
   - name: assign-metastore
     is_account_level: true
-    description: Enable Unity Catalog features on a workspace by assign a metastore to it
+    description: Enable Unity Catalog features on a workspace by assigning a metastore to it.
     flags:
       - name: workspace-id
-        description: (Optional) Workspace ID to assign a metastore to
+        description: Workspace ID to assign a metastore to
       - name: metastore-id
         description: (Optional) If there are multiple metastores in the region, specify the metastore ID to assign
       - name: default-catalog
         description: (Optional) Default catalog to assign to the workspace. If not provided, it will be hive_metastore
 
+  - name: create-ucx-catalog
+    description: Create UCX artifact catalog
+
   - name: migrate-tables
     description: |
       Trigger the `migrate-tables` workflow and, optionally, `migrate-external-hiveserde-tables-in-place-experimental`

src/databricks/labs/ucx/account/metastores.py

@@ -27,38 +27,33 @@ def show_all_metastores(self, workspace_id: str | None = None):
     def assign_metastore(
         self,
         prompts: Prompts,
-        str_workspace_id: str | None = None,
+        workspace_id: int,
+        *,
         metastore_id: str | None = None,
         default_catalog: str | None = None,
     ):
-        if not str_workspace_id:
-            workspace_choices = self._get_all_workspaces()
-            workspace_id = prompts.choice_from_dict("Please select a workspace:", workspace_choices)
-        else:
-            workspace_id = int(str_workspace_id)
-        if not metastore_id:
+        if metastore_id is None:
             # search for all matching metastores
             metastore_choices = self._get_all_metastores(self._get_region(workspace_id))
             if len(metastore_choices) == 0:
-                raise ValueError(f"No matching metastore found for workspace {workspace_id}")
+                raise ValueError(f"No matching metastore found for workspace: {workspace_id}")
             # if there are multiple matches, prompt users to select one
             if len(metastore_choices) > 1:
                 metastore_id = prompts.choice_from_dict(
                     "Multiple metastores found, please select one:", metastore_choices
                 )
             else:
                 metastore_id = list(metastore_choices.values())[0]
-        if metastore_id is not None:
-            self._ac.metastore_assignments.create(workspace_id, metastore_id)
+        self._ac.metastore_assignments.create(workspace_id, metastore_id)
         # set the default catalog using the default_namespace setting API
         if default_catalog is not None:
             self._set_default_catalog(workspace_id, default_catalog)
 
-    def _get_region(self, workspace_id: int) -> str:
+    def _get_region(self, workspace_id: int) -> str | None:
         workspace = self._ac.workspaces.get(workspace_id)
         if self._ac.config.is_aws:
-            return str(workspace.aws_region)
-        return str(workspace.location)
+            return workspace.aws_region
+        return workspace.location
 
     def _get_all_workspaces(self) -> dict[str, int]:
         output = dict[str, int]()

src/databricks/labs/ucx/cli.py

@@ -590,11 +590,27 @@ def assign_metastore(
     workspace_id: str | None = None,
     metastore_id: str | None = None,
     default_catalog: str | None = None,
+    ctx: AccountContext | None = None,
 ):
     """Assign metastore to a workspace"""
     logger.info(f"Account ID: {a.config.account_id}")
-    ctx = AccountContext(a)
-    ctx.account_metastores.assign_metastore(ctx.prompts, workspace_id, metastore_id, default_catalog)
+    ctx = ctx or AccountContext(a)
+    ctx.account_metastores.assign_metastore(
+        ctx.prompts,
+        workspace_id,
+        metastore_id=metastore_id,
+        default_catalog=default_catalog,
+    )
+
+
+@ucx.command
+def create_ucx_catalog(w: WorkspaceClient, prompts: Prompts, ctx: WorkspaceContext | None = None) -> None:
+    """Create and setup UCX artifact catalog
+
+    Amongst other things, the artifacts are used for tracking the migration progress across workspaces.
+    """
+    workspace_context = ctx or WorkspaceContext(w)
+    workspace_context.catalog_schema.create_ucx_catalog(prompts)
 
 
 @ucx.command

src/databricks/labs/ucx/config.py

@@ -11,6 +11,7 @@ class WorkspaceConfig:  # pylint: disable=too-many-instance-attributes
     __version__ = 2
 
     inventory_database: str
+    ucx_catalog: str = "ucx"  # Catalog to store UCX artifact tables (shared across workspaces)
     # Group name conversion parameters.
     workspace_group_regex: str | None = None
     workspace_group_replace: str | None = None
@@ -25,7 +26,7 @@ class WorkspaceConfig:  # pylint: disable=too-many-instance-attributes
     connect: Config | None = None
     num_threads: int | None = 10
     database_to_catalog_mapping: dict[str, str] | None = None
-    default_catalog: str | None = "ucx_default"
+    default_catalog: str | None = "ucx_default"  # DEPRECATED: Keeping to avoid errors when loading old configurations
     log_level: str | None = "INFO"
 
     # Starting path for notebooks and directories crawler

src/databricks/labs/ucx/contexts/account_cli.py

@@ -3,7 +3,6 @@
 
 from databricks.sdk import AccountClient
 
-
 from databricks.labs.ucx.account.aggregate import AccountAggregate
 from databricks.labs.ucx.account.metastores import AccountMetastores
 from databricks.labs.ucx.account.workspaces import AccountWorkspaces

src/databricks/labs/ucx/contexts/application.py

@@ -18,7 +18,7 @@
 from databricks.labs.ucx.source_code.directfs_access import DirectFsAccessCrawler
 from databricks.labs.ucx.source_code.python_libraries import PythonLibraryResolver
 from databricks.sdk import AccountClient, WorkspaceClient, core
-from databricks.sdk.errors import ResourceDoesNotExist
+from databricks.sdk.errors import NotFound
 from databricks.sdk.service import sql
 
 from databricks.labs.ucx.account.workspaces import WorkspaceInfo
@@ -314,7 +314,7 @@ def principal_locations(self):
                 eligible_locations = self.aws_acl.get_eligible_locations_principals()
             if self.is_gcp:
                 raise NotImplementedError("Not implemented for GCP.")
-        except ResourceDoesNotExist:
+        except NotFound:
             pass
         return eligible_locations
 
@@ -354,6 +354,7 @@ def catalog_schema(self):
             self.principal_acl,
             self.sql_backend,
             self.grants_crawler,
+            self.config.ucx_catalog,
         )
 
     @cached_property

src/databricks/labs/ucx/hive_metastore/catalog_schema.py

@@ -18,29 +18,49 @@
 
 
 class CatalogSchema:
+
     def __init__(
         self,
         ws: WorkspaceClient,
         table_mapping: TableMapping,
         principal_grants: PrincipalACL,
         sql_backend: SqlBackend,
         grants_crawler: GrantsCrawler,
+        ucx_catalog: str,
     ):
         self._ws = ws
         self._table_mapping = table_mapping
         self._external_locations = self._ws.external_locations.list()
         self._principal_grants = principal_grants
         self._backend = sql_backend
         self._hive_grants_crawler = grants_crawler
+        self._ucx_catalog = ucx_catalog
+
+    def create_ucx_catalog(self, prompts: Prompts, *, properties: dict[str, str] | None = None) -> None:
+        """Create the UCX catalog.
+
+        Args:
+            prompts : Prompts
+                The prompts object to use for interactive input.
+            properties : (dict[str, str] | None), default None
+                The properties to pass to the catalog. If None, no properties are passed.
+        """
+        try:
+            self._create_catalog_validate(self._ucx_catalog, prompts, properties=properties)
+        except BadRequest as e:
+            if "already exists" in str(e):
+                logger.warning(f"Catalog '{self._ucx_catalog}' already exists. Skipping.")
+                return
+            raise
 
     def create_all_catalogs_schemas(self, prompts: Prompts) -> None:
         candidate_catalogs, candidate_schemas = self._get_missing_catalogs_schemas()
         for candidate_catalog in candidate_catalogs:
             try:
-                self._create_catalog_validate(candidate_catalog, prompts)
+                self._create_catalog_validate(candidate_catalog, prompts, properties=None)
             except BadRequest as e:
                 if "already exists" in str(e):
-                    logger.warning(f"Catalog {candidate_catalog} already exists. Skipping.")
+                    logger.warning(f"Catalog '{candidate_catalog}' already exists. Skipping.")
                     continue
         for candidate_catalog, schemas in candidate_schemas.items():
             for candidate_schema in schemas:
@@ -49,7 +69,7 @@ def create_all_catalogs_schemas(self, prompts: Prompts) -> None:
                 except BadRequest as e:
                     if "already exists" in str(e):
                         logger.warning(
-                            f"Schema {candidate_schema} in catalog {candidate_catalog} " f"already exists. Skipping."
+                            f"Schema '{candidate_schema}' in catalog '{candidate_catalog}' already exists. Skipping."
                         )
                         continue
         self._apply_from_legacy_table_acls()
@@ -121,20 +141,19 @@ def _get_database_source_target_mapping(self) -> dict[str, list[SchemaInfo]]:
                 src_trg_schema_mapping[table_mapping.src_schema].append(schema)
         return src_trg_schema_mapping
 
-    def _create_catalog_validate(self, catalog, prompts: Prompts):
-        logger.info(f"Creating UC catalog: {catalog}")
-        # create catalogs
+    def _create_catalog_validate(self, catalog: str, prompts: Prompts, *, properties: dict[str, str] | None) -> None:
+        logger.info(f"Validating UC catalog: {catalog}")
         attempts = 3
         while True:
             catalog_storage = prompts.question(
-                f"Please provide storage location url for catalog:{catalog}.", default="metastore"
+                f"Please provide storage location url for catalog: {catalog}", default="metastore"
             )
             if self._validate_location(catalog_storage):
                 break
             attempts -= 1
             if attempts == 0:
                 raise NotFound(f"Failed to validate location for {catalog} catalog")
-        self._create_catalog(catalog, catalog_storage)
+        self._create_catalog(catalog, catalog_storage, properties=properties)
 
     def _list_existing(self) -> tuple[set[str], dict[str, set[str]]]:
         """generate a list of existing UC catalogs and schema."""
@@ -199,12 +218,17 @@ def _validate_location(self, location: str):
                 return True
         return False
 
-    def _create_catalog(self, catalog, catalog_storage):
+    def _create_catalog(self, catalog: str, catalog_storage: str, *, properties: dict[str, str] | None) -> None:
         logger.info(f"Creating UC catalog: {catalog}")
         if catalog_storage == "metastore":
-            self._ws.catalogs.create(catalog, comment="Created by UCX")
+            self._ws.catalogs.create(catalog, comment="Created by UCX", properties=properties)
         else:
-            self._ws.catalogs.create(catalog, storage_root=catalog_storage, comment="Created by UCX")
+            self._ws.catalogs.create(
+                catalog,
+                storage_root=catalog_storage,
+                comment="Created by UCX",
+                properties=properties,
+            )
 
     def _create_schema(self, catalog, schema):
         logger.info(f"Creating UC schema: {schema} in catalog: {catalog}")

src/databricks/labs/ucx/install.py

@@ -234,6 +234,7 @@ def _prompt_for_new_installation(self) -> WorkspaceConfig:
         inventory_database = self.prompts.question(
             "Inventory Database stored in hive_metastore", default=default_database, valid_regex=r"^\w+$"
         )
+        ucx_catalog = self.prompts.question("Catalog to store UCX artifacts in", default="ucx", valid_regex=r"^\w+$")
         log_level = self.prompts.question("Log level", default="INFO").upper()
         num_threads = int(self.prompts.question("Number of threads", default="8", valid_number=True))
         configure_groups = ConfigureGroups(self.prompts)
@@ -248,6 +249,7 @@ def _prompt_for_new_installation(self) -> WorkspaceConfig:
         )
         return WorkspaceConfig(
             inventory_database=inventory_database,
+            ucx_catalog=ucx_catalog,
             workspace_group_regex=configure_groups.workspace_group_regex,
             workspace_group_replace=configure_groups.workspace_group_replace,
             account_group_regex=configure_groups.account_group_regex,