feat: introduce aws_py_lambda macro

MODULE.bazel

@@ -8,16 +8,19 @@ module(
 
 # Lower-bound dependency versions.
 # Do not change unless the rules no longer work with the current version.
+bazel_dep(name = "aspect_rules_py", version = "0.3.0")
 bazel_dep(name = "bazel_skylib", version = "1.4.1")
-bazel_dep(name = "platforms", version = "0.0.5")
+bazel_dep(name = "platforms", version = "0.0.7")
+bazel_dep(name = "rules_oci", version = "1.4.0")
+bazel_dep(name = "rules_pkg", version = "0.9.1")
+bazel_dep(name = "rules_python", version = "0.25.0")
 
 # Development dependencies which are not exposed to users
-bazel_dep(name = "gazelle", version = "0.32.0", dev_dependency = True, repo_name = "bazel_gazelle")
-bazel_dep(name = "bazel_skylib_gazelle_plugin", version = "1.4.2", dev_dependency = True)
 bazel_dep(name = "aspect_bazel_lib", version = "1.32.1", dev_dependency = True)
+bazel_dep(name = "bazel_skylib_gazelle_plugin", version = "1.4.2", dev_dependency = True)
 bazel_dep(name = "buildifier_prebuilt", version = "6.1.2", dev_dependency = True)
-bazel_dep(name = "rules_pkg", version = "0.9.1", dev_dependency = True)
-bazel_dep(name = "rules_oci", version = "1.2.0", dev_dependency = True)
+bazel_dep(name = "container_structure_test", version = "1.16.0", dev_dependency = True)
+bazel_dep(name = "gazelle", version = "0.32.0", dev_dependency = True, repo_name = "bazel_gazelle")
 
 aws = use_extension("//aws:extensions.bzl", "aws")
 aws.toolchain(aws_cli_version = "2.13.0")
@@ -40,3 +43,22 @@ oci.pull(
 use_repo(oci, "ubuntu")
 
 register_toolchains("@aws_toolchains//:all")
+
+aws_py_lambda = use_extension(
+    "@aspect_rules_aws//aws:repositories.oci.bzl",
+    "aws_py_lambda",
+    dev_dependency = True,
+)
+use_repo(aws_py_lambda, "aws_lambda_python")
+
+pip = use_extension(
+    "@rules_python//python/extensions:pip.bzl",
+    "pip",
+    dev_dependency = True,
+)
+pip.parse(
+    hub_name = "pip",
+    python_version = "3.11",
+    requirements_lock = "//examples/python_lambda:requirements.txt",
+)
+use_repo(pip, "pip")

aws/defs.bzl

@@ -1 +1,54 @@
 "Public API re-exports"
+
+load("@rules_oci//oci:defs.bzl", "oci_image")
+load("//aws/private:py_lambda.bzl", "py_lambda_tars")
+load("@rules_python//python:defs.bzl", "py_binary")
+
+def aws_py_lambda(name, entry_point = "lambda_function.py", deps = [], base = "@aws_lambda_python"):
+    """Defines a Lambda run on the Python runtime.
+
+    See https://docs.aws.amazon.com/lambda/latest/dg/lambda-python.html
+
+    Produces an oci_image target following https://docs.aws.amazon.com/lambda/latest/dg/python-image.html
+
+    Args:
+        name: name of resulting target
+        entry_point: python source file implementing the handler
+        deps: third-party packages required at runtime
+        base: a base image that includes the AWS Runtime Interface Emulator
+
+    TODO:
+    - produce a [name].zip output following https://docs.aws.amazon.com/lambda/latest/dg/python-package.html#python-package-create-dependencies
+    """
+
+    bin_target = "_{}.bin".format(name)
+    tars_target = "_{}.tars".format(name)
+
+    # Convert //my/pkg:entry_point.py to my.pkg.entry_point.handler
+    cmd = "{}.{}.handler".format(native.package_name().replace("/", "."), entry_point.replace(".py", ""))
+
+    py_binary(
+        name = bin_target,
+        srcs = [entry_point],
+        main = entry_point,
+        deps = deps,
+    )
+
+    py_lambda_tars(
+        name = tars_target,
+        target = bin_target,
+    )
+
+    oci_image(
+        name = name,
+        base = base,
+        cmd = [cmd],
+        # Only allow building on linux, since we don't want to upload a lambda zip file
+        # with e.g. macos compiled binaries.
+        target_compatible_with = ["@platforms//os:linux"],
+        # N.B. deps layer appears first since it's larger and changes less frequently.
+        tars = [
+            "{}.deps".format(tars_target),
+            tars_target,
+        ],
+    )

aws/private/py_lambda.bzl

@@ -0,0 +1,116 @@
+"Rule to produce tar files with py_binary deps and app"
+
+# buildifier: disable=bzl-visibility
+load(
+    "@rules_pkg//pkg/private:pkg_files.bzl",
+    "add_empty_file",
+    "add_single_file",
+    "write_manifest",
+)
+
+def _short_path(file_):
+    # Remove prefixes for external and generated files.
+    # E.g.,
+    #   ../py_deps_pypi__pydantic/pydantic/__init__.py -> pydantic/__init__.py
+    short_path = file_.short_path
+    if short_path.startswith("../"):
+        second_slash = short_path.index("/", 3)
+        short_path = short_path[second_slash + 1:]
+    return short_path
+
+def _py_lambda_tar_impl(ctx):
+    deps = ctx.attr.target[DefaultInfo].default_runfiles.files
+    content_map = {}  # content handled in the manifest
+    files = []  # Files needed by rule implementation at runtime
+    args = ctx.actions.args()
+    args.add("--output", ctx.outputs.output.path)
+
+    for dep in deps.to_list():
+        short_path = _short_path(dep)
+
+        if dep.owner.workspace_name == "" and ctx.attr.kind == "app":
+            add_single_file(
+                content_map,
+                ctx.attr.prefix + "/" + dep.short_path,
+                dep,
+                ctx.label,
+            )
+        elif short_path.startswith("site-packages") and ctx.attr.kind == "deps":
+            short_path = short_path[len("site-packages"):]
+            add_single_file(
+                content_map,
+                ctx.attr.prefix + short_path,
+                dep,
+                ctx.label,
+            )
+
+    if ctx.attr.kind == "app" and ctx.attr.init_files:
+        path = ""
+        for dir in ctx.attr.init_files.split("/"):
+            path = path + "/" + dir
+            add_empty_file(
+                content_map,
+                ctx.attr.prefix + path + "/__init__.py",
+                ctx.label,
+            )
+
+    manifest_file = ctx.actions.declare_file(ctx.label.name + ".manifest")
+    files.append(manifest_file)
+    write_manifest(ctx, manifest_file, content_map)
+    args.add("--manifest", manifest_file.path)
+    args.add("--directory", "/")
+    args.set_param_file_format("flag_per_line")
+    args.use_param_file("@%s", use_always = False)
+    inputs = depset(direct = files, transitive = [deps])
+    ctx.actions.run(
+        outputs = [ctx.outputs.output],
+        inputs = inputs,
+        executable = ctx.executable._tar,
+        arguments = [args],
+        progress_message = "Creating archive...",
+        mnemonic = "PackageTar",
+    )
+
+    out = depset(direct = [ctx.outputs.output])
+    return [
+        DefaultInfo(files = out),
+        OutputGroupInfo(all_files = out),
+    ]
+
+_py_lambda_tar = rule(
+    implementation = _py_lambda_tar_impl,
+    attrs = {
+        "target": attr.label(
+            # require PyRuntimeInfo provider to be sure it's a py_binary ?
+        ),
+        "_tar": attr.label(
+            default = Label("@rules_pkg//pkg/private/tar:build_tar"),
+            cfg = "exec",
+            executable = True,
+        ),
+        "prefix": attr.string(doc = "path prefix for each entry in the tar"),
+        "init_files": attr.string(doc = "path where __init__ files will be placed"),
+        "kind": attr.string(values = ["app", "deps"]),
+        "output": attr.output(),
+    },
+)
+
+def py_lambda_tars(name, target, prefix = "/var/task", init_files = "examples/python_lambda", **kwargs):
+    _py_lambda_tar(
+        name = name,
+        kind = "app",
+        target = target,
+        prefix = prefix,
+        init_files = init_files,
+        output = name + ".app.tar",
+        **kwargs
+    )
+
+    _py_lambda_tar(
+        name = name + ".deps",
+        kind = "deps",
+        target = target,
+        prefix = prefix,
+        output = name + ".deps.tar",
+        **kwargs
+    )

aws/repositories.oci.bzl

@@ -0,0 +1,25 @@
+"Dependencies to fetch from remote docker registries"
+
+load("@rules_oci//oci:pull.bzl", "oci_pull")
+
+# As of 30 August 2023
+PY_LAMBDA_LATEST = "sha256:489d4abc8644060e2e16db2ffaaafa157359761feaf9438bf26ed88e37e43d9c"
+
+# See https://docs.aws.amazon.com/lambda/latest/dg/python-image.html#python-image-base
+def aws_py_lambda_repositories(digest = PY_LAMBDA_LATEST):
+    oci_pull(
+        name = "aws_lambda_python",
+        # tag = "3.11",
+        digest = digest,
+        platforms = ["linux/arm64/v8", "linux/amd64"],
+        image = "public.ecr.aws/lambda/python",
+    )
+
+def _aws_py_lambda_impl(_):
+    aws_py_lambda_repositories()
+
+aws_py_lambda = module_extension(
+    implementation = _aws_py_lambda_impl,
+    # TODO: allow bzlmod users to control the digest
+    # tag_classes = {"digest": digest},
+)

examples/python_lambda/BUILD.bazel

@@ -0,0 +1,28 @@
+"Example Python AWS lambda using a container"
+
+load("@aspect_rules_aws//aws:defs.bzl", "aws_py_lambda")
+load("@rules_oci//oci:defs.bzl", "oci_tarball")
+load("@rules_python//python:pip.bzl", "compile_pip_requirements")
+load("@pip//:requirements.bzl", "requirement")
+
+compile_pip_requirements(
+    name = "requirements",
+)
+
+aws_py_lambda(
+    name = "image",
+    entry_point = "lambda_function.py",
+    deps = [requirement("requests")],
+)
+
+# Manually verify the image in a local container:
+# $ bazel run //examples/python_lambda:tarball
+# $ docker run -p 9000:8080 --rm aws_lambda_hello_world:latest
+# (in another terminal)
+# $ curl "http://localhost:9000/2015-03-31/functions/function/invocations" -d '{}'
+oci_tarball(
+    name = "tarball",
+    image = ":image",
+    repo_tags = ["aws_lambda_hello_world:latest"],
+    visibility = [":__subpackages__"],
+)

examples/python_lambda/README.md

@@ -0,0 +1,7 @@
+# AWS Lambda example in Python
+
+This is a reference project for python AWS Lambda:
+
+- Generates python container image to be deployed to AWS Lambda
+- Integration test running AWS Lambda locally in docker using Runtime Interface Emulator
+- Test to verify container image file structure

examples/python_lambda/lambda_function.py

@@ -0,0 +1,9 @@
+# Copied from https://docs.aws.amazon.com/lambda/latest/dg/python-image.html#python-image-instructions
+import requests
+import sys
+
+def handler(event, context):
+    r = requests.get("https://www.example.com")
+    print(r.ok)
+
+    return 'Hello from AWS Lambda using Python' + sys.version + '!'

examples/python_lambda/requirements.in

@@ -0,0 +1,3 @@
+boto3
+testcontainers
+pytest