acryl-datahub 1.3.0.1rc4__py3-none-any.whl → 1.3.0.1rc5__py3-none-any.whl

{acryl_datahub-1.3.0.1rc4.dist-info → acryl_datahub-1.3.0.1rc5.dist-info}/RECORD

@@ -1,7 +1,7 @@
-acryl_datahub-1.3.0.1rc4.dist-info/licenses/LICENSE,sha256=9xNHpsD0uYF5ONzXsKDCuHHB-xbiCrSbueWXqrTNsxk,11365
+acryl_datahub-1.3.0.1rc5.dist-info/licenses/LICENSE,sha256=9xNHpsD0uYF5ONzXsKDCuHHB-xbiCrSbueWXqrTNsxk,11365
 datahub/__init__.py,sha256=aq_i5lVREmoLfYIqcx_pEQicO855YlhD19tWc1eZZNI,59
 datahub/__main__.py,sha256=pegIvQ9hzK7IhqVeUi1MeADSZ2QlP-D3K0OQdEg55RU,106
-datahub/_version.py,sha256=wuuyfFhYvCgNlVjy7dlOQ6sSJuFvaUAnTlhSVJzQ-fM,323
+datahub/_version.py,sha256=GFe5nZs9PKs-LLDaT1H1D9udtmOwDf_NsyGlgBGOywE,323
 datahub/entrypoints.py,sha256=VcbU6Z47b_JKW1zI-WJMYIngm05FSogKLiuvFNtyNcI,9088
 datahub/errors.py,sha256=p5rFAdAGVCk4Lqolol1YvthceadUSwpaCxLXRcyCCFQ,676
 datahub/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -241,7 +241,7 @@ datahub/ingestion/source/abs/source.py,sha256=z86K5_P_gu8kTytLOAYyQqqD2g14JGSrv1
 datahub/ingestion/source/apply/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 datahub/ingestion/source/apply/datahub_apply.py,sha256=xTD-Iq3UHhxcz61RwNuI2kJjRrnQEfZFSgvS1X6loV4,7703
 datahub/ingestion/source/aws/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-datahub/ingestion/source/aws/aws_common.py,sha256=1NgwpasOO3fpBGtZnxWEzjHFoHjEs5hpg_HkTvsbn0M,18147
+datahub/ingestion/source/aws/aws_common.py,sha256=Va9uxo5aKsAR7qIC625VpRO3XDqzNIg4SfK_eFg25Rw,23781
 datahub/ingestion/source/aws/glue.py,sha256=dUaMWcI5Ed-TzbbSrF6suT4L1vcRHoHfFCdTvAINc4w,67423
 datahub/ingestion/source/aws/platform_resource_repository.py,sha256=0eUfGy1FbaBltCSNTtXyLrkrdqTc1KkTgDJB1Gd-Ydk,853
 datahub/ingestion/source/aws/s3_boto_utils.py,sha256=rGlWAkKZpkeA1_wMvcJvSDvobvduShszowU-KcrQudg,7011
@@ -271,7 +271,7 @@ datahub/ingestion/source/bigquery_v2/bigquery_platform_resource_helper.py,sha256
 datahub/ingestion/source/bigquery_v2/bigquery_queries.py,sha256=2syDMaRpYEbtGUVejVAK5d6g8HqM54ZyEM908uLJ55o,3393
 datahub/ingestion/source/bigquery_v2/bigquery_report.py,sha256=zlTkqOmt5zxnO40rVTYHF3fclj4OVlLtqUXwW5WIIcM,7855
 datahub/ingestion/source/bigquery_v2/bigquery_schema.py,sha256=zbYb1EYnCJxgvsU8oT_76l0q_BW1exVjMWM1GAgd1nc,32600
-datahub/ingestion/source/bigquery_v2/bigquery_schema_gen.py,sha256=_NLFRRXsrxMZ8Vjg2jVL4Pg1_NGt9hzn9EWBooJZ8so,51566
+datahub/ingestion/source/bigquery_v2/bigquery_schema_gen.py,sha256=PbSCMj5ACwEu_HQNe29IHs4y1bn15_nnz6ZW1Yt17wI,51796
 datahub/ingestion/source/bigquery_v2/bigquery_test_connection.py,sha256=cATxwi5IPzj3BldRRAVcLqzSFmmYEPvqa7U0RFJbaAc,7645
 datahub/ingestion/source/bigquery_v2/common.py,sha256=IinOy-RO4UZGxSf5scaN02672BzZuNsjJZ56axti6iI,4016
 datahub/ingestion/source/bigquery_v2/lineage.py,sha256=jju14mJbAUMA_K3j2yq-TdZV202cjd5rBAsDPJGEVno,44900
