feat(ecr): add option to auto delete images upon ECR repo removal

kirintwn · kirintwn · commit 4b7ff58e5880 · 2021-08-07T16:28:30.000+08:00
diff --git a/packages/@aws-cdk/aws-ecr/README.md b/packages/@aws-cdk/aws-ecr/README.md
@@ -94,3 +94,21 @@ is important here):
 repository.addLifecycleRule({ tagPrefixList: ['prod'], maxImageCount: 9999 });
 repository.addLifecycleRule({ maxImageAge: cdk.Duration.days(30) });
 ```
+
+### Repository deletion
+
+When a repository is removed from a stack (or the stack is deleted), the ECR
+repository will be removed according to its removal policy (which by default will
+simply orphan the repository and leave it in your AWS account). If the removal
+policy is set to `RemovalPolicy.DESTROY`, the repository will be deleted as long
+as it does not contain any images.
+
+To override this and force all images to get deleted during repository deletion,
+enable the`autoDeleteImages` option.
+
+```ts
+const repository = new Repository(this, 'MyTempRepo', {
+  removalPolicy: RemovalPolicy.DESTROY,
+  autoDeleteImages: true,
+});
+```
diff --git a/packages/@aws-cdk/aws-ecr/lib/auto-delete-images-handler/index.ts b/packages/@aws-cdk/aws-ecr/lib/auto-delete-images-handler/index.ts
@@ -0,0 +1,48 @@
+// eslint-disable-next-line import/no-extraneous-dependencies
+import { ECR } from 'aws-sdk';
+
+const ecr = new ECR();
+
+export async function handler(event: AWSLambda.CloudFormationCustomResourceEvent) {
+  switch (event.RequestType) {
+    case 'Create':
+    case 'Update':
+      return;
+    case 'Delete':
+      return onDelete(event);
+  }
+}
+
+/**
+ * Recursively delete all images in the repository
+ *
+ * @param ECR.ListImagesRequest the repositoryName & nextToken if presented
+ */
+async function emptyRepository(params: ECR.ListImagesRequest) {
+  const listedImages = await ecr.listImages(params).promise();
+  const imageIds = listedImages?.imageIds ?? [];
+  const nextToken = listedImages.nextToken ?? null;
+  if (imageIds.length === 0) {
+    return;
+  }
+
+  await ecr.batchDeleteImage({
+    repositoryName: params.repositoryName,
+    imageIds,
+  }).promise();
+
+  if (nextToken) {
+    await emptyRepository({
+      ...params,
+      nextToken,
+    });
+  }
+}
+
+async function onDelete(deleteEvent: AWSLambda.CloudFormationCustomResourceDeleteEvent) {
+  const repositoryName = deleteEvent.ResourceProperties?.RepositoryName;
+  if (!repositoryName) {
+    throw new Error('No RepositoryName was provided.');
+  }
+  await emptyRepository({ repositoryName });
+}
diff --git a/packages/@aws-cdk/aws-ecr/lib/repository.ts b/packages/@aws-cdk/aws-ecr/lib/repository.ts
@@ -1,7 +1,11 @@
 import { EOL } from 'os';
+import * as path from 'path';
 import * as events from '@aws-cdk/aws-events';
 import * as iam from '@aws-cdk/aws-iam';
-import { IResource, Lazy, RemovalPolicy, Resource, Stack, Token } from '@aws-cdk/core';
+import {
+  IResource, Lazy, RemovalPolicy, Resource, Stack, Token,
+  CustomResource, CustomResourceProvider, CustomResourceProviderRuntime,
+} from '@aws-cdk/core';
 import { IConstruct, Construct } from 'constructs';
 import { CfnRepository } from './ecr.generated';
 import { LifecycleRule, TagStatus } from './lifecycle';
@@ -349,6 +353,17 @@ export interface RepositoryProps {
    */
   readonly removalPolicy?: RemovalPolicy;
 
+  /**
+   * Whether all images should be automatically deleted when the repository is
+   * removed from the stack or when the stack is deleted.
+   *
+   * Requires the `removalPolicy` to be set to `RemovalPolicy.DESTROY`.
+   *
+   * @default false
+   */
+  readonly autoDeleteImages?: boolean;
+
+
   /**
    * Enable the scan on push when creating the repository
    *
@@ -441,6 +456,7 @@ export class Repository extends RepositoryBase {
     });
   }
 
+  private static readonly AUTO_DELETE_IMAGES_RESOURCE_TYPE = 'Custom::ECRAutoDeleteImages';
 
   private static validateRepositoryName(physicalName: string) {
     const repositoryName = physicalName;
@@ -496,6 +512,12 @@ export class Repository extends RepositoryBase {
     if (props.lifecycleRules) {
       props.lifecycleRules.forEach(this.addLifecycleRule.bind(this));
     }
+    if (props.autoDeleteImages) {
+      if (props.removalPolicy !== RemovalPolicy.DESTROY) {
+        throw new Error('Cannot use \'autoDeleteImages\' property on a repository without setting removal policy to \'DESTROY\'.');
+      }
+      this.enableAutoDeleteImages();
+    }
 
     this.repositoryName = this.getResourceNameAttribute(resource.ref);
     this.repositoryArn = this.getResourceArnAttribute(resource.attrArn, {
@@ -600,6 +622,35 @@ export class Repository extends RepositoryBase {
     validateAnyRuleLast(ret);
     return ret;
   }
+
+  private enableAutoDeleteImages() {
+    const provider = CustomResourceProvider.getOrCreateProvider(this, Repository.AUTO_DELETE_IMAGES_RESOURCE_TYPE, {
+      codeDirectory: path.join(__dirname, 'auto-delete-images-handler'),
+      runtime: CustomResourceProviderRuntime.NODEJS_14_X,
+    });
+
+    // Use a iam policy to allow the custom resource to list & delete
+    // images in the repository
+    this.addToResourcePolicy(new iam.PolicyStatement({
+      actions: [
+        'ecr:BatchDeleteImage',
+        'ecr:ListImages',
+      ],
+      resources: [
+        this.repositoryArn,
+      ],
+      principals: [new iam.ArnPrincipal(provider.roleArn)],
+    }));
+
+    const customResource = new CustomResource(this, 'AutoDeleteImagesCustomResource', {
+      resourceType: Repository.AUTO_DELETE_IMAGES_RESOURCE_TYPE,
+      serviceToken: provider.serviceToken,
+      properties: {
+        RepositoryName: Lazy.any({ produce: () => this.repositoryName }),
+      },
+    });
+    customResource.node.addDependency(this);
+  }
 }
 
 function validateAnyRuleLast(rules: LifecycleRule[]) {
diff --git a/packages/@aws-cdk/aws-ecr/test/auto-delete-images-handler.test.ts b/packages/@aws-cdk/aws-ecr/test/auto-delete-images-handler.test.ts
@@ -0,0 +1,198 @@
+const mockECRClient = {
+  listImages: jest.fn().mockReturnThis(),
+  batchDeleteImage: jest.fn().mockReturnThis(),
+  promise: jest.fn(),
+};
+
+import { handler } from '../lib/auto-delete-images-handler';
+
+jest.mock('aws-sdk', () => {
+  return { ECR: jest.fn(() => mockECRClient) };
+});
+
+beforeEach(() => {
+  mockECRClient.listImages.mockReturnThis();
+  mockECRClient.batchDeleteImage.mockReturnThis();
+});
+
+afterEach(() => {
+  jest.resetAllMocks();
+});
+
+test('does nothing on create event', async () => {
+  // GIVEN
+  const event: Partial<AWSLambda.CloudFormationCustomResourceCreateEvent> = {
+    RequestType: 'Create',
+    ResourceProperties: {
+      ServiceToken: 'Foo',
+      RepositoryName: 'MyRepo',
+    },
+  };
+
+  // WHEN
+  await invokeHandler(event);
+
+  // THEN
+  expect(mockECRClient.listImages).toHaveBeenCalledTimes(0);
+  expect(mockECRClient.batchDeleteImage).toHaveBeenCalledTimes(0);
+});
+
+test('does nothing on update event', async () => {
+  // GIVEN
+  const event: Partial<AWSLambda.CloudFormationCustomResourceUpdateEvent> = {
+    RequestType: 'Update',
+    ResourceProperties: {
+      ServiceToken: 'Foo',
+      RepositoryName: 'MyRepo',
+    },
+  };
+
+  // WHEN
+  await invokeHandler(event);
+
+  // THEN
+  expect(mockECRClient.listImages).toHaveBeenCalledTimes(0);
+  expect(mockECRClient.batchDeleteImage).toHaveBeenCalledTimes(0);
+});
+
+test('deletes no images on delete event when repository has no images', async () => {
+  // GIVEN
+  mockECRClient.promise.mockResolvedValue({ imageIds: [] }); // listedImages() call
+
+  // WHEN
+  const event: Partial<AWSLambda.CloudFormationCustomResourceDeleteEvent> = {
+    RequestType: 'Delete',
+    ResourceProperties: {
+      ServiceToken: 'Foo',
+      RepositoryName: 'MyRepo',
+    },
+  };
+  await invokeHandler(event);
+
+  // THEN
+  expect(mockECRClient.listImages).toHaveBeenCalledTimes(1);
+  expect(mockECRClient.listImages).toHaveBeenCalledWith({ repositoryName: 'MyRepo' });
+  expect(mockECRClient.batchDeleteImage).toHaveBeenCalledTimes(0);
+});
+
+test('deletes all images on delete event', async () => {
+  mockECRClient.promise.mockResolvedValue({ // listedImages() call
+    imageIds: [
+      {
+        imageTag: 'tag1',
+        imageDigest: 'sha256-1',
+      },
+      {
+        imageTag: 'tag2',
+        imageDigest: 'sha256-2',
+      },
+    ],
+  });
+
+  // WHEN
+  const event: Partial<AWSLambda.CloudFormationCustomResourceDeleteEvent> = {
+    RequestType: 'Delete',
+    ResourceProperties: {
+      ServiceToken: 'Foo',
+      RepositoryName: 'MyRepo',
+    },
+  };
+  await invokeHandler(event);
+
+  // THEN
+  expect(mockECRClient.listImages).toHaveBeenCalledTimes(1);
+  expect(mockECRClient.listImages).toHaveBeenCalledWith({ repositoryName: 'MyRepo' });
+  expect(mockECRClient.batchDeleteImage).toHaveBeenCalledTimes(1);
+  expect(mockECRClient.batchDeleteImage).toHaveBeenCalledWith({
+    repositoryName: 'MyRepo',
+    imageIds: [
+      {
+        imageTag: 'tag1',
+        imageDigest: 'sha256-1',
+      },
+      {
+        imageTag: 'tag2',
+        imageDigest: 'sha256-2',
+      },
+    ],
+  });
+});
+
+test('delete event where repo has many images does recurse appropriately', async () => {
+  // GIVEN
+  mockECRClient.promise // listedImages() call
+    .mockResolvedValueOnce({
+      imageIds: [
+        {
+          imageTag: 'tag1',
+          imageDigest: 'sha256-1',
+        },
+        {
+          imageTag: 'tag2',
+          imageDigest: 'sha256-2',
+        },
+      ],
+      nextToken: 'token1',
+    })
+    .mockResolvedValueOnce(undefined) // batchDeleteImage() call
+    .mockResolvedValueOnce({ // listedImages() call
+      imageIds: [
+        {
+          imageTag: 'tag3',
+          imageDigest: 'sha256-3',
+        },
+        {
+          imageTag: 'tag4',
+          imageDigest: 'sha256-4',
+        },
+      ],
+    });
+
+  // WHEN
+  const event: Partial<AWSLambda.CloudFormationCustomResourceDeleteEvent> = {
+    RequestType: 'Delete',
+    ResourceProperties: {
+      ServiceToken: 'Foo',
+      RepositoryName: 'MyRepo',
+    },
+  };
+  await invokeHandler(event);
+
+  // THEN
+  expect(mockECRClient.listImages).toHaveBeenCalledTimes(2);
+  expect(mockECRClient.listImages).toHaveBeenCalledWith({ repositoryName: 'MyRepo' });
+  expect(mockECRClient.batchDeleteImage).toHaveBeenCalledTimes(2);
+  expect(mockECRClient.batchDeleteImage).toHaveBeenNthCalledWith(1, {
+    repositoryName: 'MyRepo',
+    imageIds: [
+      {
+        imageTag: 'tag1',
+        imageDigest: 'sha256-1',
+      },
+      {
+        imageTag: 'tag2',
+        imageDigest: 'sha256-2',
+      },
+    ],
+  });
+  expect(mockECRClient.batchDeleteImage).toHaveBeenNthCalledWith(2, {
+    repositoryName: 'MyRepo',
+    imageIds: [
+      {
+        imageTag: 'tag3',
+        imageDigest: 'sha256-3',
+      },
+      {
+        imageTag: 'tag4',
+        imageDigest: 'sha256-4',
+      },
+    ],
+  });
+});
+
+
+// helper function to get around TypeScript expecting a complete event object,
+// even though our tests only need some of the fields
+async function invokeHandler(event: Partial<AWSLambda.CloudFormationCustomResourceEvent>) {
+  return handler(event as AWSLambda.CloudFormationCustomResourceEvent);
+}
diff --git a/packages/@aws-cdk/aws-ecr/test/integ.repository-auto-delete-images.expected.json b/packages/@aws-cdk/aws-ecr/test/integ.repository-auto-delete-images.expected.json
diff --git a/packages/@aws-cdk/aws-ecr/test/integ.repository-auto-delete-images.ts b/packages/@aws-cdk/aws-ecr/test/integ.repository-auto-delete-images.ts
