sadpandajoe commented on code in PR #43695:
URL: https://github.com/apache/superset/pull/43695#discussion_r4401557255
##########
tests/unit_tests/db_engine_specs/test_databricks.py:
##########
@@ -86,9 +86,211 @@ 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",
+ "encryption": False,
+ }
+ )
+ 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 adjust_engine_params maps catalog/schema into the URL and
args."""
+ uri_str = DatabricksNativeEngineSpec.build_sqlalchemy_uri(
+ {
+ "access_token": "test-token",
+ "host": "localhost",
+ "port": 443,
+ "database": "default",
+ "encryption": False,
+ }
+ )
+ assert uri_str ==
"databricks+connector://token:test-token@localhost:443/default"
+
+ # Test default resolution when no catalog is supplied: catalog is omitted
+ # so server default applies, while URI database is preserved as schema
+ url_default, args_default =
DatabricksNativeEngineSpec.adjust_engine_params(
+ make_url(uri_str),
+ connect_args={"http_path": "/sql/1.0/warehouses/test"},
+ )
+ assert "catalog" not in url_default.query
+ assert url_default.query["schema"] == "default"
+ assert url_default.query["http_path"] == "/sql/1.0/warehouses/test"
+ assert "catalog" not in args_default
+ assert args_default["schema"] == "default"
+ assert args_default["http_path"] == "/sql/1.0/warehouses/test"
+
+ # Test fallback to "default" schema when connection has no database or
schema
+ url_no_db_schema, args_no_db_schema = (
+ DatabricksNativeEngineSpec.adjust_engine_params(
+ make_url(
+ "databricks+connector://token:test-token@localhost:443/"
+ "?catalog=main&http_path=/sql/1.0/warehouses/test"
+ ),
+ connect_args={},
+ )
+ )
+ assert url_no_db_schema.query["catalog"] == "main"
+ assert url_no_db_schema.query["schema"] == "default"
+ assert url_no_db_schema.query["http_path"] == "/sql/1.0/warehouses/test"
+ assert args_no_db_schema["catalog"] == "main"
+ assert args_no_db_schema["schema"] == "default"
+ assert args_no_db_schema["http_path"] == "/sql/1.0/warehouses/test"
+
+ # Test catalog on a connection with no configured catalog (dataset creation
+ # with the server default catalog from `get_default_catalog`): the URI
+ # database is kept as schema, so reflection targets `main.finance.<table>`
+ url_no_cat, args_no_cat = DatabricksNativeEngineSpec.adjust_engine_params(
+ make_url(
+ "databricks+connector://token:test-token@localhost:443/finance"
+ "?http_path=/sql/1.0/warehouses/test"
+ ),
+ connect_args={},
+ catalog="main",
+ schema=None,
+ )
+ assert url_no_cat.query["catalog"] == "main"
+ assert url_no_cat.query["schema"] == "finance"
+ assert args_no_cat["catalog"] == "main"
+ assert args_no_cat["schema"] == "finance"
+
+ # Test catalog override when URI already contains
?catalog=main&schema=finance
+ # Switching to a different catalog resolves the catalog's default schema
+ url_query_cat_override, args_query_cat_override = (
+ DatabricksNativeEngineSpec.adjust_engine_params(
+ make_url(
+ "databricks+connector://token:test-token@localhost:443/"
+
"?catalog=main&schema=finance&http_path=/sql/1.0/warehouses/test"
+ ),
+ connect_args={},
+ catalog="other",
+ )
+ )
+ assert url_query_cat_override.query["catalog"] == "other"
+ assert url_query_cat_override.query["schema"] == "default"
+ assert url_query_cat_override.query["http_path"] ==
"/sql/1.0/warehouses/test"
+ assert args_query_cat_override["catalog"] == "other"
+ assert args_query_cat_override["schema"] == "default"
+ assert args_query_cat_override["http_path"] == "/sql/1.0/warehouses/test"
+
+ # Test catalog matching the catalog configured in the URI with schema=None
+ # (dataset creation scenario): preserves the URI-defined schema ("finance")
+ url_cat_match, args_cat_match =
DatabricksNativeEngineSpec.adjust_engine_params(
+ make_url(
+ "databricks+connector://token:test-token@localhost:443/finance"
+ "?catalog=main&http_path=/sql/1.0/warehouses/test"
+ ),
+ connect_args={},
+ catalog="main",
+ schema=None,
+ )
+ assert url_cat_match.query["catalog"] == "main"
+ assert url_cat_match.query["schema"] == "finance"
+ assert args_cat_match["catalog"] == "main"
+ assert args_cat_match["schema"] == "finance"
+
+ # Test switching away from a catalog configured in connect_args without a
+ # schema: the saved schema belongs to the old catalog, so "default" is used
+ url_cat_switch, args_cat_switch =
DatabricksNativeEngineSpec.adjust_engine_params(
+
make_url("databricks+connector://token:test-token@localhost:443/finance"),
+ connect_args={"catalog": "main"},
+ catalog="other",
+ )
+ assert url_cat_switch.query["catalog"] == "other"
+ assert url_cat_switch.query["schema"] == "default"
+ assert args_cat_switch["catalog"] == "other"
+ assert args_cat_switch["schema"] == "default"
+
+ # Test explicit catalog and schema with http_path in connect_args
+ adjusted_url, adjusted_args =
DatabricksNativeEngineSpec.adjust_engine_params(
+ make_url(uri_str),
+ connect_args={"http_path": "/sql/1.0/warehouses/test"},
+ catalog="main",
+ schema="custom_schema",
+ )
+ assert adjusted_url.query["catalog"] == "main"
+ assert adjusted_url.query["schema"] == "custom_schema"
+ assert adjusted_url.query["http_path"] == "/sql/1.0/warehouses/test"
+ assert adjusted_args["catalog"] == "main"
+ assert adjusted_args["schema"] == "custom_schema"
+ assert adjusted_args["http_path"] == "/sql/1.0/warehouses/test"
+
+ # Test reconciling conflicting schema in connect_args with passed schema
+ reconciled_url, reconciled_args =
DatabricksNativeEngineSpec.adjust_engine_params(
+ make_url(uri_str),
+ connect_args={
+ "catalog": "connect_arg_cat",
+ "schema": "stale_schema",
+ },
+ schema="reconciled_schema",
+ )
+ assert reconciled_url.query["catalog"] == "connect_arg_cat"
+ assert reconciled_url.query["schema"] == "reconciled_schema"
+ assert reconciled_args["catalog"] == "connect_arg_cat"
+ assert reconciled_args["schema"] == "reconciled_schema"
+
+ # Test catalog from connect_args and fallback schema from URI database
+ adjusted_url_fallback, adjusted_args_fallback = (
+ DatabricksNativeEngineSpec.adjust_engine_params(
+ make_url(uri_str),
+ connect_args={"catalog": "connect_arg_cat"},
+ )
+ )
+ assert adjusted_url_fallback.query["catalog"] == "connect_arg_cat"
+ assert adjusted_url_fallback.query["schema"] == "default"
+ assert adjusted_args_fallback["catalog"] == "connect_arg_cat"
+ assert adjusted_args_fallback["schema"] == "default"
+
+
+def test_databricks_connector_dialect_registration() -> None:
+ """Test that databricks.connector dialect is registered and resolves to
DatabricksDialect."""
Review Comment:
This docstring line and the comment at line 282 are 97 characters, over the
repo's 88-column Ruff limit, and the `import sys` / `from unittest import mock`
/ `from sqlalchemy.dialects import registry` block has no blank line between
stdlib and third-party groups, so the Ruff check fails with E501 and I001.
Could the two lines be wrapped and the import groups separated?
##########
tests/unit_tests/db_engine_specs/test_databricks.py:
##########
@@ -86,9 +86,211 @@ 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",
+ "encryption": False,
+ }
+ )
+ 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 adjust_engine_params maps catalog/schema into the URL and
args."""
+ uri_str = DatabricksNativeEngineSpec.build_sqlalchemy_uri(
+ {
+ "access_token": "test-token",
+ "host": "localhost",
+ "port": 443,
+ "database": "default",
+ "encryption": False,
+ }
+ )
+ assert uri_str ==
"databricks+connector://token:test-token@localhost:443/default"
+
+ # Test default resolution when no catalog is supplied: catalog is omitted
+ # so server default applies, while URI database is preserved as schema
+ url_default, args_default =
DatabricksNativeEngineSpec.adjust_engine_params(
+ make_url(uri_str),
+ connect_args={"http_path": "/sql/1.0/warehouses/test"},
+ )
+ assert "catalog" not in url_default.query
+ assert url_default.query["schema"] == "default"
+ assert url_default.query["http_path"] == "/sql/1.0/warehouses/test"
+ assert "catalog" not in args_default
+ assert args_default["schema"] == "default"
+ assert args_default["http_path"] == "/sql/1.0/warehouses/test"
+
+ # Test fallback to "default" schema when connection has no database or
schema
+ url_no_db_schema, args_no_db_schema = (
+ DatabricksNativeEngineSpec.adjust_engine_params(
+ make_url(
+ "databricks+connector://token:test-token@localhost:443/"
+ "?catalog=main&http_path=/sql/1.0/warehouses/test"
+ ),
+ connect_args={},
+ )
+ )
+ assert url_no_db_schema.query["catalog"] == "main"
+ assert url_no_db_schema.query["schema"] == "default"
+ assert url_no_db_schema.query["http_path"] == "/sql/1.0/warehouses/test"
+ assert args_no_db_schema["catalog"] == "main"
+ assert args_no_db_schema["schema"] == "default"
+ assert args_no_db_schema["http_path"] == "/sql/1.0/warehouses/test"
+
+ # Test catalog on a connection with no configured catalog (dataset creation
+ # with the server default catalog from `get_default_catalog`): the URI
+ # database is kept as schema, so reflection targets `main.finance.<table>`
+ url_no_cat, args_no_cat = DatabricksNativeEngineSpec.adjust_engine_params(
+ make_url(
+ "databricks+connector://token:test-token@localhost:443/finance"
+ "?http_path=/sql/1.0/warehouses/test"
+ ),
+ connect_args={},
+ catalog="main",
+ schema=None,
+ )
+ assert url_no_cat.query["catalog"] == "main"
+ assert url_no_cat.query["schema"] == "finance"
+ assert args_no_cat["catalog"] == "main"
+ assert args_no_cat["schema"] == "finance"
+
+ # Test catalog override when URI already contains
?catalog=main&schema=finance
+ # Switching to a different catalog resolves the catalog's default schema
+ url_query_cat_override, args_query_cat_override = (
+ DatabricksNativeEngineSpec.adjust_engine_params(
+ make_url(
+ "databricks+connector://token:test-token@localhost:443/"
+
"?catalog=main&schema=finance&http_path=/sql/1.0/warehouses/test"
+ ),
+ connect_args={},
+ catalog="other",
+ )
+ )
+ assert url_query_cat_override.query["catalog"] == "other"
+ assert url_query_cat_override.query["schema"] == "default"
+ assert url_query_cat_override.query["http_path"] ==
"/sql/1.0/warehouses/test"
+ assert args_query_cat_override["catalog"] == "other"
+ assert args_query_cat_override["schema"] == "default"
+ assert args_query_cat_override["http_path"] == "/sql/1.0/warehouses/test"
+
+ # Test catalog matching the catalog configured in the URI with schema=None
+ # (dataset creation scenario): preserves the URI-defined schema ("finance")
+ url_cat_match, args_cat_match =
DatabricksNativeEngineSpec.adjust_engine_params(
+ make_url(
+ "databricks+connector://token:test-token@localhost:443/finance"
+ "?catalog=main&http_path=/sql/1.0/warehouses/test"
+ ),
+ connect_args={},
+ catalog="main",
+ schema=None,
+ )
+ assert url_cat_match.query["catalog"] == "main"
+ assert url_cat_match.query["schema"] == "finance"
+ assert args_cat_match["catalog"] == "main"
+ assert args_cat_match["schema"] == "finance"
+
+ # Test switching away from a catalog configured in connect_args without a
+ # schema: the saved schema belongs to the old catalog, so "default" is used
+ url_cat_switch, args_cat_switch =
DatabricksNativeEngineSpec.adjust_engine_params(
+
make_url("databricks+connector://token:test-token@localhost:443/finance"),
+ connect_args={"catalog": "main"},
+ catalog="other",
+ )
+ assert url_cat_switch.query["catalog"] == "other"
+ assert url_cat_switch.query["schema"] == "default"
+ assert args_cat_switch["catalog"] == "other"
+ assert args_cat_switch["schema"] == "default"
+
+ # Test explicit catalog and schema with http_path in connect_args
+ adjusted_url, adjusted_args =
DatabricksNativeEngineSpec.adjust_engine_params(
+ make_url(uri_str),
+ connect_args={"http_path": "/sql/1.0/warehouses/test"},
+ catalog="main",
+ schema="custom_schema",
+ )
+ assert adjusted_url.query["catalog"] == "main"
+ assert adjusted_url.query["schema"] == "custom_schema"
+ assert adjusted_url.query["http_path"] == "/sql/1.0/warehouses/test"
+ assert adjusted_args["catalog"] == "main"
+ assert adjusted_args["schema"] == "custom_schema"
+ assert adjusted_args["http_path"] == "/sql/1.0/warehouses/test"
+
+ # Test reconciling conflicting schema in connect_args with passed schema
+ reconciled_url, reconciled_args =
DatabricksNativeEngineSpec.adjust_engine_params(
Review Comment:
Every case here that sets `connect_args["schema"]` also passes an explicit
`schema=`, so the new branch that preserves a saved `connect_args` schema when
no override is given is never exercised; dropping
`adjusted_connect_args.get("schema")` from `adjust_engine_params` would leave
this suite green while silently selecting the wrong schema. Could a case call
`adjust_engine_params` with a `/legacy?schema=url_schema` URI,
`connect_args={"schema": "finance"}`, and no overrides, asserting `finance` in
both the returned query and connect args?
##########
superset/db_engine_specs/databricks.py:
##########
@@ -629,9 +657,22 @@ class
DatabricksNativeEngineSpec(DatabricksDynamicBaseEngineSpec):
"databricks+connector://token:{access_token}@{host}:{port}/{database_name}"
)
- # Note: Primary metadata is in DatabricksPythonConnectorEngineSpec which
- # consolidates all Databricks connection methods. This spec exists for
- # backwards compatibility with legacy databricks-dbapi connections.
+ metadata = {
+ "description": ("Databricks legacy connector using databricks-dbapi."),
+ "logo": "databricks.png",
+ "homepage_url": "https://www.databricks.com/";,
+ "categories": [
+ DatabaseCategory.CLOUD_DATA_WAREHOUSES,
+ DatabaseCategory.ANALYTICAL_DATABASES,
+ DatabaseCategory.HOSTED_OPEN_SOURCE,
+ ],
+ "pypi_packages": ["databricks-dbapi[sqlalchemy]"],
Review Comment:
With the supported `databricks-sqlalchemy` 2.0 dialect, a `/default` URI
plus only `http_path` still leaves the catalog unset after
`adjust_engine_params`, so reflection can build ``DESCRIBE TABLE EXTENDED
`None`.`default`.`orders` `` for a table in the connection's default catalog
(the earlier `hive_metastore` fallback was dropped). Could the Native spec
resolve the server's actual default catalog for reflection without forcing one
onto the connection, with a test that exercises the no-catalog path?
##########
superset/db_engine_specs/databricks.py:
##########
@@ -1011,6 +1020,49 @@ 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()
+
+ effective_catalog = (
+ catalog
+ or adjusted_connect_args.get("catalog")
+ or uri.query.get("catalog")
+ or cls.default_catalog
+ )
+ effective_schema = (
+ schema
+ or adjusted_connect_args.get("schema")
+ or uri.query.get("schema")
Review Comment:
When the saved URI is `/finance?http_path=…` with no configured catalog,
`get_all_schema_names(catalog="other")` still sends `schema="finance"` because
`original_catalog` is `None`; if `finance` exists only in the workspace default
catalog, the connection to `other.finance` fails before schemas load. The
caller's catalog is not necessarily the server default, so could the schema
fallback distinguish the default catalog from a genuine switch, with a test for
the latter?
--
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]
