acryl-datahub 1.0.0rc18__py3-none-any.whl → 1.3.0.1rc9__py3-none-any.whl

datahub/ingestion/sink/datahub_rest.py

@@ -3,7 +3,6 @@ import contextlib
 import dataclasses
 import functools
 import logging
-import os
 import threading
 import uuid
 from enum import auto
@@ -16,15 +15,18 @@ from datahub.configuration.common import (
     ConfigurationError,
     OperationalError,
 )
+from datahub.configuration.env_vars import (
+    get_rest_sink_default_max_threads,
+    get_rest_sink_default_mode,
+)
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.emitter.mcp_builder import mcps_from_mce
 from datahub.emitter.rest_emitter import (
     BATCH_INGEST_MAX_PAYLOAD_LENGTH,
-    DEFAULT_REST_SINK_ENDPOINT,
-    DEFAULT_REST_TRACE_MODE,
+    DEFAULT_REST_EMITTER_ENDPOINT,
     DataHubRestEmitter,
+    EmitMode,
     RestSinkEndpoint,
-    RestTraceMode,
 )
 from datahub.ingestion.api.common import RecordEnvelope, WorkUnit
 from datahub.ingestion.api.sink import (
@@ -34,7 +36,7 @@ from datahub.ingestion.api.sink import (
     WriteCallback,
 )
 from datahub.ingestion.api.workunit import MetadataWorkUnit
-from datahub.ingestion.graph.client import DatahubClientConfig
+from datahub.ingestion.graph.config import ClientMode, DatahubClientConfig
 from datahub.metadata.com.linkedin.pegasus2avro.mxe import (
     MetadataChangeEvent,
     MetadataChangeProposal,
@@ -48,9 +50,7 @@ from datahub.utilities.server_config_util import set_gms_config
 
 logger = logging.getLogger(__name__)
 
-_DEFAULT_REST_SINK_MAX_THREADS = int(
-    os.getenv("DATAHUB_REST_SINK_DEFAULT_MAX_THREADS", 15)
-)
+_DEFAULT_REST_SINK_MAX_THREADS = get_rest_sink_default_max_threads()
 
 
 class RestSinkMode(ConfigEnum):
@@ -64,14 +64,14 @@ class RestSinkMode(ConfigEnum):
 
 
 _DEFAULT_REST_SINK_MODE = pydantic.parse_obj_as(
-    RestSinkMode, os.getenv("DATAHUB_REST_SINK_DEFAULT_MODE", RestSinkMode.ASYNC_BATCH)
+    RestSinkMode, get_rest_sink_default_mode() or RestSinkMode.ASYNC_BATCH
 )
 
 
 class DatahubRestSinkConfig(DatahubClientConfig):
     mode: RestSinkMode = _DEFAULT_REST_SINK_MODE
-    endpoint: RestSinkEndpoint = DEFAULT_REST_SINK_ENDPOINT
-    default_trace_mode: RestTraceMode = DEFAULT_REST_TRACE_MODE
+    endpoint: RestSinkEndpoint = DEFAULT_REST_EMITTER_ENDPOINT
+    server_config_refresh_interval: Optional[int] = None
 
     # These only apply in async modes.
     max_threads: pydantic.PositiveInt = _DEFAULT_REST_SINK_MAX_THREADS
@@ -92,6 +92,7 @@ class DatahubRestSinkConfig(DatahubClientConfig):
 @dataclasses.dataclass
 class DataHubRestSinkReport(SinkReport):
     mode: Optional[RestSinkMode] = None
+    endpoint: Optional[RestSinkEndpoint] = None
     max_threads: Optional[int] = None
     gms_version: Optional[str] = None
     pending_requests: int = 0
@@ -134,18 +135,15 @@ class DatahubRestSink(Sink[DatahubRestSinkConfig, DataHubRestSinkReport]):
         self._emitter_thread_local = threading.local()
 
         try:
-            gms_config = self.emitter.get_server_config()
+            gms_config = self.emitter.server_config
         except Exception as exc:
             raise ConfigurationError(
                 f"💥 Failed to connect to DataHub with {repr(self.emitter)}"
             ) from exc
 
-        self.report.gms_version = (
-            gms_config.get("versions", {})
-            .get("acryldata/datahub", {})
-            .get("version", None)
-        )
+        self.report.gms_version = gms_config.service_version
         self.report.mode = self.config.mode
+        self.report.endpoint = self.config.endpoint
         self.report.max_threads = self.config.max_threads
         logger.debug("Setting env variables to override config")
         logger.debug("Setting gms config")
@@ -179,7 +177,8 @@ class DatahubRestSink(Sink[DatahubRestSinkConfig, DataHubRestSinkReport]):
             client_certificate_path=config.client_certificate_path,
             disable_ssl_verification=config.disable_ssl_verification,
             openapi_ingestion=config.endpoint == RestSinkEndpoint.OPENAPI,
-            default_trace_mode=config.default_trace_mode == RestTraceMode.ENABLED,
+            client_mode=config.client_mode,
+            datahub_component=config.datahub_component,
         )
 
     @property
@@ -190,6 +189,7 @@ class DatahubRestSink(Sink[DatahubRestSinkConfig, DataHubRestSinkReport]):
         # https://github.com/psf/requests/issues/1871#issuecomment-32751346
         thread_local = self._emitter_thread_local
         if not hasattr(thread_local, "emitter"):
+            self.config.client_mode = ClientMode.INGESTION
             thread_local.emitter = DatahubRestSink._make_emitter(self.config)
         return thread_local.emitter
 
@@ -253,9 +253,10 @@ class DatahubRestSink(Sink[DatahubRestSinkConfig, DataHubRestSinkReport]):
             MetadataChangeProposal,
             MetadataChangeProposalWrapper,
         ],
