This is an automated email from the ASF dual-hosted git repository. beto pushed a commit to branch awm-iam in repository https://gitbox.apache.org/repos/asf/superset.git
commit 4f82393da99045f78f7dfc4815a633696cff6da3 Author: Beto Dealmeida <[email protected]> AuthorDate: Thu Jan 22 18:27:21 2026 -0500 Pivot --- docker-compose-light.yml | 2 + superset/db_engine_specs/aurora.py | 73 ++----------------------- superset/db_engine_specs/postgres.py | 37 +++++++++++++ tests/unit_tests/db_engine_specs/test_aurora.py | 57 ++++++++++++------- 4 files changed, 80 insertions(+), 89 deletions(-) diff --git a/docker-compose-light.yml b/docker-compose-light.yml index b06be681af..f4dd30487e 100644 --- a/docker-compose-light.yml +++ b/docker-compose-light.yml @@ -108,6 +108,8 @@ services: extra_hosts: - "host.docker.internal:host-gateway" user: *superset-user + ports: + - "${SUPERSET_PORT:-8088}:8088" depends_on: superset-init-light: condition: service_completed_successfully diff --git a/superset/db_engine_specs/aurora.py b/superset/db_engine_specs/aurora.py index f925db4323..c7bcbc77a4 100644 --- a/superset/db_engine_specs/aurora.py +++ b/superset/db_engine_specs/aurora.py @@ -14,20 +14,8 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -from __future__ import annotations - -import logging -from typing import Any, TYPE_CHECKING - -from superset.db_engine_specs.aws_iam import AWSIAMAuthMixin from superset.db_engine_specs.mysql import MySQLEngineSpec from superset.db_engine_specs.postgres import PostgresEngineSpec -from superset.utils import json - -if TYPE_CHECKING: - from superset.models.core import Database - -logger = logging.getLogger(__name__) class AuroraMySQLDataAPI(MySQLEngineSpec): @@ -56,67 +44,14 @@ class AuroraPostgresDataAPI(PostgresEngineSpec): ) -class AuroraPostgresEngineSpec(AWSIAMAuthMixin, PostgresEngineSpec): +class AuroraPostgresEngineSpec(PostgresEngineSpec): """ - Aurora PostgreSQL engine spec with IAM authentication support. + Aurora PostgreSQL engine spec. - This engine spec extends PostgresEngineSpec with cross-account IAM - authentication capabilities. It allows Superset to connect to Aurora - PostgreSQL databases using temporary IAM credentials instead of - static database passwords. - - Configuration is provided via the database's encrypted_extra JSON: - - { - "aws_iam": { - "enabled": true, - "role_arn": "arn:aws:iam::222222222222:role/SupersetDatabaseAccess", - "external_id": "superset-prod-12345", // optional - "region": "us-east-1", - "db_username": "superset_iam_user", - "session_duration": 3600 // optional, defaults to 3600 - } - } + IAM authentication is handled by the parent PostgresEngineSpec via + the aws_iam config in encrypted_extra. """ engine = "postgresql" engine_name = "Aurora PostgreSQL" default_driver = "psycopg2" - - # Sensitive fields that should be masked when displaying encrypted_extra - encrypted_extra_sensitive_fields = { - "$.aws_iam.external_id", - "$.aws_iam.role_arn", - } - - @staticmethod - def update_params_from_encrypted_extra( - database: Database, - params: dict[str, Any], - ) -> None: - """ - Extract aws_iam config from encrypted_extra and apply IAM authentication. - - If IAM is not enabled in the configuration, falls back to standard - PostgreSQL behavior (uses credentials from the SQLAlchemy URI). - - :param database: Database model instance - :param params: Engine parameters dict to modify - """ - if not database.encrypted_extra: - return - - try: - encrypted_extra = json.loads(database.encrypted_extra) - except json.JSONDecodeError as ex: - logger.error("Failed to parse encrypted_extra: %s", ex, exc_info=True) - raise - - iam_config = encrypted_extra.get("aws_iam", {}) - - # If IAM is not enabled, fall back to standard auth - if not iam_config.get("enabled", False): - return - - logger.debug("Applying AWS IAM authentication for Aurora PostgreSQL") - AWSIAMAuthMixin._apply_iam_authentication(database, params, iam_config) diff --git a/superset/db_engine_specs/postgres.py b/superset/db_engine_specs/postgres.py index b259164e4f..28e08b8f2e 100644 --- a/superset/db_engine_specs/postgres.py +++ b/superset/db_engine_specs/postgres.py @@ -218,6 +218,12 @@ class PostgresEngineSpec(BasicParametersMixin, PostgresBaseEngineSpec): max_column_name_length = 63 try_remove_schema_from_table_name = False # pylint: disable=invalid-name + # Sensitive fields that should be masked in encrypted_extra + encrypted_extra_sensitive_fields = { + "$.aws_iam.external_id", + "$.aws_iam.role_arn", + } + column_type_mappings = ( ( re.compile(r"^double precision", re.IGNORECASE), @@ -320,6 +326,37 @@ class PostgresEngineSpec(BasicParametersMixin, PostgresBaseEngineSpec): return uri, connect_args + @staticmethod + def update_params_from_encrypted_extra( + database: Database, + params: dict[str, Any], + ) -> None: + """ + Extract sensitive parameters from encrypted_extra. + + Handles AWS IAM authentication if configured, then merges any + remaining encrypted_extra keys into params (standard behavior). + """ + if not database.encrypted_extra: + return + + try: + encrypted_extra = json.loads(database.encrypted_extra) + except json.JSONDecodeError as ex: + logger.error(ex, exc_info=True) + raise + + # Handle AWS IAM auth: pop the key so it doesn't reach create_engine() + iam_config = encrypted_extra.pop("aws_iam", None) + if iam_config and iam_config.get("enabled"): + from superset.db_engine_specs.aws_iam import AWSIAMAuthMixin + + AWSIAMAuthMixin._apply_iam_authentication(database, params, iam_config) + + # Standard behavior: merge remaining keys into params + if encrypted_extra: + params.update(encrypted_extra) + @classmethod def get_default_catalog(cls, database: Database) -> str: """ diff --git a/tests/unit_tests/db_engine_specs/test_aurora.py b/tests/unit_tests/db_engine_specs/test_aurora.py index ca0949892b..a225f3641e 100644 --- a/tests/unit_tests/db_engine_specs/test_aurora.py +++ b/tests/unit_tests/db_engine_specs/test_aurora.py @@ -33,11 +33,10 @@ def test_aurora_postgres_engine_spec_properties() -> None: assert AuroraPostgresEngineSpec.engine == "postgresql" assert AuroraPostgresEngineSpec.engine_name == "Aurora PostgreSQL" assert AuroraPostgresEngineSpec.default_driver == "psycopg2" - assert AuroraPostgresEngineSpec.supports_iam_authentication is True def test_update_params_from_encrypted_extra_without_iam() -> None: - from superset.db_engine_specs.aurora import AuroraPostgresEngineSpec + from superset.db_engine_specs.postgres import PostgresEngineSpec database = MagicMock() database.encrypted_extra = json.dumps({}) @@ -46,14 +45,14 @@ def test_update_params_from_encrypted_extra_without_iam() -> None: ) params: dict[str, Any] = {} - AuroraPostgresEngineSpec.update_params_from_encrypted_extra(database, params) + PostgresEngineSpec.update_params_from_encrypted_extra(database, params) # No modifications should be made assert params == {} def test_update_params_from_encrypted_extra_iam_disabled() -> None: - from superset.db_engine_specs.aurora import AuroraPostgresEngineSpec + from superset.db_engine_specs.postgres import PostgresEngineSpec database = MagicMock() database.encrypted_extra = json.dumps( @@ -71,15 +70,15 @@ def test_update_params_from_encrypted_extra_iam_disabled() -> None: ) params: dict[str, Any] = {} - AuroraPostgresEngineSpec.update_params_from_encrypted_extra(database, params) + PostgresEngineSpec.update_params_from_encrypted_extra(database, params) # No modifications should be made when IAM is disabled assert params == {} def test_update_params_from_encrypted_extra_with_iam() -> None: - from superset.db_engine_specs.aurora import AuroraPostgresEngineSpec from superset.db_engine_specs.aws_iam import AWSIAMAuthMixin + from superset.db_engine_specs.postgres import PostgresEngineSpec database = MagicMock() database.encrypted_extra = json.dumps( @@ -114,7 +113,7 @@ def test_update_params_from_encrypted_extra_with_iam() -> None: return_value="iam-auth-token", ), ): - AuroraPostgresEngineSpec.update_params_from_encrypted_extra(database, params) + PostgresEngineSpec.update_params_from_encrypted_extra(database, params) assert "connect_args" in params assert params["connect_args"]["password"] == "iam-auth-token" # noqa: S105 @@ -122,21 +121,43 @@ def test_update_params_from_encrypted_extra_with_iam() -> None: assert params["connect_args"]["sslmode"] == "require" +def test_update_params_merges_remaining_encrypted_extra() -> None: + from superset.db_engine_specs.postgres import PostgresEngineSpec + + database = MagicMock() + database.encrypted_extra = json.dumps( + { + "aws_iam": {"enabled": False}, + "pool_size": 10, + } + ) + database.sqlalchemy_uri_decrypted = ( + "postgresql://user:[email protected]:5432/mydb" + ) + + params: dict[str, Any] = {} + PostgresEngineSpec.update_params_from_encrypted_extra(database, params) + + # aws_iam should be consumed, pool_size should be merged + assert "aws_iam" not in params + assert params["pool_size"] == 10 + + def test_update_params_from_encrypted_extra_no_encrypted_extra() -> None: - from superset.db_engine_specs.aurora import AuroraPostgresEngineSpec + from superset.db_engine_specs.postgres import PostgresEngineSpec database = MagicMock() database.encrypted_extra = None params: dict[str, Any] = {} - AuroraPostgresEngineSpec.update_params_from_encrypted_extra(database, params) + PostgresEngineSpec.update_params_from_encrypted_extra(database, params) # No modifications should be made assert params == {} def test_update_params_from_encrypted_extra_invalid_json() -> None: - from superset.db_engine_specs.aurora import AuroraPostgresEngineSpec + from superset.db_engine_specs.postgres import PostgresEngineSpec database = MagicMock() database.encrypted_extra = "not-valid-json" @@ -144,25 +165,21 @@ def test_update_params_from_encrypted_extra_invalid_json() -> None: params: dict[str, Any] = {} with pytest.raises(json.JSONDecodeError): - AuroraPostgresEngineSpec.update_params_from_encrypted_extra(database, params) + PostgresEngineSpec.update_params_from_encrypted_extra(database, params) def test_encrypted_extra_sensitive_fields() -> None: - from superset.db_engine_specs.aurora import AuroraPostgresEngineSpec + from superset.db_engine_specs.postgres import PostgresEngineSpec # Verify sensitive fields are properly defined assert ( - "$.aws_iam.external_id" - in AuroraPostgresEngineSpec.encrypted_extra_sensitive_fields - ) - assert ( - "$.aws_iam.role_arn" - in AuroraPostgresEngineSpec.encrypted_extra_sensitive_fields + "$.aws_iam.external_id" in PostgresEngineSpec.encrypted_extra_sensitive_fields ) + assert "$.aws_iam.role_arn" in PostgresEngineSpec.encrypted_extra_sensitive_fields def test_mask_encrypted_extra() -> None: - from superset.db_engine_specs.aurora import AuroraPostgresEngineSpec + from superset.db_engine_specs.postgres import PostgresEngineSpec encrypted_extra = json.dumps( { @@ -176,7 +193,7 @@ def test_mask_encrypted_extra() -> None: } ) - masked = AuroraPostgresEngineSpec.mask_encrypted_extra(encrypted_extra) + masked = PostgresEngineSpec.mask_encrypted_extra(encrypted_extra) assert masked is not None masked_config = json.loads(masked)