@@ -456,9 +456,9 @@ datahub/ingestion/source/redshift/profile.py,sha256=H1Xtc2rXScUv4w0b2BbM7POjYEwq
 datahub/ingestion/source/redshift/query.py,sha256=HKobQ-0crARgT8Mkfe-WBqVR9ZadYCZ9DGaUoEHHHww,48234
 datahub/ingestion/source/redshift/redshift.py,sha256=RN8rao3j7nocnnD6oPcEju09-8mOZTE4vFkgy_13Az8,41293
 datahub/ingestion/source/redshift/redshift_data_reader.py,sha256=zc69jwXHdF-w8J4Hq-ZQ6BjHQ75Ij2iNDMpoRJlcmlU,1724
-datahub/ingestion/source/redshift/redshift_schema.py,sha256=7F-l_omOuKMuGE_rBWXVPG_GWXFKnCMzC4frNxZB9cs,24800
+datahub/ingestion/source/redshift/redshift_schema.py,sha256=2U8IIPRJkL-HWUeWswOzvcT1hdTBQgPMhr6tYCDuqrM,25226
 datahub/ingestion/source/redshift/report.py,sha256=aCFDFUbz5xde8b_eRIHSBiELoo9LZFtDpp2lSadiPHU,2937
-datahub/ingestion/source/redshift/usage.py,sha256=Q7R-caJovLXv33uZepMGX5Cvm4DqQSLZdiL_s-p06wU,17473
+datahub/ingestion/source/redshift/usage.py,sha256=szdg3cUw4UpZ8rXMZufIAVTq2iiI1VsmrFgjAEH8Dv4,17521
 datahub/ingestion/source/s3/__init__.py,sha256=HjqFPj11WtNFZM3kcVshlDb7kOsc19-l_3LM8PBjlJM,56
 datahub/ingestion/source/s3/config.py,sha256=lElFXgEpKDT9SVoiXvtx98wV6Gp880qP4pLQaOGJGOo,7828
 datahub/ingestion/source/s3/datalake_profiler_config.py,sha256=FfrcgK-JEF94vw-l3q6pN6FENXb-wZzW2w1VUZVkwW8,3620
@@ -523,9 +523,9 @@ datahub/ingestion/source/sql/hana.py,sha256=V6bGVLVjI1VL0deebg8VxIL8Ls-oxUvpSvX9
 datahub/ingestion/source/sql/hive.py,sha256=SPmAWlk63V-s-loBTU2hXsQA7xA4sa0iPK6pCbF-AJ8,31600
 datahub/ingestion/source/sql/hive_metastore.py,sha256=UBB7mV2eKuCxv3voi0F3tqF2MyRObSYxArAxETZfO4E,35997
 datahub/ingestion/source/sql/mariadb.py,sha256=om6QoG5UtDldt1N6AfIWp3T-HXNaaqFmpz2i0JAemfM,654
-datahub/ingestion/source/sql/mysql.py,sha256=_KhTODU7mqAoJOlrvRdPa7ihQkYLkgrZwaseQbasotM,5358
+datahub/ingestion/source/sql/mysql.py,sha256=h0kv86-8SxBTmaEhmcyqGcKoaWQ4peUdiE9sfhrTuqY,9734
 datahub/ingestion/source/sql/oracle.py,sha256=nKMM1O67SkxCgT781eENl5xXpIR8_p5joTSdAYzQwHY,29988
-datahub/ingestion/source/sql/postgres.py,sha256=blkO6bI0eDKFK8UNwUYcYtm_ObrQuWVSy5GyfdhL5dg,14274
+datahub/ingestion/source/sql/postgres.py,sha256=Vk1NVI0zJPzMj4SKJ9jfyGu4DY3bow724BDn10BaxzU,17478
 datahub/ingestion/source/sql/presto.py,sha256=58py4M3UYxkGpbBFA1o96H154eUhD2dBm1hpxxYlYYM,4256
 datahub/ingestion/source/sql/sql_common.py,sha256=EZGoeGlOYZoOrXOiKDI-S1mw-sPVV33PZQ_mPJlEvRc,57759
 datahub/ingestion/source/sql/sql_config.py,sha256=u3nGZYYl1WtaxfNsDU5bglgZ5Jq3Fxk9xei_CUIAXB0,8222
