Use API Gateway Stage Variables to Manage Lambda Functions for Multiple Deployment Environments

Nowadays, AWS Lambda and API Gateway are one of the most important and widely used services to build a serverless application. During development, we often need to create several deployment environments such as development, stage or production. Based on the environments we may want to connect to a different database from the Lambda function or trigger a completely different Lambda function or run some specific tasks. Now we can do this easily by using API Gateway stage variables and Lambda function aliases without creating environment-wise redundant API and Lambda functions.

In this article, I am going to create a simple API using API Gateway which will trigger a Lambda function. We will have two deployment environment: test and prod. Depending on the environment, the response of the API will be different.

Create a Lambda Function

Sign in to the AWS Management console, navigate to the Lambda console, and create a Lambda function, getCustomerList.

Inside the function code, use the following sample code and save it.

import json

def lambda_handler(event, context):
   # TODO implement
   return {
       'statusCode': 200,
       'body': json.dumps('Hello from Lambda prod environment!')
   }

We can assume that this is our latest stable code and we want to publish it in production. For this, we need to create a version for this Lambda function code. From Actions, select Publish new version and give it an appropriate version name. So, now we have two versions: 1 and $LATEST, both pointing to our same Lambda function code.

Again, suppose we had some change requests and we needed to update our Lambda function with the following code.

import json

def lambda_handler(event, context):
    # TODO implement
    return {
        'statusCode': 200,
        'body': json.dumps('Hello from Lambda test environment!')
    }

This time, we want to deploy this updated code in our test environment. Now, version $LATEST is pointing to the Lambda code mentioned above, whereas version 1 is pointing to our production Lambda code.

Create Lambda Function Aliases

AWS Lambda allows creating aliases for the Lambda function versions. An alias can only point to a function version, not to another alias. Unlike versions, aliases can be modified. We can update aliases to point to different versions of functions at any time.

For our Lambda function, we will create two aliases. From actions, select Create alias and create aliases: test and prod where alias test points to the $LATEST version and alias prod points to version 1. Now, depending on the environment we just need to call our Lambda function with the appropriate alias name.

Create API Gateway

Navigate to the API Gateway console, and create a new API, myDemoAPI.

From Actions, Create a resource, customers and create a GET method for this resource.

From integration request page, choose Lambda Function as the integration type, select the Lambda region in where our Lambda function is created and write getCustomerList:${stageVariables.alias} as the Lambda Function. Here, alias is an user-defined stage variable.

A popup will appear with some instructions and a CLI command to add function policy to the Lambda function.

Open AWS CLI and execute the given command on the popup by replacing the stage variable ${stageVariables.alias} with our Lambda function aliases name as below.

aws lambda add-permission
--function-name "arn:aws:lambda:ap-south-1:113510440346:function:getCustomerList:test"
--source-arn "arn:aws:execute-api:ap-south-1:51034134604:uiastwznv/*/GET/customers"
--principal apigateway.amazonaws.com
--statement-id 3aab5c51-8758-9da1-706efd9f5451
--action lambda:InvokeFunction

aws lambda add-permission
--function-name "arn:aws:lambda:ap-south-1:113510440346:function:getCustomerList:prod"
--source-arn "arn:aws:execute-api:ap-south-1:51034134604:uiastwznv/*/GET/customers"
--principal apigateway.amazonaws.com
--statement-id 3aab5c51-8758-9da1-706efd9f5451
--action lambda:InvokeFunction

If AWS CLI has not been configured before, then follow this URL to get started with it.

Deploy API

At this point, we need to deploy our API to make it publicly available. From Action, select Deploy API. Choose new deployment stage, enter test as stage name and deploy it. Similarly, create another stage for prod. Make sure that the stage names have to be the same as Lambda function aliases.

Create Stage Variables

Remember that we wrote getCustomerList:${stageVariables.alias} as the Lambda function in API integration request page, but still, we have not created any stage variable named alias. Now, navigate to Stages under MyDemoAPI and we can see the prod and test stages listed that we created earlier. Select test stage and under Stage Variables, add a new stage variable. Enter alias as name and test as value.

Similarly, create a new stage variable for prod stage and this time put alias as name and prod as value. That means, the test stage of our API points to the test alias Lambda function code and prod stage of our API points to the prod alias Lambda function code using stage variables. After that, we need to redeploy our API for both test and prod stage so that API Gateway can reflect the new changes.

Navigate to one of the API stages and select the GET method of the API to get the invoke URL. For the test stage, the invoke URL should be similar this:

https://uiastwznvw.execute-api.ap-south-1.amazonaws.com/test/customers

Open this API in the browser and we will get the response from Lambda function code with alias test. Again, open the invoke URL for the prod stage. This time, we will get the response from Lambda function code with prod alias.

https://uiastwznvw.execute-api.ap-south-1.amazonaws.com/prod/customers

The example mentioned above is one of the applications of API Gateway stage variables. Based on the requirements and application needs, it can be used in many ways, but the underline working idea is same.