rebenitez1802 commented on code in PR #45933:
URL: https://github.com/apache/superset/pull/45933#discussion_r4447649092
##########
superset/models/core.py:
##########
@@ -768,6 +768,15 @@ def _get_sqla_engine( # pylint: disable=too-many-locals
# noqa: C901
self.update_params_from_encrypted_extra(engine_kwargs)
+ # ``adjust_engine_params`` applies the engine spec's enforced
connection
+ # arguments, but ``encrypted_extra`` is merged afterwards and can
replace
+ # ``connect_args`` wholesale. Re-apply the enforced arguments last so
they
+ # stay authoritative regardless of what ``encrypted_extra`` provides.
+ if enforced_connect_args :=
self.db_engine_spec.enforce_uri_query_params.get(
+ sqlalchemy_url.get_driver_name(), {}
+ ):
Review Comment:
nit: worth calling out a side effect here.
`DorisEngineSpec.adjust_engine_params` skips `super()`, and pydoris'
`DorisDialect` subclasses `MySQLDialect_mysqldb` (driver `mysqldb`), so
`doris://` connections now get `local_infile=0` too (before: `{}`, or
`local_infile=1` if set via `extra.engine_params.connect_args`). That's a good
change, but it would be nice to mention it in the PR description and/or cover
it with a test so it's clearly intended.
##########
tests/unit_tests/models/core_test.py:
##########
@@ -678,6 +678,34 @@ def
test_get_sqla_engine_honors_adjusted_connect_args(mocker: MockerFixture) ->
assert kwargs["connect_args"].get("local_infile") == 0
+def test_get_sqla_engine_enforced_connect_args_outlast_encrypted_extra(
+ mocker: MockerFixture,
+) -> None:
+ """
+ ``encrypted_extra`` is merged into ``engine_kwargs`` after
+ ``adjust_engine_params`` has applied the engine spec's enforced
+ ``connect_args``. The enforced values must stay authoritative, so an
+ ``encrypted_extra`` that also provides ``connect_args`` cannot replace
them.
+ Exercised via MySQL, which enforces ``local_infile=0``.
+ """
+ from superset.models.core import Database
+
+ create_engine_mock = mocker.patch(
+ "superset.models.core.create_engine",
+ return_value=create_engine("sqlite://"),
+ )
+
+ database = Database(
+ database_name="my_db",
+ sqlalchemy_uri="mysql://u:p@h/db",
+ encrypted_extra=json.dumps({"connect_args": {"local_infile": 1}}),
+ )
+ database._get_sqla_engine(nullpool=False)
+
+ _, kwargs = create_engine_mock.call_args
+ assert kwargs["connect_args"].get("local_infile") == 0
Review Comment:
nit: asserting that an unrelated key survives would lock in the
merge-not-replace behavior:
```suggestion
encrypted_extra=json.dumps(
{"connect_args": {"local_infile": 1, "ssl": {"ca": "/x"}}}
),
)
database._get_sqla_engine(nullpool=False)
_, kwargs = create_engine_mock.call_args
assert kwargs["connect_args"].get("local_infile") == 0
assert kwargs["connect_args"].get("ssl") == {"ca": "/x"}
```
##########
superset/models/core.py:
##########
@@ -768,6 +768,15 @@ def _get_sqla_engine( # pylint: disable=too-many-locals
# noqa: C901
self.update_params_from_encrypted_extra(engine_kwargs)
+ # ``adjust_engine_params`` applies the engine spec's enforced
connection
+ # arguments, but ``encrypted_extra`` is merged afterwards and can
replace
+ # ``connect_args`` wholesale. Re-apply the enforced arguments last so
they
+ # stay authoritative regardless of what ``encrypted_extra`` provides.
Review Comment:
nit: this is only true for the base `adjust_engine_params`; Doris, StarRocks
and SingleStore override it without calling `super()`, so they never get the
enforced args from there. Suggested wording:
```suggestion
# The base ``adjust_engine_params`` applies the engine spec's
enforced
# connection arguments, but overrides that skip ``super()`` (e.g.
Doris)
# never apply them, and ``encrypted_extra`` is merged afterwards and
can
# replace ``connect_args`` wholesale. Re-apply the enforced
arguments last
# so they stay authoritative for every spec.
```
##########
superset/models/core.py:
##########
@@ -768,6 +768,15 @@ def _get_sqla_engine( # pylint: disable=too-many-locals
# noqa: C901
self.update_params_from_encrypted_extra(engine_kwargs)
+ # ``adjust_engine_params`` applies the engine spec's enforced
connection
+ # arguments, but ``encrypted_extra`` is merged afterwards and can
replace
+ # ``connect_args`` wholesale. Re-apply the enforced arguments last so
they
+ # stay authoritative regardless of what ``encrypted_extra`` provides.
+ if enforced_connect_args :=
self.db_engine_spec.enforce_uri_query_params.get(
+ sqlalchemy_url.get_driver_name(), {}
+ ):
+ engine_kwargs.setdefault("connect_args",
{}).update(enforced_connect_args)
Review Comment:
nit: if `encrypted_extra` stores a non-dict `connect_args` (e.g. `null`),
this raises a raw `AttributeError` before the `try` that maps driver errors via
`get_dbapi_mapped_exception` (on master it surfaced as a mapped `TypeError`
from `create_engine`). A list of pairs also used to work and now fails.
Building a new dict handles both:
```suggestion
engine_kwargs["connect_args"] = {
**dict(engine_kwargs.get("connect_args") or {}),
**enforced_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]