@@ -535,7 +535,7 @@ datahub/ingestion/source/sql/sql_report.py,sha256=gw-OPHSExp_b6DRjvwqE1U6Bpkwekx
 datahub/ingestion/source/sql/sql_types.py,sha256=AVeBBXw8aKB1_jw6Wtg58miu-YUfN_-7ZcXwSF-ESgA,16021
 datahub/ingestion/source/sql/sql_utils.py,sha256=q-Bsk6WxlsRtrw9RXBxvqI3zuaMTC_F25T2VrCziR9I,8418
 datahub/ingestion/source/sql/sqlalchemy_data_reader.py,sha256=FvHZ4JEK3aR2DYOBZiT_ZsAy12RjTu4t_KIR_92B11k,2644
-datahub/ingestion/source/sql/sqlalchemy_uri.py,sha256=u0ZvgdJjXZdo_vl7YIQfYuuWbGwpnH6OSozI2e8ZV4I,858
+datahub/ingestion/source/sql/sqlalchemy_uri.py,sha256=wE_GX2TtkFEvscC_Epy5hUfBxtE6mUQoVPb7fUea0jk,1882
 datahub/ingestion/source/sql/sqlalchemy_uri_mapper.py,sha256=KOpbmDIE2h1hyYEsbVHJi2B7FlsyUMTXZx4diyzltQg,1826
 datahub/ingestion/source/sql/teradata.py,sha256=YydlPGndFGZcpvlmim3T-1yaAmsFt08TZVOTo1R3GLo,66871
 datahub/ingestion/source/sql/trino.py,sha256=o5hm84iwRHO59TD2LaEqYgF2LYIcSUIKmlgu1VudGBY,19254
@@ -646,12 +646,12 @@ datahub/lite/lite_registry.py,sha256=bpH0kasP-LtwwUFNA2QsOIehfekAYfJtN-AkQLmSWnw
 datahub/lite/lite_server.py,sha256=p9Oa2nNs65mqcssSIVOr7VOzWqfVstz6ZQEdT4f82S0,1949
 datahub/lite/lite_util.py,sha256=G0LQHKkyEb1pc_q183g6hflShclGx7kikgMaOxtVVcs,4545
 datahub/metadata/__init__.py,sha256=AjhXPjI6cnpdcrBRrE5gOWo15vv2TTl2ctU4UAnUN7A,238
-datahub/metadata/_internal_schema_classes.py,sha256=iJEKSTxDX1uH8YgUwxt0_HSd2AryK0-vUTBeN-7wG0Q,1076970
-datahub/metadata/schema.avsc,sha256=sn-1LEnMogcM9WhP7f8PWa_I7mgLALVeN75XtuCd4hY,812831
+datahub/metadata/_internal_schema_classes.py,sha256=M3j1TDrK43RYlhKUsaqeHgjbw5ERkAzKxdC4pD4dLEg,1077060
+datahub/metadata/schema.avsc,sha256=mUP8XLRosJg0WlBymK0-EAlGqrC-j7bQIDbOVtODrR8,775591
 datahub/metadata/schema_classes.py,sha256=tPT8iHCak4IsZi_oL0nirbPpI8ETTPTZzapqLRpeKU4,1326
 datahub/metadata/urns.py,sha256=nfrCTExR-k2P9w272WVtWSN3xW1VUJngPwP3xnvULjU,1217
 datahub/metadata/_urns/__init__.py,sha256=cOF3GHMDgPhmbLKbN02NPpuLGHSu0qNgQyBRv08eqF0,243
-datahub/metadata/_urns/urn_defs.py,sha256=q1vWaLDYps2cfLh5IYsWUrbFYx_N2txLAFIl9vKdq-M,143257
+datahub/metadata/_urns/urn_defs.py,sha256=_LgqKLHrmHHxpvrP-93NMJSLEnoFI8q72lkX17mK1XA,143257
 datahub/metadata/com/__init__.py,sha256=gsAIuTxzfJdI7a9ybZlgMIHMAYksM1SxGxXjtySgKSc,202
 datahub/metadata/com/linkedin/__init__.py,sha256=gsAIuTxzfJdI7a9ybZlgMIHMAYksM1SxGxXjtySgKSc,202
 datahub/metadata/com/linkedin/events/__init__.py,sha256=s_dR0plZF-rOxxIbE8ojekJqwiHzl2WYR-Z3kW6kKS0,298