+        emit_mode: EmitMode,
     ) -> None:
         # TODO: Add timing metrics
-        self.emitter.emit(record)
+        self.emitter.emit(record, emit_mode=emit_mode)
 
     def _emit_batch_wrapper(
         self,
@@ -270,8 +271,10 @@ class DatahubRestSink(Sink[DatahubRestSinkConfig, DataHubRestSinkReport]):
         ],
     ) -> None:
         events: List[Union[MetadataChangeProposal, MetadataChangeProposalWrapper]] = []
+
         for record in records:
             event = record[0]
+
             if isinstance(event, MetadataChangeEvent):
                 # Unpack MCEs into MCPs.
                 mcps = mcps_from_mce(event)
@@ -279,7 +282,7 @@ class DatahubRestSink(Sink[DatahubRestSinkConfig, DataHubRestSinkReport]):
             else:
                 events.append(event)
 
-        chunks = self.emitter.emit_mcps(events)
+        chunks = self.emitter.emit_mcps(events, emit_mode=EmitMode.ASYNC)
         self.report.async_batches_prepared += 1
         if chunks > 1:
             self.report.async_batches_split += chunks
@@ -310,6 +313,7 @@ class DatahubRestSink(Sink[DatahubRestSinkConfig, DataHubRestSinkReport]):
                     partition_key,
                     self._emit_wrapper,
                     record,
+                    EmitMode.ASYNC,
                     done_callback=functools.partial(
                         self._write_done_callback, record_envelope, write_callback
                     ),
@@ -321,6 +325,7 @@ class DatahubRestSink(Sink[DatahubRestSinkConfig, DataHubRestSinkReport]):
                 self.executor.submit(
                     partition_key,
                     record,
+                    EmitMode.ASYNC,
                     done_callback=functools.partial(
                         self._write_done_callback, record_envelope, write_callback
                     ),
@@ -329,7 +334,7 @@ class DatahubRestSink(Sink[DatahubRestSinkConfig, DataHubRestSinkReport]):
             else:
                 # execute synchronously
                 try:
-                    self._emit_wrapper(record)
+                    self._emit_wrapper(record, emit_mode=EmitMode.SYNC_PRIMARY)
                     write_callback.on_success(record_envelope, success_metadata={})
                 except Exception as e:
                     write_callback.on_failure(record_envelope, e, failure_metadata={})
@@ -341,11 +346,14 @@ class DatahubRestSink(Sink[DatahubRestSinkConfig, DataHubRestSinkReport]):
         ],
     ) -> None:
         return self.write_record_async(
-            RecordEnvelope(item, metadata={}),
-            NoopWriteCallback(),
+            RecordEnvelope(item, metadata={}), NoopWriteCallback()
         )
 
     def close(self):
