AWS CDK - AWS Lambda Asset

my_stack.py

# from pathlib import Path  # Optional

# Using CDK v2
from aws_cdk import Duration
from aws_cdk.aws_lambda import Function, Runtime

from aws_cdk_lambda_asset.zip_asset_code import ZipAssetCode


# NOTE: Only example usage below (needs some modification)

# work_dir = Path(__file__).parents[1]

# noinspection PyTypeChecker
py_runtime: Runtime = Runtime.PYTHON_3_10

# Defines an AWS Lambda resource
fn = Function(
    self,
    # optional syntax, requires Python 3.8+
    fn_name := 'my-function-name',
    function_name=fn_name,  # Optional
    runtime=py_runtime,
    code=ZipAssetCode(
        include=['path/to/code'],  # Example: ['code']
        file_name='my-lambda.zip'
        # runtime=runtime,   # Optional
        # work_dir=work_dir, # Optional
        # test=False, # Optional
        # reqs_filename='requirements.txt', # Optional
    ),
    handler='path.to.handler.func',  # Example: 'code.handlers.my_handler'
    timeout=Duration.minutes(2),  # Optional
    memory_size=512,  # Optional
    environment={...}
)

zip_asset_code.py

"""
# Credits

https://gitlab.com/josef.stach/aws-cdk-lambda-asset/

# Note

Create a `aws_cdk_lambda_asset` package, and put this module in there.

"""

import platform
import uuid
from filecmp import cmp, dircmp
from glob import glob
from os import chdir, remove, system
from os.path import isdir, isfile
from pathlib import Path
from shutil import copy, copytree, make_archive, move, rmtree
from typing import List

import docker
import requests
from aws_cdk.aws_lambda import Architecture, AssetCode, Runtime, RuntimeFamily

# Project Root Directory, which contains the `cdk.json` file.
# Iterate backwards, until we find such a directory.
PROJECT_DIR = Path(__file__)
for p in PROJECT_DIR.parents:
    if (p / 'cdk.json').is_file():
        PROJECT_DIR = p
        break


def is_linux() -> bool:
    """
    :return: True if running on Linux. False otherwise.
    """
    return platform.system().lower() == 'linux'


# Check if two directories have the same contents
# Credits: https://stackoverflow.com/a/37790329/10237506
def same_folders(dcmp):
    if dcmp.diff_files or dcmp.left_only or dcmp.right_only:
        return False
    for sub_dcmp in dcmp.subdirs.values():
        if not same_folders(sub_dcmp):
            return False
    return True


class ZipAssetCode(AssetCode):
    """
    CDK AssetCode which builds lambda function and produces a ZIP file with dependencies.
    Lambda function is built either in Docker or natively when running on Linux.
    """

    def __init__(
        self,
        include: List[str],
        work_dir: Path = PROJECT_DIR,
        file_name: str | None = None,
        runtime: Runtime = Runtime.PYTHON_3_10,
        architecture: Architecture = Architecture.X86_64,
        reqs_filename='requirements.txt',
        test=False,
        always_build=False,
    ) -> None:
        """
        :param include: List of packages to include in the lambda archive. Examples: ['src', 'code']
        :param work_dir: Path to root directory, containing `requirements.txt`
        :param file_name: Lambda ZIP archive name. Example: 'my_lambda.zip'
        :param runtime: Python version of Lambda runtime. Example: PYTHON_3_10 (3.10)
        :param architecture: Lambda architecture (defaults to x86_64)
        :param reqs_filename: Requirements filename (defaults to requirements.txt)
        :param test: True if being run from a test environment
        :param always_build: True to always build (i.e. install dependencies with
          Docker), even if the `reqs_filename` has not changed.

        """
        if file_name is None:
            file_name = str(uuid.uuid4())[:8]

        if test:
            # any dummy file here
            asset_path = work_dir.joinpath(file_name or '')

        else:
            asset_path = LambdaPackaging(
                arch=architecture,
                runtime=runtime,
                include_paths=include,
                reqs_filename=reqs_filename,
                work_dir=work_dir,
                out_file=file_name,
                always_build=always_build,
            ).package()

        super().__init__(asset_path.as_posix())

    @property
    def is_inline(self) -> bool:
        return False


