Improved automation for MANAGED table migration and continued building tables migration component (#295)

william-conti · nfx · web-flow · commit 2e8a880dc197 · 2023-09-28T17:07:58.000Z
Fixes #106 --------- Co-authored-by: Serge Smertin <259697+nfx@users.noreply.github.com>
diff --git a/src/databricks/labs/ucx/config.py b/src/databricks/labs/ucx/config.py
@@ -205,6 +205,8 @@ class WorkspaceConfig(_Config["WorkspaceConfig"]):
     connect: ConnectConfig | None = None
     num_threads: int | None = 10
     log_level: str | None = "INFO"
+    database_to_catalog_mapping: dict[str, str] = None
+    default_catalog: str = "ucx_default"
 
     # Starting path for notebooks and directories crawler
     workspace_start_path: str = "/"
@@ -220,6 +222,8 @@ def from_dict(cls, raw: dict):
             warehouse_id=raw.get("warehouse_id", None),
             num_threads=raw.get("num_threads", 10),
             log_level=raw.get("log_level", "INFO"),
+            database_to_catalog_mapping=raw.get("database_to_catalog_mapping", None),
+            default_catalog=raw.get("default_catalog", "main"),
         )
 
     def to_workspace_client(self) -> WorkspaceClient:
diff --git a/src/databricks/labs/ucx/hive_metastore/tables.py b/src/databricks/labs/ucx/hive_metastore/tables.py
@@ -3,6 +3,8 @@
 from dataclasses import dataclass
 from functools import partial
 
+from databricks.sdk import WorkspaceClient
+
 from databricks.labs.ucx.framework.crawlers import CrawlerBase, SqlBackend
 from databricks.labs.ucx.framework.parallel import ThreadedExecution
 from databricks.labs.ucx.mixins.sql import Row
@@ -35,39 +37,38 @@ def key(self) -> str:
     def kind(self) -> str:
         return "VIEW" if self.view_text is not None else "TABLE"
 
-    def _sql_alter(self, catalog):
-        return (
-            f"ALTER {self.kind} {self.key} SET"
-            f" TBLPROPERTIES ('upgraded_to' = '{catalog}.{self.database}.{self.name}');"
-        )
-
     def _sql_external(self, catalog):
-        # TODO: https://github.com/databricks/ucx/issues/106
-        return (
-            f"CREATE TABLE IF NOT EXISTS {catalog}.{self.database}.{self.name}"
-            f" LIKE {self.key} COPY LOCATION;" + self._sql_alter(catalog)
-        )
+        return f"SYNC TABLE {catalog}.{self.database}.{self.name} FROM {self.key};"
 
     def _sql_managed(self, catalog):
         if not self.is_delta:
             msg = f"{self.key} is not DELTA: {self.table_format}"
             raise ValueError(msg)
-        return (
-            f"CREATE TABLE IF NOT EXISTS {catalog}.{self.database}.{self.name}"
-            f" DEEP CLONE {self.key};" + self._sql_alter(catalog)
-        )
+        return f"CREATE TABLE IF NOT EXISTS {catalog}.{self.database}.{self.name} DEEP CLONE {self.key};"
 
     def _sql_view(self, catalog):
         return f"CREATE VIEW IF NOT EXISTS {catalog}.{self.database}.{self.name} AS {self.view_text};"
 
     def uc_create_sql(self, catalog):
         if self.kind == "VIEW":
             return self._sql_view(catalog)
-        elif self.location is not None:
+        elif self.object_type == "EXTERNAL":
             return self._sql_external(catalog)
         else:
             return self._sql_managed(catalog)
 
+    def sql_alter_to(self, catalog):
+        return (
+            f"ALTER {self.kind} {self.key} SET"
+            f" TBLPROPERTIES ('upgraded_to' = '{catalog}.{self.database}.{self.name}');"
+        )
+
+    def sql_alter_from(self, catalog):
+        return (
+            f"ALTER {self.kind} {catalog}.{self.database}.{self.name} SET"
+            f" TBLPROPERTIES ('upgraded_from' = '{self.key}');"
+        )
+
 
 class TablesCrawler(CrawlerBase):
     def __init__(self, backend: SqlBackend, schema):
@@ -143,3 +144,52 @@ def _describe(self, catalog: str, database: str, table: str) -> Table | None:
         except Exception as e:
             logger.error(f"Couldn't fetch information for table {full_name} : {e}")
             return None
