Handling Custom Python Errors in AWS Step Functions

AWS Step Functions are an amazing way to orchestrate Lambda Functions for serverless applications. Handling and catching errors is really interesting and one of the main advantages of using Step Functions.

AWS Step Functions State Machine

I created a state machine consisting of 3 step functions. The start function is a Lambda Function, called GetBook, that essentially throws a custom Python exception, called BookNotFoundException, when the isbn provided as input is book not found.

Here is the state machine in code and a visual representation that shows the path taken when the custom exception occurs.

{
  "Comment": "Handling Custom Python Errors in AWS Step Functions.",
  "StartAt": "GetBook",
  "States": {
    "GetBook": {
      "Type": "Task",
      "Resource": "{arn}",
      "Next": "SendEmail",
      "Catch": [
          {
             "ErrorEquals": ["BookNotFoundException"],
             "ResultPath": "$.error-info",
             "Next": "CustomerSupport"
          }
       ]
    },
    "CustomerSupport": {
      "Type": "Pass",
      "Result": "Customer support contacted!",
      "End": true
    },
    "SendEmail": {
      "Type": "Pass",
      "Result": "Email sent!",
      "End": true
    }
  }
}

AWS Step Function for Serverless Applications

The part of the state machine that interests me the most is the Catch section of the GetBook task.

"Catch": [
    {
        "ErrorEquals": ["BookNotFoundException"],
        "ResultPath": "$.error-info",
        "Next": "CustomerSupport"
    }
  ]

If the step function throws a custom error, called BookNotFoundException, it will proceed to the CustomerSupport task.

Notice the ResultPath attribute. It passes the error information in a custom attribute, called error-info, to CustomerSupport while also passing along all the input that was provided as input to the GetBook task.

This is the input passed to the CustomerSupport Step Function.

{
  "isbn": "book not found",
  "error-info": {
    "Error": "BookNotFoundException",
    "Cause": "{\"errorType\": \"BookNotFoundException\", \"stackTrace\": [[\"/var/task/lambda_function.py\", 14, \"lambda_handler\", \"raise BookNotFoundException()\"]]}"
  }
}

Lambda Function

The GetBook task in the state machine is associated with a GetBook Lambda Function written in Python 3. As mentioned, if you provide an isbn of book not found as input to this function, it will raise a custom BookNotFoundException.

Here is the input to simulate a custom Python Exception.

{
  "isbn": "book not found"
}

Here is the actual Lambda Function.

import logging

logger = logging.getLogger()
logger.setLevel(logging.INFO)

class BookNotFoundException(Exception): pass

def lambda_handler(event, context):
    isbn = event["isbn"]
    
    # simulate exception
    if isbn == "book not found":
        logger.info("Intentionally raising BookNotFoundException.")
        raise BookNotFoundException()

    return {
        "book": {
            "isbn": isbn,
            "title": "my favorite book",
            "author": "my favorite author"
        }
    }

Lots of interesting stuff happening here with Step Functions that I want to keep in mind for future development.

Contents