Assigning Different Lambda Functions to Stages in AWS CDK Stack

Question

I am working on a cdk deploy script and I have it somewhat working but am lost/not having luck setting the different stages and applying different lambdas to the api resources.

So I have

    // Construct lambdas - prod 
    const lambdaBacklogGet = new lambdajs.NodejsFunction(this, "name", {
      nodeModules: ['axios'],
      entry: './src/path/index.js',
      handler: 'handler',
      runtime: lambda.Runtime.NODEJS_12_X,
      timeout: cdk.Duration.seconds(20),
      role: webformRole
    });

   // Construct lambdas - dev 
    const devLambdaBacklogGet = new lambdajs.NodejsFunction(this, "name-dev", {
      nodeModules: ['axios'],
      entry: './src/path/index.js',
      handler: 'handler',
      runtime: lambda.Runtime.NODEJS_12_X,
      timeout: cdk.Duration.seconds(20),
      role: webformRole
    });

    // then I Construct API 
    const api = new apiGateway.RestApi(this, "name-api", {
      defaultCorsPreflightOptions: {
        allowOrigins: apiGateway.Cors.ALL_ORIGINS,
        allowHeaders: apiGateway.Cors.DEFAULT_HEADERS,
        allowMethods: apiGateway.Cors.ALL_METHODS,
      },
      description: "API for Something",
      deploy: true,
    });

    // Default response parameters, to allow CORS
    const corsResponseParams = {
      "method.response.header.Access-Control-Allow-Origin": true,
    };
    const corsIntegrationResponseParams = {
      "method.response.header.Access-Control-Allow-Origin": "'*'",
    };

    // REST API - /adding prod endpoint
    const backlogResourceAPI = api.root.addResource("backlog");
    const issuesBacklogGetIntegration = new apiGateway.LambdaIntegration(lambdaBacklogGet, {});
    backlogResourceAPI.addMethod("GET", issuesBacklogGetIntegration, {});

//at the end I do this
    // Then create an explicit Deployment construct
    const deployment  = new apiGateway.Deployment(this, 'my_deployment', { api });

    // And different stages
    const [devStage, testStage, prodStage] = ['dev', 'test'].map(item => 
      new apiGateway.Stage(this, `${item}_stage`, { deployment, stageName: item }));

    //api.deploymentStage = prodStage
    api.deploymentStage = devStage

the last part I know isn't setup correctly. But basically I want one API Gateway and I want the prod/dev stages to have identical resources/structure. However prod stage should pull specific prod lambdas and dev should call dev lambdas.

Is there any way to set that up through the CDK?

Thanks, Tim

Answer 1

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