Create Stateful Serverless Workflows with AWS Step Functions and JSONata

As an avid user of AWS Step Functions, I've been pleased by several excellent releases over the past few years, including Distributed Map, Express Workflows, Intrinsic functions, TestState, redrive, service integrations, and so many others. Those are all fantastic releases, but in my humble opinion, none of them are as big of a deal as the introduction of JSONata expressions. AWS announced this game-changing feature a week before re:Invent 2024.

JSONata

JSONata is a language for querying and transforming JSON data developed by Andrew Coleman of IBM. I had not heard of or worked with JSONata before its inclusion in Step Functions, but it seems like a perfect fit.

I haven't used it yet outside of Step Functions, but JSONata seems quite remarkable in its own right, and I can imagine using it in other contexts. Check out the docs at https://docs.jsonata.org/overview.

JSONata in Step Functions

JSONata simplifies how data is passed from step to step over the existing JSONPath approach. This is a very good thing because the developer experience of Step Functions hasn't always been the greatest. Even frequent users of this service get lost in the flow of InputPath, Parameters, Task Result, ResultSelector, ResultPath, and OutputPath. Wait, which one do I use to add a value to the state input? AWS released the Data flow simulator to help navigate this flow, but in some ways, all that did was drive home the point that this can be complex and hard to use.

Passing Data

When we use JSONata with Step Functions, we do away with the above flow and instead work with Parameters, Output, and one new directive, Assign. Because JSONata lets us create complex expressions, we can use Output to achieve what used to require ResultSelector, ResultPath, and OutputPath, a significant simplification.

The Assign Directive

As good as that is, Assign pushes JSONata support to the top of an impressive feature set. Assign lets us create scoped variables that can be read in later steps! If you've ever had to pass a value from step 2 to step 7 to be read, you're probably out of your seat and cheering. Note that these variables will be scoped when you use Assign in steps within a Map or Parallel state. The Map or Parallel state can use Assign to make values available later in the state machine execution.

Comparing JSONPath with JSONata

Let's look at a simple example. I refactored the state machine I used in this 2019 blogpost. It demonstrates a simple workflow that moves through different states, each implemented by a Lambda Function.

Five years ago, we had no choice but to pass data along from state to state. We could refer back to the original input using $$.Execution.Input, but anything derived by a step must be passed from step to step (or stored in a database). So while my simple state machine example works fine, it introduces an opportunity for weird bugs. I'm passing the CaseId from step to step. It should be understood I'm working on the same case all the way through. What if I pass "001" to the "Assign Case" step and then return a case number of "ABC123"? My step definitions may create a "garbage-in, garbage-out" machine.

{
  "Next": "Assign Case",
  "Type": "Task",
  "OutputPath": "$.Payload",
  "Resource": "arn:aws:states:::lambda:invoke",
  "Parameters": {
    "FunctionName": "arn:aws:lambda:us-east-1:123456890:function:assign-case",
    "Payload.$": "$"
  }
}

{
  "Next": "Work On Case",
  "Type": "Task",
  "OutputPath": "$.Payload",
  "Resource": "arn:aws:states:::lambda:invoke",
  "Parameters": {
    "FunctionName": "arn:aws:lambda:us-east-1:123456890:function:assign-case",
    "Payload.$": "$"
  }
}

It's easy enough to convert these to JSONata steps and start seeing the value of Assign. We set QueryLanguage to JSONata and now we can use the Assign directive.

{
  "QueryLanguage": "JSONata",
  "Next": "Assign Case",
  "Type": "Task",
  "Arguments": {
    "FunctionName": "arn:aws:lambda:us-east-1:1234567890:function:open-case",
    "Payload": "{% $states.input %}"
  },
  "Assign": {
    "caseId": "{% $states.input.inputCaseID %}",
    "message": "{% $states.result.Payload.message %}",
    "status": "{% $states.result.Payload.status %}"
  },
  "Resource": "arn:aws:states:::lambda:invoke"
}

{
  "QueryLanguage": "JSONata",
  "Next": "Work On Case",
  "Type": "Task",
  "Arguments": {
    "FunctionName": "arn:aws:lambda:us-east-1:1234567890:function:assign-case",
    "Payload": {
      "caseId": "{% $caseId %}",
      "message": "{% $message %}",
      "status": "{% $status %}"
    }
  },
  "Assign": {
    "message": "{% $states.result.Payload.message %}",
    "status": "{% $states.result.Payload.status %}"
  },
  "Resource": "arn:aws:states:::lambda:invoke"
}

The Assign Case step now sets the message and status properties, but the caseId remains constant throughout the state machine execution.

JSONata in the AWS Console

Still not convinced? The console experience should close the deal for you. We get a new tab for every JSONata step: "Variables" that shows us any variables the step used and any it assigned.

But that's not all. We also get to view the history of each variable and the assignments made during the state machine execution.

This is impressively useful for debugging complex workflows. I can't think of another tool that records the state of each variable and lets me view it after execution. Just think of how many console.log statements this will save you!

Migrating to JSONata

Making the switch is relatively easy, although more complex workflows require more effort. I have had success doing this by hand. I spent a few hours converting a very complex and hard-to-debug workflow to use JSONata and immediately reaped the benefits of greater observability and maintainability. Although that isn't code I can share, readers might be interested in the diff that upgraded my sample project.

But we don't need to do it by hand! The console now includes a tool we can use to help with the migration.

You get completion right in the tool.

You can also build out powerful Assign expressions here.

It's possible to mix JSONPath and JSONata expressions in the same state machine execution. I wouldn't intentionally do this, but it's useful to be able to refactor one state at a time and test to ensure the workflow is still functional.

What Else Does JSONata Offer?

I feel like we've barely scratched the surface. Predicate queries look very powerful, and now we can do this in an Assign step instead of writing an ES6 reduce callback. JSONata also features an impressive list of built-in functions.

One Last Thing

Think this is cool, but not fully sold on Step Functions? Let me convince you.