Added databricks labs ucx skip --schema ... --table ... command to mark table/schema for skipping in the table migration process (#680)

closes #672

Author: FastLee

labs.yml

@@ -26,6 +26,14 @@ commands:
       {{range .}}{{.user_name}}\t{{.database}}\t{{.warehouse_id}}
       {{end}}
 
+  - name: skip
+    description: Create a skip comment on a schema or a table
+    flags:
+      - name: schema
+        description: Schema Name to Skip
+      - name: table
+        description: Table Name to Skip
+
   - name: sync-workspace-info
     is_account_level: true
     description: upload workspace config to all workspaces in the account where ucx is installed

src/databricks/labs/ucx/cli.py

@@ -4,6 +4,7 @@
 import webbrowser
 
 from databricks.sdk import WorkspaceClient
+from databricks.sdk.errors import NotFound
 
 from databricks.labs.ucx.account import AccountWorkspaces, WorkspaceInfo
 from databricks.labs.ucx.config import AccountConfig, ConnectConfig
@@ -40,6 +41,31 @@ def list_installations():
     print(json.dumps(all_users))
 
 
+def skip(schema: str, table: str | None = None):
+    logger.info("Running skip command")
+    if not schema:
+        logger.error("--Schema is a required parameter.")
+        return None
+    ws = WorkspaceClient()
+    installation_manager = WorkspaceInstaller(ws)
+    logger.info("Fetching installation config.")
+    try:
+        warehouse_id = installation_manager._current_config.warehouse_id
+        sql_backend = StatementExecutionBackend(ws, warehouse_id)
+    except NotFound:
+        logger.error(
+            "Couldn't find UCX configuration in the user's home folder. "
+            "Make sure the current user has configured and installed UCX."
+        )
+        return None
+
+    mapping = TableMapping(ws)
+    if table:
+        mapping.skip_table(sql_backend, schema, table)
+    else:
+        mapping.skip_schema(sql_backend, schema)
+
+
 def sync_workspace_info():
     workspaces = AccountWorkspaces(AccountConfig(connect=ConnectConfig()))
     workspaces.sync_workspace_info()
@@ -71,6 +97,7 @@ def create_table_mapping():
     "sync-workspace-info": sync_workspace_info,
     "manual-workspace-info": manual_workspace_info,
     "create-table-mapping": create_table_mapping,
+    "skip": skip,
 }
 
 
@@ -86,7 +113,6 @@ def main(raw):
         log_level = "info"
     databricks_logger = logging.getLogger("databricks")
     databricks_logger.setLevel(log_level.upper())
-
     kwargs = {k.replace("-", "_"): v for k, v in flags.items()}
     MAPPING[command](**kwargs)
 

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

@@ -1,17 +1,21 @@
 import csv
 import dataclasses
 import io
+import logging
 import re
 from dataclasses import dataclass
 
 from databricks.sdk import WorkspaceClient
-from databricks.sdk.errors import NotFound
+from databricks.sdk.errors import BadRequest, NotFound
 from databricks.sdk.service.workspace import ImportFormat
 
 from databricks.labs.ucx.account import WorkspaceInfo
+from databricks.labs.ucx.framework.crawlers import StatementExecutionBackend
 from databricks.labs.ucx.hive_metastore import TablesCrawler
 from databricks.labs.ucx.hive_metastore.tables import Table
 
+logger = logging.getLogger(__name__)
+
 
 @dataclass
 class Rule:
@@ -35,6 +39,8 @@ def initial(cls, workspace_name: str, catalog_name: str, table: Table) -> "Rule"
 
 
 class TableMapping:
+    UCX_SKIP_PROPERTY = "databricks.labs.ucx.skip"
+
     def __init__(self, ws: WorkspaceClient, folder: str | None = None):
         if not folder:
             folder = f"/Users/{ws.current_user.me().user_name}/.ucx"
@@ -76,3 +82,27 @@ def load(self) -> list[Rule]:
         except NotFound:
             msg = "Please run: databricks labs ucx table-mapping"
             raise ValueError(msg) from None
