Fixed escaping SQL object names (#836)

Author: larsgeorge-db

src/databricks/labs/ucx/framework/utils.py

@@ -0,0 +1,25 @@
+import string
+
+_allowed_object_chars = set(string.ascii_letters + string.digits + '_')
+
+
+def escape_sql_identifier(path: str, optional: bool | None = True) -> str:
+    """
+    Escapes the path components to make them SQL safe.
+
+    Args:
+        path (str): The dot-separated path of a catalog object.
+        optional (bool): If `True` then do not escape if no special characters are present.
+
+    Returns:
+         str: The path with all parts escaped in backticks.
+    """
+    parts = path.split(".", maxsplit=2)
+    escaped = []
+    for part in parts:
+        if not part.startswith("`") and not part.endswith("`"):
+            part = part.strip("`")
+            if not optional or not set(part) <= _allowed_object_chars:
+                part = f"`{part}`"
+        escaped.append(part)
+    return ".".join(escaped)

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

@@ -8,6 +8,7 @@
 from databricks.sdk.service.catalog import SchemaInfo, TableInfo
 
 from databricks.labs.ucx.framework.crawlers import CrawlerBase
+from databricks.labs.ucx.framework.utils import escape_sql_identifier
 from databricks.labs.ucx.hive_metastore.tables import TablesCrawler
 from databricks.labs.ucx.hive_metastore.udfs import UdfsCrawler
 
@@ -95,13 +96,13 @@ def hive_grant_sql(self) -> list[str]:
 
     def hive_revoke_sql(self) -> str:
         object_type, object_key = self.this_type_and_key()
-        return f"REVOKE {self.action_type} ON {object_type} {object_key} FROM `{self.principal}`"
+        return f"REVOKE {self.action_type} ON {object_type} {escape_sql_identifier(object_key)} FROM `{self.principal}`"
 
     def _set_owner_sql(self, object_type, object_key):
-        return f"ALTER {object_type} {object_key} OWNER TO `{self.principal}`"
+        return f"ALTER {object_type} {escape_sql_identifier(object_key)} OWNER TO `{self.principal}`"
 
     def _apply_grant_sql(self, action_type, object_type, object_key):
-        return f"GRANT {action_type} ON {object_type} {object_key} TO `{self.principal}`"
+        return f"GRANT {action_type} ON {object_type} {escape_sql_identifier(object_key)} TO `{self.principal}`"
 
     def _uc_action(self, action_type):
         def inner(object_type, object_key):
@@ -155,7 +156,7 @@ def snapshot(self) -> Iterable[Grant]:
         return self._snapshot(partial(self._try_load), partial(self._crawl))
 
     def _try_load(self):
-        for row in self._fetch(f"SELECT * FROM {self._full_name}"):
+        for row in self._fetch(f"SELECT * FROM {escape_sql_identifier(self._full_name)}"):
             yield Grant(*row)
 
     def _crawl(self) -> Iterable[Grant]:
@@ -189,7 +190,7 @@ def _crawl(self) -> Iterable[Grant]:
         tasks.append(partial(self._grants, catalog=catalog, any_file=True))
         tasks.append(partial(self._grants, catalog=catalog, anonymous_function=True))
         # scan all databases, even empty ones
-        for row in self._fetch(f"SHOW DATABASES FROM {catalog}"):
+        for row in self._fetch(f"SHOW DATABASES FROM {escape_sql_identifier(catalog)}"):
             tasks.append(partial(self._grants, catalog=catalog, database=row.databaseName))
         for table in self._tc.snapshot():
             fn = partial(self._grants, catalog=catalog, database=table.database)
@@ -273,7 +274,7 @@ def _grants(
                 "ANY_FILE": "ANY FILE",
                 "ANONYMOUS_FUNCTION": "ANONYMOUS FUNCTION",
             }
-            for row in self._fetch(f"SHOW GRANTS ON {on_type} {key}"):
+            for row in self._fetch(f"SHOW GRANTS ON {on_type} {escape_sql_identifier(key)}"):
                 (principal, action_type, object_type, _) = row
                 object_type = object_type_normalization.get(object_type, object_type)
                 if on_type != object_type:

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

@@ -10,6 +10,7 @@
 from databricks.sdk.service.workspace import ImportFormat
 
 from databricks.labs.ucx.framework.crawlers import CrawlerBase, SqlBackend
+from databricks.labs.ucx.framework.utils import escape_sql_identifier
 from databricks.labs.ucx.mixins.sql import Row
 
 logger = logging.getLogger(__name__)
@@ -110,7 +111,7 @@ def _add_jdbc_location(self, external_locations, location, table):
     def _external_location_list(self) -> Iterable[ExternalLocation]:
         tables = list(
             self._backend.fetch(
-                f"SELECT location, storage_properties FROM {self._schema}.tables WHERE location IS NOT NULL"
+                f"SELECT location, storage_properties FROM {escape_sql_identifier(self._schema)}.tables WHERE location IS NOT NULL"
             )
         )
         mounts = Mounts(self._backend, self._ws, self._schema).snapshot()
@@ -120,7 +121,9 @@ def snapshot(self) -> Iterable[ExternalLocation]:
         return self._snapshot(self._try_fetch, self._external_location_list)
 
     def _try_fetch(self) -> Iterable[ExternalLocation]:
-        for row in self._fetch(f"SELECT * FROM {self._schema}.{self._table}"):
+        for row in self._fetch(
+            f"SELECT * FROM {escape_sql_identifier(self._schema)}.{escape_sql_identifier(self._table)}"
+        ):
             yield ExternalLocation(*row)
 
     def _get_ext_location_definitions(self, missing_locations: list[ExternalLocation]) -> list:
@@ -227,5 +230,7 @@ def snapshot(self) -> Iterable[Mount]:
         return self._snapshot(self._try_fetch, self._list_mounts)
 
     def _try_fetch(self) -> Iterable[Mount]:
-        for row in self._fetch(f"SELECT * FROM {self._schema}.{self._table}"):
+        for row in self._fetch(
+            f"SELECT * FROM {escape_sql_identifier(self._schema)}.{escape_sql_identifier(self._table)}"
+        ):
             yield Mount(*row)

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

@@ -13,6 +13,7 @@
 
 from databricks.labs.ucx.account import WorkspaceInfo
 from databricks.labs.ucx.framework.crawlers import SqlBackend
+from databricks.labs.ucx.framework.utils import escape_sql_identifier
 from databricks.labs.ucx.hive_metastore import TablesCrawler
 from databricks.labs.ucx.hive_metastore.tables import Table
 
@@ -99,7 +100,7 @@ def skip_table(self, schema: str, table: str):
         # Marks a table to be skipped in the migration process by applying a table property
         try:
             self._backend.execute(
-                f"ALTER TABLE `{schema}`.`{table}` SET TBLPROPERTIES('{self.UCX_SKIP_PROPERTY}' = true)"
+                f"ALTER TABLE {escape_sql_identifier(schema)}.{escape_sql_identifier(table)} SET TBLPROPERTIES('{self.UCX_SKIP_PROPERTY}' = true)"
             )
         except NotFound as nf:
             if "[TABLE_OR_VIEW_NOT_FOUND]" in str(nf) or "[DELTA_TABLE_NOT_FOUND]" in str(nf):
@@ -112,7 +113,9 @@ def skip_table(self, schema: str, table: str):
     def skip_schema(self, schema: str):
         # Marks a schema to be skipped in the migration process by applying a table property
         try:
-            self._backend.execute(f"ALTER SCHEMA `{schema}` SET DBPROPERTIES('{self.UCX_SKIP_PROPERTY}' = true)")
+            self._backend.execute(
+                f"ALTER SCHEMA {escape_sql_identifier(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.")
@@ -156,7 +159,7 @@ def _get_databases_in_scope(self, databases: set[str]):
 
     def _get_database_in_scope_task(self, database: str) -> str | None:
         describe = {}
-        for value in self._backend.fetch(f"DESCRIBE SCHEMA EXTENDED {database}"):
+        for value in self._backend.fetch(f"DESCRIBE SCHEMA EXTENDED {escape_sql_identifier(database)}"):
             describe[value["database_description_item"]] = value["database_description_value"]
         if self.UCX_SKIP_PROPERTY in TablesCrawler.parse_database_props(describe.get("Properties", "").lower()):
             logger.info(f"Database {database} is marked to be skipped")
@@ -170,7 +173,9 @@ def _get_table_in_scope_task(self, table_to_migrate: TableToMigrate) -> TableToM
         if self._exists_in_uc(table, rule.as_uc_table_key):
             logger.info(f"The intended target for {table.key}, {rule.as_uc_table_key}, already exists.")
             return None
-        result = self._backend.fetch(f"SHOW TBLPROPERTIES `{table.database}`.`{table.name}`")
+        result = self._backend.fetch(
+            f"SHOW TBLPROPERTIES {escape_sql_identifier(table.database)}.{escape_sql_identifier(table.name)}"
+        )
         for value in result:
             if value["key"] == self.UCX_SKIP_PROPERTY:
                 logger.info(f"{table.key} is marked to be skipped")

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

@@ -8,6 +8,7 @@
 from databricks.labs.blueprint.parallel import Threads
 
 from databricks.labs.ucx.framework.crawlers import CrawlerBase, SqlBackend
+from databricks.labs.ucx.framework.utils import escape_sql_identifier
 from databricks.labs.ucx.mixins.sql import Row
 
 logger = logging.getLogger(__name__)
@@ -96,16 +97,16 @@ def is_databricks_dataset(self) -> bool:
         return False
 
     def sql_migrate_external(self, target_table_key):
-        return f"SYNC TABLE {target_table_key} FROM {self.key};"
+        return f"SYNC TABLE {escape_sql_identifier(target_table_key)} FROM {escape_sql_identifier(self.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};"
+        return f"CREATE TABLE IF NOT EXISTS {escape_sql_identifier(target_table_key)} DEEP CLONE {escape_sql_identifier(self.key)};"
 
     def sql_migrate_view(self, target_table_key):
-        return f"CREATE VIEW IF NOT EXISTS {target_table_key} AS {self.view_text};"
+        return f"CREATE VIEW IF NOT EXISTS {escape_sql_identifier(target_table_key)} AS {self.view_text};"
 
 
 @dataclass
@@ -163,7 +164,7 @@ def parse_database_props(tbl_props: str) -> dict:
 
     def _try_load(self) -> Iterable[Table]:
         """Tries to load table information from the database or throws TABLE_OR_VIEW_NOT_FOUND error"""
-        for row in self._fetch(f"SELECT * FROM {self._full_name}"):
+        for row in self._fetch(f"SELECT * FROM {escape_sql_identifier(self._full_name)}"):
             yield Table(*row)
 
     def _crawl(self) -> Iterable[Table]:
@@ -184,7 +185,9 @@ def _crawl(self) -> Iterable[Table]:
         catalog = "hive_metastore"
         for (database,) in self._all_databases():
             logger.debug(f"[{catalog}.{database}] listing tables")
-            for _, table, _is_tmp in self._fetch(f"SHOW TABLES FROM {catalog}.{database}"):
+            for _, table, _is_tmp in self._fetch(
+                f"SHOW TABLES FROM {escape_sql_identifier(catalog)}.{escape_sql_identifier(database)}"
+            ):
                 tasks.append(partial(self._describe, catalog, database, table))
         catalog_tables, errors = Threads.gather(f"listing tables in {catalog}", tasks)
         if len(errors) > 0:
@@ -209,7 +212,7 @@ def _describe(self, catalog: str, database: str, table: str) -> Table | None:
         try:
             logger.debug(f"[{full_name}] fetching table metadata")
             describe = {}
-            for key, value, _ in self._fetch(f"DESCRIBE TABLE EXTENDED {full_name}"):
+            for key, value, _ in self._fetch(f"DESCRIBE TABLE EXTENDED {escape_sql_identifier(full_name)}"):
                 describe[key] = value
             return Table(
                 catalog=catalog.lower(),

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

@@ -6,6 +6,7 @@
 from databricks.labs.blueprint.parallel import Threads
 
 from databricks.labs.ucx.framework.crawlers import CrawlerBase, SqlBackend
+from databricks.labs.ucx.framework.utils import escape_sql_identifier
 from databricks.labs.ucx.mixins.sql import Row
 
 logger = logging.getLogger(__name__)
@@ -54,7 +55,7 @@ def snapshot(self) -> list[Udf]:
 
     def _try_load(self) -> Iterable[Udf]:
         """Tries to load udf information from the database or throws TABLE_OR_VIEW_NOT_FOUND error"""
-        for row in self._fetch(f"SELECT * FROM {self._full_name}"):
+        for row in self._fetch(f"SELECT * FROM {escape_sql_identifier(self._full_name)}"):
             yield Udf(*row)
 
     def _crawl(self) -> Iterable[Udf]:
@@ -63,10 +64,12 @@ def _crawl(self) -> Iterable[Udf]:
         catalog = "hive_metastore"
         # need to set the current catalog otherwise "SHOW USER FUNCTIONS FROM" is raising error:
         # "target schema <database> is not in the current catalog"
-        self._exec(f"USE CATALOG {catalog};")
+        self._exec(f"USE CATALOG {escape_sql_identifier(catalog)};")
         for (database,) in self._all_databases():
             logger.debug(f"[{catalog}.{database}] listing udfs")
-            for (udf,) in self._fetch(f"SHOW USER FUNCTIONS FROM {catalog}.{database};"):
+            for (udf,) in self._fetch(
+                f"SHOW USER FUNCTIONS FROM {escape_sql_identifier(catalog)}.{escape_sql_identifier(database)};"
+            ):
                 if udf.startswith(f"{catalog}.{database}"):
                     udf_name = udf[udf.rfind(".") + 1 :]  # remove catalog and database info from the name
                     tasks.append(partial(self._describe, catalog, database, udf_name))
@@ -83,7 +86,7 @@ def _describe(self, catalog: str, database: str, udf: str) -> Udf | None:
         try:
             logger.debug(f"[{full_name}] fetching udf metadata")
             describe = {}
-            for key_value in self._fetch(f"DESCRIBE FUNCTION EXTENDED {full_name}"):
+            for key_value in self._fetch(f"DESCRIBE FUNCTION EXTENDED {escape_sql_identifier(full_name)}"):
                 if ":" in key_value:  # skip free text configs that don't have a key
                     key, value = key_value.split(":")
                     describe[key] = value.strip()

tests/unit/framework/test_utils.py

@@ -0,0 +1,25 @@
+from databricks.labs.ucx.framework.utils import escape_sql_identifier
+
+
+def test_escaped_path():
+    # test with optional escaping
+    assert "a" == escape_sql_identifier("a")
+    assert "a.b" == escape_sql_identifier("a.b")
+    assert "a.b.c" == escape_sql_identifier("a.b.c")
+    assert "a" == escape_sql_identifier("a")
+    assert "a.b" == escape_sql_identifier("a.b")
+    assert "`a`.`b`.`c`" == escape_sql_identifier("`a`.`b`.`c`")
+    assert "`a.b`.`c`" == escape_sql_identifier("`a.b`.`c`")
+    assert "`a-b`.c.d" == escape_sql_identifier("a-b.c.d")
+    assert "a.`b-c`.d" == escape_sql_identifier("a.b-c.d")
+    assert "a.b.`c-d`" == escape_sql_identifier("a.b.c-d")
+    assert "`✨`.`🍰`.`✨`" == escape_sql_identifier("✨.🍰.✨")
+    assert "a.b.`c.d`" == escape_sql_identifier("a.b.c.d")
+    # test with escaping enforced
+    assert "`a`" == escape_sql_identifier("a", False)
+    assert "`a`.`b`" == escape_sql_identifier("a.b", False)
+    assert "`a`.`b`.`c`" == escape_sql_identifier("a.b.c", False)
+    assert "`a-b`.`c`.`d`" == escape_sql_identifier("a-b.c.d", False)
+    assert "`a`.`b-c`.`d`" == escape_sql_identifier("a.b-c.d", False)
+    assert "`a`.`b`.`c-d`" == escape_sql_identifier("a.b.c-d", False)
+    assert "`a`.`b`.`c.d`" == escape_sql_identifier("a.b.c.d", False)

tests/unit/hive_metastore/test_mapping.py

@@ -131,12 +131,10 @@ def test_skip_happy_path(mocker, caplog):
     sbe = mocker.patch("databricks.labs.ucx.framework.crawlers.StatementExecutionBackend.__init__")
     mapping = TableMapping(ws, sbe)
     mapping.skip_table(schema="schema", table="table")
-    sbe.execute.assert_called_with(
-        f"ALTER TABLE `schema`.`table` SET TBLPROPERTIES('{mapping.UCX_SKIP_PROPERTY}' = true)"
-    )
+    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(schema="schema")
-    sbe.execute.assert_called_with(f"ALTER SCHEMA `schema` SET DBPROPERTIES('{mapping.UCX_SKIP_PROPERTY}' = true)")
+    sbe.execute.assert_called_with(f"ALTER SCHEMA schema SET DBPROPERTIES('{mapping.UCX_SKIP_PROPERTY}' = true)")
     assert len(caplog.records) == 0
 
 
@@ -174,13 +172,13 @@ def test_skip_tables_marked_for_skipping_or_upgraded():
             ["test_schema2"],
             ["test_schema3"],
         ],
-        "SHOW TBLPROPERTIES `test_schema1`.`test_table1`": [
+        "SHOW TBLPROPERTIES test_schema1.test_table1": [
             {"key": "upgraded_to", "value": "fake_dest"},
         ],
-        "SHOW TBLPROPERTIES `test_schema1`.`test_view1`": [
+        "SHOW TBLPROPERTIES test_schema1.test_view1": [
             {"key": "databricks.labs.ucx.skip", "value": "true"},
         ],
-        "SHOW TBLPROPERTIES `test_schema1`.`test_table2`": [
+        "SHOW TBLPROPERTIES test_schema1.test_table2": [
             {"key": "upgraded_to", "value": "fake_dest"},
         ],
         "DESCRIBE SCHEMA EXTENDED test_schema1": [],
@@ -280,7 +278,7 @@ def test_skip_tables_marked_for_skipping_or_upgraded():
 def test_table_with_no_target_reverted():
     errors = {}
     rows = {
-        "SHOW TBLPROPERTIES `schema1`.`table1`": [
+        "SHOW TBLPROPERTIES schema1.table1": [
             {"key": "upgraded_to", "value": "non.existing.table"},
         ],
     }
@@ -429,8 +427,8 @@ def test_skipping_rules_database_skipped():
     ]
     table_mapping.get_tables_to_migrate(tables_crawler)
 
-    assert "SHOW TBLPROPERTIES `schema1`.`table1`" in backend.queries
-    assert "SHOW TBLPROPERTIES `schema2`.`table2`" not in backend.queries
+    assert "SHOW TBLPROPERTIES schema1.table1" in backend.queries
+    assert "SHOW TBLPROPERTIES schema2.table2" not in backend.queries
 
 
 def test_skip_missing_table_in_snapshot():