fix(ingestion): refresh RDS IAM auth token per connection for MySQL (open-metadata#28730)

* fix(ingestion): refresh RDS IAM auth token per connection for MySQL

RDS IAM auth tokens expire after ~15 minutes. The token was generated
once in get_connection_url_common and baked into the SQLAlchemy URL, so
a single engine reused one frozen token for every pooled connection and
never refreshed it. With max_overflow=-1 and multi-threaded extraction,
any connection opened after the 15-minute TTL — common on large catalogs
and during profiling — authenticated with a stale token and failed mid
run with "Access denied".

Add RdsIamAuthTokenManager (caches the token, parses its expiry from the
presigned-URL params, refreshes before it lapses) and wire a do_connect
event listener in the MySQL connection handler that injects a fresh token
on every new connection instead of embedding it in the URL. SSL is forced
on for IAM (PyMySQL requires it) while preserving any existing SSL config.

Scoped to the MySQL connector. The shared builders.py IAM path is still
token-frozen for the other RDS connectors (Postgres, Redshift, Greenplum,
Timescale); the reusable token manager lives in aws_client.py so they can
adopt it next.

* fix: Thread-safety in RdsIamAuthTokenManager, username not URL-encoded

* fix: base_url drops databaseSchema + connectionOptions

* fix(ingestion): enforce required TLS for MySQL RDS IAM connections

The do_connect listener set cparams["ssl"] = {} when no SSL config was
present. An empty dict is falsy, so PyMySQL only enables TLS in PREFERRED
mode (self.ssl=True, _ssl_required=False), which silently falls back to
plaintext if the server doesn't offer TLS. RDS IAM auth mandates TLS, so
this is the wrong guarantee.

Inject {"check_hostname": True} instead: a truthy dict makes PyMySQL treat
SSL as required (_ssl_required=True) and verifies the RDS server cert.
Explicitly provided ssl config is still preserved.

Update test_listener_enables_ssl_required_by_pymysql_for_iam to assert a
truthy value, and add test_injected_ssl_makes_pymysql_require_tls which
drives the value through real PyMySQL and asserts conn.ssl is True and
conn._ssl_required is True, so the test verifies actual TLS-required
behavior rather than the mock value.

---------

Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>

Author: harshsoni2024

ingestion/src/metadata/clients/aws_client.py

@@ -13,9 +13,11 @@
 """
 
 import datetime
+import threading
 from enum import Enum
 from functools import partial
 from typing import Any, Callable, Dict, Optional, Type, TypeVar  # noqa: UP035
+from urllib.parse import parse_qs, urlparse
 
 import boto3
 import botocore.session
@@ -263,3 +265,88 @@ def get_redshift_serverless_client(self):
 
     def get_mwaa_client(self):
         return self.get_client(AWSServices.MWAA.value)
+
+
+RDS_IAM_TOKEN_DEFAULT_TTL = datetime.timedelta(minutes=15)
+RDS_IAM_TOKEN_REFRESH_THRESHOLD = datetime.timedelta(minutes=5)
+
+
+class RdsIamAuthTokenManager:
+    """
+    Manages the lifecycle of an AWS RDS IAM authentication token.
+
+    RDS IAM tokens are short-lived (~15 minutes) presigned URLs. A long ingestion
+    that opens new pooled connections after the token expires would otherwise
+    authenticate with a stale token and fail. This manager caches the current
+    token, derives its expiry from the presigned URL, and regenerates it shortly
+    before it lapses so every connection receives a valid token.
+    """
+
+    def __init__(
+        self,
+        host: str,
+        port: str,
+        username: str,
+        aws_config: AWSCredentials,
+        refresh_threshold: datetime.timedelta = RDS_IAM_TOKEN_REFRESH_THRESHOLD,
+    ):
+        self.host = host
+        self.port = port
+        self.username = username
+        self.aws_config = aws_config
+        self.refresh_threshold = refresh_threshold
+        self._token: Optional[str] = None  # noqa: UP045
+        self._expires_at: Optional[datetime.datetime] = None  # noqa: UP045
+        self._lock = threading.Lock()
+
+    def get_token(self) -> str:
+        """Return a valid token, refreshing if needed.
+
+        The check-and-refresh is serialized: the engine shares one manager across
+        all worker threads (each calls ``engine.connect()`` from the ``do_connect``
+        listener), so without the lock multiple threads could refresh concurrently
+        and observe a token paired with a stale expiry.
+        """
+        with self._lock:
+            if self._needs_refresh():
+                self._refresh_token()
+            if self._token is None:
+                raise RuntimeError("Failed to generate RDS IAM authentication token")
+            return self._token
+
+    def _needs_refresh(self) -> bool:
+        needs_refresh = True
+        if self._token is not None and self._expires_at is not None:
+            time_left = self._expires_at - datetime.datetime.now(datetime.timezone.utc)
+            needs_refresh = time_left <= self.refresh_threshold
+        return needs_refresh
+
+    def _refresh_token(self) -> None:
+        logger.debug(f"Generating RDS IAM auth token for {self.username}@{self.host}")
+        rds_client = AWSClient(config=self.aws_config).get_rds_client()
+        token = rds_client.generate_db_auth_token(
+            DBHostname=self.host,
+            Port=self.port,
+            DBUsername=self.username,
+            Region=self.aws_config.awsRegion,
+        )
+        self._token = token
+        self._expires_at = self._parse_token_expiry(token)
+
+    def _parse_token_expiry(self, token: str) -> datetime.datetime:
+        """Derive token expiry from the presigned URL's X-Amz-Date / X-Amz-Expires.
+
+        Falls back to a conservative default TTL if the token can't be parsed so a
+        malformed token still triggers periodic refresh rather than never expiring.
+        """
+        now = datetime.datetime.now(datetime.timezone.utc)
+        expires_at = now + RDS_IAM_TOKEN_DEFAULT_TTL
+        try:
+            query_params = parse_qs(urlparse(token).query)
+            amz_date = query_params["X-Amz-Date"][0]
+            amz_expires = int(query_params["X-Amz-Expires"][0])
+            issued_at = datetime.datetime.strptime(amz_date, "%Y%m%dT%H%M%SZ").replace(tzinfo=datetime.timezone.utc)
+            expires_at = issued_at + datetime.timedelta(seconds=amz_expires)
+        except (KeyError, ValueError, IndexError) as exc:
+            logger.warning(f"Could not parse RDS IAM token expiry, using default TTL: {exc}")
+        return expires_at

ingestion/src/metadata/ingestion/source/database/mysql/connection.py

@@ -13,10 +13,12 @@
 Source connection handler
 """
 
-from typing import Optional
+from typing import Any, Dict, Optional, cast  # noqa: UP035
 
 from sqlalchemy.engine import Engine
+from sqlalchemy.event import listen
 
+from metadata.clients.aws_client import RdsIamAuthTokenManager
 from metadata.generated.schema.entity.automations.workflow import (
     Workflow as AutomationWorkflow,
 )
@@ -29,6 +31,9 @@
 from metadata.generated.schema.entity.services.connections.database.common.gcpCloudSqlConfig import (
     GcpCloudsqlConfigurationSource,
 )
+from metadata.generated.schema.entity.services.connections.database.common.iamAuthConfig import (
+    IamAuthConfigurationSource,
+)
 from metadata.generated.schema.entity.services.connections.database.mysqlConnection import (
     MysqlConnection as MySQLConnectionConfig,
 )
@@ -72,12 +77,63 @@ def _get_client(self) -> Engine:
         if isinstance(connection.authType, GcpCloudsqlConfigurationSource):
             return self._get_cloudsql_engine(connection)
 
+        if isinstance(connection.authType, IamAuthConfigurationSource):
+            return self._get_iam_engine(connection)
+
         return create_generic_db_connection(
             connection=connection,
             get_connection_url_fn=get_connection_url_common,
             get_connection_args_fn=get_connection_args_common,
         )
 
+    def _get_iam_engine(self, connection: MySQLConnectionConfig) -> Engine:
+        """Build an engine that refreshes the RDS IAM token per connection.
+
+        RDS IAM tokens expire after ~15 minutes. Rather than baking a single token
+        into the connection URL (which would go stale for connections opened later
+        in a long ingestion), a ``do_connect`` listener injects a freshly minted
+        token on every new connection.
+        """
+        auth_type = cast("IamAuthConfigurationSource", connection.authType)
+        if auth_type.awsConfig is None:
+            raise ValueError("awsConfig is required for MySQL RDS IAM authentication")
+
+        host, port = connection.hostPort.split(":")
+        token_manager = RdsIamAuthTokenManager(
+            host=host,
+            port=port,
+            username=connection.username,
+            aws_config=auth_type.awsConfig,
+        )
+        engine = create_generic_db_connection(
+            connection=connection,
+            get_connection_url_fn=self._build_iam_url,
+            get_connection_args_fn=get_connection_args_common,
+        )
+
+        def inject_iam_token(_dialect, _conn_rec, _cargs, cparams: Dict[str, Any]):  # noqa: UP006
+            cparams["password"] = token_manager.get_token()
+            # RDS IAM auth requires TLS. A truthy ssl dict makes PyMySQL treat SSL
+            # as required (an empty dict only yields PREFERRED, which can silently
+            # fall back to plaintext). check_hostname also verifies the RDS cert.
+            # Any explicitly provided ssl config is preserved.
+            if "ssl" not in cparams:
+                cparams["ssl"] = {"check_hostname": True}
+
+        listen(engine, "do_connect", inject_iam_token)
+        return engine
+
+    @staticmethod
+    def _build_iam_url(connection: MySQLConnectionConfig) -> str:
+        """Build the connection URL exactly like ``get_connection_url_common`` but
+        without a password/token: the ``do_connect`` listener injects a fresh RDS
+        IAM token per connection. Reusing the common helper with the IAM auth
+        neutralized keeps databaseSchema and connectionOptions handling in sync.
+        """
+        url_connection = connection.model_copy()
+        url_connection.authType = BasicAuth(password="")  # type: ignore
+        return get_connection_url_common(url_connection)
+
     def _get_cloudsql_engine(self, connection: MySQLConnectionConfig) -> Engine:
         try:
             from google.cloud.sql.connectors import Connector  # noqa: PLC0415

ingestion/tests/unit/connections/test_iam_token_refresh.py

@@ -0,0 +1,127 @@
+#  Copyright 2025 Collate
+#  Licensed under the Collate Community License, Version 1.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  https://github.com/open-metadata/OpenMetadata/blob/main/ingestion/LICENSE
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+"""
+Characterization tests for the SHARED RDS IAM auth path in ``builders.py``.
+
+RDS IAM auth tokens expire after ~15 minutes. ``get_connection_url_common`` mints
+the token once and bakes it into the connection URL string, so a single engine
+reuses one frozen token for every pooled connection and cannot refresh it.
+
+The MySQL connector now works around this at the connector level (it builds the
+engine without a token in the URL and attaches a ``do_connect`` listener that
+injects a fresh token per connection) — see
+``tests/unit/source/database/test_mysql_iam.py``.
+
+These tests pin the behaviour of the SHARED helper, which is still token-frozen
+and is the remaining gap for the other RDS connectors that go through it
+(Postgres, Redshift, Greenplum, Timescale). They are expected to FAIL — and
+should be flipped to assert refreshed behaviour — once the shared path is fixed
+too. ``test_iam_token_is_used_for_authentication`` is the invariant that holds
+either way.
+"""
+
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from metadata.generated.schema.entity.services.connections.database.common.iamAuthConfig import (
+    IamAuthConfigurationSource,
+)
+from metadata.generated.schema.entity.services.connections.database.mysqlConnection import (
+    MysqlConnection,
+)
+from metadata.generated.schema.security.credentials.awsCredentials import (
+    AWSCredentials,
+)
+from metadata.ingestion.connections import builders
+
+IAM_TOKEN = "iam-token-v1"
+HOST = "myrds.abc.us-east-1.rds.amazonaws.com"
+PORT = "3306"
+USERNAME = "iam_user"
+REGION = "us-east-1"
+
+
+def _iam_connection() -> MysqlConnection:
+    return MysqlConnection(
+        username=USERNAME,
+        hostPort=f"{HOST}:{PORT}",
+        authType=IamAuthConfigurationSource(awsConfig=AWSCredentials(awsRegion=REGION)),
+    )
+
+
+@pytest.fixture
+def mock_rds_token():
+    """Patch AWSClient so generate_db_auth_token returns a fixed token and is counted."""
+    with patch.object(builders, "AWSClient") as mock_aws_client:
+        rds_client = MagicMock()
+        rds_client.generate_db_auth_token.return_value = IAM_TOKEN
+        mock_aws_client.return_value.get_rds_client.return_value = rds_client
+        yield rds_client
+
+
+class TestRdsIamTokenRefresh:
+    def test_iam_token_is_used_for_authentication(self, mock_rds_token):
+        """Invariant: the generated IAM token authenticates the connection.
+
+        Holds before and after the B1 fix.
+        """
+        url = builders.get_connection_url_common(_iam_connection())
+
+        assert IAM_TOKEN in url
+        assert mock_rds_token.generate_db_auth_token.call_args.kwargs == {
+            "DBHostname": HOST,
+            "Port": PORT,
+            "DBUsername": USERNAME,
+            "Region": REGION,
+        }
+
+    def test_iam_token_is_generated_once_and_frozen_in_url(self, mock_rds_token):
+        """Documents the remaining shared-path gap: token baked into the URL string.
+
+        ``get_connection_url_common`` embeds the token in the URL, so every engine
+        built from it reuses that single token for its whole lifetime — it is never
+        refreshed and goes stale after the ~15 min RDS IAM token TTL. The MySQL
+        connector bypasses this path; the connectors that still rely on it do not.
+
+        EXPECTED TO FAIL once the shared path is fixed to mint a token per
+        connection rather than embedding it in the URL.
+        """
+        conn = _iam_connection()
+
+        first_url = builders.get_connection_url_common(conn)
+        second_url = builders.get_connection_url_common(conn)
+
+        assert IAM_TOKEN in first_url, "token is embedded directly in the URL (the bug)"
+        assert first_url == second_url, "same frozen token is reused, never refreshed"
+
+    def test_no_do_connect_listener_is_registered_for_iam(self, mock_rds_token):
+        """Documents the shared-path gap: no per-connection token injection hook.
+
+        The shared ``create_generic_db_connection`` does not attach a ``do_connect``
+        listener for IAM, so an engine built purely through it cannot refresh the
+        baked-in token once it expires. (The MySQL connector adds this listener
+        itself; see test_mysql_iam.py.)
+
+        EXPECTED TO FAIL once the shared path grows its own IAM listener.
+        """
+        with patch.object(builders, "listen") as mock_listen:
+            builders.create_generic_db_connection(
+                connection=_iam_connection(),
+                get_connection_url_fn=builders.get_connection_url_common,
+                get_connection_args_fn=builders.get_connection_args_common,
+            )
+
+        do_connect_registrations = [call for call in mock_listen.call_args_list if "do_connect" in call.args]
+        assert not do_connect_registrations, (
+            "no do_connect listener is wired for IAM today — this is the gap B1 "
+            "describes; once fixed, invert this assertion"
+        )