Pattern 2 cleanup

kazmer97 · kazmer97 · commit a96ce7f326e6 · 2025-10-14T11:31:40.000+01:00
diff --git a/patterns/pattern-2/template.yaml b/patterns/pattern-2/template.yaml
@@ -188,8 +188,6 @@ Resources:
 
   Pattern2ECRRepository:
     Type: AWS::ECR::Repository
-    DeletionPolicy: Retain
-    UpdateReplacePolicy: Retain
     Properties:
       ImageScanningConfiguration:
         ScanOnPush: true
@@ -284,6 +282,12 @@ Resources:
                   - logs:CreateLogStream
                   - logs:PutLogEvents
                 Resource: "*"
+              - Effect: Allow
+                Action:
+                  - ecr:ListImages
+                  - ecr:BatchDeleteImage
+                Resource:
+                  - !GetAtt Pattern2ECRRepository.Arn
           # Used by custom resource helper poller
           # https://github.com/aws-cloudformation/custom-resource-helper
         - PolicyName: CustomResourcePoller
@@ -351,6 +355,15 @@ Resources:
       BuildProjectName: !Ref Pattern2DockerBuildProject
       # Force rebuild when source changes (filename includes content hash)
       CodeLocation: !Sub "arn:${AWS::Partition}:s3:::${ArtifactBucket}/${ArtifactPrefix}/${Pattern2SourceZipfile}"
+
+  Pattern2EcrRepositoryCleanup:
+    Type: Custom::ECRRepositoryCleanup
+    DependsOn:
+      - Pattern2DockerBuildRun
+      - Pattern2ECRRepository
+    Properties:
+      ServiceToken: !GetAtt Pattern2CodeBuildTrigger.Arn
+      RepositoryName: !Ref Pattern2ECRRepository
   # Shared IAM policy for Lambda functions to pull container images from ECR
   LambdaECRAccessPolicy:
     Type: AWS::IAM::ManagedPolicy
diff --git a/src/lambda/start_codebuild/index.py b/src/lambda/start_codebuild/index.py
@@ -1,12 +1,14 @@
 # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 # SPDX-License-Identifier: Apache-2.0
-"""CodeBuild Starter Lambda Function"""
+"""CodeBuild starter and ECR cleanup Lambda function used by Pattern 2 deployments."""
 import logging
 from os import getenv
 import json
+from typing import List
 
 import boto3
 from botocore.config import Config as BotoCoreConfig
+from botocore.exceptions import ClientError
 from crhelper import CfnResource
 
 
@@ -24,7 +26,8 @@
     CLIENT_CONFIG = BotoCoreConfig(
         retries={"mode": "adaptive", "max_attempts": 5},
     )
-    CLIENT = boto3.client("codebuild", config=CLIENT_CONFIG)
+    CODEBUILD_CLIENT = boto3.client("codebuild", config=CLIENT_CONFIG)
+    ECR_CLIENT = boto3.client("ecr", config=CLIENT_CONFIG)
 except Exception as init_exception:  # pylint: disable=broad-except
     HELPER.init_failure(init_exception)
 
@@ -39,7 +42,7 @@ def create_or_update(event, _):
     if resource_type == "Custom::CodeBuildRun":
         try:
             project_name = resource_properties["BuildProjectName"]
-            response = CLIENT.start_build(projectName=project_name)
+            response = CODEBUILD_CLIENT.start_build(projectName=project_name)
             build_id = response["build"]["id"]
             HELPER.Data["build_id"] = build_id
         except Exception as exception:  # pylint: disable=broad-except
@@ -48,6 +51,12 @@ def create_or_update(event, _):
 
         return
 
+    if resource_type == "Custom::ECRRepositoryCleanup":
+        repository_name = resource_properties["RepositoryName"]
+        LOGGER.info("registered ECR cleanup resource for repository %s", repository_name)
+        HELPER.Data["repository_name"] = repository_name
+        return
+
     raise ValueError(f"invalid resource type: {resource_type}")
 
 
@@ -61,7 +70,7 @@ def poll_create_or_update(event, _):
     if resource_type == "Custom::CodeBuildRun":
         try:
             build_id = helper_data["build_id"]