@@ -782,7 +782,7 @@ datahub/metadata/schemas/DataHubIngestionSourceInfo.avsc,sha256=4wac7sluRIq-0ZjO
 datahub/metadata/schemas/DataHubIngestionSourceKey.avsc,sha256=TGmm9WEGTaABs7kt5Uc-N-kbc5Sd-2sQwx-JpfAptvw,545
 datahub/metadata/schemas/DataHubOpenAPISchemaKey.avsc,sha256=q6ZyMoxInwmrkrXkUgMe-i-WZzAxbjcvJ-EI99SnEp8,599
 datahub/metadata/schemas/DataHubPageModuleKey.avsc,sha256=NyFN8cVO6s6rtgoLGJJGfcPfpGr5PfmZlIhM6ajldfQ,460
-datahub/metadata/schemas/DataHubPageModuleProperties.avsc,sha256=53Fj4ztBJqo9QMWuza2Kdtfpr2nTOTW0XuuXW77ugB8,10347
+datahub/metadata/schemas/DataHubPageModuleProperties.avsc,sha256=hbIEkpjQxVxSZOQmTVIjOGnGTNeaA0r0oYNKPpwuddg,10443
 datahub/metadata/schemas/DataHubPageTemplateKey.avsc,sha256=0sVqwL97Rp8YHPytp2RqUP5hIW048hmT2hPNP5k6arc,472
 datahub/metadata/schemas/DataHubPageTemplateProperties.avsc,sha256=FyNcZIniQy9m6yN9DT4XsPkDrxUsU7tRTqmfdGoEtMU,8565
 datahub/metadata/schemas/DataHubPersonaInfo.avsc,sha256=OUvbTgPQsBtzkDDb9pxHXpQ6A7dkL77ZnCXZ-MLEG14,227
@@ -1128,8 +1128,8 @@ datahub_provider/operators/datahub_assertion_operator.py,sha256=uvTQ-jk2F0sbqqxp
 datahub_provider/operators/datahub_assertion_sensor.py,sha256=lCBj_3x1cf5GMNpHdfkpHuyHfVxsm6ff5x2Z5iizcAo,140
 datahub_provider/operators/datahub_operation_operator.py,sha256=aevDp2FzX7FxGlXrR0khoHNbxbhKR2qPEX5e8O2Jyzw,174
 datahub_provider/operators/datahub_operation_sensor.py,sha256=8fcdVBCEPgqy1etTXgLoiHoJrRt_nzFZQMdSzHqSG7M,168
-acryl_datahub-1.3.0.1rc4.dist-info/METADATA,sha256=cJb6Hz3UpjAA15FINI4riQLYt7C5Q-aUsoWz2pelBp0,184504
-acryl_datahub-1.3.0.1rc4.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-acryl_datahub-1.3.0.1rc4.dist-info/entry_points.txt,sha256=pzsBoTx-D-iTcmpX8oCGCyzlHP2112EygUMzZWz56M8,10105
-acryl_datahub-1.3.0.1rc4.dist-info/top_level.txt,sha256=iLjSrLK5ox1YVYcglRUkcvfZPvKlobBWx7CTUXx8_GI,25
-acryl_datahub-1.3.0.1rc4.dist-info/RECORD,,
+acryl_datahub-1.3.0.1rc5.dist-info/METADATA,sha256=ckjfgTlPEUgZH1sYvZm3MLupBZvBbsC-zokD3Q2ekno,184688
+acryl_datahub-1.3.0.1rc5.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+acryl_datahub-1.3.0.1rc5.dist-info/entry_points.txt,sha256=pzsBoTx-D-iTcmpX8oCGCyzlHP2112EygUMzZWz56M8,10105
+acryl_datahub-1.3.0.1rc5.dist-info/top_level.txt,sha256=iLjSrLK5ox1YVYcglRUkcvfZPvKlobBWx7CTUXx8_GI,25
+acryl_datahub-1.3.0.1rc5.dist-info/RECORD,,

datahub/_version.py

@@ -1,6 +1,6 @@
 # Published at https://pypi.org/project/acryl-datahub/.
 __package_name__ = "acryl-datahub"
-__version__ = "1.3.0.1rc4"
+__version__ = "1.3.0.1rc5"
 
 
 def is_dev_mode() -> bool:

datahub/ingestion/source/aws/aws_common.py

@@ -3,11 +3,13 @@ from datetime import datetime, timedelta, timezone
 from enum import Enum
 from http import HTTPStatus
 from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Tuple, Union
+from urllib.parse import parse_qs, urlparse
 
 import boto3
 import requests
 from boto3.session import Session
 from botocore.config import DEFAULT_TIMEOUT, Config
