sadpandajoe commented on code in PR #43695:
URL: https://github.com/apache/superset/pull/43695#discussion_r4268240319
##########
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:
This docstring line is over the 88-column Ruff limit, so the lint check will
flag it. Could it be shortened?
```suggestion
"""Test that adjust_engine_params maps catalog/schema into the URL and
args."""
```
##########
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:
With a saved URI like
`databricks+connector://…?catalog=main&schema=finance&http_path=…`, calling
this with `catalog="other"` and no schema pops `schema` from `connect_args` but
`update_query_dict` only merges, so `schema=finance` stays in the URL and the
dialect still opens the connection in `other.finance`. If that schema doesn't
exist in `other`, listing schemas for the other catalog fails at connect time
instead of returning its schemas. Could the stale `schema` query key be removed
when the catalog is overridden without a schema, with a test that starts from a
URL carrying `?schema=` (the current catalog-override case only covers the
`/finance` database path)?
--
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]
