Automating AWS Budget Notifications

AWS Lambda offers a powerful way to execute code in response to various events within the AWS ecosystem. One practical use case is to trigger a Lambda function when your AWS Budgets threshold alarm is breached. This enables you to automate actions like sending notifications, invoking other AWS services, or implementing custom logic to manage your costs proactively.

Terraform

# Terraform backend
# Replace to your own configuration:
# https://github.com/hands-on-cloud/cloudformation-templates/tree/master/terraform-remote-state-infrastructure
terraform {
  backend "s3" {
    bucket  = "hands-on-cloud-terraform-remote-state-s3"
    key     = "budget-lambda-trigger-demo.tfstate"
    region  = "us-west-2"
    encrypt = "true"
    dynamodb_table = "hands-on-cloud-terraform-remote-state-dynamodb"
  }
}

provider "aws" {
  region = "us-east-1"
}
locals {
  prefix = "budget-lambda-trigger-demo"
}
data "archive_file" "lambda_zip" {
  type        = "zip"
  source_file = "${path.module}/lambda_function.py"
  output_path = "${path.module}/lambda_function.zip"
}
resource "aws_lambda_function" "budget_alert_lambda" {
  filename         = data.archive_file.lambda_zip.output_path
  function_name    = "${local.prefix}-lambda"
  role             = aws_iam_role.lambda_exec_role.arn
  handler          = "lambda_function.lambda_handler"
  source_code_hash = data.archive_file.lambda_zip.output_base64sha256
  runtime          = "python3.11"
}
resource "aws_iam_role" "lambda_exec_role" {
  name = "${local.prefix}-lambda-exec-role"
  assume_role_policy = jsonencode({
    Version = "2012-10-17"
    Statement = [
      {
        Action = "sts:AssumeRole"
        Effect = "Allow"
        Principal = {
          Service = "lambda.amazonaws.com"
        }
      },
    ]
  })
}
resource "aws_sns_topic" "budget_alert_topic" {
  name = "${local.prefix}-budget-alert-topic"
}
resource "aws_sns_topic_subscription" "budget_alert_subscription" {
  topic_arn = aws_sns_topic.budget_alert_topic.arn
  protocol  = "lambda"
  endpoint  = aws_lambda_function.budget_alert_lambda.arn
}
resource "aws_lambda_permission" "allow_sns_to_call_lambda" {
  statement_id  = "AllowExecutionFromSNS"
  action        = "lambda:InvokeFunction"
  function_name = aws_lambda_function.budget_alert_lambda.function_name
  principal     = "sns.amazonaws.com"
  source_arn    = aws_sns_topic.budget_alert_topic.arn
}
resource "aws_budgets_budget" "account" {
  name              = "${local.prefix}-budget-aws-monthly"
  budget_type       = "COST"
  limit_amount      = "10"
  limit_unit        = "USD"
  time_period_end   = "2087-06-15_00:00"
  time_period_start = "2024-01-01_00:00"
  time_unit         = "MONTHLY"
  notification {
    comparison_operator        = "GREATER_THAN"
    threshold                  = 90
    threshold_type             = "PERCENTAGE"
    notification_type          = "FORECASTED"
    subscriber_sns_topic_arns = [ aws_sns_topic.budget_alert_topic.arn ]
  }
}

Lambda function code

import json
def lambda_handler(event, context):
    message = json.loads(event['Records'][0]['Sns']['Message'])
    print(f"Budget Alert: {message}")
    return {
        'statusCode': 200,
        'body': json.dumps('Budget alert processed.')
    }

CloudFormation Stack

