cartography-cncf#1321: Refactor get_launch_template_versions to skip …

…race condition. (cartography-cncf#1332)

See cartography-cncf#1321

cartography/intel/aws/ec2/launch_templates.py

@@ -3,6 +3,7 @@
 
 import boto3
 import neo4j
+from botocore.exceptions import ClientError
 
 from .util import get_botocore_config
 from cartography.client.core.tx import load
@@ -17,13 +18,30 @@
 
 @timeit
 @aws_handle_regions
-def get_launch_templates(boto3_session: boto3.session.Session, region: str) -> list[dict[str, Any]]:
+def get_launch_templates(
+    boto3_session: boto3.session.Session,
+    region: str,
+) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
     client = boto3_session.client('ec2', region_name=region, config=get_botocore_config())
     paginator = client.get_paginator('describe_launch_templates')
     templates: list[dict[str, Any]] = []
+    template_versions: list[dict[str, Any]] = []
     for page in paginator.paginate():
-        templates.extend(page['LaunchTemplates'])
-    return templates
+        paginated_templates = page['LaunchTemplates']
+        for template in paginated_templates:
+            template_id = template['LaunchTemplateId']
+            try:
+                versions = get_launch_template_versions_by_template(boto3_session, template_id, region)
+            except ClientError as e:
+                logger.warning(
+                    f"Failed to get launch template versions for {template_id}: {e}",
+                    exc_info=True,
+                )
+                versions = []
+            # Using a key not defined in latest boto3 documentation
+            template_versions.extend(versions)
+        templates.extend(paginated_templates)
+    return templates, template_versions
 
 
 def transform_launch_templates(templates: list[dict[str, Any]]) -> list[dict[str, Any]]:
@@ -55,17 +73,16 @@ def load_launch_templates(
 
 @timeit
 @aws_handle_regions
-def get_launch_template_versions(
+def get_launch_template_versions_by_template(
         boto3_session: boto3.session.Session,
-        templates: list[dict[str, Any]],
+        template: str,
         region: str,
 ) -> list[dict[str, Any]]:
     client = boto3_session.client('ec2', region_name=region, config=get_botocore_config())
     v_paginator = client.get_paginator('describe_launch_template_versions')
     template_versions = []
-    for template in templates:
-        for versions in v_paginator.paginate(LaunchTemplateId=template['LaunchTemplateId']):
-            template_versions.extend(versions['LaunchTemplateVersions'])
+    for versions in v_paginator.paginate(LaunchTemplateId=template):
+        template_versions.extend(versions['LaunchTemplateVersions'])
     return template_versions
 
 
@@ -136,11 +153,9 @@ def sync_ec2_launch_templates(
 ) -> None:
     for region in regions:
         logger.info(f"Syncing launch templates for region '{region}' in account '{current_aws_account_id}'.")
-        templates = get_launch_templates(boto3_session, region)
+        templates, versions = get_launch_templates(boto3_session, region)
         templates = transform_launch_templates(templates)
         load_launch_templates(neo4j_session, templates, region, current_aws_account_id, update_tag)
-
-        versions = get_launch_template_versions(boto3_session, templates, region)
         versions = transform_launch_template_versions(versions)
         load_launch_template_versions(neo4j_session, versions, region, current_aws_account_id, update_tag)
 

cartography/util.py

@@ -225,7 +225,7 @@ def aws_paginate(
     return items
 
 
-AWSGetFunc = TypeVar('AWSGetFunc', bound=Callable[..., List])
+AWSGetFunc = TypeVar('AWSGetFunc', bound=Callable[..., Iterable])
 
 # fix for AWS TooManyRequestsException
 # https://github.com/lyft/cartography/issues/297

test-requirements.txt

@@ -1,4 +1,5 @@
 backoff>=2.1.2
+moto
 pre-commit
 pytest>=6.2.4
 pytest-mock

tests/integration/cartography/intel/aws/ec2/test_launch_templates.py

@@ -1,3 +1,11 @@
+from unittest.mock import patch
+
+import boto3
+from botocore.exceptions import ClientError
+from moto import mock_aws
+
+import cartography.intel.aws.ec2.launch_templates
+from cartography.intel.aws.ec2.launch_templates import get_launch_templates
 from cartography.intel.aws.ec2.launch_templates import load_launch_template_versions
 from cartography.intel.aws.ec2.launch_templates import load_launch_templates
 from cartography.intel.aws.ec2.launch_templates import transform_launch_template_versions
@@ -11,6 +19,48 @@
 TEST_UPDATE_TAG = 123456789
 
 
+@mock_aws(config={'core': {'reset_boto3_session': True, 'mock_credentials': True}})
+@patch.object(
+    cartography.intel.aws.ec2.launch_templates,
+    'get_launch_template_versions_by_template',
+)
+def test_get_launch_template_throws_exception(mock_get_template_versions, *args):
+    # Arrange
+    template_data = {
+        "ImageId": "ami-abc123",
+        "TagSpecifications": [
+            {
+                "ResourceType": "instance", "Tags": [
+                    {"Key": "eks:cluster-name", "Value": "eks-cluster-example"},
+                    {"Key": "eks:nodegroup-name", "Value": "private-node-group-example"},
+                ],
+            },
+        ],
+        "SecurityGroupIds": ["sg-1234"],
+    }
+    client = boto3.client('ec2', region_name=TEST_REGION)
+    mock_template = client.create_launch_template(
+        LaunchTemplateName='eks-00000000-0000-0000-0000-000000000000',
+        LaunchTemplateData=template_data,
+    )
+    template_id = mock_template['LaunchTemplate']['LaunchTemplateId']
+    error_response = {
+        "Error": {
+            "Code": "InvalidLaunchTemplateId.NotFound",
+            "Message": f"The specified launch template, with template ID {template_id}, does not exist.",
+        },
+    }
+    mock_get_template_versions.side_effect = ClientError(error_response, "DescribeLaunchTemplateVersions")
+    session = boto3.Session(region_name=TEST_REGION)
+    # Act: get the launch template versions
+
+    templates, versions = get_launch_templates(session, TEST_REGION)
+
+    # Assert: the launch template versions are as expected
+    assert len(templates) == 1
+    assert len(versions) == 0
+
+
 def test_load_launch_templates(neo4j_session, *args):
     # Arrange: an AWSAccount must exist
     neo4j_session.run(