+        # Execute pre-shutdown callbacks first (handled by parent class)
+        super().close()
+
+        # Then perform sink-specific shutdown
         with self.report.main_thread_blocking_timer:
             self.executor.shutdown()
 

datahub/ingestion/sink/file.py

@@ -79,6 +79,7 @@ class FileSink(Sink[FileSinkConfig, SinkReport]):
             write_callback.on_success(record_envelope, {})
 
     def close(self):
+        super().close()
         self.file.write("\n]")
         self.file.close()
 

datahub/ingestion/source/abs/config.py

@@ -151,7 +151,7 @@ class DataLakeSourceConfig(
             raise ValueError("platform must not be empty")
         return platform
 
-    @pydantic.root_validator()
+    @pydantic.root_validator(skip_on_failure=True)
     def ensure_profiling_pattern_is_passed_to_profiling(
         cls, values: Dict[str, Any]
     ) -> Dict[str, Any]:

datahub/ingestion/source/abs/datalake_profiler_config.py

@@ -72,7 +72,7 @@ class DataLakeProfilerConfig(ConfigModel):
         description="Whether to profile for the sample values for all columns.",
     )
 
-    @pydantic.root_validator()
+    @pydantic.root_validator(skip_on_failure=True)
     def ensure_field_level_settings_are_normalized(
         cls: "DataLakeProfilerConfig", values: Dict[str, Any]
     ) -> Dict[str, Any]:

datahub/ingestion/source/abs/source.py

@@ -44,7 +44,11 @@ from datahub.ingestion.source.azure.abs_utils import (
     get_key_prefix,
     strip_abs_prefix,
 )
