Assigning Different Lambda Functions to Stages in AWS CDK Stack

While it is possible to use a Stage Variable when adding a Lambda integration to an API gateway in the AWS Console it does not seem possible when using CDK at the time of writing this answer.

It is possible to use a Stage Variable in a path in an S3 integration but for a Lambda integration, the best I could manage was to read the Stage Variables from the event and use logic in the Lambda to depending upon which stage was used.

The code for my stack (in Python):

from aws_cdk import core
from aws_cdk import aws_lambda, aws_s3
from my_cdk.api_gateway import ApiGateway


class MyCdkStack(core.Stack):

    def __init__(
        self,
        scope: core.Construct,
        construct_id: str,
        **kwargs
    ) -> None:
        super().__init__(scope, construct_id, **kwargs)

        s3_bucket = aws_s3.Bucket(
            self,
            "s3_bucket",
            bucket_name="dan-dev-so-test-01"
        )

        default_lambda = aws_lambda.Function(
            self,
            "default_lambda",
            code=aws_lambda.Code.asset("./my_cdk/lambdas/"),
            handler="test_lambda.handler",
            runtime=aws_lambda.Runtime.PYTHON_3_8,
        )

        test_lambda = aws_lambda.Function(
            self,
            "test lambda",
            code=aws_lambda.Code.asset("./my_cdk/lambdas/"),
            handler="test_lambda.handler",
            runtime=aws_lambda.Runtime.PYTHON_3_8,
        )

        ApiGateway(
            self,
            "Api Gateway construct",
            default_lambda=default_lambda,
            test_lambda=test_lambda,
            s3_bucket=s3_bucket

        )

The code for api_gateway.py

from aws_cdk import core
from aws_cdk import aws_apigateway, aws_lambda, aws_iam, aws_s3


class ApiGateway(core.Construct):

    def __init__(
        self,
        scope: core.Construct,
        construct_id: str,
        default_lambda: aws_lambda.Function,
        test_lambda: aws_lambda.Function,
        s3_bucket: aws_s3.Bucket
    ) -> None:
        super().__init__(scope, construct_id)

        api = aws_apigateway.LambdaRestApi(
            self,
            "api",
            proxy=False,
            handler=default_lambda,
            deploy_options=aws_apigateway.StageOptions(
                stage_name='v1',
                variables={
                    'stage_name': 'v1'
                }
            )
        )

        v2_deployment = aws_apigateway.Deployment(
            self,
            "deployment2",
            api=api,
            description="v2"
        )

        aws_apigateway.Stage(
            self,
            "stage 2",
            deployment=v2_deployment,
            stage_name="v2",
            variables={
                'stage_name': 'v2'
            }
        )

        test_lambda.add_permission(
            "permission",
            principal=aws_iam.ServicePrincipal('apigateway.amazonaws.com'),
            source_arn=api.arn_for_execute_api(
                stage='v2', method='GET', path='/lambdas'
            )
        )
        lambda_endpoint = api.root.add_resource("lambdas")
        s3_endpoint = api.root.add_resource("test_s3")
        lambda_endpoint.add_method(
            "GET",
            aws_apigateway.LambdaIntegration(test_lambda)
        )

        response_200 = aws_apigateway.IntegrationResponse(
            status_code="200"
        )

        method_200 =aws_apigateway.MethodResponse(
            status_code="200"
        )

        credential_role = aws_iam.Role(
            self,
            "s3_role",
            assumed_by=aws_iam.ServicePrincipal("apigateway.amazonaws.com"),
            managed_policies=[aws_iam.ManagedPolicy.from_aws_managed_policy_name("AmazonS3FullAccess")]
        )

        s3_endpoint_integration = aws_apigateway.AwsIntegration(
            service="s3",
            integration_http_method="GET",
            options=aws_apigateway.IntegrationOptions(
                credentials_role=credential_role,
                integration_responses=[response_200]
            ),
            path=f"{s3_bucket.bucket_name}/${{stageVariables.stage_name}}/file.json",


        )

        s3_endpoint.add_method(
            "GET",
            s3_endpoint_integration,
            method_responses=[method_200]
        )

The code for the Lambda:

import json


def handler(event, context):
    body = {"Error": "Not found"}
    if event['stageVariables'].get('stage_name') == 'v1':
        body = {"hi": "From V1"}
    elif event['stageVariables'].get('stage_name') == 'v2':
        body = {"hi": "From V2"}
    return {
        "statusCode": 200,
        "body": json.dumps(body)
    }

Given the API is called using the /v1/lambdas endpoint it will return:

{"hi": "From V1"}

Given the API is called using the /v2/lambdas endpoint it will return:

{"hi": "From V2"}

Given there is a file under the path v1/file.json in the S3 bucket the file contents will be returned using the endpoint v1/test_s3

Given there is a file under the path v2/file.json in the S3 bucket the file contents will be returned using the endpoint v2/test_s3