FrancescoCastaldi commented on code in PR #43695:
URL: https://github.com/apache/superset/pull/43695#discussion_r4269747506
##########
superset/db_engine_specs/databricks.py:
##########
@@ -1011,6 +1019,55 @@ def build_sqlalchemy_uri( # type: ignore
query=query,
).render_as_string(hide_password=False)
+ @classmethod
+ def adjust_engine_params(
+ cls,
+ uri: URL,
+ connect_args: dict[str, Any],
+ catalog: str | None = None,
+ schema: str | None = None,
+ ) -> tuple[URL, dict[str, Any]]:
+ adjusted_connect_args = connect_args.copy()
+
+ is_catalog_overridden = catalog is not None
+ effective_catalog = (
+ catalog
+ or adjusted_connect_args.get("catalog")
+ or uri.query.get("catalog")
+ )
+ effective_schema = (
+ schema
+ or (adjusted_connect_args.get("schema") if not
is_catalog_overridden else None)
+ or (uri.query.get("schema") if not is_catalog_overridden else None)
+ or (uri.database if uri.database and not is_catalog_overridden
else None)
+ )
+ effective_http_path = (
+ adjusted_connect_args.get("http_path")
+ or uri.query.get("http_path")
+ )
+
+ query_params: dict[str, str] = {}
+ if effective_catalog:
+ query_params["catalog"] = effective_catalog
+ adjusted_connect_args["catalog"] = effective_catalog
+ else:
+ adjusted_connect_args.pop("catalog", None)
+
+ if effective_schema:
+ query_params["schema"] = effective_schema
+ adjusted_connect_args["schema"] = effective_schema
+ else:
+ adjusted_connect_args.pop("schema", None)
+
+ if effective_http_path:
+ query_params["http_path"] = effective_http_path
+ adjusted_connect_args["http_path"] = effective_http_path
+
+ if query_params:
+ uri = uri.update_query_dict(query_params)
Review Comment:
Good catch @sadpandajoe! Because `update_query_dict` only merges keys, any
existing `schema` key in `uri.query` persisted when switching catalogs.
In commit `1c2d9f83b5`, `adjust_engine_params` now mutates query parameters
via `query_dict = dict(uri.query)` and assigns via `uri.set(query=query_dict)`.
When `effective_schema` is None (such as during a catalog override without an
explicit schema), `query_dict.pop("schema", None)` explicitly purges `schema`
from the URL query dict as well as from `connect_args`.
I also added a dedicated unit test starting from a URI carrying
`?catalog=main&schema=finance&http_path=...` verifying that overriding with
`catalog="other"` drops `schema` completely from both the URL query and
`connect_args`.
##########
tests/unit_tests/db_engine_specs/test_databricks.py:
##########
@@ -74,9 +74,129 @@ def test_dialect_supports_installed_sqlalchemy(mocker:
MockerFixture) -> None:
assert str(text("SELECT :value").compile(engine)) == "SELECT :value"
finally:
engine.dispose()
+
+ native_uri_str = DatabricksNativeEngineSpec.build_sqlalchemy_uri(
+ {
+ "access_token": "test-token",
+ "host": "localhost",
+ "port": 443,
+ "database": "default",
+ }
+ )
+ native_url, native_connect_args =
DatabricksNativeEngineSpec.adjust_engine_params(
+ make_url(native_uri_str),
+ connect_args={
+ "http_path": "/sql/1.0/warehouses/test",
+ "catalog": "hive_metastore",
+ },
+ )
+ connector_engine = create_engine(native_url,
connect_args=native_connect_args)
+ try:
+ assert isinstance(connector_engine.dialect, dialect.DatabricksDialect)
+ _, connector_dialect_args =
connector_engine.dialect.create_connect_args(
+ connector_engine.url
+ )
+ assert connector_dialect_args["http_path"] ==
"/sql/1.0/warehouses/test"
+ assert connector_dialect_args["catalog"] == "hive_metastore"
+ assert connector_dialect_args["schema"] == "default"
+ assert native_connect_args["http_path"] == "/sql/1.0/warehouses/test"
+ assert native_connect_args["catalog"] == "hive_metastore"
+ assert native_connect_args["schema"] == "default"
+ assert connector_engine.dialect.catalog == "hive_metastore"
+ assert connector_engine.dialect.schema == "default"
+ finally:
+ connector_engine.dispose()
socket.assert_not_called()
+def test_databricks_native_adjust_engine_params() -> None:
+ """Test that DatabricksNativeEngineSpec translates catalog/schema and
syncs connect_args."""
Review Comment:
Applied the suggestion directly! In commit `1c2d9f83b5`, the docstring in
`test_databricks.py` is now shortened to 79 characters, well within the
88-column Ruff limit.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
