feat(lambda-python): add support for custom build image

packages/@aws-cdk/aws-lambda-python/README.md

@@ -36,6 +36,32 @@ new PythonFunction(this, 'MyFunction', {
 });
 ```
 
+Custom Docker images for bundling dependencies can be specified by specifying additional `buildImageOptions`. If using a custom Docker image, please ensure that the dependencies are stored at `/var/dependencies` within the Docker image for them to be bundled into the Lambda asset.
+
+A different bundling Docker image `Dockerfile.build` can be specified as:
+
+```ts
+new PythonFunction(this, 'MyFunction', {
+   ...
+   dockerBuildImageOptions: {
+     file: "Dockerfile.build",
+   },
+});
+```
+
+All bundling images are passed in the `IMAGE` Docker build arg that specifies the correct AWS SAM build image based on the runtime of the function. Additional build args can be specified as:
+
+```ts
+new PythonFunction(this, 'MyFunction', {
+  ...
+  dockerBuildImageOptions: {
+    buildArgs: {
+      HTTPS_PROXY: 'https://127.0.0.1:3001',
+    },
+  },
+});
+```
+
 All other properties of `lambda.Function` are supported, see also the [AWS Lambda construct library](https://github.com/aws/aws-cdk/tree/master/packages/%40aws-cdk/aws-lambda).
 
 ## Module Dependencies
@@ -44,9 +70,9 @@ If `requirements.txt` or `Pipfile` exists at the entry path, the construct will
 all required modules in a [Lambda compatible Docker container](https://gallery.ecr.aws/sam/build-python3.7)
 according to the `runtime`.
 
-Python bundles are only recreated and published when a file in a source directory has changed. 
+Python bundles are only recreated and published when a file in a source directory has changed.
 Therefore (and as a general best-practice), it is highly recommended to commit a lockfile with a
-list of all transitive dependencies and their exact versions. 
+list of all transitive dependencies and their exact versions.
 This will ensure that when any dependency version is updated, the bundle asset is recreated and uploaded.
 
 To that end, we recommend using [`pipenv`] or [`poetry`] which has lockfile support.

packages/@aws-cdk/aws-lambda-python/lib/bundling.ts

@@ -71,13 +71,23 @@ export interface BundlingOptions {
    * @default - based on `assetHashType`
    */
   readonly assetHash?: string;
+
+  /**
+   * Bundling Docker image options to use. If no options are provided, the default bundling image
+   * will be used. This is useful for specifying a custom Docker image for bundling. Additionally,
+   * the correct AWS SAM build image based on the runtime of the function will be passed as the build arg
+   * `IMAGE` to the Docker image.
+   *
+   * @default - uses default bundling.
+   */
+  readonly buildImageOptions?: cdk.DockerBuildOptions;
 }
 
 /**
  * Produce bundled Lambda asset code
  */
 export function bundle(options: BundlingOptions): lambda.Code {
-  const { entry, runtime, outputPathSuffix } = options;
+  const { entry, runtime, outputPathSuffix, buildImageOptions } = options;
 
   const stagedir = cdk.FileSystem.mkdtemp('python-bundling-');
   const hasDeps = stageDependencies(entry, stagedir);
@@ -97,12 +107,20 @@ export function bundle(options: BundlingOptions): lambda.Code {
 
   // copy Dockerfile to workdir
   fs.copyFileSync(path.join(__dirname, dockerfile), path.join(stagedir, dockerfile));
+  // if custom build Dockerfile is provided, copy it to workdir
+  if (buildImageOptions?.file) {
+    fs.copyFileSync(path.join(entry, buildImageOptions.file), path.join(stagedir, buildImageOptions.file));
+  }
+
+  const buildArgs = {
+    IMAGE: runtime.bundlingImage.image,
+    ...buildImageOptions?.buildArgs,
+  };
 
   const image = cdk.DockerImage.fromBuild(stagedir, {
-    buildArgs: {
-      IMAGE: runtime.bundlingImage.image,
-    },
     file: dockerfile,
+    ...buildImageOptions,
+    buildArgs,
   });
 
   return lambda.Code.fromAsset(entry, {

packages/@aws-cdk/aws-lambda-python/lib/function.ts

@@ -1,7 +1,7 @@
 import * as fs from 'fs';
 import * as path from 'path';
 import * as lambda from '@aws-cdk/aws-lambda';
-import { AssetHashType } from '@aws-cdk/core';
+import { AssetHashType, DockerBuildOptions } from '@aws-cdk/core';
 import { bundle } from './bundling';
 
 // keep this import separate from other imports to reduce chance for merge conflicts with v2-main
@@ -77,6 +77,12 @@ export interface PythonFunctionProps extends lambda.FunctionOptions {
    * @default - based on `assetHashType`
    */
   readonly assetHash?: string;
+
+  /** Custom build options for the bundling Docker image.
+   *
+   * @default - uses default bundling image and options.
+   */
+  readonly dockerBuildImageOptions?: DockerBuildOptions;
 }
 
 /**
@@ -112,6 +118,7 @@ export class PythonFunction extends lambda.Function {
         outputPathSuffix: '.',
         assetHashType: props.assetHashType,
         assetHash: props.assetHash,
+        buildImageOptions: props.dockerBuildImageOptions,
       }),
       handler: `${index.slice(0, -3)}.${handler}`,
     });

packages/@aws-cdk/aws-lambda-python/lib/layer.ts

@@ -1,5 +1,6 @@
 import * as path from 'path';
 import * as lambda from '@aws-cdk/aws-lambda';
+import { DockerBuildOptions } from '@aws-cdk/core';
 import { bundle } from './bundling';
 
 // keep this import separate from other imports to reduce chance for merge conflicts with v2-main
@@ -21,6 +22,12 @@ export interface PythonLayerVersionProps extends lambda.LayerVersionOptions {
    * @default - All runtimes are supported.
    */
   readonly compatibleRuntimes?: lambda.Runtime[];
+
+  /** Custom build options for the bundling Docker image.
+   *
+   * @default - uses default bundling image and options.
+   */
+  readonly dockerBuildImageOptions?: DockerBuildOptions;
 }
 
 /**
@@ -50,6 +57,7 @@ export class PythonLayerVersion extends lambda.LayerVersion {
         entry,
         runtime,
         outputPathSuffix: 'python',
+        buildImageOptions: props.dockerBuildImageOptions,
       }),
     });
   }

packages/@aws-cdk/aws-lambda-python/test/bundling.test.ts

@@ -1,7 +1,7 @@
 import * as fs from 'fs';
 import * as path from 'path';
 import { Code, Runtime } from '@aws-cdk/aws-lambda';
-import { FileSystem } from '@aws-cdk/core';
+import { DockerImage, FileSystem } from '@aws-cdk/core';
 import { stageDependencies, bundle } from '../lib/bundling';
 
 jest.spyOn(Code, 'fromAsset');
@@ -19,8 +19,17 @@ jest.mock('child_process', () => ({
   }),
 }));
 
+// Mock DockerImage.fromAsset() to avoid building the image
+let fromBuildMock: jest.SpyInstance<DockerImage>;
 beforeEach(() => {
   jest.clearAllMocks();
+
+  fromBuildMock = jest.spyOn(DockerImage, 'fromBuild').mockReturnValue({
+    image: 'built-image',
+    cp: () => 'dest-path',
+    run: () => {},
+    toJSON: () => 'built-image',
+  });
 });
 
 test('Bundling a function without dependencies', () => {
@@ -139,3 +148,29 @@ describe('Dependency detection', () => {
     expect(stageDependencies(sourcedir, '/dummy')).toEqual(false);
   });
 });
+
+test('Bundling Docker with custom bundling image', () => {
+  const entry = path.join(__dirname, 'lambda-handler-custom-build-docker-image');
+  bundle({
+    entry,
+    runtime: Runtime.PYTHON_3_7,
+    outputPathSuffix: '.',
+    buildImageOptions: {
+      buildArgs: {
+        HELLO: 'WORLD',
+        IMAGE: Runtime.PYTHON_3_7.bundlingImage.image,
+      },
+      file: 'Dockerfile.build',
+    },
+  });
+
+  expect(fromBuildMock).toHaveBeenCalledWith(expect.stringContaining('python-bundling'),
+    expect.objectContaining({
+      buildArgs: expect.objectContaining({
+        HELLO: 'WORLD',
+        IMAGE: Runtime.PYTHON_3_7.bundlingImage.image,
+      }),
+      file: expect.stringContaining('Dockerfile.build'),
+    }),
+  );
+});

packages/@aws-cdk/aws-lambda-python/test/integ.function.custom-build.expected.json

@@ -0,0 +1,113 @@
+{
+  "Resources": {
+    "myhandlerServiceRole77891068": {
+      "Type": "AWS::IAM::Role",
+      "Properties": {
+        "AssumeRolePolicyDocument": {
+          "Statement": [
+            {
+              "Action": "sts:AssumeRole",
+              "Effect": "Allow",
+              "Principal": {
+                "Service": "lambda.amazonaws.com"
+              }
+            }
+          ],
+          "Version": "2012-10-17"
+        },
+        "ManagedPolicyArns": [
+          {
+            "Fn::Join": [
+              "",
+              [
+                "arn:",
+                {
+                  "Ref": "AWS::Partition"
+                },
+                ":iam::aws:policy/service-role/AWSLambdaBasicExecutionRole"
+              ]
+            ]
+          }
+        ]
+      }
+    },
+    "myhandlerD202FA8E": {
+      "Type": "AWS::Lambda::Function",
+      "Properties": {
+        "Code": {
+          "S3Bucket": {
+            "Ref": "AssetParameters8fc4fd4f1abe3a706b6b352735cbfd7feb66aca750b63d920c6f018e9d665b9dS3Bucket0D7ED7BA"
+          },
+          "S3Key": {
+            "Fn::Join": [
+              "",
+              [
+                {
+                  "Fn::Select": [
+                    0,
+                    {
+                      "Fn::Split": [
+                        "||",
+                        {
+                          "Ref": "AssetParameters8fc4fd4f1abe3a706b6b352735cbfd7feb66aca750b63d920c6f018e9d665b9dS3VersionKey5CD58274"
+                        }
+                      ]
+                    }
+                  ]
+                },
+                {
+                  "Fn::Select": [
+                    1,
+                    {
+                      "Fn::Split": [
+                        "||",
+                        {
+                          "Ref": "AssetParameters8fc4fd4f1abe3a706b6b352735cbfd7feb66aca750b63d920c6f018e9d665b9dS3VersionKey5CD58274"
+                        }
+                      ]
+                    }
+                  ]
+                }
+              ]
+            ]
+          }
+        },
+        "Role": {
+          "Fn::GetAtt": [
+            "myhandlerServiceRole77891068",
+            "Arn"
+          ]
+        },
+        "Handler": "index.handler",
+        "Runtime": "python3.8"
+      },
+      "DependsOn": [
+        "myhandlerServiceRole77891068"
+      ]
+    }
+  },
+  "Parameters": {
+    "AssetParameters8fc4fd4f1abe3a706b6b352735cbfd7feb66aca750b63d920c6f018e9d665b9dS3Bucket0D7ED7BA": {
+      "Type": "String",
+      "Description": "S3 bucket for asset \"8fc4fd4f1abe3a706b6b352735cbfd7feb66aca750b63d920c6f018e9d665b9d\""
+    },
+    "AssetParameters8fc4fd4f1abe3a706b6b352735cbfd7feb66aca750b63d920c6f018e9d665b9dS3VersionKey5CD58274": {
+      "Type": "String",
+      "Description": "S3 key for asset version \"8fc4fd4f1abe3a706b6b352735cbfd7feb66aca750b63d920c6f018e9d665b9d\""
+    },
+    "AssetParameters8fc4fd4f1abe3a706b6b352735cbfd7feb66aca750b63d920c6f018e9d665b9dArtifactHash3FD1FF77": {
+      "Type": "String",
+      "Description": "Artifact hash for asset \"8fc4fd4f1abe3a706b6b352735cbfd7feb66aca750b63d920c6f018e9d665b9d\""
+    }
+  },
+  "Outputs": {
+    "FunctionArn": {
+      "Value": {
+        "Fn::GetAtt": [
+          "myhandlerD202FA8E",
+          "Arn"
+        ]
+      }
+    }
+  }
+}

packages/@aws-cdk/aws-lambda-python/test/integ.function.custom-build.ts

@@ -0,0 +1,30 @@
+import * as path from 'path';
+import { Runtime } from '@aws-cdk/aws-lambda';
+import { App, CfnOutput, Stack, StackProps } from '@aws-cdk/core';
+import { Construct } from 'constructs';
+import * as lambda from '../lib';
+
+/*
+ * Stack verification steps:
+ * * aws lambda invoke --function-name <deployed fn name> --invocation-type Event --payload '"OK"' response.json
+ */
+
+class TestStack extends Stack {
+  constructor(scope: Construct, id: string, props?: StackProps) {
+    super(scope, id, props);
+
+    const fn = new lambda.PythonFunction(this, 'my_handler', {
+      entry: path.join(__dirname, 'lambda-handler-custom-build-docker-image'),
+      runtime: Runtime.PYTHON_3_8,
+      dockerBuildImageOptions: { file: 'Dockerfile.build' },
+    });
+
+    new CfnOutput(this, 'FunctionArn', {
+      value: fn.functionArn,
+    });
+  }
+}
+
+const app = new App();
+new TestStack(app, 'cdk-integ-lambda-python');
+app.synth();