iop
/
superset
4
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

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

This commit is contained in:
Vitor Avila 2024-12-02 14:00:00 -03:00 committed by GitHub
parent 06fb330569
commit d66ac9f3f4
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 20 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
superset/db_engine_specs/databricks.py
View File

@ -464,8 +464,10 @@ class DatabricksNativeEngineSpec(DatabricksDynamicBaseEngineSpec):
) -> list[str]: ) -> list[str]:
prequeries = [] prequeries = []
if catalog: if catalog:
catalog = f"`{catalog}`" if not catalog.startswith("`") else catalog
prequeries.append(f"USE CATALOG {catalog}") prequeries.append(f"USE CATALOG {catalog}")
if schema: if schema:
schema = f"`{schema}`" if not schema.startswith("`") else schema
prequeries.append(f"USE SCHEMA {schema}") prequeries.append(f"USE SCHEMA {schema}")
return prequeries return prequeries

22
tests/unit_tests/db_engine_specs/test_databricks.py
View File

@ -257,14 +257,28 @@ def test_get_prequeries(mocker: MockerFixture) -> None:
assert DatabricksNativeEngineSpec.get_prequeries(database) == [] assert DatabricksNativeEngineSpec.get_prequeries(database) == []
assert DatabricksNativeEngineSpec.get_prequeries(database, schema="test") == [ assert DatabricksNativeEngineSpec.get_prequeries(database, schema="test") == [
"USE SCHEMA test", "USE SCHEMA `test`",
] ]
assert DatabricksNativeEngineSpec.get_prequeries(database, catalog="test") == [ assert DatabricksNativeEngineSpec.get_prequeries(database, catalog="test") == [
"USE CATALOG test", "USE CATALOG `test`",
] ]
assert DatabricksNativeEngineSpec.get_prequeries( assert DatabricksNativeEngineSpec.get_prequeries(
database, catalog="foo", schema="bar" database, catalog="foo", schema="bar"
) == [ ) == [
"USE CATALOG foo", "USE CATALOG `foo`",
"USE SCHEMA bar", "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`",
] ]