Added functionality to determine migration method based on DBFS Root (#759)

Author: FastLee

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

@@ -134,6 +134,9 @@ def get_tables_to_migrate(self, tables_crawler: TablesCrawler):
             if rule.src_schema not in databases_in_scope:
                 logger.info(f"Table {rule.as_hms_table_key} is in a database that was marked to be skipped")
                 continue
+            if crawled_tables_keys[rule.as_hms_table_key].is_databricks_dataset:
+                logger.info(f"Table {rule.as_hms_table_key} is a db demo dataset and will not be upgraded")
+                continue
             tasks.append(
                 partial(self._get_table_in_scope_task, TableToMigrate(crawled_tables_keys[rule.as_hms_table_key], rule))
             )

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

@@ -39,24 +39,25 @@ def _migrate_table(self, src_table: Table, rule: Rule):
         if self._table_already_upgraded(rule.as_uc_table_key):
             logger.info(f"Table {src_table.key} already upgraded to {rule.as_uc_table_key}")
             return True
-        if src_table.object_type == "MANAGED":
-            return self._migrate_managed_table(src_table, rule)
+        if src_table.kind == "TABLE" and src_table.table_format == "DELTA" and src_table.is_dbfs_root:
+            return self._migrate_dbfs_root_table(src_table, rule)
+        if src_table.kind == "TABLE" and src_table.is_format_supported_for_sync:
+            return self._migrate_external_table(src_table, rule)
         if src_table.kind == "VIEW":
             return self._migrate_view(src_table, rule)
-        if src_table.object_type == "EXTERNAL":
-            return self._migrate_external_table(src_table, rule)
+        logger.info(f"Table {src_table.key} is not supported for migration")
         return True
 
     def _migrate_external_table(self, src_table: Table, rule: Rule):
         target_table_key = rule.as_uc_table_key
-        table_migrate_sql = src_table.uc_create_sql(target_table_key)
+        table_migrate_sql = src_table.sql_migrate_external(target_table_key)
         logger.debug(f"Migrating external table {src_table.key} to using SQL query: {table_migrate_sql}")
         self._backend.execute(table_migrate_sql)
         return True
 
-    def _migrate_managed_table(self, src_table: Table, rule: Rule):
+    def _migrate_dbfs_root_table(self, src_table: Table, rule: Rule):
         target_table_key = rule.as_uc_table_key
-        table_migrate_sql = src_table.uc_create_sql(target_table_key)
+        table_migrate_sql = src_table.sql_migrate_dbfs(target_table_key)
         logger.debug(f"Migrating managed table {src_table.key} to using SQL query: {table_migrate_sql}")
         self._backend.execute(table_migrate_sql)
         self._backend.execute(src_table.sql_alter_to(rule.as_uc_table_key))
@@ -65,16 +66,13 @@ def _migrate_managed_table(self, src_table: Table, rule: Rule):
 
     def _migrate_view(self, src_table: Table, rule: Rule):
         target_table_key = rule.as_uc_table_key
-        table_migrate_sql = src_table.uc_create_sql(target_table_key)
+        table_migrate_sql = src_table.sql_migrate_view(target_table_key)
         logger.debug(f"Migrating view {src_table.key} to using SQL query: {table_migrate_sql}")
         self._backend.execute(table_migrate_sql)
         self._backend.execute(src_table.sql_alter_to(rule.as_uc_table_key))
         self._backend.execute(src_table.sql_alter_from(rule.as_uc_table_key))
         return True
 
-        msg = f"Table {src_table.key} is a {src_table.object_type} and is not supported for migration yet"
-        logger.info(msg)
-
     def _init_seen_tables(self):
         for catalog in self._ws.catalogs.list():
             for schema in self._ws.schemas.list(catalog_name=catalog.name):

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

@@ -40,7 +40,7 @@ def _crawl(self) -> Iterable[TableSize]:
         for table in self._tables_crawler.snapshot():
             if not table.kind == "TABLE":
                 continue
-            if not table.is_dbfs_root():
+            if not table.is_dbfs_root:
                 continue
             size_in_bytes = self.get_table_size(table.key)
             yield TableSize(

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

@@ -35,6 +35,9 @@ class Table:
     DBFS_ROOT_PREFIX_EXCEPTIONS: typing.ClassVar[list[str]] = [
         "/dbfs/mnt",
         "dbfs:/mnt",
+    ]
+
+    DBFS_DATABRICKS_DATASETS_PREFIXES: typing.ClassVar[list[str]] = [
         "/dbfs/databricks-datasets",
         "dbfs:/databricks-datasets",
     ]
@@ -53,14 +56,6 @@ def key(self) -> str:
     def kind(self) -> str:
         return "VIEW" if self.view_text is not None else "TABLE"
 
-    def uc_create_sql(self, target_table_key):
-        if self.kind == "VIEW":
-            return self._sql_migrate_view(target_table_key)
-        elif self.object_type == "EXTERNAL":
-            return self._sql_migrate_external(target_table_key)
-        else:
-            return self._sql_migrate_managed(target_table_key)
-
     def sql_alter_to(self, target_table_key):
         return f"ALTER {self.kind} {self.key} SET TBLPROPERTIES ('upgraded_to' = '{target_table_key}');"
 
@@ -70,27 +65,46 @@ def sql_alter_from(self, target_table_key):
     def sql_unset_upgraded_to(self):
         return f"ALTER {self.kind} {self.key} UNSET TBLPROPERTIES IF EXISTS('upgraded_to');"
 
+    @property
     def is_dbfs_root(self) -> bool:
         if not self.location:
             return False
-        for exception in self.DBFS_ROOT_PREFIX_EXCEPTIONS:
-            if self.location.startswith(exception):
-                return False
         for prefix in self.DBFS_ROOT_PREFIXES:
             if self.location.startswith(prefix):
+                for exception in self.DBFS_ROOT_PREFIX_EXCEPTIONS:
+                    if self.location.startswith(exception):
+                        return False
+                for db_datasets in self.DBFS_DATABRICKS_DATASETS_PREFIXES:
+                    if self.location.startswith(db_datasets):
+                        return False
+                return True
+        return False
+
+    @property
+    def is_format_supported_for_sync(self) -> bool:
+        if self.table_format is None:
+            return False
+        return self.table_format.upper() in ("DELTA", "PARQUET", "CSV", "JSON", "ORC", "TEXT")
+
+    @property
+    def is_databricks_dataset(self) -> bool:
+        if not self.location:
+            return False
+        for db_datasets in self.DBFS_DATABRICKS_DATASETS_PREFIXES:
+            if self.location.startswith(db_datasets):
                 return True
         return False
 
-    def _sql_migrate_external(self, target_table_key):
+    def sql_migrate_external(self, target_table_key):
         return f"SYNC TABLE {target_table_key} FROM {self.key};"
 
-    def _sql_migrate_managed(self, target_table_key):
+    def sql_migrate_dbfs(self, target_table_key):
         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 {target_table_key} DEEP CLONE {self.key};"
 
-    def _sql_migrate_view(self, target_table_key):
+    def sql_migrate_view(self, target_table_key):
         return f"CREATE VIEW IF NOT EXISTS {target_table_key} AS {self.view_text};"
 
 

tests/integration/hive_metastore/test_migrate.py

@@ -198,6 +198,9 @@ def test_mapping_skips_tables_databases(ws, sql_backend, inventory_schema, make_
     src_schema1 = make_schema(catalog_name="hive_metastore")
     src_schema2 = make_schema(catalog_name="hive_metastore")
     table_to_migrate = make_table(schema_name=src_schema1.name)
+    table_databricks_dataset = make_table(
+        schema_name=src_schema1.name, external_csv="dbfs:/databricks-datasets/adult/adult.data"
+    )
     table_to_skip = make_table(schema_name=src_schema1.name)
     table_in_skipped_database = make_table(schema_name=src_schema2.name)
     all_tables = [table_to_migrate, table_to_skip, table_in_skipped_database]
@@ -224,6 +227,14 @@ def test_mapping_skips_tables_databases(ws, sql_backend, inventory_schema, make_
             table_to_skip.name,
             table_to_skip.name,
         ),
+        Rule(
+            "workspace",
+            dst_catalog.name,
+            src_schema1.name,
+            dst_schema1.name,
+            table_databricks_dataset.name,
+            table_databricks_dataset.name,
+        ),
         Rule(
             "workspace",
             dst_catalog.name,

tests/unit/hive_metastore/test_table_migrate.py

@@ -23,7 +23,7 @@
 logger = logging.getLogger(__name__)
 
 
-def test_migrate_managed_tables_should_produce_proper_queries():
+def test_migrate_dbfs_root_tables_should_produce_proper_queries():
     errors = {}
     rows = {}
     backend = MockBackend(fails_on_first=errors, rows=rows)
@@ -32,20 +32,38 @@ def test_migrate_managed_tables_should_produce_proper_queries():
     table_mapping = create_autospec(TableMapping)
     table_mapping.get_tables_to_migrate.return_value = [
         TableToMigrate(
-            Table("hive_metastore", "db1_src", "managed_src", "MANAGED", "DELTA"),
-            Rule("workspace", "ucx_default", "db1_src", "db1_dst", "managed_src", "managed_dst"),
-        )
+            Table("hive_metastore", "db1_src", "managed_dbfs", "MANAGED", "DELTA", "dbfs:/some_location"),
+            Rule("workspace", "ucx_default", "db1_src", "db1_dst", "managed_dbfs", "managed_dbfs"),
+        ),
+        TableToMigrate(
+            Table("hive_metastore", "db1_src", "managed_mnt", "MANAGED", "DELTA", "s3:/mnt/location"),
+            Rule("workspace", "ucx_default", "db1_src", "db1_dst", "managed_mnt", "managed_mnt"),
+        ),
+        TableToMigrate(
+            Table("hive_metastore", "db1_src", "managed_other", "MANAGED", "DELTA", "s3:/location"),
+            Rule("workspace", "ucx_default", "db1_src", "db1_dst", "managed_other", "managed_other"),
+        ),
     ]
     table_migrate = TablesMigrate(table_crawler, client, backend, table_mapping)
     table_migrate.migrate_tables()
 
-    assert (list(backend.queries)) == [
-        "CREATE TABLE IF NOT EXISTS ucx_default.db1_dst.managed_dst DEEP CLONE hive_metastore.db1_src.managed_src;",
-        "ALTER TABLE hive_metastore.db1_src.managed_src "
-        "SET TBLPROPERTIES ('upgraded_to' = 'ucx_default.db1_dst.managed_dst');",
-        "ALTER TABLE ucx_default.db1_dst.managed_dst "
-        "SET TBLPROPERTIES ('upgraded_from' = 'hive_metastore.db1_src.managed_src');",
-    ]
+    assert (
+        "CREATE TABLE IF NOT EXISTS ucx_default.db1_dst.managed_dbfs DEEP CLONE hive_metastore.db1_src.managed_dbfs;"
+    ) in list(backend.queries)
+    assert "SYNC TABLE ucx_default.db1_dst.managed_mnt FROM hive_metastore.db1_src.managed_mnt;" in list(
+        backend.queries
+    )
+    assert (
+        "ALTER TABLE hive_metastore.db1_src.managed_dbfs "
+        "SET TBLPROPERTIES ('upgraded_to' = 'ucx_default.db1_dst.managed_dbfs');"
+    ) in list(backend.queries)
+    assert (
+        "ALTER TABLE ucx_default.db1_dst.managed_dbfs "
+        "SET TBLPROPERTIES ('upgraded_from' = 'hive_metastore.db1_src.managed_dbfs');"
+    ) in list(backend.queries)
+    assert "SYNC TABLE ucx_default.db1_dst.managed_other FROM hive_metastore.db1_src.managed_other;" in list(
+        backend.queries
+    )
 
 
 def test_migrate_external_tables_should_produce_proper_queries():
@@ -85,13 +103,15 @@ def test_migrate_view_should_produce_proper_queries():
     table_migrate = TablesMigrate(table_crawler, client, backend, table_mapping)
     table_migrate.migrate_tables()
 
-    assert (list(backend.queries)) == [
-        "CREATE VIEW IF NOT EXISTS ucx_default.db1_dst.view_dst AS SELECT * FROM table;",
+    assert "CREATE VIEW IF NOT EXISTS ucx_default.db1_dst.view_dst AS SELECT * FROM table;" in list(backend.queries)
+    assert (
         "ALTER VIEW hive_metastore.db1_src.view_src "
-        "SET TBLPROPERTIES ('upgraded_to' = 'ucx_default.db1_dst.view_dst');",
+        "SET TBLPROPERTIES ('upgraded_to' = 'ucx_default.db1_dst.view_dst');"
+    ) in list(backend.queries)
+    assert (
         "ALTER VIEW ucx_default.db1_dst.view_dst "
-        "SET TBLPROPERTIES ('upgraded_from' = 'hive_metastore.db1_src.view_src');",
-    ]
+        "SET TBLPROPERTIES ('upgraded_from' = 'hive_metastore.db1_src.view_src');"
+    ) in list(backend.queries)
 
 
 def get_table_migrate(backend: SqlBackend) -> TablesMigrate:

tests/unit/hive_metastore/test_tables.py

@@ -41,23 +41,42 @@ def test_sql_managed_non_delta():
     with pytest.raises(ValueError):
         Table(
             catalog="catalog", database="db", name="table", object_type="type", table_format="PARQUET"
-        )._sql_migrate_managed("catalog")
+        ).sql_migrate_dbfs("catalog")
 
 
 @pytest.mark.parametrize(
     "table,target,query",
     [
         (
-            Table(catalog="catalog", database="db", name="managed_table", object_type="..", table_format="DELTA"),
+            Table(
+                catalog="catalog",
+                database="db",
+                name="managed_table",
+                object_type="MANAGED",
+                table_format="DELTA",
+                location="dbfs:/location/table",
+            ),
             "new_catalog.db.managed_table",
             "CREATE TABLE IF NOT EXISTS new_catalog.db.managed_table DEEP CLONE catalog.db.managed_table;",
         ),
+        (
+            Table(
+                catalog="catalog",
+                database="db",
+                name="managed_table",
+                object_type="MANAGED",
+                table_format="DELTA",
+                location="dbfs:/mnt/location/table",
+            ),
+            "new_catalog.db.managed_table",
+            "SYNC TABLE new_catalog.db.managed_table FROM catalog.db.managed_table;",
+        ),
         (
             Table(
                 catalog="catalog",
                 database="db",
                 name="view",
-                object_type="..",
+                object_type="VIEW",
                 table_format="DELTA",
                 view_text="SELECT * FROM table",
             ),
@@ -79,7 +98,12 @@ def test_sql_managed_non_delta():
     ],
 )
 def test_uc_sql(table, target, query):
-    assert table.uc_create_sql(target) == query
+    if table.kind == "VIEW":
+        assert table.sql_migrate_view(target) == query
+    if table.kind == "TABLE" and table.is_dbfs_root:
+        assert table.sql_migrate_dbfs(target) == query
+    if table.kind == "TABLE" and not table.is_dbfs_root:
+        assert table.sql_migrate_external(target) == query
 
 
 def test_tables_crawler_inventory_table():
@@ -113,15 +137,51 @@ def test_tables_returning_error_when_describing():
 
 
 def test_is_dbfs_root():
-    assert Table("a", "b", "c", "MANAGED", "DELTA", location="dbfs:/somelocation/tablename").is_dbfs_root()
-    assert Table("a", "b", "c", "MANAGED", "DELTA", location="/dbfs/somelocation/tablename").is_dbfs_root()
-    assert not Table("a", "b", "c", "MANAGED", "DELTA", location="dbfs:/mnt/somelocation/tablename").is_dbfs_root()
-    assert not Table("a", "b", "c", "MANAGED", "DELTA", location="/dbfs/mnt/somelocation/tablename").is_dbfs_root()
+    assert Table("a", "b", "c", "MANAGED", "DELTA", location="dbfs:/somelocation/tablename").is_dbfs_root
+    assert Table("a", "b", "c", "MANAGED", "DELTA", location="/dbfs/somelocation/tablename").is_dbfs_root
+    assert not Table("a", "b", "c", "MANAGED", "DELTA", location="dbfs:/mnt/somelocation/tablename").is_dbfs_root
+    assert not Table("a", "b", "c", "MANAGED", "DELTA", location="/dbfs/mnt/somelocation/tablename").is_dbfs_root
     assert not Table(
         "a", "b", "c", "MANAGED", "DELTA", location="dbfs:/databricks-datasets/somelocation/tablename"
-    ).is_dbfs_root()
+    ).is_dbfs_root
+    assert not Table(
+        "a", "b", "c", "MANAGED", "DELTA", location="/dbfs/databricks-datasets/somelocation/tablename"
+    ).is_dbfs_root
+    assert not Table("a", "b", "c", "MANAGED", "DELTA", location="s3:/somelocation/tablename").is_dbfs_root
+    assert not Table("a", "b", "c", "MANAGED", "DELTA", location="adls:/somelocation/tablename").is_dbfs_root
+
+
+def test_is_db_dataset():
+    assert not Table("a", "b", "c", "MANAGED", "DELTA", location="dbfs:/somelocation/tablename").is_databricks_dataset
+    assert not Table("a", "b", "c", "MANAGED", "DELTA", location="/dbfs/somelocation/tablename").is_databricks_dataset
+    assert not Table(
+        "a", "b", "c", "MANAGED", "DELTA", location="dbfs:/mnt/somelocation/tablename"
+    ).is_databricks_dataset
     assert not Table(
+        "a", "b", "c", "MANAGED", "DELTA", location="/dbfs/mnt/somelocation/tablename"
+    ).is_databricks_dataset
+    assert Table(
+        "a", "b", "c", "MANAGED", "DELTA", location="dbfs:/databricks-datasets/somelocation/tablename"
+    ).is_databricks_dataset
+    assert Table(
         "a", "b", "c", "MANAGED", "DELTA", location="/dbfs/databricks-datasets/somelocation/tablename"
-    ).is_dbfs_root()
-    assert not Table("a", "b", "c", "MANAGED", "DELTA", location="s3:/somelocation/tablename").is_dbfs_root()
-    assert not Table("a", "b", "c", "MANAGED", "DELTA", location="adls:/somelocation/tablename").is_dbfs_root()
+    ).is_databricks_dataset
+    assert not Table("a", "b", "c", "MANAGED", "DELTA", location="s3:/somelocation/tablename").is_databricks_dataset
+    assert not Table("a", "b", "c", "MANAGED", "DELTA", location="adls:/somelocation/tablename").is_databricks_dataset
+
+
+def test_is_supported_for_sync():
+    assert Table(
+        "a", "b", "c", "EXTERNAL", "DELTA", location="dbfs:/somelocation/tablename"
+    ).is_format_supported_for_sync
+    assert Table("a", "b", "c", "EXTERNAL", "CSV", location="dbfs:/somelocation/tablename").is_format_supported_for_sync
+    assert Table(
+        "a", "b", "c", "EXTERNAL", "TEXT", location="dbfs:/somelocation/tablename"
+    ).is_format_supported_for_sync
+    assert Table("a", "b", "c", "EXTERNAL", "ORC", location="dbfs:/somelocation/tablename").is_format_supported_for_sync
+    assert Table(
+        "a", "b", "c", "EXTERNAL", "JSON", location="dbfs:/somelocation/tablename"
+    ).is_format_supported_for_sync
+    assert not (
+        Table("a", "b", "c", "EXTERNAL", "AVRO", location="dbfs:/somelocation/tablename").is_format_supported_for_sync
+    )