+
+
+class TablesMigrate:
+    def __init__(
+        self,
+        tc: TablesCrawler,
+        ws: WorkspaceClient,
+        backend: SqlBackend,
+        inventory_database: str,
+        default_catalog=None,
+        database_to_catalog_mapping: dict[str, str] | None = None,
+    ):
+        self._tc = tc
+        self._backend = backend
+        self._ws = ws
+        self._inventory_database = inventory_database
+        self._database_to_catalog_mapping = database_to_catalog_mapping
+        self._seen_tables = {}
+        self._default_catalog = self._init_default_catalog(default_catalog)
+
+    @staticmethod
+    def _init_default_catalog(default_catalog):
+        if default_catalog:
+            return default_catalog
+        else:
+            return "ucx_default"  # TODO : Fetch current workspace name and append it to the default catalog.
+
+    def migrate_tables(self):
+        tasks = []
+        for table in self._tc.snapshot():
+            target_catalog = self._default_catalog
+            if self._database_to_catalog_mapping:
+                target_catalog = self._database_to_catalog_mapping[table.database]
+            tasks.append(partial(self._migrate_table, target_catalog, table))
+        ThreadedExecution.gather("migrate tables", tasks)
+
+    def _migrate_table(self, target_catalog, table):
+        try:
+            sql = table.uc_create_sql(target_catalog)
+            logger.debug(f"Migrating table {table.key} to using SQL query: {sql}")
+
+            if table.object_type == "MANAGED":
+                self._backend.execute(sql)
+                self._backend.execute(table.sql_alter_to(target_catalog))
+                self._backend.execute(table.sql_alter_from(target_catalog))
+            else:
+                logger.info(f"Table {table.key} is a {table.object_type} and is not supported for migration yet ")
+        except Exception as e:
+            logger.error(f"Could not create table {table.name} because: {e}")
diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py
@@ -82,12 +82,17 @@ def test_catalog_fixture(make_catalog):
 
 @pytest.fixture
 def make_schema(sql_exec, make_random):
-    def create(*, catalog="hive_metastore"):
-        name = f"{catalog}.ucx_S{make_random(4)}".lower()
-        sql_exec(f"CREATE SCHEMA {name}")
-        return name
+    def create(*, catalog="hive_metastore", schema):
+        if schema is None:
+            schema = f"{catalog}.ucx_S{make_random(4)}".lower()
+        else:
+            schema = f"{catalog}.{schema}"
+        sql_exec(f"CREATE SCHEMA {schema}")
+        return schema
 
-    yield from factory("schema", create, lambda name: sql_exec(f"DROP SCHEMA IF EXISTS {name} CASCADE"))  # noqa: F405
+    yield from factory(  # noqa: F405
+        "schema", create, lambda schema_name: sql_exec(f"DROP SCHEMA IF EXISTS {schema_name} CASCADE")
+    )
 
 
 def test_schema_fixture(make_schema):