+
+    def skip_table(self, backend: StatementExecutionBackend, schema: str, table: str):
+        # Marks a table to be skipped in the migration process by applying a table property
+        try:
+            backend.execute(f"ALTER TABLE `{schema}`.`{table}` SET TBLPROPERTIES('{self.UCX_SKIP_PROPERTY}' = true)")
+        except NotFound as nf:
+            if "[TABLE_OR_VIEW_NOT_FOUND]" in str(nf):
+                logger.error(f"Failed to apply skip marker for Table {schema}.{table}. Table not found.")
+            else:
+                logger.error(nf)
+        except BadRequest as br:
+            logger.error(br)
+
+    def skip_schema(self, backend: StatementExecutionBackend, schema: str):
+        # Marks a schema to be skipped in the migration process by applying a table property
+        try:
+            backend.execute(f"ALTER SCHEMA `{schema}` SET DBPROPERTIES('{self.UCX_SKIP_PROPERTY}' = true)")
+        except NotFound as nf:
+            if "[SCHEMA_NOT_FOUND]" in str(nf):
+                logger.error(f"Failed to apply skip marker for Schema {schema}. Schema not found.")
+            else:
+                logger.error(nf)
+        except BadRequest as br:
+            logger.error(br)

tests/unit/hive_metastore/test_mapping.py

@@ -99,3 +99,35 @@ def test_load_mapping():
             dst_table="bar",
         )
     ] == rules
+
+
+def test_skip_happy_path(mocker, caplog):
+    ws = mocker.patch("databricks.sdk.WorkspaceClient.__init__")
+    sbe = mocker.patch("databricks.labs.ucx.framework.crawlers.StatementExecutionBackend.__init__")
+    mapping = TableMapping(ws)
+    mapping.skip_table(sbe, schema="schema", table="table")
+    sbe.execute.assert_called_with(
+        f"ALTER TABLE `schema`.`table` SET TBLPROPERTIES('{mapping.UCX_SKIP_PROPERTY}' = true)"
+    )
+    assert len(caplog.records) == 0
+    mapping.skip_schema(sbe, schema="schema")
+    sbe.execute.assert_called_with(f"ALTER SCHEMA `schema` SET DBPROPERTIES('{mapping.UCX_SKIP_PROPERTY}' = true)")
+    assert len(caplog.records) == 0
+
+
+def test_skip_missing_schema(mocker, caplog):
+    ws = mocker.patch("databricks.sdk.WorkspaceClient.__init__")
+    sbe = mocker.patch("databricks.labs.ucx.framework.crawlers.StatementExecutionBackend.__init__")
+    sbe.execute.side_effect = NotFound("[SCHEMA_NOT_FOUND]")
+    mapping = TableMapping(ws)
+    mapping.skip_schema(sbe, schema="schema")
+    assert [rec.message for rec in caplog.records if "schema not found" in rec.message.lower()]
+
+
+def test_skip_missing_table(mocker, caplog):
+    ws = mocker.patch("databricks.sdk.WorkspaceClient.__init__")
+    sbe = mocker.patch("databricks.labs.ucx.framework.crawlers.StatementExecutionBackend.__init__")
+    sbe.execute.side_effect = NotFound("[TABLE_OR_VIEW_NOT_FOUND]")
+    mapping = TableMapping(ws)
+    mapping.skip_table(sbe, schema="schema", table="table")
+    assert [rec.message for rec in caplog.records if "table not found" in rec.message.lower()]

tests/unit/test_cli.py

@@ -0,0 +1,34 @@
+import pytest
+from databricks.sdk.errors import NotFound
+from databricks.sdk.service import iam
+from databricks.sdk.service.iam import User
+
+from databricks.labs.ucx.cli import skip
+
+
+@pytest.fixture
+def ws(mocker):
+    ws = mocker.patch("databricks.sdk.WorkspaceClient.__init__")
+    ws.current_user.me = lambda: iam.User(user_name="[email protected]", groups=[iam.ComplexValue(display="admins")])
+    ws.return_value = None
+
+
+def test_skip_no_schema(mocker, caplog):
+    ws = mocker.patch("databricks.sdk.WorkspaceClient.__init__")
+    ws.users.list.return_value = [User(user_name="foo")]
+    ws.workspace.download.side_effect = NotFound(...)
+    skip(schema=None, table="table")
+    assert [rec.message for rec in caplog.records if "schema" in rec.message.lower()]
+
+
+def test_skip_no_ucx(caplog, mocker):
+    mocker.patch("databricks.sdk.WorkspaceClient.__init__", return_value=None)
+    mocker.patch("databricks.labs.ucx.install.WorkspaceInstaller.__init__", return_value=None)
+    mocker.patch("databricks.labs.ucx.install.WorkspaceInstaller._current_config", return_value="foo")
+    mocker.patch(
+        "databricks.labs.ucx.framework.crawlers.StatementExecutionBackend.__init__",
+        return_value=None,
+        side_effect=NotFound("..."),
+    )
+    skip(schema="schema", table="table")
+    assert [rec.message for rec in caplog.records if "UCX configuration" in rec.message]