FrancescoCastaldi commented on code in PR #43695:
URL: https://github.com/apache/superset/pull/43695#discussion_r4422534351
##########
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:
Good catch. The reason we cannot pass `catalog` during the handshake is that
older Databricks clusters and clusters without Unity Catalog reject connections
when `catalog` is set as an engine connection parameter.
To address the reflection issue where `databricks-sqlalchemy` constructs
`DESCRIBE TABLE EXTENDED \`None\`.\`default\`.\`orders\`` when
`dialect.catalog` is `None`, we have updated
`DatabricksNativeEngineSpec.get_columns` and `get_table_names` to dynamically
inspect the connection: if `dialect.catalog` is `None`, it resolves the
server's default catalog via `SHOW CATALOGS` / `SELECT current_catalog()` and
sets `dialect.catalog` on the engine before reflection executes.
This guarantees that:
1. Clusters without multi-catalog support can connect without handshake
failures.
2. Reflection uses the server's active default catalog without falling back
to `None`.
Added test
`test_databricks_native_get_columns_resolves_no_catalog_reflection` verifying
this dynamic resolution.
##########
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 a connection has an explicit catalog configured, any switch away from
it invalidates and resets the schema to `"default"`.
When no catalog is configured on the connection, the catalog initially
passed by the caller comes from `get_default_catalog` (matching the server
default catalog, e.g. `"main"`), so the schema from the URI/connection still
applies. To prevent this from suppressing resets when a user genuinely switches
catalogs, `adjust_engine_params` now checks against `server_default_catalog`:
any catalog selected that differs from the server default catalog is treated as
a genuine switch and resets the schema to `"default"`.
Tested in `test_databricks_native_adjust_engine_params` for both
unconfigured connections with server default catalog and genuine switches to
other catalogs.
##########
tests/unit_tests/db_engine_specs/test_lint_metadata.py:
##########
@@ -0,0 +1,248 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND,
+# either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import io
+from pathlib import Path
+from textwrap import dedent
+
+import pytest
+
+from superset.db_engine_specs.lint_metadata import (
+ get_all_engine_specs_ast,
+ main,
+)
+
+
+def _write_spec(directory: Path, filename: str, code: str) -> Path:
+ """Write a Python source file into *directory* and return the path."""
+ spec_file = directory / filename
+ spec_file.write_text(dedent(code), encoding="utf-8")
+ return spec_file
+
+
+def test_get_all_engine_specs_ast_annotated_attributes(
+ tmp_path: Path,
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ """
+ Engine specs with type-annotated engine, engine_name, and metadata
+ (inline initializers) are correctly discovered, including when named
+ with the *BaseEngineSpec suffix.
+ """
+ _write_spec(
+ tmp_path,
+ "_tmp_typed_inline.py",
+ """\
+ from typing import Any
+ from superset.db_engine_specs.base import BaseEngineSpec
+
+ class TypedInlineBaseEngineSpec(BaseEngineSpec):
+ engine: str = "typed_engine"
+ engine_name: str = "Typed Engine"
+ metadata: dict[str, Any] = {
+ "description": "A typed engine spec",
+ }
+ """,
+ )
+
+ monkeypatch.setattr(
+ "superset.db_engine_specs.lint_metadata.SPECS_DIR",
+ str(tmp_path),
+ )
+
+ specs = get_all_engine_specs_ast()
+ matching = [s for s in specs if s["class_name"] ==
"TypedInlineBaseEngineSpec"]
+ assert len(matching) == 1
+ spec = matching[0]
+ assert spec["engine_name"] == "Typed Engine"
+ assert spec["metadata"] == {"description": "A typed engine spec"}
+
+
+def test_get_all_engine_specs_ast_declaration_then_assignment(
+ tmp_path: Path,
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ """
+ Regression: a bare annotation `engine: str` followed by a plain
+ assignment `engine = "value"` must still be discovered, including
+ when named with the *BaseEngineSpec suffix (the scanner must not break
+ on the bare annotation or skip the spec as a true base class).
+ """
+ _write_spec(
+ tmp_path,
+ "_tmp_typed_decl.py",
+ """\
+ from typing import Any
+ from superset.db_engine_specs.base import BaseEngineSpec
+
+ class DeclThenAssignBaseEngineSpec(BaseEngineSpec):
+ engine: str
+ engine = "decl_then_assign"
+ engine_name: str
+ engine_name = "DeclThenAssign"
+ metadata: dict[str, Any]
+ metadata = {
+ "description": "Engine with declaration-then-assignment",
+ }
+ """,
+ )
+
+ monkeypatch.setattr(
+ "superset.db_engine_specs.lint_metadata.SPECS_DIR",
+ str(tmp_path),
+ )
+
+ specs = get_all_engine_specs_ast()
+ matching = [s for s in specs if s["class_name"] ==
"DeclThenAssignBaseEngineSpec"]
+ found_names = [s["class_name"] for s in specs]
+ assert len(matching) == 1, (
+ f"DeclThenAssignBaseEngineSpec not found; got: {found_names}"
+ )
+ spec = matching[0]
+ assert spec["engine_name"] == "DeclThenAssign"
+ assert spec["metadata"] == {
+ "description": "Engine with declaration-then-assignment"
+ }
+
+
+def test_get_all_engine_specs_ast_set_and_unary_minus(
+ tmp_path: Path,
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ """
+ Regression: ast.Set and unary-minus in metadata values are correctly
+ evaluated and preserved in the AST scan.
+ """
+ _write_spec(
+ tmp_path,
+ "_tmp_complex_values.py",
+ """\
+ from typing import Any
+ from superset.db_engine_specs.base import BaseEngineSpec
+
+ class ComplexValuesBaseEngineSpec(BaseEngineSpec):
+ engine: str = "complex_engine"
+ engine_name: str = "Complex Engine"
+ metadata: dict[str, Any] = {
Review Comment:
Updated `ComplexValuesBaseEngineSpec` fixture in
`tests/unit_tests/db_engine_specs/test_lint_metadata.py` with `"pair":
("value", -1)` to explicitly assert `ast.Tuple` evaluation in
`test_get_all_engine_specs_ast_set_and_unary_minus`.
##########
superset/db_engine_specs/clickhouse.py:
##########
@@ -424,6 +426,7 @@ class ClickHouseConnectEngineSpec(BasicParametersMixin,
ClickHouseEngineSpec):
"logo": "clickhouse.svg",
"logo_dark": "clickhouse-dark.svg",
"homepage_url": "https://clickhouse.com/";,
+ "sqlalchemy_docs_url":
"https://clickhouse.com/docs/en/integrations/sql-clients/superset";,
Review Comment:
Updated both `sqlalchemy_docs_url` and `docs_url` in
`ClickHouseConnectEngineSpec` to the live documentation guide:
`https://clickhouse.com/docs/integrations/connectors/data-visualization/superset-and-clickhouse`.
##########
tests/unit_tests/db_engine_specs/test_databricks.py:
##########
@@ -86,9 +86,227 @@ 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"
+
+ # Test preserving connect_args schema over URI query schema when no
override
+ # is given
+ legacy_url = (
+
"databricks+connector://token:abc@my_hostname:1234/legacy?schema=url_schema"
+ )
+ adj_url_saved, adj_args_saved =
DatabricksNativeEngineSpec.adjust_engine_params(
+ make_url(legacy_url),
+ connect_args={"schema": "finance"},
+ )
+ assert adj_url_saved.query["schema"] == "finance"
+ assert adj_args_saved["schema"] == "finance"
Review Comment:
Added test case in `test_databricks_native_adjust_engine_params` exercising
`make_url("databricks+connector://token:abc@my_hostname:1234/legacy?catalog=main&schema=finance")`
with empty `connect_args`, verifying that `uri.query["schema"]` is preserved
and passed to both the returned URL and `connect_args`.
--
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]
