diff --git a/superset/db_engine_specs/databricks.py b/superset/db_engine_specs/databricks.py index 88d6407a3ee10..1395ca3aa287c 100644 --- 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 index 025784fe2c568..0c9e3843d6e06 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`", ]