diff --git a/tests/integration/hive_metastore/test_migrate.py b/tests/integration/hive_metastore/test_migrate.py
@@ -0,0 +1,64 @@
+import logging
+import os
+
+import pytest
+
+from databricks.labs.ucx.framework.crawlers import StatementExecutionBackend
+from databricks.labs.ucx.hive_metastore import TablesCrawler
+from databricks.labs.ucx.hive_metastore.tables import TablesMigrate
+
+logger = logging.getLogger(__name__)
+
+
+def test_migrate_managed_tables(ws, make_catalog, make_schema, make_table):
+    target_catalog = make_catalog()
+    schema_a = make_schema(catalog="hive_metastore")
+    _, target_schema = schema_a.split(".")
+
+    make_schema(catalog=target_catalog, schema_name=target_schema)
+
+    managed_table = make_table(schema=schema_a)
+
+    logger.info(f"target catalog={target_catalog}, managed_table={managed_table}")
+
+    inventory_schema = make_schema(catalog="hive_metastore")
+    _, inventory_schema = inventory_schema.split(".")
+
+    backend = StatementExecutionBackend(ws, os.environ["TEST_DEFAULT_WAREHOUSE_ID"])
+    crawler = TablesCrawler(backend, inventory_schema)
+    tm = TablesMigrate(crawler, ws, backend, target_catalog, inventory_schema)
+    tm.migrate_tables()
+
+    target_tables = list(backend.fetch(f"SHOW TABLES IN {target_catalog}.{target_schema}"))
+    assert len(target_tables) == 1
+
+    _, _, managed_table_name = managed_table.split(".")
+    target_table_properties = ws.tables.get(f"{target_catalog}.{target_schema}.{managed_table_name}").properties
+
+    assert target_table_properties["upgraded_from"] == managed_table
+
+
+@pytest.mark.skip(reason="Needs Storage credential + External Location in place")
+def test_migrate_external_table(ws, make_catalog, make_schema, make_table):
+    target_catalog = make_catalog()
+    schema_a = make_schema(catalog="hive_metastore")
+    _, target_schema = schema_a.split(".")
+
+    make_schema(catalog=target_catalog, schema_name=target_schema)
+
+    external_table = make_table(schema=schema_a, external=True)
+
+    logger.info(f"target catalog={target_catalog}, external_table={external_table} ")
+
+    inventory_schema = make_schema(catalog="hive_metastore")
+    _, inventory_schema = inventory_schema.split(".")
+
+    backend = StatementExecutionBackend(ws, os.environ["TEST_DEFAULT_WAREHOUSE_ID"])
+
+    backend = StatementExecutionBackend(ws, os.environ["TEST_DEFAULT_WAREHOUSE_ID"])
+    crawler = TablesCrawler(backend, inventory_schema)
+    tm = TablesMigrate(crawler, ws, backend, target_catalog, inventory_schema)
+    tm.migrate_tables()
+
+    target_tables = list(backend.fetch(f"SHOW TABLES IN {target_catalog}.{target_schema}"))
+    assert len(target_tables) == 1
diff --git a/tests/unit/hive_metastore/test_migrate.py b/tests/unit/hive_metastore/test_migrate.py
@@ -0,0 +1,99 @@
+import logging
+from unittest.mock import MagicMock
+
+import pytest
+
+from databricks.labs.ucx.hive_metastore.tables import TablesCrawler, TablesMigrate
+
+from ..framework.mocks import MockBackend
+
+logger = logging.getLogger(__name__)
+
+
+def test_migrate_managed_tables_should_produce_proper_queries():
+    errors = {}
+    rows = {
+        "SELECT": [
+            (
+                "hive_metastore",
+                "db1",
+                "managed",
+                "MANAGED",
+                "DELTA",
+                None,
+                None,
+            ),
+        ]
+    }
+    backend = MockBackend(fails_on_first=errors, rows=rows)
+    tc = TablesCrawler(backend, "inventory_database")
+    client = MagicMock()
+    tm = TablesMigrate(tc, client, backend, "")
+    tm.migrate_tables()
+
+    assert (list(backend.queries)) == [
+        "SELECT * FROM hive_metastore.inventory_database.tables",
+        "CREATE TABLE IF NOT EXISTS ucx_default.db1.managed DEEP CLONE hive_metastore.db1.managed;",
+        "ALTER TABLE hive_metastore.db1.managed SET TBLPROPERTIES ('upgraded_to' = 'ucx_default.db1.managed');",
+        "ALTER TABLE ucx_default.db1.managed SET TBLPROPERTIES ('upgraded_from' = 'hive_metastore.db1.managed');",
+    ]
+
+
+@pytest.mark.skip(reason="Not implemented yet")
+def test_migrate_managed_tables_should_do_nothing_if_upgrade_tag_is_present():
+    errors = {}
+    rows = {
+        "SELECT": [
+            ("hive_metastore", "db1", "managed", "MANAGED", "DELTA", None, None, "[upgraded_to=target]"),
+        ]
+    }
+    backend = MockBackend(fails_on_first=errors, rows=rows)
+    tc = TablesCrawler(backend, "inventory_database")
+    client = MagicMock()
+    tm = TablesMigrate(tc, client, backend, "")
+    tm.migrate_tables()
+
+    assert (list(backend.queries)) == ["SELECT * FROM hive_metastore.inventory_database.tables"]
+
+
+def test_migrate_tables_should_migrate_tables_to_default_catalog_if_not_found_in_mapping():
+    errors = {}
+    rows = {
+        "SELECT": [
+            ("hive_metastore", "db1", "managed", "MANAGED", "DELTA", None, None),
+        ]
+    }
+    backend = MockBackend(fails_on_first=errors, rows=rows)
+    tc = TablesCrawler(backend, "inventory_database")
+    client = MagicMock()
+    database_to_catalog_mapping = {"db1": "catalog_1", "db2": "catalog_2"}
+    tm = TablesMigrate(tc, client, backend, "", database_to_catalog_mapping=database_to_catalog_mapping)
+    tm.migrate_tables()
+
+    assert (list(backend.queries)) == [
+        "SELECT * FROM hive_metastore.inventory_database.tables",
+        "CREATE TABLE IF NOT EXISTS catalog_1.db1.managed DEEP CLONE hive_metastore.db1.managed;",
+        "ALTER TABLE hive_metastore.db1.managed SET TBLPROPERTIES ('upgraded_to' = 'catalog_1.db1.managed');",
+        "ALTER TABLE catalog_1.db1.managed SET TBLPROPERTIES ('upgraded_from' = 'hive_metastore.db1.managed');",
+    ]
+
+
+def test_migrate_tables_should_migrate_tables_to_default_catalog_if_specified():
+    errors = {}
+    rows = {
+        "SELECT": [
+            ("hive_metastore", "db1", "managed", "MANAGED", "DELTA", None, None),
+        ]
+    }
+    backend = MockBackend(fails_on_first=errors, rows=rows)
+    tc = TablesCrawler(backend, "inventory_database")
+    client = MagicMock()
+    tm = TablesMigrate(tc, client, backend, "", default_catalog="test_catalog")
+    tm.migrate_tables()
+
+    assert (list(backend.queries)) == [
+        "SELECT * FROM hive_metastore.inventory_database.tables",
+        "CREATE TABLE IF NOT EXISTS test_catalog.db1.managed DEEP CLONE hive_metastore.db1.managed;",
+        "ALTER TABLE hive_metastore.db1.managed SET TBLPROPERTIES ('upgraded_to' = 'test_catalog.db1.managed');",
+        "ALTER TABLE test_catalog.db1.managed SET TBLPROPERTIES ('upgraded_from' = 'hive_metastore.db1.managed');",
+    ]
diff --git a/tests/unit/hive_metastore/test_tables.py b/tests/unit/hive_metastore/test_tables.py
@@ -49,9 +49,7 @@ def test_sql_managed_non_delta():
     [
         (
             Table(catalog="catalog", database="db", name="managed_table", object_type="..", table_format="DELTA"),
-            "CREATE TABLE IF NOT EXISTS new_catalog.db.managed_table DEEP CLONE "
-            "catalog.db.managed_table;ALTER TABLE catalog.db.managed_table SET "
-            "TBLPROPERTIES ('upgraded_to' = 'new_catalog.db.managed_table');",
+            "CREATE TABLE IF NOT EXISTS new_catalog.db.managed_table DEEP CLONE catalog.db.managed_table;",
         ),
         (
             Table(
@@ -69,14 +67,11 @@ def test_sql_managed_non_delta():
                 catalog="catalog",
                 database="db",
                 name="external_table",
-                object_type="..",
+                object_type="EXTERNAL",
                 table_format="DELTA",
                 location="s3a://foo/bar",
             ),
-            "CREATE TABLE IF NOT EXISTS new_catalog.db.external_table LIKE "
-            "catalog.db.external_table COPY LOCATION;ALTER TABLE "
-            "catalog.db.external_table SET TBLPROPERTIES ('upgraded_to' = "
-            "'new_catalog.db.external_table');",
+            "SYNC TABLE new_catalog.db.external_table FROM catalog.db.external_table;",
         ),
     ],
 )
diff --git a/tests/unit/test_install.py b/tests/unit/test_install.py
@@ -49,7 +49,8 @@ def not_found(_):
 
     ws.workspace.upload.assert_called_with(
         "/Users/me@example.com/.ucx/config.yml",
-        b"""groups:
+        b"""default_catalog: ucx_default
+groups:
   backup_group_prefix: '42'
   selected:
   - '42'
@@ -106,7 +107,8 @@ def mock_question(text: str, *, default: str | None = None) -> str:
 
     ws.workspace.upload.assert_called_with(
         "/Users/me@example.com/.ucx/config.yml",
-        b"""groups:
+        b"""default_catalog: ucx_default
+groups:
   auto: true
   backup_group_prefix: '42'
 inventory_database: '42'
@@ -146,7 +148,8 @@ def mock_question(text: str, *, default: str | None = None) -> str:
 
     ws.workspace.upload.assert_called_with(
         "/Users/me@example.com/.ucx/config.yml",
-        b"""groups:
+        b"""default_catalog: ucx_default
+groups:
   backup_group_prefix: '42'
   selected:
   - g1
