When building APIs using AWS API Gateway backed by AWS Lambda functions, there’s a common pattern for creating separate development and production environments. Multiple stages (e.g., dev and prod) are created for the API. The name of the Lambda function (or its alias) to call for serving requests is stored in a stage variable. While building an AWS CDK stack, I found implementing this pattern surprisingly difficult, and it took me quite a while to find a solution. I’m sharing the solution I found here to make it easier for others.

Here are the variables I will use in my code to allow you to substitute your own values and adapt the code to your environment:

const region = "us-east-1";
const awsAccountId = "0123456789";
const apiConstructName = "YourRestApi";
const functionName = "your-lambda-function";
const devFunctionAlias = "your-lambda-function:dev";
const prodFunctionAlias = "your-lambda-function:prod";
const stageVariableName = "LambdaName";

First, you need to create the REST API itself:

const api = new apigateway.RestApi(this, apiConstructName, {
    // ...specify options here
});

Then you should give the API permission to call the Lambdas:

const credentialsRole = new iam.Role(this, `${props.apiConstructName}Role`, {
    assumedBy: new iam.ServicePrincipal("apigateway.amazonaws.com"),
});

credentialsRole.addToPolicy(
    new iam.PolicyStatement({
        actions: ["lambda:InvokeFunction"],
        resources: [
            // Add for each function you need to call:
            `arn:aws:lambda:${region}:${awsAccountId}:function:${functionName}`,
            // Add this if you want to use aliases:
            `arn:aws:lambda:${REGION}:${awsAccountId}:function:${functionName}:*`,
        ],
        effect: iam.Effect.ALLOW,
    })
);

Then, for each function, you also need to create an integration, which will allow you to attach it to a method. This is the tricky part. Instead of a LambdaIntegration, we create a more general AwsIntegration, which allows us to refer to a stage variable in the path. LambdaIntegration requires you to pass in an IFunction instance, and even if you use Function.fromFunctionArn(), including a stage variable reference in the ARN leads to a synthesis-time error.

const integration = new apigateway.AwsIntegration({
    proxy: true,
    service: "lambda",
    path: `2015-03-31/functions/arn:aws:lambda:${REGION}:${awsAccountId}:function:\${stageVariables.${stageVariableName}}/invocations`,
    options: {
        credentialsRole: credentialsRole,
    },
});

After this, you can use integration anywhere you need to call this function:

// /items
const itemsResource = api.root.addResource("items");
// POST /sam
const postItemMethod = itemsResource.addMethod("POST", integration);

And finally, of course, you need to configure the API Stages to use the stage variables:

const deployment = new apigateway.Deployment(this, `${props.apiConstructName}Deployment`, { api });

const devStage = new apigateway.Stage(this, `${props.apiConstructName}DevStage`, {
    stageName: 'dev',
    deployment: deployment,
    variables: {
        [stageVariableName]: devFunctionAlias,
    },
});

const prodStage = new apigateway.Stage(this, `${props.apiConstructName}ProdStage`, {
    stageName: 'prod',
    deployment: deployment,
    variables: {
        [stageVariableName]: prodFunctionAlias,
    },
});

api.deploymentStage = prodStage;

Hope this helps – happy hacking!