+from botocore.exceptions import ClientError, NoCredentialsError
 from botocore.utils import fix_s3_host
 from pydantic.fields import Field
 
@@ -465,6 +467,165 @@ class AwsConnectionConfig(ConfigModel):
     def get_lakeformation_client(self) -> "LakeFormationClient":
         return self.get_session().client("lakeformation", config=self._aws_config())
 
+    def get_rds_client(self):
+        """Get an RDS client for generating IAM auth tokens."""
+        return self.get_session().client("rds", config=self._aws_config())
+
+
+def generate_rds_iam_token(
+    endpoint: str,
+    username: str,
+    port: int,
+    aws_config: AwsConnectionConfig,
+) -> str:
+    """
+    Generate an AWS RDS IAM authentication token.
+
+    boto3's generate_db_auth_token() returns a presigned URL in the format:
+    "hostname:port/?Action=connect&DBUser=username&X-Amz-Date=...&X-Amz-Expires=..."
+
+    This token should be used as-is by pymysql/psycopg2 drivers.
+
+    Args:
+        endpoint: RDS endpoint hostname
+        username: Database username for IAM authentication
+        port: Database port (5432 for PostgreSQL, 3306 for MySQL)
+        aws_config: AwsConnectionConfig for session management and credentials
+
+    Returns:
+        Authentication token (presigned URL format)
+
+    Raises:
+        ValueError: If AWS credentials are not found or token generation fails
+
+    """
+    try:
+        client = aws_config.get_rds_client()
+        token = client.generate_db_auth_token(
+            DBHostname=endpoint, Port=port, DBUsername=username
+        )
+        logger.debug(f"Generated RDS IAM token for {username}@{endpoint}:{port}")
+        return token
+    except NoCredentialsError as e:
+        raise ValueError("AWS credentials not found") from e
+    except ClientError as e:
+        raise ValueError(f"Failed to generate RDS IAM token: {e}") from e
+
+
+class RDSIAMTokenManager:
+    """
+    Manages RDS IAM token lifecycle with automatic refresh.
+
+    RDS IAM tokens include expiration information in the URL parameters.
+    This manager parses the token expiry and refreshes before expiration
+    to ensure uninterrupted database access.
+    """
+
+    def __init__(
+        self,
+        endpoint: str,
+        username: str,
+        port: int,
+        aws_config: AwsConnectionConfig,
+        refresh_threshold_minutes: int = 5,
+    ):
+        """
+        Initialize the token manager.
+
+        Args:
+            endpoint: RDS endpoint hostname
+            username: Database username for IAM authentication
+            port: Database port
+            aws_config: AwsConnectionConfig for session management and credentials
+            refresh_threshold_minutes: Refresh token when this many minutes remain before expiry
+        """
+        self.endpoint = endpoint
+        self.username = username
+        self.port = port
+        self.aws_config = aws_config
+        self.refresh_threshold = timedelta(minutes=refresh_threshold_minutes)
+
+        self._current_token: Optional[str] = None
+        self._token_expires_at: Optional[datetime] = None
+
+    def get_token(self) -> str:
+        """
+        Get current token, refreshing if necessary.
+
+        Returns:
+            Valid authentication token
+
+        Raises:
+            RuntimeError: If token generation or refresh fails
+        """
+        if self._needs_refresh():
+            self._refresh_token()
+
+        assert self._current_token is not None
+        return self._current_token
+
+    def _needs_refresh(self) -> bool:
+        """Check if token needs to be refreshed."""
+        if self._current_token is None or self._token_expires_at is None:
+            return True
+
+        time_until_expiry = self._token_expires_at - datetime.now(timezone.utc)
+        return time_until_expiry <= self.refresh_threshold
+
+    def _parse_token_expiry(self, token: str) -> datetime:
+        """
+        Parse token expiry from X-Amz-Date and X-Amz-Expires URL parameters.
+
+        Args:
+            token: RDS IAM authentication token (presigned URL)
+
+        Returns:
+            Expiration datetime in UTC
+
+        Raises:
+            ValueError: If token URL format is invalid or missing required parameters
+        """
+        try:
+            parsed_url = urlparse(token)
+            query_params = parse_qs(parsed_url.query)
+
+            # Extract X-Amz-Date (ISO 8601 format: YYYYMMDDTHHMMSSZ)
+            amz_date_list = query_params.get("X-Amz-Date")
+            if not amz_date_list:
+                raise ValueError("Missing X-Amz-Date parameter in RDS IAM token")
+            amz_date_str = amz_date_list[0]
+
+            # Extract X-Amz-Expires (duration in seconds)
+            amz_expires_list = query_params.get("X-Amz-Expires")
+            if not amz_expires_list:
+                raise ValueError("Missing X-Amz-Expires parameter in RDS IAM token")
+            amz_expires_seconds = int(amz_expires_list[0])
+
+            # Parse X-Amz-Date to datetime
+            token_issued_at = datetime.strptime(amz_date_str, "%Y%m%dT%H%M%SZ").replace(
+                tzinfo=timezone.utc
+            )
+
+            # Calculate expiration
+            return token_issued_at + timedelta(seconds=amz_expires_seconds)
+
+        except (ValueError, KeyError, IndexError) as e:
+            raise ValueError(
+                f"Failed to parse RDS IAM token expiry: {e}. Token format may be invalid."
+            ) from e
+
+    def _refresh_token(self) -> None:
+        """Generate and store a new token with parsed expiry."""
+        logger.info("Refreshing RDS IAM authentication token")
+        self._current_token = generate_rds_iam_token(
+            endpoint=self.endpoint,
+            username=self.username,
+            port=self.port,
+            aws_config=self.aws_config,
+        )
+        self._token_expires_at = self._parse_token_expiry(self._current_token)
+        logger.debug(f"Token will expire at {self._token_expires_at}")
+
 
 class AwsSourceConfig(EnvConfigMixin, AwsConnectionConfig):
     """

datahub/ingestion/source/bigquery_v2/bigquery_schema_gen.py

@@ -449,10 +449,12 @@ class BigQuerySchemaGenerator:
                 ):
                     yield wu
             except Exception as e:
-                if self.config.is_profiling_enabled():
-                    action_mesage = "Does your service account have bigquery.tables.list, bigquery.routines.get, bigquery.routines.list permission, bigquery.tables.getData permission?"
+                # If configuration indicates we need table data access (for profiling or use_tables_list_query_v2),
+                # include bigquery.tables.getData in the error message since that's likely the missing permission
+                if self.config.have_table_data_read_permission:
+                    action_mesage = "Does your service account have bigquery.tables.list, bigquery.routines.get, bigquery.routines.list, bigquery.tables.getData permissions?"
                 else:
-                    action_mesage = "Does your service account have bigquery.tables.list, bigquery.routines.get, bigquery.routines.list permission?"
+                    action_mesage = "Does your service account have bigquery.tables.list, bigquery.routines.get, bigquery.routines.list permissions?"
 
                 self.report.failure(
                     title="Unable to get tables for dataset",

datahub/ingestion/source/redshift/redshift_schema.py

@@ -15,6 +15,7 @@ from datahub.ingestion.source.sql.sql_generic import BaseColumn, BaseTable
 from datahub.metadata.com.linkedin.pegasus2avro.schema import SchemaField
 from datahub.sql_parsing.sqlglot_lineage import SqlParsingResult
 from datahub.utilities.hive_schema_to_avro import get_schema_fields_for_hive_column
+from datahub.utilities.perf_timer import PerfTimer
 
 logger: logging.Logger = logging.getLogger(__name__)
 
@@ -243,9 +244,13 @@ class RedshiftDataDictionary:
         conn: redshift_connector.Connection, query: str
     ) -> redshift_connector.Cursor:
         cursor: redshift_connector.Cursor = conn.cursor()
-
-        logger.debug(f"Query : {query}")
-        cursor.execute(query)
+        with PerfTimer() as timer:
+            query_hash_id = hash(query)
+            logger.info(f"Executing query [{query_hash_id}]\n{query}")
+            cursor.execute(query)
+            logger.info(
+                f"Time taken query [{query_hash_id}: {timer.elapsed_seconds():.3f} seconds"
+            )
         return cursor
 
     @staticmethod
@@ -545,8 +550,7 @@ class RedshiftDataDictionary:
         conn: redshift_connector.Connection,
         query: str,
     ) -> Iterable[LineageRow]:
-        cursor = conn.cursor()
-        cursor.execute(query)
+        cursor = RedshiftDataDictionary.get_query_result(conn=conn, query=query)
         field_names = [i[0] for i in cursor.description]
 
         rows = cursor.fetchmany()
@@ -603,9 +607,7 @@ class RedshiftDataDictionary:
         conn: redshift_connector.Connection,
         query: str,
     ) -> Iterable[TempTableRow]:
-        cursor = conn.cursor()
-
-        cursor.execute(query)
+        cursor = RedshiftDataDictionary.get_query_result(conn=conn, query=query)
 
         field_names = [i[0] for i in cursor.description]
 
@@ -662,8 +664,9 @@ class RedshiftDataDictionary:
     def get_outbound_datashares(
         conn: redshift_connector.Connection,
     ) -> Iterable[OutboundDatashare]:
-        cursor = conn.cursor()
-        cursor.execute(RedshiftCommonQuery.list_outbound_datashares())
+        cursor = RedshiftDataDictionary.get_query_result(
+            conn=conn, query=RedshiftCommonQuery.list_outbound_datashares()
+        )
         for item in cursor.fetchall():
             yield OutboundDatashare(
                 share_name=item[1],
@@ -678,8 +681,10 @@ class RedshiftDataDictionary:
         conn: redshift_connector.Connection,
         database: str,
     ) -> Optional[InboundDatashare]:
-        cursor = conn.cursor()
-        cursor.execute(RedshiftCommonQuery.get_inbound_datashare(database))
+        cursor = RedshiftDataDictionary.get_query_result(
+            conn=conn,
+            query=RedshiftCommonQuery.get_inbound_datashare(database),
+        )
         item = cursor.fetchone()
         if item:
             return InboundDatashare(

datahub/ingestion/source/redshift/usage.py

@@ -25,6 +25,7 @@ from datahub.ingestion.source.redshift.query import (
     RedshiftServerlessQuery,
 )
 from datahub.ingestion.source.redshift.redshift_schema import (
+    RedshiftDataDictionary,
     RedshiftTable,
     RedshiftView,
 )
@@ -263,8 +264,7 @@ class RedshiftUsageExtractor:
         connection: redshift_connector.Connection,
         all_tables: Dict[str, Dict[str, List[Union[RedshiftView, RedshiftTable]]]],
     ) -> Iterable[RedshiftAccessEvent]:
-        cursor = connection.cursor()
-        cursor.execute(query)
+        cursor = RedshiftDataDictionary.get_query_result(conn=connection, query=query)
         results = cursor.fetchmany()
         field_names = [i[0] for i in cursor.description]
         while results:

datahub/ingestion/source/sql/mysql.py

@@ -1,14 +1,17 @@
 # This import verifies that the dependencies are available.
-
-from typing import List
+import logging
+from typing import TYPE_CHECKING, Any, List, Optional
 
 import pymysql  # noqa: F401
 from pydantic.fields import Field
-from sqlalchemy import util
+from sqlalchemy import create_engine, event, inspect, util
 from sqlalchemy.dialects.mysql import BIT, base
 from sqlalchemy.dialects.mysql.enumerated import SET
 from sqlalchemy.engine.reflection import Inspector
 
+if TYPE_CHECKING:
+    from sqlalchemy.engine import Engine
+
 from datahub.configuration.common import AllowDenyPattern, HiddenFromDocs
 from datahub.ingestion.api.decorators import (
     SourceCapability,
@@ -18,11 +21,16 @@ from datahub.ingestion.api.decorators import (
     platform_name,
     support_status,
 )
+from datahub.ingestion.source.aws.aws_common import (
+    AwsConnectionConfig,
+    RDSIAMTokenManager,
+)
 from datahub.ingestion.source.sql.sql_common import (
     make_sqlalchemy_type,
     register_custom_type,
 )
 from datahub.ingestion.source.sql.sql_config import SQLAlchemyConnectionConfig
+from datahub.ingestion.source.sql.sqlalchemy_uri import parse_host_port
 from datahub.ingestion.source.sql.stored_procedures.base import (
     BaseProcedure,
 )
@@ -31,6 +39,9 @@ from datahub.ingestion.source.sql.two_tier_sql_source import (
     TwoTierSQLAlchemySource,
 )
 from datahub.metadata.schema_classes import BytesTypeClass
+from datahub.utilities.str_enum import StrEnum
+
+logger = logging.getLogger(__name__)
 
 SET.__repr__ = util.generic_repr  # type:ignore
 
@@ -54,11 +65,33 @@ base.ischema_names["polygon"] = POLYGON
 base.ischema_names["decimal128"] = DECIMAL128
 
 
+class MySQLAuthMode(StrEnum):
+    """Authentication mode for MySQL connection."""
+
+    PASSWORD = "PASSWORD"
+    AWS_IAM = "AWS_IAM"
+
+
 class MySQLConnectionConfig(SQLAlchemyConnectionConfig):
     # defaults
     host_port: str = Field(default="localhost:3306", description="MySQL host URL.")
     scheme: HiddenFromDocs[str] = "mysql+pymysql"
 
+    # Authentication configuration
+    auth_mode: MySQLAuthMode = Field(
+        default=MySQLAuthMode.PASSWORD,
+        description="Authentication mode to use for the MySQL connection. "
+        "Options are 'PASSWORD' (default) for standard username/password authentication, "
+        "or 'AWS_IAM' for AWS RDS IAM authentication.",
+    )
+    aws_config: AwsConnectionConfig = Field(
+        default_factory=AwsConnectionConfig,
+        description="AWS configuration for RDS IAM authentication (only used when auth_mode is AWS_IAM). "
+        "Provides full control over AWS credentials, region, profiles, role assumption, retry logic, and proxy settings. "
+        "If not explicitly configured, boto3 will automatically use the default credential chain and region from "
+        "environment variables (AWS_DEFAULT_REGION, AWS_REGION), AWS config files (~/.aws/config), or IAM role metadata.",
+    )
+
 
 class MySQLConfig(MySQLConnectionConfig, TwoTierSQLAlchemyConfig):
     def get_identifier(self, *, schema: str, table: str) -> str:
@@ -91,9 +124,27 @@ class MySQLSource(TwoTierSQLAlchemySource):
     Table, row, and column statistics via optional SQL profiling
     """
 
