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]:
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

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, 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`",
]