Crawl Table ACLs in all databases (#122)

This PR enumerates all databases for `GrantsCrawler`.

---------

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

Author: 3 people

examples/migration_config.yml

@@ -4,8 +4,10 @@ inventory:
     database: default
     name: uc_migration_inventory
 
+tacl:
+  databases: [ "default" ]
 
-with_table_acls: False
+warehouse_id: None
 
 groups:
   selected: [ "analyst" ]

notebooks/toolkit.py

@@ -36,7 +36,8 @@
 # COMMAND ----------
 
 from databricks.labs.ucx.toolkits.group_migration import GroupMigrationToolkit
-from databricks.labs.ucx.config import MigrationConfig, InventoryConfig, GroupsConfig, InventoryTable
+from databricks.labs.ucx.config import MigrationConfig, InventoryConfig, GroupsConfig, InventoryTable, TaclConfig
+from databricks.labs.ucx.toolkits.table_acls import TaclToolkit
 
 # COMMAND ----------
 
@@ -47,17 +48,28 @@
 # COMMAND ----------
 
 config = MigrationConfig(
-    with_table_acls=False,
     inventory=InventoryConfig(table=InventoryTable(catalog="main", database="default", name="ucx_migration_inventory")),
     groups=GroupsConfig(
         # use this option to select specific groups manually
         selected=["groupA", "groupB"],
         # use this option to select all groups automatically
         # auto=True
     ),
+    tacl=TaclConfig(
+        # use this option to select specific databases manually
+        databases=["default"],
+        # use this option to select all databases automatically
+        # auto=True
+    ),
     log_level="DEBUG",
 )
 toolkit = GroupMigrationToolkit(config)
+tacltoolkit = TaclToolkit(
+    toolkit._ws,
+    config.inventory.table.catalog,
+    config.inventory.table.schema,
+    databases=config.tacl.databases,
+)
 
 # COMMAND ----------
 
@@ -152,6 +164,16 @@
 
 # COMMAND ----------
 
+# MAGIC %md
+# MAGIC
+# MAGIC ## Inventorize Table ACL's
+
+# COMMAND ----------
+
+tacltoolkit.grants_snapshot()
+
+# COMMAND ----------
+
 # MAGIC %md
 # MAGIC
 # MAGIC ## Cleanup the inventory table

pyproject.toml

@@ -187,4 +187,4 @@ exclude_lines = [
     "no cov",
     "if __name__ == .__main__.:",
     "if TYPE_CHECKING:",
-]
+]

src/databricks/labs/ucx/config.py

@@ -90,10 +90,28 @@ def from_dict(cls, raw: dict):
         return cls(**raw)
 
 
+@dataclass
+class TaclConfig:
+    databases: list[str] | None = None
+    auto: bool | None = None
+
+    def __post_init__(self):
+        if not self.databases and self.auto is None:
+            msg = "Either selected or auto must be set"
+            raise ValueError(msg)
+        if self.databases and self.auto is False:
+            msg = "No selected groups provided, but auto-collection is disabled"
+            raise ValueError(msg)
+
+    @classmethod
+    def from_dict(cls, raw: dict):
+        return cls(**raw)
+
+
 @dataclass
 class MigrationConfig:
     inventory: InventoryConfig
-    with_table_acls: bool
+    tacl: TaclConfig
     groups: GroupsConfig
     connect: ConnectConfig | None = None
     num_threads: int | None = 4
@@ -102,9 +120,6 @@ class MigrationConfig:
     def __post_init__(self):
         if self.connect is None:
             self.connect = ConnectConfig()
-        if self.with_table_acls:
-            msg = "Table ACLS are not yet implemented"
-            raise NotImplementedError(msg)
 
     def as_dict(self) -> dict:
         from dataclasses import fields, is_dataclass
