#1206: add async to aws:s3 and aws:ecr

cartography/intel/aws/ecr.py

@@ -1,4 +1,5 @@
 import logging
+from typing import Any
 from typing import Dict
 from typing import List
 
@@ -9,6 +10,8 @@
 from cartography.util import batch
 from cartography.util import run_cleanup_job
 from cartography.util import timeit
+from cartography.util import to_asynchronous
+from cartography.util import to_synchronous
 
 logger = logging.getLogger(__name__)
 
@@ -139,6 +142,21 @@ def cleanup(neo4j_session: neo4j.Session, common_job_parameters: Dict) -> None:
     run_cleanup_job('aws_import_ecr_cleanup.json', neo4j_session, common_job_parameters)
 
 
+def _get_image_data(boto3_session: boto3.session.Session, region: str, repositories: List[Dict]) -> Dict:
+    '''
+    Given a list of repositories, get the image data for each repository,
+     return as a mapping from repositoryUri to image object
+    '''
+    image_data = {}
+
+    async def async_get_images(repo: Dict[str, Any]) -> None:
+        repo_image_obj = await to_asynchronous(get_ecr_repository_images, boto3_session, region, repo['repositoryName'])
+        image_data[repo['repositoryUri']] = repo_image_obj
+    to_synchronous(*[async_get_images(repo) for repo in repositories])
+
+    return image_data
+
+
 @timeit
 def sync(
     neo4j_session: neo4j.Session, boto3_session: boto3.session.Session, regions: List[str], current_aws_account_id: str,
@@ -148,9 +166,7 @@ def sync(
         logger.info("Syncing ECR for region '%s' in account '%s'.", region, current_aws_account_id)
         image_data = {}
         repositories = get_ecr_repositories(boto3_session, region)
-        for repo in repositories:
-            repo_image_obj = get_ecr_repository_images(boto3_session, region, repo['repositoryName'])
-            image_data[repo['repositoryUri']] = repo_image_obj
+        image_data = _get_image_data(boto3_session, region, repositories)
         load_ecr_repositories(neo4j_session, repositories, region, current_aws_account_id, update_tag)
         repo_images_list = transform_ecr_repository_images(image_data)
         load_ecr_repository_images(neo4j_session, repo_images_list, region, update_tag)

cartography/intel/aws/s3.py

@@ -1,3 +1,4 @@
+import asyncio
 import hashlib
 import json
 import logging
@@ -20,6 +21,8 @@
 from cartography.util import run_analysis_job
 from cartography.util import run_cleanup_job
 from cartography.util import timeit
+from cartography.util import to_asynchronous
+from cartography.util import to_synchronous
 
 logger = logging.getLogger(__name__)
 stat_handler = get_stats_client(__name__)
@@ -55,20 +58,33 @@ def get_s3_bucket_details(
     # a local store for s3 clients so that we may re-use clients for an AWS region
     s3_regional_clients: Dict[Any, Any] = {}
 
-    for bucket in bucket_data['Buckets']:
+    BucketDetail = Tuple[str, Dict[str, Any], Dict[str, Any], Dict[str, Any], Dict[str, Any], Dict[str, Any]]
+
+    async def _get_bucket_detail(bucket: Dict[str, Any]) -> BucketDetail:
         # Note: bucket['Region'] is sometimes None because
         # client.get_bucket_location() does not return a location constraint for buckets
         # in us-east-1 region
         client = s3_regional_clients.get(bucket['Region'])
         if not client:
             client = boto3_session.client('s3', bucket['Region'])
             s3_regional_clients[bucket['Region']] = client
-        acl = get_acl(bucket, client)
-        policy = get_policy(bucket, client)
-        encryption = get_encryption(bucket, client)
-        versioning = get_versioning(bucket, client)
-        public_access_block = get_public_access_block(bucket, client)
-        yield bucket['Name'], acl, policy, encryption, versioning, public_access_block
+        (
+            acl,
+            policy,
+            encryption,
+            versioning,
+            public_access_block,
+        ) = await asyncio.gather(
+            to_asynchronous(get_acl, bucket, client),
+            to_asynchronous(get_policy, bucket, client),
+            to_asynchronous(get_encryption, bucket, client),
+            to_asynchronous(get_versioning, bucket, client),
+            to_asynchronous(get_public_access_block, bucket, client),
+        )
+        return bucket['Name'], acl, policy, encryption, versioning, public_access_block
+
+    bucket_details = to_synchronous(*[_get_bucket_detail(bucket) for bucket in bucket_data['Buckets']])
+    yield from bucket_details
 
 
 @timeit

cartography/util.py

@@ -1,9 +1,12 @@
+import asyncio
 import logging
 import re
 import sys
+from functools import partial
 from functools import wraps
 from string import Template
 from typing import Any
+from typing import Awaitable
 from typing import BinaryIO
 from typing import Callable
 from typing import cast
@@ -25,6 +28,7 @@
 from cartography.stats import get_stats_client
 from cartography.stats import ScopedStatsClient
 
+
 if sys.version_info >= (3, 7):
     from importlib.resources import open_binary, read_text
 else:
@@ -141,6 +145,7 @@ def load_resource_binary(package: str, resource_name: str) -> BinaryIO:
     return open_binary(package, resource_name)
 
 
+R = TypeVar('R')
 F = TypeVar('F', bound=Callable[..., Any])
 
 
@@ -297,3 +302,47 @@ def batch(items: Iterable, size: int = DEFAULT_BATCH_SIZE) -> List[List]:
         items[i: i + size]
         for i in range(0, len(items), size)
     ]
+
+
+def to_asynchronous(func: Callable[..., R], *args: Any, **kwargs: Any) -> Awaitable[R]:
+    '''
+    Returns a Future that will run a function in the default threadpool.
+    Helper until we start using pytohn 3.9's asyncio.to_thread
+
+    Calls are also wrapped within a backoff decorator to handle throttling errors.
+
+    example:
+    future = to_async(my_func, my_arg, my_arg2)
+    to_sync(future)
+
+    NOTE: to use this in a Jupyter notebook, you need to do:
+    # import nest_asyncio
+    # nest_asyncio.apply()
+    '''
+    CartographyThrottlingException = type('CartographyThrottlingException', (Exception,), {})
+    throttling_error_codes = ['LimitExceededException', 'Throttling']
+
+    @wraps(func)
+    def wrapper(*args: Any, **kwargs: Any) -> R:
+        try:
+            return func(*args, **kwargs)
+        except botocore.exceptions.ClientError as error:
+            if error.response['Error']['Code'] in throttling_error_codes:
+                raise CartographyThrottlingException from error
+            raise
+
+    # don't use @backoff as decorator, to preserve typing
+    wrapped = backoff.on_exception(backoff.expo, CartographyThrottlingException)(wrapper)
+    call = partial(wrapped, *args, **kwargs)
+    return asyncio.get_event_loop().run_in_executor(None, call)
+
+
+def to_synchronous(*awaitables: Awaitable[Any]) -> Any:
+    '''
+    Waits for the Awaitable(s) to complete and returns their result(s).
+    See https://docs.python.org/3.8/library/asyncio-task.html#asyncio-awaitables
+
+    example:
+    result = to_sync(my_async_func(my_arg), another_async(my_arg2)))
+    '''
+    return asyncio.get_event_loop().run_until_complete(asyncio.gather(*awaitables))