rebenitez1802 commented on code in PR #45843:
URL: https://github.com/apache/superset/pull/45843#discussion_r4414440739
##########
superset/db_engine_specs/databricks.py:
##########
@@ -897,18 +902,54 @@ def get_parameters_from_uri( # type: ignore
"access_token": url.password,
"host": url.host,
"port": url.port,
- "http_path_field": query["http_path"],
- "default_catalog": query["catalog"],
- "default_schema": query["schema"],
+ "http_path_field": query.get("http_path", ""),
+ "default_catalog": query.get("catalog", ""),
+ "default_schema": query.get("schema", url.database or ""),
"encryption": encryption,
}
+ @classmethod
+ def get_parameters_from_database(
+ cls, database: Database, encrypted_extra: dict[str, Any]
+ ) -> dict[str, Any]:
+ """Expose legacy connect_args in the native connector's field-based
form."""
+ parameters = super().get_parameters_from_database(database,
encrypted_extra)
+ query = database.url_object.query
+ for field, key in (
+ ("http_path_field", "http_path"),
+ ("default_catalog", "catalog"),
+ ("default_schema", "schema"),
+ ):
+ if key not in query and key in database.connect_args:
+ parameters[field] = database.connect_args[key]
+ return parameters
+
+ @classmethod
+ def get_connection_identity(cls, uri: URL, connect_args: dict[str, Any])
-> URL:
+ """
+ Describe legacy connect_args and native query settings identically.
+
+ Mirrors ``adjust_engine_params``: query settings take precedence over
+ connect_args, and the URI path is the schema when neither sets one.
+ """
+ query: dict[str, str | tuple[str, ...]] = {
+ key: str(connect_args[key])
+ for key in ("http_path", "catalog", "schema")
+ if key in connect_args
+ }
+ query.update(uri.query)
+ if "schema" not in query and uri.database:
+ query["schema"] = uri.database
+ return uri._replace(database=None).set(query=query)
+
@classmethod
def get_default_catalog(
cls,
database: Database,
) -> str | None:
- return database.url_object.query.get("catalog")
+ return database.url_object.query.get("catalog") or
database.connect_args.get(
+ "catalog"
+ )
Review Comment:
🔴 High — Converted rows lose their default catalog
After conversion a row resolves to `DatabricksPythonConnectorEngineSpec`,
and this returns `None` unless the URI query or `connect_args` sets a catalog.
Most 6.0 rows only set `http_path`. The legacy
`DatabricksNativeEngineSpec.get_default_catalog` fell back to `SHOW CATALOGS` /
`current_catalog()` (e.g. `hive_metastore`), and migration `4081be5b6b74`
(`upgrade_catalog_perms(engines={"databricks"})`) already rewrote these rows'
perms and dataset catalogs to that value (`[db].[hive_metastore].[schema]`).
Once the row is converted:
- users with only `schema_access` no longer match their perms in SQL Lab;
- `UpdateDatabaseCommand` sees `None` as the original catalog, so the first
save that adds a catalog rewrites every dataset and query catalog through
`_update_catalog_attribute`.
Repro: a row with `extra={"engine_params": {"connect_args": {"http_path":
"/sql/1.0/x"}}}` returns `hive_metastore` on the legacy scheme and `None` after
`b8f3a2c91d64`. The suggestion uses the migration's marker so that native rows
stay unchanged. With it applied, the focused suite (engine spec, migration,
`update_test`, `validate_test`) passes and ruff/mypy report nothing new.
```suggestion
if catalog := database.url_object.query.get(
"catalog"
) or database.connect_args.get("catalog"):
return catalog
# Rows converted from ``databricks+connector://`` had their datasets
and
# permissions qualified with the legacy spec's runtime default
catalog.
if
database.get_extra().get("migrated_from_databricks_connector_uri"):
return DatabricksNativeEngineSpec.get_default_catalog(database)
return None
```
##########
tests/unit_tests/commands/databases/validate_test.py:
##########
@@ -894,3 +894,80 @@ def test_ssh_tunnel_missing_message_is_interpolated(
assert missing_field_messages
assert all("%(missing)s" not in msg for msg in missing_field_messages)
assert any("server_port" in msg for msg in missing_field_messages)
+
+
[email protected](
+ "rebuilt_uri, expected_uri",
+ [
+ # unchanged native-form fields: the stored token is rehydrated
+ (
+ "databricks://token:XXXXXXXXXX@host:443"
+ "?catalog=main&http_path=%2Flegacy&schema=default",
+ "databricks://token:real-token@host:443/default",
+ ),
+ # a different warehouse keeps the masked token
+ (
+ "databricks://token:XXXXXXXXXX@host:443"
+ "?catalog=main&http_path=%2Fother&schema=default",
+ "databricks://token:XXXXXXXXXX@host:443"
+ "?catalog=main&http_path=%2Fother&schema=default",
+ ),
+ # a newly typed token is used as typed
+ (
+ "databricks://token:new-token@host:443"
+ "?catalog=main&http_path=%2Flegacy&schema=default",
+ "databricks://token:new-token@host:443"
+ "?catalog=main&http_path=%2Flegacy&schema=default",
+ ),
+ ],
+)
+def test_validate_converted_databricks_connection_token(
+ mocker: MockerFixture, rebuilt_uri: str, expected_uri: str
+) -> None:
+ """
+ A converted Databricks connection keeps legacy settings in ``connect_args``
+ while the native form rebuilds them into the URI query. Validation must
+ still test an unchanged connection with its stored token.
+ """
+ extra = json.dumps(
+ {"engine_params": {"connect_args": {"http_path": "/legacy", "catalog":
"main"}}}
+ )
+ existing = mocker.MagicMock()
+ existing.safe_sqlalchemy_uri.return_value = (
+ "databricks://token:XXXXXXXXXX@host:443/default"
+ )
+ existing.sqlalchemy_uri_decrypted =
"databricks://token:real-token@host:443/default"
+ existing.encrypted_extra = "{}"
+ existing.extra = extra
+ existing.ssh_tunnel = None
+
+ database = mocker.MagicMock()
+ with database.get_sqla_engine() as engine:
+ engine.dialect.do_ping.return_value = True
+
+ DatabaseDAO = mocker.patch( # noqa: N806
+ "superset.commands.database.validate.DatabaseDAO"
+ )
+ DatabaseDAO.find_by_id.return_value = existing
+ DatabaseDAO.validate_update_uniqueness.return_value = True
+ DatabaseDAO.build_db_for_connection_test.return_value = database
+
+ mocker.patch(
+ "superset.commands.database.validate.get_engine_spec",
+ return_value=mocker.MagicMock(
+ validate_parameters=mocker.MagicMock(return_value=[]),
Review Comment:
🟡 Medium — An unchanged save of a converted row without a catalog fails
validation
This mock (and `catalog: "main"` in every fixture) hides a gap.
`DatabricksPythonConnectorEngineSpec.required_parameters` still includes
`default_catalog` (`databricks.py:843`, outside the diff, so no applyable
suggestion). Saving an unchanged converted row with no catalog therefore fails
in the modal's pre-save `validate_parameters` with "One or more parameters are
missing: default_catalog". If the user types a catalog to get past it, the
connection identity changes and they must re-enter the token, which contradicts
"an unchanged connection keeps its stored token". Before this PR the form was
blank, so this is an incomplete fix, not a regression, but it breaks the main
edit flow for the same rows as the catalog comment above.
Fix: make the catalog optional (`build_sqlalchemy_uri` already leaves out an
empty catalog):
```python
required_parameters =
DatabricksDynamicBaseEngineSpec.required_parameters | {
"default_schema",
"http_path_field",
}
```
Also add a test for a no-catalog row that calls the real
`validate_parameters` instead of mocking it.
--
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]