[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="AssumeRoleSession",
+                    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

@@ -35,6 +35,10 @@
     is_access_denied_exception,
     is_server_error,
 )
+from port_ocean.utils.async_iterators import stream_async_iterators_tasks
+import asyncio
+
+MAX_CONCURRENT_TASKS = 50
 
 
 async def _handle_global_resource_resync(
@@ -69,19 +73,37 @@ async def _handle_global_resource_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 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:
-                    continue
+    semaphore = asyncio.BoundedSemaphore(MAX_CONCURRENT_TASKS)
+
+    async def handle_account(
+        credentials: AwsCredentials,
+    ) -> ASYNC_GENERATOR_RESYNC_TYPE:
+        """Function to handle fetching resources for a single account."""
+        async with semaphore:
+            errors, regions = [], []
+            if is_global:
+                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 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)
+
+    tasks = [handle_account(credentials) 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 +116,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