This is an automated email from the ASF dual-hosted git repository. vavila pushed a commit to branch fix/databricks/escape-catalog-and-schema in repository https://gitbox.apache.org/repos/asf/superset.git
commit 6ba3b79831d0c97323252c06192098e502ea0e47 Author: Vitor Avila <[email protected]> AuthorDate: Fri Nov 29 17:32:07 2024 -0300 fix(Databricks): Escape catalog and schema names in pre-queries --- superset/db_engine_specs/databricks.py | 2 ++ .../unit_tests/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 88d6407a3e..1395ca3aa2 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 025784fe2c..0c9e3843d6 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`", ]