AWSTemplateFormatVersion: '2010-09-09'
Description: An AWS CloudFormation stack to trigger a Lambda function on AWS budget alerts.
Resources:
  AccountBudget:
    Type: AWS::Budgets::Budget
    Properties:
      Budget:
        BudgetName: budget-lambda-trigger-demo-budget-aws-monthly
        BudgetLimit:
          Amount: 10
          Unit: USD
        TimeUnit: MONTHLY
        BudgetType: COST
        TimePeriod:
          Start: '1609459200' # Equivalent to 2024-01-01_00:00 in epoch seconds
          End: '3666288000'   # Equivalent to 2087-06-15_00:00 in epoch seconds
      NotificationsWithSubscribers:
        - Notification:
            NotificationType: FORECASTED
            ComparisonOperator: GREATER_THAN
            Threshold: 90
            ThresholdType: PERCENTAGE
          Subscribers:
            - SubscriptionType: SNS
              Address: !Ref BudgetAlertTopic
  BudgetAlertLambda:
    Type: AWS::Lambda::Function
    Properties:
      Handler: index.lambda_handler
      Role: !GetAtt LambdaExecutionRole.Arn
      Runtime: python3.8
      Code:
        ZipFile: |
          import json
          def lambda_handler(event, context):
              message = json.loads(event['Records'][0]['Sns']['Message'])
              print(f"Budget Alert: {message}")
              return {
                  'statusCode': 200,
                  'body': json.dumps('Budget alert processed.')
              }
      Timeout: 10
  LambdaExecutionRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Principal:
              Service: lambda.amazonaws.com
            Action: sts:AssumeRole
      Policies:
        - PolicyName: LambdaLoggingPolicy
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Effect: Allow
                Action:
                  - logs:CreateLogGroup
                  - logs:CreateLogStream
                  - logs:PutLogEvents
                Resource: arn:aws:logs:*:*:*
  BudgetAlertTopic:
    Type: AWS::SNS::Topic
  LambdaInvokePermission:
    Type: AWS::Lambda::Permission
    Properties:
      FunctionName: !Ref BudgetAlertLambda
      Action: lambda:InvokeFunction
      Principal: sns.amazonaws.com
      SourceArn: !Ref BudgetAlertTopic
  BudgetAlertSubscription:
    Type: AWS::SNS::Subscription
    Properties:
      Protocol: lambda
      TopicArn: !Ref BudgetAlertTopic
      Endpoint: !GetAtt BudgetAlertLambda.Arn

Manual Configuration

Here’s a straightforward guide to setting this up:

Create an AWS Budget:

Navigate to the AWS Budgets console.
Click on “Create budget” and choose the type of budget you need (e.g., Cost budget, Usage budget).
Define your budget details, including the threshold percentage at which you want to trigger the alarm.

Set Up an SNS Topic:

Go to the Amazon SNS console.
Create a new topic for budget notifications.
Note the ARN (Amazon Resource Name) of the SNS topic; you’ll need it for setting up notifications and linking to your Lambda function.

Configure Budget Notifications:

In the AWS Budgets console, find your budget and navigate to the “Alerts” section.
Create a new alert. Select “Actual” for alert type, and specify the threshold based on your preference.
Under “Notification,” choose “Select an SNS topic,” and specify the ARN of the SNS topic you created earlier.

Create a Lambda Function:

Head over to the AWS Lambda console.
Click “Create function.” Give it a name, choose a runtime (e.g., Python, Node.js), and define the necessary execution role.
Implement your custom logic within the Lambda function. This could be sending an email, logging a message, or any action you deem necessary upon breaching the budget threshold.

Subscribe Lambda to the SNS Topic:

In the Amazon SNS console, find your topic and create a new subscription.
Choose “AWS Lambda” as the protocol and select your Lambda function.
This links your Lambda function to the SNS topic, allowing it to be triggered whenever a message is published to the topic.

Test Your Setup:

It’s crucial to verify that everything works as intended.
You might simulate the budget alarm (if possible) or publish a test message to your SNS topic to ensure the Lambda function is triggered correctly.

Monitor and Adjust:

Keep an eye on the execution of your Lambda function and the notifications generated by AWS Budgets.
Adjust your budget, the threshold, or the Lambda function logic as necessary based on your observations and evolving requirements.

This approach allows you to automate responses to budget alarms, helping manage costs effectively without constant manual oversight. Remember, this is just the foundation. Depending on your needs, you can expand this setup to include more complex logic, integrate with other AWS services, or handle different types of budget alerts.