feat: allow inclusion of image copy to pipeline

Author: Hi-Fi

src/index.ts

@@ -5,6 +5,8 @@
 import * as child_process from 'child_process';
 import * as path from 'path';
 import {
+  aws_codepipeline as codepipeline,
+  aws_codepipeline_actions as codepipeline_actions,
   aws_ec2 as ec2,
   aws_iam as iam,
   aws_lambda as lambda,
@@ -14,7 +16,6 @@ import {
   Token,
   DockerImage,
 } from 'aws-cdk-lib';
-import { PolicyStatement, AddToPrincipalPolicyResult } from 'aws-cdk-lib/aws-iam';
 import { Construct } from 'constructs';
 import { shouldUsePrebuiltLambda } from './config';
 
@@ -77,6 +78,11 @@ export interface ECRDeploymentProps {
    * The environment variable to set
    */
   readonly environment?: { [key: string]: string };
+
+  /**
+   * Stage to include lambda to. If this is set, lambda is invoked in pipeline instead of custom resource.
+   */
+  readonly stage?: codepipeline.IStage;
 }
 
 export interface IImageName {
@@ -200,7 +206,9 @@ export class ECRDeployment extends Construct {
   constructor(scope: Construct, id: string, props: ECRDeploymentProps) {
     super(scope, id);
     const memoryLimit = props.memoryLimit ?? 512;
-    this.handler = new lambda.SingletonFunction(this, 'CustomResourceHandler', {
+    const lambdaId = props.stage ? 'ImageCopyHandler' : 'CustomResourceHandler';
+    const lambdaPurpose = props.stage ? 'ImageCopy' : 'Custom::CDKECRDeployment';
+    this.handler = new lambda.SingletonFunction(this, lambdaId, {
       uuid: this.renderSingletonUuid(memoryLimit),
       code: lambda.Code.fromAsset(path.join(__dirname, '../lambda'), {
         bundling: {
@@ -254,7 +262,7 @@ export class ECRDeployment extends Construct {
       runtime: lambda.Runtime.GO_1_X,
       handler: 'main',
       environment: props.environment,
-      lambdaPurpose: 'Custom::CDKECRDeployment',
+      lambdaPurpose,
       timeout: Duration.minutes(15),
       role: props.role,
       memorySize: memoryLimit,
@@ -298,19 +306,33 @@ export class ECRDeployment extends Construct {
     props.src.creds?.secretManager?.secret.grantRead(handlerRole);
     props.dest.creds?.secretManager?.secret.grantRead(handlerRole);
 
-    new CustomResource(this, 'CustomResource', {
-      serviceToken: this.handler.functionArn,
-      resourceType: 'Custom::CDKBucketDeployment',
-      properties: {
-        SrcImage: props.src.uri,
-        SrcCreds: formatCredentials(props.src.creds),
-        DestImage: props.dest.uri,
-        DestCreds: formatCredentials(props.dest.creds),
-      },
-    });
+    if (props.stage) {
+      props.stage.addAction(
+        new codepipeline_actions.LambdaInvokeAction({
+          lambda: this.handler,
+          actionName: lambdaPurpose,
+          userParameters: {
+            SrcImage: props.src.uri,
+            SrcCreds: formatCredentials(props.src.creds),
+            DestImage: props.dest.uri,
+            DestCreds: formatCredentials(props.dest.creds),
+          },
+        }));
+    } else {
+      new CustomResource(this, 'CustomResource', {
+        serviceToken: this.handler.functionArn,
+        resourceType: lambdaPurpose,
+        properties: {
+          SrcImage: props.src.uri,
+          SrcCreds: formatCredentials(props.src.creds),
+          DestImage: props.dest.uri,
+          DestCreds: formatCredentials(props.dest.creds),
+        },
+      });
+    }
   }
 
-  public addToPrincipalPolicy(statement: PolicyStatement): AddToPrincipalPolicyResult {
+  public addToPrincipalPolicy(statement: iam.PolicyStatement): iam.AddToPrincipalPolicyResult {
     const handlerRole = this.handler.role;
     if (!handlerRole) { throw new Error('lambda.SingletonFunction should have created a Role'); }