-            response = CLIENT.batch_get_builds(ids=[build_id])
+            response = CODEBUILD_CLIENT.batch_get_builds(ids=[build_id])
             LOGGER.info(response)
 
             builds = response["builds"]
@@ -86,13 +95,79 @@ def poll_create_or_update(event, _):
             LOGGER.error("build poller - exception: %s", exception)
             raise
 
+    if resource_type == "Custom::ECRRepositoryCleanup":
+        LOGGER.info("ECR cleanup resource create/update completed")
+        return True
+
     raise RuntimeError(f"Invalid resource type: {resource_type}")
 
 
 @HELPER.delete
-def delete_no_op(event, _):
+def delete_resource(event, _):
     """Delete Resource"""
-    LOGGER.info("delete event ignored: %s", event)
+    resource_type = event["ResourceType"]
+
+    if resource_type == "Custom::CodeBuildRun":
+        LOGGER.info("delete event ignored for CodeBuild custom resource: %s", event)
+        return
+
+    if resource_type == "Custom::ECRRepositoryCleanup":
+        repository_name = event["ResourceProperties"]["RepositoryName"]
+        LOGGER.info("starting cleanup for repository %s", repository_name)
+        try:
+            _delete_all_ecr_images(repository_name)
+        except ClientError as error:
+            if error.response["Error"]["Code"] == "RepositoryNotFoundException":
+                LOGGER.info("repository %s already deleted", repository_name)
+                return
+            LOGGER.error(
+                "failed to purge repository %s - error: %s",
+                repository_name,
+                error,
+            )
+            raise
+        except Exception as unknown_exception:  # pylint: disable=broad-except
+            LOGGER.error(
+                "unexpected error while cleaning repository %s - error: %s",
+                repository_name,
+                unknown_exception,
+            )
+            raise
+
+        LOGGER.info("cleanup for repository %s completed", repository_name)
+        return
+
+    LOGGER.warning("received delete for unsupported resource type: %s", resource_type)
+
+
+def _delete_all_ecr_images(repository_name: str) -> None:
+    """Delete every image (tagged and untagged) from the ECR repository."""
+    paginator = ECR_CLIENT.get_paginator("list_images")
+    images_to_delete: List[dict] = []
+
+    for page in paginator.paginate(repositoryName=repository_name):
+        image_ids = page.get("imageIds", [])
+        if not image_ids:
+            continue
+        images_to_delete.extend(image_ids)
+        LOGGER.debug(
+            "queued %s images for deletion from repository %s",
+            len(image_ids),
+            repository_name,
+        )
+
+    if not images_to_delete:
+        LOGGER.info("no images found in repository %s", repository_name)
+        return
+
+    for chunk_start in range(0, len(images_to_delete), 100):
+        chunk = images_to_delete[chunk_start : chunk_start + 100]
+        LOGGER.debug(
+            "deleting %s images from repository %s",
+            len(chunk),
+            repository_name,
+        )
+        ECR_CLIENT.batch_delete_image(repositoryName=repository_name, imageIds=chunk)
 
 
 def handler(event, context):