-from datahub.ingestion.source.data_lake_common.data_lake_utils import ContainerWUCreator
+from datahub.ingestion.source.common.subtypes import SourceCapabilityModifier
+from datahub.ingestion.source.data_lake_common.data_lake_utils import (
+    ContainerWUCreator,
+    add_partition_columns_to_schema,
+)
 from datahub.ingestion.source.schema_inference import avro, csv_tsv, json, parquet
 from datahub.ingestion.source.state.stale_entity_removal_handler import (
     StaleEntityRemovalHandler,
@@ -53,10 +57,7 @@ from datahub.ingestion.source.state.stateful_ingestion_base import (
     StatefulIngestionSourceBase,
 )
 from datahub.metadata.com.linkedin.pegasus2avro.schema import (
-    SchemaField,
-    SchemaFieldDataType,
     SchemaMetadata,
-    StringTypeClass,
 )
 from datahub.metadata.schema_classes import (
     DataPlatformInstanceClass,
@@ -128,6 +129,14 @@ class TableData:
 @support_status(SupportStatus.INCUBATING)
 @capability(SourceCapability.DATA_PROFILING, "Optionally enabled via configuration")
 @capability(SourceCapability.TAGS, "Can extract ABS object/container tags if enabled")
+@capability(
+    SourceCapability.CONTAINERS,
+    "Extract ABS containers and folders",
+    subtype_modifier=[
+        SourceCapabilityModifier.FOLDER,
+        SourceCapabilityModifier.ABS_CONTAINER,
+    ],
+)
 class ABSSource(StatefulIngestionSourceBase):
     source_config: DataLakeSourceConfig
     report: DataLakeSourceReport
@@ -223,36 +232,12 @@ class ABSSource(StatefulIngestionSourceBase):
         fields = sorted(fields, key=lambda f: f.fieldPath)
 
         if self.source_config.add_partition_columns_to_schema:
-            self.add_partition_columns_to_schema(
+            add_partition_columns_to_schema(
                 fields=fields, path_spec=path_spec, full_path=table_data.full_path
             )
 
         return fields
 
-    def add_partition_columns_to_schema(
-        self, path_spec: PathSpec, full_path: str, fields: List[SchemaField]
-    ) -> None:
-        vars = path_spec.get_named_vars(full_path)
-        if vars is not None and "partition" in vars:
-            for partition in vars["partition"].values():
-                partition_arr = partition.split("=")
-                if len(partition_arr) != 2:
-                    logger.debug(
-                        f"Could not derive partition key from partition field {partition}"
-                    )
-                    continue
-                partition_key = partition_arr[0]
-                fields.append(
-                    SchemaField(
-                        fieldPath=f"{partition_key}",
-                        nativeDataType="string",
-                        type=SchemaFieldDataType(StringTypeClass()),
-                        isPartitioningKey=True,
-                        nullable=True,
-                        recursive=False,
-                    )
-                )
-
     def _create_table_operation_aspect(self, table_data: TableData) -> OperationClass:
         reported_time = int(time.time() * 1000)
 
@@ -533,7 +518,7 @@ class ABSSource(StatefulIngestionSourceBase):
             )
             path_spec.sample_files = False
             for obj in container_client.list_blobs(
-                prefix=f"{prefix}", results_per_page=PAGE_SIZE
+                name_starts_with=f"{prefix}", results_per_page=PAGE_SIZE
             ):
                 abs_path = self.create_abs_path(obj.name)
                 logger.debug(f"Path: {abs_path}")

datahub/ingestion/source/apply/datahub_apply.py

@@ -18,6 +18,7 @@ from datahub.ingestion.api.source import MetadataWorkUnitProcessor, Source, Sour
 from datahub.ingestion.api.source_helpers import auto_workunit_reporter
 from datahub.ingestion.api.workunit import MetadataWorkUnit
 from datahub.ingestion.graph.client import DataHubGraph, get_default_graph
+from datahub.ingestion.graph.config import ClientMode
 from datahub.metadata.schema_classes import (
     DomainsClass,
     GlossaryTermAssociationClass,
@@ -48,7 +49,7 @@ def apply_association_to_container(
     """
     urns: List[str] = [container_urn]
     if not graph:
-        graph = get_default_graph()
+        graph = get_default_graph(ClientMode.INGESTION)
     logger.info(f"Using {graph}")
     urns.extend(
         graph.get_urns_by_filter(
@@ -95,7 +96,7 @@ def apply_association_to_container(
 class DomainApplyConfig(ConfigModel):
     assets: List[str] = Field(
         default_factory=list,
-        description="List of assets to apply domain hierarchichaly. Currently only containers and datasets are supported",
+        description="List of assets to apply domain hierarchically. Currently only containers and datasets are supported",
     )
     domain_urn: str = Field(default="")
 
@@ -103,7 +104,7 @@ class DomainApplyConfig(ConfigModel):
 class TagApplyConfig(ConfigModel):
     assets: List[str] = Field(
         default_factory=list,
-        description="List of assets to apply tag hierarchichaly. Currently only containers and datasets are supported",
+        description="List of assets to apply tag hierarchically. Currently only containers and datasets are supported",
     )
     tag_urn: str = Field(default="")
 
@@ -111,7 +112,7 @@ class TagApplyConfig(ConfigModel):
 class TermApplyConfig(ConfigModel):
     assets: List[str] = Field(
         default_factory=list,
-        description="List of assets to apply term hierarchichaly. Currently only containers and datasets are supported",
+        description="List of assets to apply term hierarchically. Currently only containers and datasets are supported",
     )
     term_urn: str = Field(default="")
 
@@ -119,7 +120,7 @@ class TermApplyConfig(ConfigModel):
 class OwnerApplyConfig(ConfigModel):
     assets: List[str] = Field(
         default_factory=list,
-        description="List of assets to apply owner hierarchichaly. Currently only containers and datasets are supported",
+        description="List of assets to apply owner hierarchically. Currently only containers and datasets are supported",
     )
     owner_urn: str = Field(default="")
 

datahub/ingestion/source/aws/aws_common.py

@@ -1,14 +1,15 @@
 import logging
-import os
 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
 
@@ -17,6 +18,16 @@ from datahub.configuration.common import (
     ConfigModel,
     PermissiveConfigModel,
 )
+from datahub.configuration.env_vars import (
+    get_aws_app_runner_service_id,
+    get_aws_execution_env,
+    get_aws_lambda_function_name,
+    get_aws_role_arn,
+    get_aws_web_identity_token_file,
+    get_ecs_container_metadata_uri,
+    get_ecs_container_metadata_uri_v4,
+    get_elastic_beanstalk_environment_name,
+)
 from datahub.configuration.source_common import EnvConfigMixin
 
 logger = logging.getLogger(__name__)
@@ -24,6 +35,7 @@ logger = logging.getLogger(__name__)
 if TYPE_CHECKING:
     from mypy_boto3_dynamodb import DynamoDBClient
     from mypy_boto3_glue import GlueClient
+    from mypy_boto3_lakeformation import LakeFormationClient
     from mypy_boto3_s3 import S3Client, S3ServiceResource
     from mypy_boto3_sagemaker import SageMakerClient
     from mypy_boto3_sts import STSClient
@@ -99,27 +111,25 @@ def detect_aws_environment() -> AwsEnvironment:
     Order matters as some environments may have multiple indicators.
     """
     # Check Lambda first as it's most specific
-    if os.getenv("AWS_LAMBDA_FUNCTION_NAME"):
-        if os.getenv("AWS_EXECUTION_ENV", "").startswith("CloudFormation"):
+    if get_aws_lambda_function_name():
+        if (get_aws_execution_env() or "").startswith("CloudFormation"):
             return AwsEnvironment.CLOUD_FORMATION
         return AwsEnvironment.LAMBDA
 
     # Check EKS (IRSA)
-    if os.getenv("AWS_WEB_IDENTITY_TOKEN_FILE") and os.getenv("AWS_ROLE_ARN"):
+    if get_aws_web_identity_token_file() and get_aws_role_arn():
         return AwsEnvironment.EKS
 
     # Check App Runner
-    if os.getenv("AWS_APP_RUNNER_SERVICE_ID"):
+    if get_aws_app_runner_service_id():
         return AwsEnvironment.APP_RUNNER
 
     # Check ECS
-    if os.getenv("ECS_CONTAINER_METADATA_URI_V4") or os.getenv(
-        "ECS_CONTAINER_METADATA_URI"
-    ):
+    if get_ecs_container_metadata_uri_v4() or get_ecs_container_metadata_uri():
         return AwsEnvironment.ECS
 
     # Check Elastic Beanstalk
-    if os.getenv("ELASTIC_BEANSTALK_ENVIRONMENT_NAME"):
+    if get_elastic_beanstalk_environment_name():
         return AwsEnvironment.BEANSTALK
 
     if is_running_on_ec2():
@@ -154,7 +164,7 @@ def get_instance_role_arn() -> Optional[str]:
 def get_lambda_role_arn() -> Optional[str]:
     """Get the Lambda function's role ARN"""
     try:
-        function_name = os.getenv("AWS_LAMBDA_FUNCTION_NAME")
+        function_name = get_aws_lambda_function_name()
         if not function_name:
             return None
 
@@ -180,7 +190,7 @@ def get_current_identity() -> Tuple[Optional[str], Optional[str]]:
         return role_arn, AwsServicePrincipal.LAMBDA.value
 
     elif env == AwsEnvironment.EKS:
-        role_arn = os.getenv("AWS_ROLE_ARN")
+        role_arn = get_aws_role_arn()
         return role_arn, AwsServicePrincipal.EKS.value
 
     elif env == AwsEnvironment.APP_RUNNER:
@@ -193,8 +203,8 @@ def get_current_identity() -> Tuple[Optional[str], Optional[str]]:
 
     elif env == AwsEnvironment.ECS:
         try:
-            metadata_uri = os.getenv("ECS_CONTAINER_METADATA_URI_V4") or os.getenv(
-                "ECS_CONTAINER_METADATA_URI"
+            metadata_uri = (
+                get_ecs_container_metadata_uri_v4() or get_ecs_container_metadata_uri()
             )
             if metadata_uri:
                 response = requests.get(f"{metadata_uri}/task", timeout=1)
@@ -454,6 +464,168 @@ class AwsConnectionConfig(ConfigModel):
     def get_sagemaker_client(self) -> "SageMakerClient":
         return self.get_session().client("sagemaker", config=self._aws_config())
 
+    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):
     """