@@ -126,7 +141,7 @@ def inner(x):
     def from_dict(cls, raw: dict) -> "MigrationConfig":
         return cls(
             inventory=InventoryConfig.from_dict(raw.get("inventory", {})),
-            with_table_acls=raw.get("with_table_acls", False),
+            tacl=TaclConfig.from_dict(raw.get("tacl", {})),
             groups=GroupsConfig.from_dict(raw.get("groups", {})),
             connect=ConnectConfig.from_dict(raw.get("connect", {})),
             num_threads=raw.get("num_threads", 4),

src/databricks/labs/ucx/tacl/_internal.py

@@ -143,15 +143,25 @@ def _snapshot(self, klass, fetcher, loader) -> list[any]:
         Returns:
         list[any]: A list of data records, either fetched or loaded.
         """
+        loaded = False
+        trigger_load = ValueError("trigger records load")
         while True:
             try:
                 logger.debug(f"[{self._full_name}] fetching {self._table} inventory")
-                return list(fetcher())
+                cached_results = list(fetcher())
+                if len(cached_results) == 0 and loaded:
+                    return cached_results
+                if len(cached_results) == 0 and not loaded:
+                    raise trigger_load
+                return cached_results
             except Exception as e:
-                if "TABLE_OR_VIEW_NOT_FOUND" not in str(e):
+                if not (e == trigger_load or "TABLE_OR_VIEW_NOT_FOUND" in str(e)):
                     raise e
-                logger.debug(f"[{self._full_name}] {self._table} inventory not found, crawling")
-                self._append_records(klass, loader())
+            logger.debug(f"[{self._full_name}] crawling new batch for {self._table}")
+            loaded_records = list(loader())
+            if len(loaded_records) > 0:
+                self._append_records(klass, loaded_records)
+            loaded = True
 
     @staticmethod
     def _row_to_sql(row, fields):

src/databricks/labs/ucx/toolkits/table_acls.py

@@ -1,3 +1,5 @@
+import logging
+
 from databricks.sdk import WorkspaceClient
 
 from databricks.labs.ucx.tacl._internal import (
@@ -8,17 +10,38 @@
 from databricks.labs.ucx.tacl.grants import GrantsCrawler
 from databricks.labs.ucx.tacl.tables import TablesCrawler
 
+logger = logging.getLogger(__name__)
+
 
 class TaclToolkit:
-    def __init__(self, ws: WorkspaceClient, inventory_catalog, inventory_schema, warehouse_id=None):
+    def __init__(
+        self,
+        ws: WorkspaceClient,
+        inventory_catalog,
+        inventory_schema,
+        warehouse_id=None,
+        databases=None,
+    ):
         self._tc = TablesCrawler(self._backend(ws, warehouse_id), inventory_catalog, inventory_schema)
         self._gc = GrantsCrawler(self._tc)
 
-    def database_snapshot(self, schema):
-        return self._tc.snapshot("hive_metastore", schema)
+        self._databases = (
+            databases if databases else [database.as_dict()["databaseName"] for database in self._tc._all_databases()]
+        )
+
+    def database_snapshot(self):
+        tables = []
+        for db in self._databases:
+            for t in self._tc.snapshot("hive_metastore", db):
+                tables.append(t)
+        return tables
 
-    def grants_snapshot(self, schema):
-        return self._gc.snapshot("hive_metastore", schema)
+    def grants_snapshot(self):
+        grants = []
+        for db in self._databases:
+            for grant in self._gc.snapshot("hive_metastore", db):
+                grants.append(grant)
+        return grants
 
     @staticmethod
     def _backend(ws: WorkspaceClient, warehouse_id: str | None = None) -> SqlBackend:

tests/integration/test_tacls.py

@@ -8,16 +8,22 @@
 logger = logging.getLogger(__name__)
 
 
-def test_describe_all_tables(ws: WorkspaceClient, make_catalog, make_schema, make_table):
+def test_describe_all_tables_in_databases(ws: WorkspaceClient, make_catalog, make_schema, make_table):
     warehouse_id = os.environ["TEST_DEFAULT_WAREHOUSE_ID"]
 
     logger.info("setting up fixtures")
-    schema = make_schema(catalog="hive_metastore")
-    managed_table = make_table(schema=schema)
-    external_table = make_table(schema=schema, external=True)
-    tmp_table = make_table(schema=schema, ctas="SELECT 2+2 AS four")
-    view = make_table(schema=schema, ctas="SELECT 2+2 AS four", view=True)
-    non_delta = make_table(schema=schema, non_detla=True)
+
+    schema_a = make_schema(catalog="hive_metastore")
+
+    schema_b = make_schema(catalog="hive_metastore")
+
+    schema_c = make_schema(catalog="hive_metastore")
+
+    managed_table = make_table(schema=schema_a)
+    external_table = make_table(schema=schema_b, external=True)
+    tmp_table = make_table(schema=schema_a, ctas="SELECT 2+2 AS four")
+    view = make_table(schema=schema_b, ctas="SELECT 2+2 AS four", view=True)
+    non_delta = make_table(schema=schema_a, non_detla=True)
 
     logger.info(
         f"managed_table={managed_table}, "
@@ -28,10 +34,13 @@ def test_describe_all_tables(ws: WorkspaceClient, make_catalog, make_schema, mak
 
     inventory_schema = make_schema(catalog=make_catalog())
     inventory_catalog, inventory_schema = inventory_schema.split(".")
-    tak = TaclToolkit(ws, inventory_catalog, inventory_schema, warehouse_id)
+
+    databases = [schema_a.split(".")[1], schema_b.split(".")[1], schema_c.split(".")[1]]
+
+    tak = TaclToolkit(ws, inventory_catalog, inventory_schema, warehouse_id=warehouse_id, databases=databases)
 
     all_tables = {}
-    for t in tak.database_snapshot(schema.split(".")[1]):
+    for t in tak.database_snapshot():
         all_tables[t.key] = t
 
     assert len(all_tables) == 5
@@ -43,28 +52,33 @@ def test_describe_all_tables(ws: WorkspaceClient, make_catalog, make_schema, mak
     assert all_tables[view].view_text == "SELECT 2+2 AS four"
 
 
-def test_all_grants_in_database(ws: WorkspaceClient, sql_exec, make_catalog, make_schema, make_table, make_group):
+def test_all_grants_in_databases(ws: WorkspaceClient, sql_exec, make_catalog, make_schema, make_table, make_group):
     warehouse_id = os.environ["TEST_DEFAULT_WAREHOUSE_ID"]
 
     group_a = make_group()
     group_b = make_group()
-    schema = make_schema()
-    table = make_table(schema=schema, external=True)
+    schema_a = make_schema()
+    schema_b = make_schema()
+    table_a = make_table(schema=schema_a)
+    table_b = make_table(schema=schema_b)
 
     sql_exec(f"GRANT USAGE ON SCHEMA default TO `{group_a.display_name}`")
     sql_exec(f"GRANT USAGE ON SCHEMA default TO `{group_b.display_name}`")
-    sql_exec(f"GRANT SELECT ON TABLE {table} TO `{group_a.display_name}`")
-    sql_exec(f"GRANT MODIFY ON SCHEMA {schema} TO `{group_b.display_name}`")
+    sql_exec(f"GRANT SELECT ON TABLE {table_a} TO `{group_a.display_name}`")
+    sql_exec(f"GRANT SELECT ON TABLE {table_b} TO `{group_b.display_name}`")
+    sql_exec(f"GRANT MODIFY ON SCHEMA {schema_b} TO `{group_b.display_name}`")
 
     inventory_schema = make_schema(catalog=make_catalog())
     inventory_catalog, inventory_schema = inventory_schema.split(".")
-    tak = TaclToolkit(ws, inventory_catalog, inventory_schema, warehouse_id)
+
+    tak = TaclToolkit(ws, inventory_catalog, inventory_schema, warehouse_id=warehouse_id)
 
     all_grants = {}
-    for grant in tak.grants_snapshot(schema.split(".")[1]):
-        logger.info(f"grant:\n{grant}\n  hive: {grant.hive_grant_sql()}\n  uc: {grant.uc_grant_sql()}")
+    for grant in tak.grants_snapshot():
+        logging.info(f"grant:\n{grant}\n  hive: {grant.hive_grant_sql()}\n  uc: {grant.uc_grant_sql()}")
         all_grants[f"{grant.principal}.{grant.object_key}"] = grant.action_type
 
-    assert len(all_grants) >= 2, "must have at least two grants"
-    assert all_grants[f"{group_a.display_name}.{table}"] == "SELECT"
-    assert all_grants[f"{group_b.display_name}.{schema}"] == "MODIFY"
+    assert len(all_grants) >= 3, "must have at least three grants"
+    assert all_grants[f"{group_a.display_name}.{table_a}"] == "SELECT"
+    assert all_grants[f"{group_b.display_name}.{table_b}"] == "SELECT"
+    assert all_grants[f"{group_b.display_name}.{schema_b}"] == "MODIFY"

tests/unit/test_config.py

@@ -3,14 +3,14 @@
 from functools import partial
 from pathlib import Path
 
-import pytest
 import yaml
 
 from databricks.labs.ucx.config import (
     GroupsConfig,
     InventoryConfig,
     InventoryTable,
     MigrationConfig,
+    TaclConfig,
 )
 
 
@@ -19,12 +19,9 @@ def test_initialization():
         MigrationConfig,
         inventory=InventoryConfig(table=InventoryTable(catalog="catalog", database="database", name="name")),
         groups=GroupsConfig(auto=True),
+        tacl=TaclConfig(databases=["default"]),
     )
-
-    with pytest.raises(NotImplementedError):
-        mc(with_table_acls=True)
-
-    mc(with_table_acls=False)
+    mc()
 
 
 # path context manager
@@ -54,9 +51,10 @@ def test_reader(tmp_path: Path):
             MigrationConfig,
             inventory=InventoryConfig(table=InventoryTable(catalog="catalog", database="database", name="name")),
             groups=GroupsConfig(auto=True),
+            tacl=TaclConfig(databases=["default"]),
         )
 
-        config: MigrationConfig = mc(with_table_acls=False)
+        config: MigrationConfig = mc()
         config_file = tmp_path / "config.yml"
 
         as_dict = config.as_dict()

tests/unit/test_crawler_base.py

@@ -13,6 +13,13 @@ class Foo:
     second: bool
 
 
+@dataclass
+class Bar:
+    first: str
+    second: bool
+    third: float
+
+
 def test_invalid():
     with pytest.raises(ValueError):
         CrawlerBase(MockBackend(), "a.a.a", "b", "c")
@@ -42,6 +49,22 @@ def fetcher():
     assert insert == b.queries[0]
 
 
+def test_snapshot_appends_incorrect_type():
+    b = MockBackend()
+    cb = CrawlerBase(b, "a", "b", "c")
+    runs = []
+
+    def fetcher():
+        if len(runs) == 0:
+            runs.append(1)
+            msg = "TABLE_OR_VIEW_NOT_FOUND"
+            raise RuntimeError(msg)
+        return []
+
+    with pytest.raises(ValueError):
+        cb._snapshot(Bar, fetcher=fetcher, loader=lambda: [Bar(first="first", second=True, third=3.14)])
+
+
 def test_snapshot_appends_to_new_table():
     b = MockBackend(fails_on_first={"INSERT INTO a.b.c": "TABLE_OR_VIEW_NOT_FOUND ..."})
     cb = CrawlerBase(b, "a", "b", "c")
@@ -62,3 +85,19 @@ def fetcher():
     assert insert == b.queries[0]
     assert create == b.queries[1]
     assert insert == b.queries[2]
+
+
+def test_snapshot_wrong_error():
+    b = MockBackend(fails_on_first={"INSERT INTO a.b.c": "TABLE_NOT_FOUND ..."})
+    cb = CrawlerBase(b, "a", "b", "c")
+    runs = []
+
+    def fetcher():
+        if len(runs) == 0:
+            runs.append(1)
+            msg = "TABLE_OR_VIEW_NOT_FOUND"
+            raise RuntimeError(msg)
+        return []
+
+    with pytest.raises(RuntimeError):
+        cb._snapshot(Foo, fetcher=fetcher, loader=lambda: [Foo(first="first", second=True)])