Using API Gateway stage variables to manage Lambda functions

Ed Lima, Cloud Support Engineer

There’s a new feature on Amazon API Gateway called stage variables. Stage variables act like environment variables and can be used to change the behavior of your API Gateway methods for each deployment stage; for example, making it possible to reach a different back end depending on which stage the API is running on. This blog post will demonstrate how to use stage variables with two different AWS Lambda functions.

For this example we will use the sample functions from the Lambda Walkthrough. Sign in to the AWS Management console, open the Lambda console, and create the required functions (make sure you’re using the appropriate IAM execution role:

GetHelloWorld

console.log('Loading event');

exports.handler = function(event, context) {
  context.done(null, {"Hello":"World"});  // SUCCESS with message
};

GetHelloWorldWithName

console.log('Loading event');
            
exports.handler = function(event, context) {
  var name = (event.name === undefined ? 'No-Name' : event.name);
  context.done(null, {"Hello":name}); // SUCCESS with message
};

In the API Gateway console, create a new API called LambdaVar:

In the root resource, create a new GET method. In Integration type for the new method, choose Lambda Function, then select your Lambda Region, and type ${stageVariables.lbfunc} in the Lambda Function field. This tells API Gateway to read the value for this field from a stage variable at runtime:

The console detects the stage variable and displays the Add Permission to Lambda Function message:

Next, you manually give permissions to your Lambda functions, using the AWS CLI. This enables API Gateway to execute the functions. The CLI command must be issued with credentials that have permission to call the “add-permission” action of the Lambda APIs. The output from the AWS CLI will contain the policy statement that was set on the Lambda function resource policies.

aws lambda add-permission --function-name arn:aws:lambda:us-west-2:XXXXXXXXXXXXX:function:GetHelloWithName --source-arn arn:aws:execute-api:us-west-2:XXXXXXXXXXXXX:y91j2l4bnd/*/GET/ --principal apigateway.amazonaws.com --statement-id 95486b16-7d8a-4aca-9322-5f883ab702a6 --action lambda:InvokeFunction

# expected output
{
    "Statement": "{\"Condition\":{\"ArnLike\":{\"AWS:SourceArn\":\"arn:aws:execute-api:us-west-2: XXXXXXXXXXXX:y91j2l4bnd/*/GET/\"}},\"Action\":[\"lambda:InvokeFunction\"],\"Resource\":\"arn:aws:lambda:us-west-2:XXXXXXXXXXXX:function:GetHelloWithName\",\"Effect\":\"Allow\",\"Principal\":{\"Service\":\"apigateway.amazonaws.com\"},\"Sid\":\"95486b16-7d8a-4aca-9322-5f883ab702a6\"}"
}

aws lambda add-permission --function-name arn:aws:lambda:us-west-2:XXXXXXXXXXXXX:function:GetHelloWorld --source-arn arn:aws:execute-api:us-west-2:XXXXXXXXXXXXX:y91j2l4bnd/*/GET/ --principal apigateway.amazonaws.com --statement-id 95486b16-7d8a-4aca-9322-5f883ab702a6 --action lambda:InvokeFunction

# expected output
{
    "Statement": "{\"Condition\":{\"ArnLike\":{\"AWS:SourceArn\":\"arn:aws:execute-api:us-west-2: XXXXXXXXXXXXX:y91j2l4bnd/*/GET/\"}},\"Action\":[\"lambda:InvokeFunction\"],\"Resource\":\"arn:aws:lambda:us-west-2: XXXXXXXXXXXXX:function:GetHelloWorld\",\"Effect\":\"Allow\",\"Principal\":{\"Service\":\"apigateway.amazonaws.com\"},\"Sid\":\"95486b16-7d8a-4aca-9322-5f883ab702a6\"}"
}

Back in the console, you can now create your first stage. Choose Deploy API. In Stage name, type dev. In Stage description, type a description for your new stage, and then choose Deploy.

After the API deploys, on the Stage Editor page, choose the Stage Variables tab and add the stage variable from your API configuration, lbfunc. As you can see from the screenshot the value assigned to the new stage variable is the name of the Lambda function we want to invoke:

The second Lambda function, GetHelloWithName, can also receive a name parameter. You can configure the API to read the incoming parameter from the query string and pass it to the JSON body for the Lambda function by using mapping templates. To do this, go back to the GET method and choose Integration Request. Under Mapping Templates, add the following mapping template for the application/json content type:

{ "name": "$input.params('name')" }

To apply the change, deploy the API to a new stage called prod:

Next, set up the stage variable in the new deployment stage to point to the second Lambda function, GetHelloWithName:

Now you are ready to test!

The dev stage invoke URL directs you to the GetHelloWorld Lambda function:

The prod stage invoke URL with the appropriate query string directs you to the GetHelloWithName Lambda function and returns a value:

If you try to use the query string on the first stage, the query string will simply be ignored because the Lambda function is not configured to handle the parameter:

There it is: a nice way to optimize your Amazon API Gateway resources by using a single method with 2 different stages that use 2 different Lambda functions.

Alternatively, you can mix and match static names with stage variables in the integrations. For example, instead of having 2 different Lambda functions you can have a single Lambda function with multiple versions and aliases. Then, in the integration setup, you can simply use the stage variables to point to the correct alias. For instance, using an alias to one of the Lambda functions above, add the following: GetHelloWithName:${stageVariables.lambdaAlias} as Integration Type:

In your stage, add the lambdaAlias stage variable accordingly:

This variable will refer to the Lambda alias of your function:

As you can see, the new stage variables feature enables you to dynamically access different back ends, using fewer configuration steps and resources/methods in your API Gateway. The variables add even more flexibility to stages when deploying your API, which can enable different use cases in your environments.