fix(Databricks): Escape catalog and schema names in pre-queries (#31199)

apache · Dec 2, 2024 · d66ac9f · d66ac9f
1 parent 06fb330
commit d66ac9f
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 4 deletions.
diff --git a/superset/db_engine_specs/databricks.py b/superset/db_engine_specs/databricks.py
@@ -464,8 +464,10 @@ def get_prequeries(
     ) -> list[str]:
         prequeries = []
         if catalog:
+            catalog = f"`{catalog}`" if not catalog.startswith("`") else catalog
             prequeries.append(f"USE CATALOG {catalog}")
         if schema:
+            schema = f"`{schema}`" if not schema.startswith("`") else schema
             prequeries.append(f"USE SCHEMA {schema}")
         return prequeries
 

diff --git a/tests/unit_tests/db_engine_specs/test_databricks.py b/tests/unit_tests/db_engine_specs/test_databricks.py
@@ -257,14 +257,28 @@ def test_get_prequeries(mocker: MockerFixture) -> None:
 
     assert DatabricksNativeEngineSpec.get_prequeries(database) == []
     assert DatabricksNativeEngineSpec.get_prequeries(database, schema="test") == [
-        "USE SCHEMA test",
+        "USE SCHEMA `test`",
     ]
     assert DatabricksNativeEngineSpec.get_prequeries(database, catalog="test") == [
-        "USE CATALOG test",
+        "USE CATALOG `test`",
     ]
     assert DatabricksNativeEngineSpec.get_prequeries(
         database, catalog="foo", schema="bar"
     ) == [
-        "USE CATALOG foo",
-        "USE SCHEMA bar",
+        "USE CATALOG `foo`",
+        "USE SCHEMA `bar`",
+    ]
+
+    assert DatabricksNativeEngineSpec.get_prequeries(
+        database, catalog="with-hyphen", schema="hyphen-again"
+    ) == [
+        "USE CATALOG `with-hyphen`",
+        "USE SCHEMA `hyphen-again`",
+    ]
+
+    assert DatabricksNativeEngineSpec.get_prequeries(
+        database, catalog="`escaped-hyphen`", schema="`hyphen-escaped`"
+    ) == [
+        "USE CATALOG `escaped-hyphen`",
+        "USE SCHEMA `hyphen-escaped`",
     ]