-    def __init__(self, config, ctx):
+    config: MySQLConfig
+
+    def __init__(self, config: MySQLConfig, ctx: Any):
         super().__init__(config, ctx, self.get_platform())
 
+        self._rds_iam_token_manager: Optional[RDSIAMTokenManager] = None
+        if config.auth_mode == MySQLAuthMode.AWS_IAM:
+            hostname, port = parse_host_port(config.host_port, default_port=3306)
+            if port is None:
+                raise ValueError("Port must be specified for RDS IAM authentication")
+
+            if not config.username:
+                raise ValueError("username is required for RDS IAM authentication")
+
+            self._rds_iam_token_manager = RDSIAMTokenManager(
+                endpoint=hostname,
+                username=config.username,
+                port=port,
+                aws_config=config.aws_config,
+            )
+
     def get_platform(self):
         return "mysql"
 
@@ -102,6 +153,52 @@ class MySQLSource(TwoTierSQLAlchemySource):
         config = MySQLConfig.parse_obj(config_dict)
         return cls(config, ctx)
 
+    def _setup_rds_iam_event_listener(
+        self, engine: "Engine", database_name: Optional[str] = None
+    ) -> None:
+        """Setup SQLAlchemy event listener to inject RDS IAM tokens."""
+        if not (
+            self.config.auth_mode == MySQLAuthMode.AWS_IAM
+            and self._rds_iam_token_manager
+        ):
+            return
+
+        def do_connect_listener(_dialect, _conn_rec, _cargs, cparams):
+            if not self._rds_iam_token_manager:
+                raise RuntimeError("RDS IAM Token Manager is not initialized")
+            cparams["password"] = self._rds_iam_token_manager.get_token()
+            # PyMySQL requires SSL to be enabled for RDS IAM authentication.
+            # Preserve any existing SSL configuration, otherwise enable with default settings.
+            # The {"ssl": True} dict is a workaround to make PyMySQL recognize that SSL
+            # should be enabled, since the library requires a truthy value in the ssl parameter.
+            # See https://pymysql.readthedocs.io/en/latest/modules/connections.html#pymysql.connections.Connection
+            cparams["ssl"] = cparams.get("ssl") or {"ssl": True}
+
+        event.listen(engine, "do_connect", do_connect_listener)  # type: ignore[misc]
+
+    def get_inspectors(self):
+        url = self.config.get_sql_alchemy_url()
+        logger.debug(f"sql_alchemy_url={url}")
+
+        engine = create_engine(url, **self.config.options)
+        self._setup_rds_iam_event_listener(engine)
+
+        with engine.connect() as conn:
+            inspector = inspect(conn)
+            if self.config.database and self.config.database != "":
+                databases = [self.config.database]
+            else:
+                databases = inspector.get_schema_names()
+            for db in databases:
+                if self.config.database_pattern.allowed(db):
+                    url = self.config.get_sql_alchemy_url(current_db=db)
+                    db_engine = create_engine(url, **self.config.options)
+                    self._setup_rds_iam_event_listener(db_engine, database_name=db)
+
+                    with db_engine.connect() as conn:
+                        inspector = inspect(conn)
+                        yield inspector
+
     def add_profile_metadata(self, inspector: Inspector) -> None:
         if not self.config.is_profiling_enabled():
             return