fix(ingest/aws): Fix for the aws credential caching (#15278)

Author: treff7es

metadata-ingestion/src/datahub/ingestion/source/aws/aws_common.py

@@ -269,6 +269,7 @@ class AwsConnectionConfig(ConfigModel):
     """
 
     _credentials_expiration: Optional[datetime] = None
+    _cached_credentials: Optional[dict] = None
 
     aws_access_key_id: Optional[str] = Field(
         default=None,
@@ -353,10 +354,25 @@ def get_session(self) -> Session:
             )
         else:
             # Use boto3's credential autodetection
-            session = Session(region_name=self.aws_region)
-
             target_roles = self._normalized_aws_roles()
             if target_roles:
+                # If we have cached credentials that are still valid, use them
+                if (
+                    self._cached_credentials is not None
+                    and not self._should_refresh_credentials()
+                ):
+                    logger.debug("Using cached assumed role credentials")
+                    return Session(
+                        aws_access_key_id=self._cached_credentials["AccessKeyId"],
+                        aws_secret_access_key=self._cached_credentials[
+                            "SecretAccessKey"
+                        ],
+                        aws_session_token=self._cached_credentials["SessionToken"],
+                        region_name=self.aws_region,
+                    )
+
+                # Need to assume role (either first time or credentials expired)
+                session = Session(region_name=self.aws_region)
                 current_role_arn, credential_source = get_current_identity()
 
                 # Only assume role if:
@@ -368,7 +384,10 @@ def get_session(self) -> Session:
 
                 if should_assume_role:
                     env = detect_aws_environment()
-                    logger.debug(f"Assuming role(s) from {env.value} environment")
+                    role_arns = [role.RoleArn for role in target_roles]
+                    logger.debug(
+                        f"Assuming {role_arns} role(s) from {env.value} environment"
+                    )
 
                     current_credentials = session.get_credentials()
                     if current_credentials is None:
@@ -381,14 +400,15 @@ def get_session(self) -> Session:
                     }
 
                     for role in target_roles:
-                        if self._should_refresh_credentials():
-                            credentials = assume_role(
-                                role=role,
-                                aws_region=self.aws_region,
-                                credentials=credentials,
-                            )
-                            if isinstance(credentials["Expiration"], datetime):
-                                self._credentials_expiration = credentials["Expiration"]
+                        credentials = assume_role(
+                            role=role,
+                            aws_region=self.aws_region,
+                            credentials=credentials,
+                        )
+
+                    if isinstance(credentials["Expiration"], datetime):
+                        self._credentials_expiration = credentials["Expiration"]
+                    self._cached_credentials = credentials
 
                     session = Session(
                         aws_access_key_id=credentials["AccessKeyId"],
@@ -398,12 +418,17 @@ def get_session(self) -> Session:
                     )
                 else:
                     logger.debug(f"Using existing role from {credential_source}")
+            else:
+                session = Session(region_name=self.aws_region)
 
         return session
 
     def _should_refresh_credentials(self) -> bool:
         if self._credentials_expiration is None:
             return True
+        # Refresh credentials when less than 5 minutes remain before expiration.
+        # This buffer helps avoid race conditions where credentials expire between
+        # the cache check and actual AWS API usage.
         remaining_time = self._credentials_expiration - datetime.now(timezone.utc)
         return remaining_time < timedelta(minutes=5)