diff --git a/tests/unit/lambda/test_start_codebuild_cleanup.py b/tests/unit/lambda/test_start_codebuild_cleanup.py
@@ -0,0 +1,143 @@
+import importlib
+import sys
+import types
+import unittest
+from unittest.mock import MagicMock
+
+
+class ClientError(Exception):
+    """Lightweight stand-in for botocore.exceptions.ClientError."""
+
+    def __init__(self, error_response, operation_name):
+        super().__init__(error_response)
+        self.response = error_response
+        self.operation_name = operation_name
+
+
+fake_boto3 = types.ModuleType("boto3")
+
+
+def _fake_client(service_name, config=None):  # pylint: disable=unused-argument
+    return MagicMock(name=f"{service_name}_client")
+
+
+fake_boto3.client = _fake_client
+sys.modules.setdefault("boto3", fake_boto3)
+
+fake_botocore = types.ModuleType("botocore")
+fake_botocore_config = types.ModuleType("botocore.config")
+fake_botocore_config.Config = MagicMock  # type: ignore[attr-defined]
+
+fake_botocore_exceptions = types.ModuleType("botocore.exceptions")
+fake_botocore_exceptions.ClientError = ClientError
+
+fake_botocore.config = fake_botocore_config
+fake_botocore.exceptions = fake_botocore_exceptions
+
+sys.modules.setdefault("botocore", fake_botocore)
+sys.modules.setdefault("botocore.config", fake_botocore_config)
+sys.modules.setdefault("botocore.exceptions", fake_botocore_exceptions)
+
+fake_crhelper = types.ModuleType("crhelper")
+
+
+class _FakeCfnResource:
+    def __init__(self, *args, **kwargs):
+        self.Data = {}
+
+    def create(self, func):
+        return func
+
+    def update(self, func):
+        return func
+
+    def poll_create(self, func):
+        return func
+
+    def poll_update(self, func):
+        return func
+
+    def delete(self, func):
+        return func
+
+    def init_failure(self, exception):
+        raise exception
+
+    def __call__(self, event, context):
+        return None
+
+
+fake_crhelper.CfnResource = _FakeCfnResource
+sys.modules.setdefault("crhelper", fake_crhelper)
+
+start_codebuild = importlib.import_module("src.lambda.start_codebuild.index")
+start_codebuild.ClientError = ClientError  # ensure same reference in tests
+
+
+class StartCodebuildCleanupTests(unittest.TestCase):
+    def setUp(self):
+        self.repo_name = "pattern-2"
+        self.ecr_client = MagicMock(name="ecr_client")
+        self.paginator = MagicMock(name="list_images_paginator")
+        self.ecr_client.get_paginator.return_value = self.paginator
+        start_codebuild.ECR_CLIENT = self.ecr_client
+
+    def test_delete_all_ecr_images_handles_pagination_and_chunking(self):
+        images_page_one = [{"imageDigest": f"sha256:{i:064d}"} for i in range(100)]
+        images_page_two = [{"imageDigest": f"sha256:{100 + i:064d}"} for i in range(50)]
+        self.paginator.paginate.return_value = [
+            {"imageIds": images_page_one},
+            {"imageIds": images_page_two},
+        ]
+
+        start_codebuild._delete_all_ecr_images(self.repo_name)  # pylint: disable=protected-access
+
+        self.ecr_client.batch_delete_image.assert_any_call(
+            repositoryName=self.repo_name, imageIds=images_page_one
+        )
+        self.ecr_client.batch_delete_image.assert_any_call(
+            repositoryName=self.repo_name, imageIds=images_page_two
+        )
+
+    def test_delete_all_ecr_images_skips_when_repository_empty(self):
+        self.paginator.paginate.return_value = [{"imageIds": []}]
+
+        start_codebuild._delete_all_ecr_images(self.repo_name)  # pylint: disable=protected-access
+
+        self.ecr_client.batch_delete_image.assert_not_called()
+
+    def test_delete_resource_handles_repository_not_found(self):
+        cleanup_event = {
+            "ResourceType": "Custom::ECRRepositoryCleanup",
+            "ResourceProperties": {"RepositoryName": self.repo_name},
+        }
+        self.paginator.paginate.side_effect = ClientError(
+            {"Error": {"Code": "RepositoryNotFoundException"}}, "ListImages"
+        )
+
+        start_codebuild.delete_resource(cleanup_event, None)
+
+        self.ecr_client.batch_delete_image.assert_not_called()
+
+    def test_delete_resource_raises_unexpected_ecr_error(self):
+        cleanup_event = {
+            "ResourceType": "Custom::ECRRepositoryCleanup",
+            "ResourceProperties": {"RepositoryName": self.repo_name},
+        }
+        self.paginator.paginate.side_effect = ClientError(
+            {"Error": {"Code": "AccessDeniedException"}}, "ListImages"
+        )
+
+        with self.assertRaises(ClientError):
+            start_codebuild.delete_resource(cleanup_event, None)
+
+    def test_delete_resource_ignores_non_cleanup_types(self):
+        event = {"ResourceType": "Custom::CodeBuildRun", "ResourceProperties": {}}
+
+        start_codebuild.delete_resource(event, None)
+
+        self.ecr_client.batch_delete_image.assert_not_called()
+
+
+if __name__ == "__main__":
+    unittest.main()
