Add Support for migrating Schema/Catalog ACL for Interactive cluster (#1413)

HariGS-DB · web-flow · commit fd645fd39e35 · 2024-04-22T23:28:01.000+02:00
## Changes - Identified all the database acl grants from the PrincipalACL class and filters for only database grants - creates the mapping from hive metastore schema to UC schema and catalog using the Table Mapping - replaces the hive metastore action with equivalent UC action for both schema and catalog - handles both aws and azure - Note: This PR doesn't cover external location permission which is part of the same issue. I will submit separate PR for that ### Linked issues  Fixes Partially #1192 and #1193 ### Functionality - [ ] added relevant user documentation - [ ] added new CLI command - [ ] modified existing command: `databricks labs ucx ...` - [ ] added a new workflow - [ ] modified existing workflow: `...` - [ ] added a new table - [ ] modified existing table: `...` ### Tests  - [X] manually tested - [X] added unit tests - [X] added integration tests - [ ] verified on staging environment (screenshot attached)
diff --git a/src/databricks/labs/ucx/contexts/application.py b/src/databricks/labs/ucx/contexts/application.py
@@ -303,7 +303,7 @@ def table_mapping(self):
 
     @cached_property
     def catalog_schema(self):
-        return CatalogSchema(self.workspace_client, self.table_mapping)
+        return CatalogSchema(self.workspace_client, self.table_mapping, self.principal_acl, self.sql_backend)
 
     @cached_property
     def languages(self):
diff --git a/src/databricks/labs/ucx/hive_metastore/catalog_schema.py b/src/databricks/labs/ucx/hive_metastore/catalog_schema.py
@@ -1,7 +1,10 @@
 import logging
 from pathlib import PurePath
+import dataclasses
 
 from databricks.labs.blueprint.tui import Prompts
+from databricks.labs.lsql.backends import SqlBackend
+from databricks.labs.ucx.hive_metastore.grants import PrincipalACL, Grant
 from databricks.sdk import WorkspaceClient
 from databricks.sdk.errors import NotFound
 
@@ -11,10 +14,14 @@
 
 
 class CatalogSchema:
-    def __init__(self, ws: WorkspaceClient, table_mapping: TableMapping):
+    def __init__(
+        self, ws: WorkspaceClient, table_mapping: TableMapping, principal_grants: PrincipalACL, sql_backend: SqlBackend
+    ):
         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
 
     def create_all_catalogs_schemas(self, prompts: Prompts):
         candidate_catalogs, candidate_schemas = self._get_missing_catalogs_schemas()
@@ -23,6 +30,53 @@ def create_all_catalogs_schemas(self, prompts: Prompts):
         for candidate_catalog, schemas in candidate_schemas.items():
             for candidate_schema in schemas:
                 self._create_schema(candidate_catalog, candidate_schema)
+        self._update_principal_acl()
+
+    def _update_principal_acl(self):
+        grants = self._get_catalog_schema_grants()
+        for grant in grants:
+            acl_migrate_sql = grant.uc_grant_sql()
+            if acl_migrate_sql is None:
+                logger.warning(f"Cannot identify UC grant for {grant.this_type_and_key()}. Skipping.")
+                continue
+            logger.debug(f"Migrating acls on {grant.this_type_and_key()} using SQL query: {acl_migrate_sql}")
+            self._backend.execute(acl_migrate_sql)
+
+    def _get_catalog_schema_grants(self):
+        catalog_grants: set[Grant] = set()
+        new_grants = []
+        src_trg_schema_mapping = self._get_database_source_target_mapping()
+        grants = self._principal_grants.get_interactive_cluster_grants()
+        # filter on grants to only get database level grants
+        database_grants = [grant for grant in grants if grant.table is None and grant.view is None]
+        for db_grant in database_grants:
+            new_grants.append(
+                dataclasses.replace(
+                    db_grant,
+                    # replace source database with taget UC database
+                    database=src_trg_schema_mapping[db_grant.database]['target_schema'],
+                    # replace hive_metastore with target UC catalog
+                    catalog=src_trg_schema_mapping[db_grant.database]['target_catalog'],
+                )
+            )
+        for grant in new_grants:
+            catalog_grants.add(dataclasses.replace(grant, database=None))
+        new_grants.extend(catalog_grants)
+        return new_grants
+
+    def _get_database_source_target_mapping(self) -> dict[str, dict]:
+        """generate a dictionary of source database in hive_metastore and its
+        mapping of target UC catalog and schema from the table mappings."""
+        src_trg_schema_mapping: dict[str, dict] = {}
+        table_mappings = self._table_mapping.load()
+        for mappings in table_mappings:
+            if mappings.src_schema not in src_trg_schema_mapping:
+                src_trg_schema_mapping[mappings.src_schema] = {
+                    'target_catalog': mappings.catalog_name,
+                    'target_schema': mappings.dst_schema,
+                }
+                continue
+        return src_trg_schema_mapping
 
     def _create_catalog_validate(self, catalog, prompts: Prompts):
         logger.info(f"Creating UC catalog: {catalog}")
diff --git a/src/databricks/labs/ucx/hive_metastore/grants.py b/src/databricks/labs/ucx/hive_metastore/grants.py
@@ -166,6 +166,7 @@ def uc_grant_sql(self, object_type: str | None = None, object_key: str | None =
             ("DATABASE", "OWN"): self._set_owner_sql,
             ("DATABASE", "READ_METADATA"): self._uc_action("BROWSE"),
             ("CATALOG", "OWN"): self._set_owner_sql,
+            ("CATALOG", "USAGE"): self._uc_action("USE CATALOG"),
         }
         make_query = hive_to_uc.get((object_type, self.action_type), None)
         if make_query is None:
@@ -530,8 +531,6 @@ def get_interactive_cluster_grants(self) -> list[Grant]:
                 continue
             cluster_usage = self._get_grants(locations, principals, tables, mounts)
             grants.update(cluster_usage)
-            catalog_grants = [Grant(principal, "USE", "hive_metastore") for principal in principals]
-            grants.update(catalog_grants)
         return list(grants)
 
     def _get_privilege(self, table: Table, locations: dict[str, str], mounts: list[Mount]):
@@ -557,7 +556,7 @@ def _get_privilege(self, table: Table, locations: dict[str, str], mounts: list[M
     def _get_database_grants(self, tables: list[Table], principals: list[str]) -> list[Grant]:
         databases = {table.database for table in tables}
         return [
-            Grant(principal, "USE", "hive_metastore", database) for database in databases for principal in principals
+            Grant(principal, "USAGE", "hive_metastore", database) for database in databases for principal in principals
         ]
 
     def _get_grants(
diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py
@@ -665,3 +665,24 @@ def prepare_tables_for_migration(
     installation_ctx.save_mounts()
     installation_ctx.with_dummy_grants_and_tacls()
     return tables, dst_schema
+
+
+@pytest.fixture
+def prepared_principal_acl(runtime_ctx, env_or_skip, make_mounted_location, make_catalog, make_schema):
+    src_schema = make_schema(catalog_name="hive_metastore")
+    src_external_table = runtime_ctx.make_table(
+        catalog_name=src_schema.catalog_name,
+        schema_name=src_schema.name,
+        external_csv=make_mounted_location,
+    )
+    dst_catalog = make_catalog()
+    dst_schema = make_schema(catalog_name=dst_catalog.name, name=src_schema.name)
+    rules = [Rule.from_src_dst(src_external_table, dst_schema)]
+    runtime_ctx.with_table_mapping_rules(rules)
+    runtime_ctx.with_dummy_resource_permission()
+    return (
+        runtime_ctx,
+        f"{dst_catalog.name}.{dst_schema.name}.{src_external_table.name}",
+        f"{dst_catalog.name}.{dst_schema.name}",
+        dst_catalog.name,
+    )
diff --git a/tests/integration/hive_metastore/__init__.py b/tests/integration/hive_metastore/__init__.py
@@ -0,0 +1,12 @@
+def get_azure_spark_conf():
+    return {
+        "spark.databricks.cluster.profile": "singleNode",
+        "spark.master": "local[*]",
+        "fs.azure.account.auth.type.labsazurethings.dfs.core.windows.net": "OAuth",
+        "fs.azure.account.oauth.provider.type.labsazurethings.dfs.core.windows.net": "org.apache.hadoop.fs"
+        ".azurebfs.oauth2.ClientCredsTokenProvider",
+        "fs.azure.account.oauth2.client.id.labsazurethings.dfs.core.windows.net": "dummy_application_id",
+        "fs.azure.account.oauth2.client.secret.labsazurethings.dfs.core.windows.net": "dummy",
+        "fs.azure.account.oauth2.client.endpoint.labsazurethings.dfs.core.windows.net": "https://login"
+        ".microsoftonline.com/directory_12345/oauth2/token",
+    }
diff --git a/tests/integration/hive_metastore/test_catalog_schema.py b/tests/integration/hive_metastore/test_catalog_schema.py
@@ -0,0 +1,72 @@
+import logging
+from datetime import timedelta
+import pytest
+
+from databricks.labs.blueprint.tui import MockPrompts
+
+from databricks.sdk.errors import NotFound
+from databricks.sdk.retries import retried
+from databricks.sdk.service.compute import DataSecurityMode, AwsAttributes
+from databricks.sdk.service.catalog import Privilege, SecurableType, PrivilegeAssignment
+from databricks.sdk.service.iam import PermissionLevel
+
+from . import get_azure_spark_conf
+
+logger = logging.getLogger(__name__)
+_SPARK_CONF = get_azure_spark_conf()
+
+
+@retried(on=[NotFound], timeout=timedelta(minutes=2))
+def test_create_catalog_schema_with_principal_acl_azure(
+    ws, make_user, prepared_principal_acl, make_cluster_permissions, make_cluster
+):
+    if not ws.config.is_azure:
+        pytest.skip("only works in azure test env")
+    ctx, _, schema_name, catalog_name = prepared_principal_acl
+
+    cluster = make_cluster(single_node=True, spark_conf=_SPARK_CONF, data_security_mode=DataSecurityMode.NONE)
+    user = make_user()
+    make_cluster_permissions(
+        object_id=cluster.cluster_id,
+        permission_level=PermissionLevel.CAN_ATTACH_TO,
+        user_name=user.user_name,
+    )
+    catalog_schema = ctx.catalog_schema
+    mock_prompts = MockPrompts({"Please provide storage location url for catalog: *": ""})
+    catalog_schema.create_all_catalogs_schemas(mock_prompts)
+
+    schema_grants = ws.grants.get(SecurableType.SCHEMA, schema_name)
+    catalog_grants = ws.grants.get(SecurableType.CATALOG, catalog_name)
+    schema_grant = PrivilegeAssignment(user.user_name, [Privilege.USE_SCHEMA])
+    catalog_grant = PrivilegeAssignment(user.user_name, [Privilege.USE_CATALOG])
+    assert schema_grant in schema_grants.privilege_assignments
+    assert catalog_grant in catalog_grants.privilege_assignments
+
+
+@retried(on=[NotFound], timeout=timedelta(minutes=3))
+def test_create_catalog_schema_with_principal_acl_aws(
+    ws, make_user, prepared_principal_acl, make_cluster_permissions, make_cluster, env_or_skip
+):
+    ctx, _, schema_name, catalog_name = prepared_principal_acl
+
+    cluster = make_cluster(
+        single_node=True,
+        data_security_mode=DataSecurityMode.NONE,
+        aws_attributes=AwsAttributes(instance_profile_arn=env_or_skip("TEST_WILDCARD_INSTANCE_PROFILE")),
+    )
+    user = make_user()
+    make_cluster_permissions(
+        object_id=cluster.cluster_id,
+        permission_level=PermissionLevel.CAN_ATTACH_TO,
+        user_name=user.user_name,
+    )
+    catalog_schema = ctx.catalog_schema
+    mock_prompts = MockPrompts({"Please provide storage location url for catalog: *": ""})
+    catalog_schema.create_all_catalogs_schemas(mock_prompts)
+
+    schema_grants = ws.grants.get(SecurableType.SCHEMA, schema_name)
+    catalog_grants = ws.grants.get(SecurableType.CATALOG, catalog_name)
+    schema_grant = PrivilegeAssignment(user.user_name, [Privilege.USE_SCHEMA])
+    catalog_grant = PrivilegeAssignment(user.user_name, [Privilege.USE_CATALOG])
+    assert schema_grant in schema_grants.privilege_assignments
+    assert catalog_grant in catalog_grants.privilege_assignments
diff --git a/tests/integration/hive_metastore/test_migrate.py b/tests/integration/hive_metastore/test_migrate.py
@@ -10,20 +10,11 @@
 
 from databricks.labs.ucx.hive_metastore.mapping import Rule, TableMapping
 from databricks.labs.ucx.hive_metastore.tables import AclMigrationWhat, Table, What
+from . import get_azure_spark_conf
 
 
 logger = logging.getLogger(__name__)
-_SPARK_CONF = {
-    "spark.databricks.cluster.profile": "singleNode",
-    "spark.master": "local[*]",
-    "fs.azure.account.auth.type.labsazurethings.dfs.core.windows.net": "OAuth",
-    "fs.azure.account.oauth.provider.type.labsazurethings.dfs.core.windows.net": "org.apache.hadoop.fs"
-    ".azurebfs.oauth2.ClientCredsTokenProvider",
-    "fs.azure.account.oauth2.client.id.labsazurethings.dfs.core.windows.net": "dummy_application_id",
-    "fs.azure.account.oauth2.client.secret.labsazurethings.dfs.core.windows.net": "dummy",
-    "fs.azure.account.oauth2.client.endpoint.labsazurethings.dfs.core.windows.net": "https://login"
-    ".microsoftonline.com/directory_12345/oauth2/token",
-}
+_SPARK_CONF = get_azure_spark_conf()
 
 
 @retried(on=[NotFound], timeout=timedelta(minutes=2))
@@ -425,31 +416,13 @@ def test_migrate_managed_tables_with_acl(ws, sql_backend, runtime_ctx, make_cata
     assert target_table_grants.privilege_assignments[0].privileges == [Privilege.MODIFY, Privilege.SELECT]
 
 
-@pytest.fixture
-def prepared_principal_acl(runtime_ctx, env_or_skip, make_mounted_location, make_catalog, make_schema):
-    src_schema = make_schema(catalog_name="hive_metastore")
-    src_external_table = runtime_ctx.make_table(
-        catalog_name=src_schema.catalog_name,
-        schema_name=src_schema.name,
-        external_csv=make_mounted_location,
-    )
-    dst_catalog = make_catalog()
-    dst_schema = make_schema(catalog_name=dst_catalog.name, name=src_schema.name)
-    rules = [Rule.from_src_dst(src_external_table, dst_schema)]
-    runtime_ctx.with_table_mapping_rules(rules)
-    return (
-        runtime_ctx,
-        f"{dst_catalog.name}.{dst_schema.name}.{src_external_table.name}",
-    )
-
-
 @retried(on=[NotFound], timeout=timedelta(minutes=2))
 def test_migrate_external_tables_with_principal_acl_azure(
     ws, make_user, prepared_principal_acl, make_cluster_permissions, make_cluster, make_ucx_group
 ):
     if not ws.config.is_azure:
         pytest.skip("only works in azure test env")
-    ctx, table_full_name = prepared_principal_acl
+    ctx, table_full_name, _, _ = prepared_principal_acl
     cluster = make_cluster(single_node=True, spark_conf=_SPARK_CONF, data_security_mode=DataSecurityMode.NONE)
     ctx.with_dummy_resource_permission()
     table_migrate = ctx.tables_migrator
@@ -490,7 +463,7 @@ def test_migrate_external_tables_with_principal_acl_azure(
 def test_migrate_external_tables_with_principal_acl_aws(
     ws, make_user, prepared_principal_acl, make_cluster_permissions, make_cluster, env_or_skip
 ):
-    ctx, table_full_name = prepared_principal_acl
+    ctx, table_full_name, _, _ = prepared_principal_acl
     ctx.with_dummy_resource_permission()
     cluster = make_cluster(
         single_node=True,
diff --git a/tests/unit/hive_metastore/test_catalog_schema.py b/tests/unit/hive_metastore/test_catalog_schema.py
@@ -9,14 +9,16 @@
 from databricks.sdk.service.catalog import CatalogInfo, ExternalLocationInfo, SchemaInfo
 
 from databricks.labs.ucx.hive_metastore.catalog_schema import CatalogSchema
+from databricks.labs.ucx.hive_metastore.grants import PrincipalACL, Grant
 from databricks.labs.ucx.hive_metastore.mapping import TableMapping
 
 
-def prepare_test(ws) -> CatalogSchema:
+def prepare_test(ws, backend: MockBackend | None = None) -> CatalogSchema:
     ws.catalogs.list.return_value = [CatalogInfo(name="catalog1")]
     ws.schemas.list.return_value = [SchemaInfo(name="schema1")]
     ws.external_locations.list.return_value = [ExternalLocationInfo(url="s3://foo/bar")]
-    backend = MockBackend()
+    if backend is None:
+        backend = MockBackend()
     installation = MockInstallation(
         {
             'mapping.csv': [
@@ -48,16 +50,26 @@ def prepare_test(ws) -> CatalogSchema:
         }
     )
     table_mapping = TableMapping(installation, ws, backend)
-
-    return CatalogSchema(ws, table_mapping)
+    principal_acl = create_autospec(PrincipalACL)
+    grants = [
+        Grant('user1', 'SELECT', 'catalog1', 'schema3', 'table'),
+        Grant('user1', 'MODIFY', 'catalog2', 'schema2', 'table'),
+        Grant('user1', 'SELECY', 'catalog2', 'schema2', 'table2'),
+        Grant('user1', 'USAGE', 'hive_metastore', 'schema3'),
+        Grant('user1', 'USAGE', 'hive_metastore', 'schema2'),
+    ]
+    principal_acl.get_interactive_cluster_grants.return_value = grants
+    return CatalogSchema(ws, table_mapping, principal_acl, backend)
 
 
 def test_create():
     ws = create_autospec(WorkspaceClient)
     mock_prompts = MockPrompts({"Please provide storage location url for catalog: *": "s3://foo/bar"})
 
     catalog_schema = prepare_test(ws)
-    catalog_schema.create_all_catalogs_schemas(mock_prompts)
+    catalog_schema.create_all_catalogs_schemas(
+        mock_prompts,
+    )
     ws.catalogs.create.assert_called_once_with("catalog2", storage_root="s3://foo/bar", comment="Created by UCX")
     ws.schemas.create.assert_any_call("schema2", "catalog2", comment="Created by UCX")
     ws.schemas.create.assert_any_call("schema3", "catalog1", comment="Created by UCX")
@@ -89,3 +101,22 @@ def test_no_catalog_storage():
     catalog_schema = prepare_test(ws)
     catalog_schema.create_all_catalogs_schemas(mock_prompts)
     ws.catalogs.create.assert_called_once_with("catalog2", comment="Created by UCX")
+
+
+def test_catalog_schema_acl():
+    ws = create_autospec(WorkspaceClient)
+    backend = MockBackend()
+    mock_prompts = MockPrompts({"Please provide storage location url for catalog: *": ""})
+    catalog_schema = prepare_test(ws, backend)
+    catalog_schema.create_all_catalogs_schemas(
+        mock_prompts,
+    )
+    queries = [
+        'GRANT USE SCHEMA ON DATABASE catalog1.schema3 TO `user1`',
+        'GRANT USE SCHEMA ON DATABASE catalog2.schema2 TO `user1`',
+        'GRANT USE CATALOG ON CATALOG catalog1 TO `user1`',
+        'GRANT USE CATALOG ON CATALOG catalog2 TO `user1`',
+    ]
+    assert len(backend.queries) == 4
+    for query in queries:
+        assert query in backend.queries
diff --git a/tests/unit/hive_metastore/test_principal_grants.py b/tests/unit/hive_metastore/test_principal_grants.py
