[Integration][AWS] - Fix ExpiredTokenException

integrations/aws/CHANGELOG.md

@@ -6,12 +6,24 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 <!-- towncrier release notes start -->
+
+## 0.2.43 (2024-09-18)
+
+
+### Improvements
+
+- Improved support for parralel fetching of aws account resources
+- Fixed ExpiredTokenException by replacing event-based caching with a time-dependent caching mechanism. The new approach reassumes the role and refreshes session credentials when 80% of the session duration has been used, ensuring credentials are refreshed before expiry.
+
+
 ## 0.2.42 (2024-09-24)
 
+
 ### Bug Fixes
 
 - Fixes an issue where `is_access_denied_exception` could raise an `AttributeError` if `e.response` is `None`.
 
+
 ## 0.2.41 (2024-09-22)
 
 

integrations/aws/aws/session_manager.py

@@ -15,6 +15,9 @@ class AccountNotFoundError(OceanAbortException):
     pass
 
 
+ASSUME_ROLE_DURATION_SECONDS = 3600  # 1 hour
+
+
 class SessionManager:
     def __init__(self) -> None:
         """
@@ -93,7 +96,9 @@ async def _get_organization_session(self) -> aioboto3.Session | None:
         async with application_session.client("sts") as sts_client:
             try:
                 organizations_client = await sts_client.assume_role(
-                    RoleArn=organization_role_arn, RoleSessionName="AssumeRoleSession"
+                    RoleArn=organization_role_arn,
+                    RoleSessionName="OceanOrgAssumeRoleSession",
+                    DurationSeconds=ASSUME_ROLE_DURATION_SECONDS,
                 )
 
                 credentials = organizations_client["Credentials"]
@@ -150,7 +155,8 @@ async def _assume_role_and_update_credentials(
         try:
             account_role = await sts_client.assume_role(
                 RoleArn=f'arn:aws:iam::{account["Id"]}:role/{self._get_account_read_role_name()}',
-                RoleSessionName="AssumeRoleSession",
+                RoleSessionName="OceanMemberAssumeRoleSession",
+                DurationSeconds=ASSUME_ROLE_DURATION_SECONDS,
             )
             raw_credentials = account_role["Credentials"]
             credentials = AwsCredentials(

integrations/aws/main.py

@@ -34,7 +34,9 @@
     ResourceKindsWithSpecialHandling,
     is_access_denied_exception,
     is_server_error,
+    semaphore,
 )
+from port_ocean.utils.async_iterators import stream_async_iterators_tasks
 
 
 async def _handle_global_resource_resync(
@@ -65,23 +67,44 @@ async def _handle_global_resource_resync(
                     raise e
 
 
-@ocean.on_resync()
-async def resync_all(kind: str) -> ASYNC_GENERATOR_RESYNC_TYPE:
-    if kind in iter(ResourceKindsWithSpecialHandling):
-        return
-    await update_available_access_credentials()
-    is_global = is_global_resource(kind)
-    async for credentials in get_accounts():
-        if is_global:
+async def resync_resources_for_account(
+    credentials: AwsCredentials, kind: str
+) -> ASYNC_GENERATOR_RESYNC_TYPE:
+    """Function to handle fetching resources for a single account."""
+
+    async with semaphore:  # limit the number of concurrent tasks
+        errors, regions = [], []
+
+        if is_global_resource(kind):
             async for batch in _handle_global_resource_resync(kind, credentials):
                 yield batch
         else:
             async for session in credentials.create_session_for_each_region():
                 try:
                     async for batch in resync_cloudcontrol(kind, session):
                         yield batch
-                except Exception:
+                except Exception as exc:
+                    regions.append(session.region_name)
+                    errors.append(exc)
                     continue
+        if errors:
+            message = f"Failed to fetch {kind} for these regions {regions} with {len(errors)} errors in account {credentials.account_id}"
+            raise ExceptionGroup(message, errors)
+
+
+@ocean.on_resync()
+async def resync_all(kind: str) -> ASYNC_GENERATOR_RESYNC_TYPE:
+    if kind in iter(ResourceKindsWithSpecialHandling):
+        return
+
+    await update_available_access_credentials()
+    tasks = [
+        resync_resources_for_account(credentials, kind)
+        async for credentials in get_accounts()
+    ]
+    if tasks:
+        async for batch in stream_async_iterators_tasks(*tasks):
+            yield batch
 
 
 @ocean.on_resync(kind=ResourceKindsWithSpecialHandling.ACCOUNT)
@@ -94,76 +117,102 @@ async def resync_account(kind: str) -> ASYNC_GENERATOR_RESYNC_TYPE:
 @ocean.on_resync(kind=ResourceKindsWithSpecialHandling.ELASTICACHE_CLUSTER)
 async def resync_elasticache(kind: str) -> ASYNC_GENERATOR_RESYNC_TYPE:
     await update_available_access_credentials()
-    async for session in get_sessions():
-        async for batch in resync_custom_kind(
+
+    tasks = [
+        resync_custom_kind(
             kind,
             session,
             "elasticache",
             "describe_cache_clusters",
             "CacheClusters",
             "Marker",
-        ):
+        )
+        async for session in get_sessions()
+    ]
+    if tasks:
+        async for batch in stream_async_iterators_tasks(*tasks):
             yield batch
 
 
 @ocean.on_resync(kind=ResourceKindsWithSpecialHandling.ELBV2_LOAD_BALANCER)
 async def resync_elv2_load_balancer(kind: str) -> ASYNC_GENERATOR_RESYNC_TYPE:
     await update_available_access_credentials()
-    async for session in get_sessions():
-        async for batch in resync_custom_kind(
+
+    tasks = [
+        resync_custom_kind(
             kind,
             session,
             "elbv2",
             "describe_load_balancers",
             "LoadBalancers",
             "Marker",
-        ):
+        )
+        async for session in get_sessions()
+    ]
+
+    if tasks:
+        async for batch in stream_async_iterators_tasks(*tasks):
             yield batch
 
 
 @ocean.on_resync(kind=ResourceKindsWithSpecialHandling.ACM_CERTIFICATE)
 async def resync_acm(kind: str) -> ASYNC_GENERATOR_RESYNC_TYPE:
     await update_available_access_credentials()
-    async for session in get_sessions():
-        async for batch in resync_custom_kind(
+
+    tasks = [
+        resync_custom_kind(
             kind,
             session,
             "acm",
             "list_certificates",
             "CertificateSummaryList",
             "NextToken",
-        ):
+        )
+        async for session in get_sessions()
+    ]
+
+    if tasks:
+        async for batch in stream_async_iterators_tasks(*tasks):
             yield batch
 
 
 @ocean.on_resync(kind=ResourceKindsWithSpecialHandling.AMI_IMAGE)
 async def resync_ami(kind: str) -> ASYNC_GENERATOR_RESYNC_TYPE:
     await update_available_access_credentials()
-    async for session in get_sessions():
-        async for batch in resync_custom_kind(
+    tasks = [
+        resync_custom_kind(
             kind,
             session,
             "ec2",
             "describe_images",
             "Images",
             "NextToken",
             {"Owners": ["self"]},
-        ):
+        )
+        async for session in get_sessions()
+    ]
+    if tasks:
+        async for batch in stream_async_iterators_tasks(*tasks):
             yield batch
 
 
 @ocean.on_resync(kind=ResourceKindsWithSpecialHandling.CLOUDFORMATION_STACK)
 async def resync_cloudformation(kind: str) -> ASYNC_GENERATOR_RESYNC_TYPE:
     await update_available_access_credentials()
-    async for session in get_sessions():
-        async for batch in resync_custom_kind(
+    tasks = [
+        resync_custom_kind(
             kind,
             session,
             "cloudformation",
             "describe_stacks",
             "Stacks",
             "NextToken",
-        ):
+        )
+        async for session in get_sessions()
+    ]
+
+    if tasks:
+        async for batch in stream_async_iterators_tasks(*tasks):
             yield batch