class LambdaPackaging:
    """
    EXCLUDE_DEPENDENCIES - List of libraries already included in the lambda runtime environment. No need to package these.
    EXCLUDE_FILES - List of files not required and therefore safe to be removed to save space.
    """

    EXCLUDE_DEPENDENCIES = {
        'bin',
        'boto3',
        'botocore',
        'dateutil',  # python-dateutil
        'docutils',
        'jmespath',
        'pip',
        's3transfer',
        'setuptools',
        'six.py',
        'urllib3',
    }

    EXCLUDE_FILES = {
        '*.dist-info',
        '__pycache__',
        '*.pyc',
        '*.pyo',
    }

    # Configuration for `pip install` within a Docker container
    TRUSTED_HOST = 'my-internal-pypi.com'
    CERT_NAME = 'my-custom-cert.pem'

    def __init__(
        self,
        *,
        arch: Architecture,
        runtime: Runtime,
        include_paths: List[str],
        out_file: str,
        reqs_filename: str,
        work_dir: Path,
        always_build=False,
    ) -> None:
        assert runtime.family is RuntimeFamily.PYTHON, 'Requires a PYTHON Runtime!'

        self._include_paths = include_paths
        self._zip_file = out_file.replace('.zip', '')
        self.arch = arch
        self.py_version = runtime.name[6:]
        self.work_dir = work_dir
        self.build_dir = build_dir = work_dir / '.build'
        self.requirements_dir = build_dir / 'requirements'
        self.requirements_txt = work_dir / reqs_filename
        self.reqs_filename = reqs_filename
        self.cert_file = work_dir / self.CERT_NAME

        if always_build:
            self.force_build = True
        else:
            try:
                self.force_build = not cmp(
                    self.build_dir / self.reqs_filename, self.requirements_txt
                )
            except FileNotFoundError:
                self.force_build = True

    @property
    def path(self) -> Path:
        return self.work_dir.joinpath(self._zip_file + '.zip').resolve()

    def package(self) -> Path:
        try:
            chdir(self.work_dir.as_posix())
            print(f'Working directory: {Path.cwd()}')

            if self.force_build:
                print(f'Build directory: {self.build_dir}')
                self._prepare_build()
                self._build_lambda()
            else:
                print(
                    f"Skipping build, as '{self.reqs_filename}' has not changed.\n"
                    f'    Enable `always_build` to override this behavior.'
                )

            self._package_lambda()
            return self.path

        except requests.exceptions.ConnectionError:
            raise Exception('Could not connect to Docker daemon.')

        except Exception as ex:
            raise Exception('Error during build.', str(ex))

    def _prepare_build(self) -> None:
        rmtree(self.build_dir, ignore_errors=True)
        self.requirements_dir.mkdir(parents=True)

        # Needed when building in Docker
        copy(self.requirements_txt, self.requirements_dir)
        copy(self.cert_file, self.requirements_dir)

        # print(f'Exporting poetry dependencies: {self.requirements_txt}')
        # result = system(f'poetry export --without-hashes --format requirements.txt --output {self.requirements_txt}')
        # if result != 0:
        #     raise EnvironmentError('Version of your poetry is not compatible - please update to 1.0.0b1 or newer')

    def _build_lambda(self) -> None:
        if is_linux():
            self._build_natively()
        else:
            self._build_in_docker()
        self._remove_bundled_files()

    def _build_in_docker(self) -> None:
        """
        Build lambda dependencies in a container as-close-as-possible to the actual runtime environment.
        """
        tag = f'python:{self.py_version}-{self.arch.name}'
        print(f'({tag}) Installing dependencies [running in Docker]...')
        client = docker.from_env()
        client.containers.run(
            image=f'public.ecr.aws/lambda/{tag}',
            entrypoint='/bin/sh',
            command=f"-c 'python -m pip config set global.cert /var/task/{self.CERT_NAME} && "
            f"python -m pip install "
            f"--trusted-host {self.TRUSTED_HOST} "
            f"--target /var/task/ --requirement /var/task/{self.reqs_filename} && "
            "find /var/task -name \\*.so -exec strip \\{{\\}} \\;'",
            remove=True,
            volumes={
                self.requirements_dir.as_posix(): {'bind': '/var/task', 'mode': 'rw'}
            },
            user=0,
        )

    def _build_natively(self) -> None:
        """
        Build lambda dependencies natively on linux. Should be the same architecture though.
        """
        print('Installing dependencies [running on Linux]...')
        if (
            system(
                f"/bin/sh -c 'pip3 config set global.cert {self.CERT_NAME} && "
                f"pip3 install --trusted-host {self.TRUSTED_HOST} --target {self.requirements_dir} "
                f"--requirement {self.requirements_txt} && "
                f"find {self.requirements_dir} -name \\*.so -exec strip \\{{\\}} \\;'"
            )
            != 0
        ):
            raise Exception(
                'Error running build in Docker. Make sure Docker daemon is running on your machine.'
            )

    def _package_lambda(self) -> None:
        build = has_changes = self.force_build

        if build:
            print(
                f'Moving required dependencies to the build directory: {self.build_dir}'
            )
            for req_dir in self.requirements_dir.glob('*'):
                move(str(req_dir), str(self.build_dir))
            rmtree(self.requirements_dir, ignore_errors=True)

        print('Copying \'include\' resources:')
        for include_path in self._include_paths:
            dst = self.build_dir / include_path
            print(f'    -  {(Path.cwd() / include_path).resolve()}')

            # if needed, check if 'include' folder contents have changed (since last run)
            if not has_changes:
                include_has_changes = not same_folders(dircmp(include_path, dst))
                if include_has_changes:
                    has_changes = True
                else:
                    # don't need to copy over directory, as it's the same
                    continue

            if dst.exists():
                rmtree(dst)

            copytree(include_path, dst)

        # Create zip archive, only if files/folders in build folder have changed
        zip_file_path = (self.work_dir / self._zip_file).resolve()
        zip_file_with_ext = zip_file_path.with_suffix('.zip')

        if has_changes or not zip_file_with_ext.exists():
            print(f'Packaging application into {zip_file_with_ext}')
            make_archive(
                str(zip_file_path), 'zip', root_dir=str(self.build_dir), verbose=True
            )

    def _remove_bundled_files(self) -> None:
        """
        Remove caches and dependencies already bundled in the lambda runtime environment.
        """
        print('Removing dependencies bundled in lambda runtime and caches:')
        for pattern in self.EXCLUDE_DEPENDENCIES.union(self.EXCLUDE_FILES):
            pattern = str(self.requirements_dir / '**' / pattern)
            print(f'    -  {pattern}')
            files = glob(pattern, recursive=True)
            for file_path in files:
                try:
                    if isdir(file_path):
                        rmtree(file_path)
                    if isfile(file_path):
                        remove(file_path)
                except OSError:
                    print(f'Error while deleting file: {file_path}')