From d66ac9f3f46eed9fe870c378a34df4a7174717d9 Mon Sep 17 00:00:00 2001
From: Vitor Avila <96086495+Vitor-Avila@users.noreply.github.com>
Date: Mon, 2 Dec 2024 14:00:00 -0300
Subject: [PATCH] fix(Databricks): Escape catalog and schema names in
 pre-queries (#31199)

---
 superset/db_engine_specs/databricks.py        |  2 ++
 .../db_engine_specs/test_databricks.py        | 22 +++++++++++++++----
 2 files changed, 20 insertions(+), 4 deletions(-)

diff --git a/superset/db_engine_specs/databricks.py b/superset/db_engine_specs/databricks.py
index 88d6407a3..1395ca3aa 100644
--- a/superset/db_engine_specs/databricks.py
+++ b/superset/db_engine_specs/databricks.py
@@ -464,8 +464,10 @@ class DatabricksNativeEngineSpec(DatabricksDynamicBaseEngineSpec):
     ) -> 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
index 025784fe2..0c9e3843d 100644
--- 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`",
     ]