chore: Mask specific Druid/Presto/Trino fields

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."

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``,

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.

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"}}
+    )

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"}}
+    )

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"},
+        }
+    )