This is an automated email from the ASF dual-hosted git repository. Vitor-Avila pushed a commit to branch chore/mask-trino-presto-druid in repository https://gitbox.apache.org/repos/asf/superset.git
commit 9eefb287f349c32391f0cc628843fcaef0c42f87 Author: Vitor Avila <[email protected]> AuthorDate: Mon Jul 27 21:13:25 2026 -0300 chore: Mask specific Druid/Presto/Trino fields --- superset/db_engine_specs/druid.py | 5 ++ superset/db_engine_specs/presto.py | 6 ++ superset/db_engine_specs/trino.py | 5 ++ tests/unit_tests/db_engine_specs/test_druid.py | 47 +++++++++++++++ tests/unit_tests/db_engine_specs/test_presto.py | 65 ++++++++++++++++++++ tests/unit_tests/db_engine_specs/test_trino.py | 79 +++++++++++++++++++++++++ 6 files changed, 207 insertions(+) diff --git a/superset/db_engine_specs/druid.py b/superset/db_engine_specs/druid.py index 61845909c2a..992b997f1b6 100644 --- a/superset/db_engine_specs/druid.py +++ b/superset/db_engine_specs/druid.py @@ -53,6 +53,11 @@ class DruidEngineSpec(BaseEngineSpec): type_probe_needs_row = True requires_column_value_normalization = True + encrypted_extra_sensitive_fields = { + "$.connect_args.jwt": "JWT Token", + "$.connect_args.password": "Password", + } + metadata = { "description": ( "Apache Druid is a high performance real-time analytics database." diff --git a/superset/db_engine_specs/presto.py b/superset/db_engine_specs/presto.py index 3f100a17134..cc16484ae0b 100644 --- a/superset/db_engine_specs/presto.py +++ b/superset/db_engine_specs/presto.py @@ -166,6 +166,12 @@ class PrestoBaseEngineSpec(BaseEngineSpec, metaclass=ABCMeta): supports_dynamic_schema = True supports_catalog = supports_dynamic_catalog = supports_cross_catalog_queries = True + + encrypted_extra_sensitive_fields = { + "$.auth_params.password": "Password", + "$.auth_params.token": "JWT Token", + "$.connect_args.requests_kwargs.jwt": "JWT Token", + } # Not set here: GROUPING SETS support is opted in per-concrete-engine # (``PrestoEngineSpec``, ``TrinoEngineSpec``) rather than on this shared # base, since Hive-family descendants (``HiveEngineSpec``, ``SparkEngineSpec``, diff --git a/superset/db_engine_specs/trino.py b/superset/db_engine_specs/trino.py index 1ad98d06bd1..78c18858b1d 100644 --- a/superset/db_engine_specs/trino.py +++ b/superset/db_engine_specs/trino.py @@ -73,6 +73,11 @@ class TrinoEngineSpec(PrestoBaseEngineSpec): allows_alias_to_source_column = False supports_grouping_sets = True + encrypted_extra_sensitive_fields = { + **PrestoBaseEngineSpec.encrypted_extra_sensitive_fields, + "$.oauth2_client_info.secret": "OAuth2 client secret", + } + # The full set of columns Trino's "<table>$partitions" exposes for an # Iceberg table. The real partition keys are nested in the "partition" ROW, # so none of these are user partition columns. diff --git a/tests/unit_tests/db_engine_specs/test_druid.py b/tests/unit_tests/db_engine_specs/test_druid.py index 6a55c1e0b9a..85c5fa28638 100644 --- a/tests/unit_tests/db_engine_specs/test_druid.py +++ b/tests/unit_tests/db_engine_specs/test_druid.py @@ -211,3 +211,50 @@ def test_non_string_cursor_type_unaffected_by_druid_spec() -> None: col = result_set.columns[0] assert col["type"] == "INT" + + +def test_mask_encrypted_extra() -> None: + """ + Only the credentials inside `connect_args` should be masked, not the whole object. + """ + from superset.db_engine_specs.druid import DruidEngineSpec + from superset.utils import json + + config = json.dumps( + { + "connect_args": { + "scheme": "https", + "jwt": "my-secret-token", + "password": "my-password", + }, + } + ) + + assert DruidEngineSpec.mask_encrypted_extra(config) == json.dumps( + { + "connect_args": { + "scheme": "https", + "jwt": "XXXXXXXXXX", + "password": "XXXXXXXXXX", + }, + } + ) + + +def test_unmask_encrypted_extra() -> None: + """ + Masked credentials are reused from the previous value; edited ones are kept. + """ + from superset.db_engine_specs.druid import DruidEngineSpec + from superset.utils import json + + old = json.dumps( + {"connect_args": {"scheme": "https", "jwt": "old-token", "password": "old"}} + ) + new = json.dumps( + {"connect_args": {"scheme": "http", "jwt": "XXXXXXXXXX", "password": "new"}} + ) + + assert DruidEngineSpec.unmask_encrypted_extra(old, new) == json.dumps( + {"connect_args": {"scheme": "http", "jwt": "old-token", "password": "new"}} + ) diff --git a/tests/unit_tests/db_engine_specs/test_presto.py b/tests/unit_tests/db_engine_specs/test_presto.py index 5a22747a0a2..e9a56ba3d15 100644 --- a/tests/unit_tests/db_engine_specs/test_presto.py +++ b/tests/unit_tests/db_engine_specs/test_presto.py @@ -476,3 +476,68 @@ def test_partition_query_escapes_single_quote_in_filter_value( # by injected SQL) must NOT appear anywhere in the output — that would # mean the payload broke out of the literal. assert "'2024-01-01' UNION SELECT" not in sql + + +def test_mask_encrypted_extra() -> None: + """ + The sensitive `auth_params` values are masked, while `auth_method` and + non-sensitive fields such as `username` stay visible. + """ + from superset.db_engine_specs.presto import PrestoEngineSpec + from superset.utils import json + + config = json.dumps( + { + "auth_method": "basic", + "auth_params": {"username": "alice", "password": "my-password"}, + } + ) + + assert PrestoEngineSpec.mask_encrypted_extra(config) == json.dumps( + { + "auth_method": "basic", + "auth_params": {"username": "alice", "password": "XXXXXXXXXX"}, + } + ) + + +def test_mask_encrypted_extra_jwt_in_connect_args() -> None: + """ + A JWT passed via `connect_args.requests_kwargs` is masked without touching + the surrounding connection settings. + """ + from superset.db_engine_specs.presto import PrestoEngineSpec + from superset.utils import json + + config = json.dumps( + { + "connect_args": { + "protocol": "https", + "requests_kwargs": {"jwt": "my-secret-token"}, + }, + } + ) + + assert PrestoEngineSpec.mask_encrypted_extra(config) == json.dumps( + { + "connect_args": { + "protocol": "https", + "requests_kwargs": {"jwt": "XXXXXXXXXX"}, + }, + } + ) + + +def test_unmask_encrypted_extra() -> None: + """ + Masked credentials are reused from the previous value; edited ones are kept. + """ + from superset.db_engine_specs.presto import PrestoEngineSpec + from superset.utils import json + + old = json.dumps({"auth_method": "jwt", "auth_params": {"token": "old-token"}}) + new = json.dumps({"auth_method": "jwt", "auth_params": {"token": "XXXXXXXXXX"}}) + + assert PrestoEngineSpec.unmask_encrypted_extra(old, new) == json.dumps( + {"auth_method": "jwt", "auth_params": {"token": "old-token"}} + ) diff --git a/tests/unit_tests/db_engine_specs/test_trino.py b/tests/unit_tests/db_engine_specs/test_trino.py index 64b2ffa94a5..b7cb70d5374 100644 --- a/tests/unit_tests/db_engine_specs/test_trino.py +++ b/tests/unit_tests/db_engine_specs/test_trino.py @@ -1650,3 +1650,82 @@ def test_handle_boolean_filter() -> None: str(result_computed.compile(compile_kwargs={"literal_binds": True})) == "(expiration = 1) = true" ) + + +def test_mask_encrypted_extra() -> None: + """ + All `auth_params` values and the OAuth2 client secret are masked, while + `auth_method` and other non-sensitive fields stay visible. + """ + from superset.db_engine_specs.trino import TrinoEngineSpec + + config = json.dumps( + { + "auth_method": "jwt", + "auth_params": {"token": "my-secret-token"}, + "oauth2_client_info": {"id": "client-id", "secret": "my-secret"}, + } + ) + + assert TrinoEngineSpec.mask_encrypted_extra(config) == json.dumps( + { + "auth_method": "jwt", + "auth_params": {"token": "XXXXXXXXXX"}, + "oauth2_client_info": {"id": "client-id", "secret": "XXXXXXXXXX"}, + } + ) + + +def test_mask_encrypted_extra_jwt_in_connect_args() -> None: + """ + A JWT passed via `connect_args.requests_kwargs` is masked without touching + the surrounding connection settings. + """ + from superset.db_engine_specs.trino import TrinoEngineSpec + + config = json.dumps( + { + "connect_args": { + "protocol": "https", + "requests_kwargs": {"jwt": "my-secret-token"}, + }, + } + ) + + assert TrinoEngineSpec.mask_encrypted_extra(config) == json.dumps( + { + "connect_args": { + "protocol": "https", + "requests_kwargs": {"jwt": "XXXXXXXXXX"}, + }, + } + ) + + +def test_unmask_encrypted_extra() -> None: + """ + Masked credentials are reused from the previous value; edited ones are kept. + """ + from superset.db_engine_specs.trino import TrinoEngineSpec + + old = json.dumps( + { + "auth_method": "basic", + "auth_params": {"username": "alice", "password": "old-password"}, + } + ) + # `username` is not masked on read, so it comes back in cleartext; only the + # masked `password` is revealed from the previous value. + new = json.dumps( + { + "auth_method": "basic", + "auth_params": {"username": "alice", "password": "XXXXXXXXXX"}, + } + ) + + assert TrinoEngineSpec.unmask_encrypted_extra(old, new) == json.dumps( + { + "auth_method": "basic", + "auth_params": {"username": "alice", "password": "old-password"}, + } + )
