AWS Step FunctionsのJSONを生成するPythonプログラムを作ってみた

はじめに

最近、現在のプロジェクトのワークフローツールとしてAWS Step Functionsを採用することになりました。
Step Functionsが何かについては上のリンクを見ていただければわかるかと思います。まだ使い始めたところですが、Step Functionsを使うことでLambda等をつなぐワークフローのロジックを比較的低い学習コストで実現できそうで、これから本格的に使っていこうと思っています。

問題点

現在のプロジェクトのAWSインフラはすべてCloudFormationを使って構築しているので、Step FunctionsもCloudFormationで構築、管理したいと考えているのですが、ここに少し問題があります。
CloudFormationはインフラをJSON/YAMLで管理するものですが、現在のプロジェクトでは直接JSON/YAMLを書くのではなく、Pythonの troposphere というライブラリを使っています。
理由はそこそこ複雑なインフラなのでJSON/YAMLを直接編集するのは非現実的だからです。
Step FunctionsもStateMachineをJSONで記述するので、JSONを直接編集することなく、 troposphereで管理したいと思っていました。しかし、一応StateMachineは定義されているものの、肝心の中身の定義に関してはJSONのストリングをそのまま突っ込むという仕様なので結局のところ自分でJSONを編集するしかなく、あまりありがたみがありません。
ということで、StateMachineの中身のJSONの定義はAmazon States Languageという仕様で定義されているのでそれに則ってPythonのクラスを定義してそれをdictで吐き出してjson.dumpsすればいいのではと思い実際に作ってみました。

Step FunctionsのJSONを生成するPythonプログラムの実装

こちらのGistに実装を公開しています。
https://gist.github.com/Y4suyuki/a2e9951f8acb9fd07bff5b521c7aa8c0

一旦、Amazon States Languageの仕様に沿ってPythonクラスを定義してそれをdictで吐き出すメソッドがあるだけなので、変数のValidationなどはないですが、とりあえず当初意図したものができました。
AWS Step FunctionsのドキュメントのStateMachineの例で示すと元のJSON

{
  "Comment": "A Hello World example of the Amazon States Language using a Pass state",
  "StartAt": "HelloWorld",
  "States": {
    "HelloWorld": {
      "Type": "Pass",
      "Result": "Hello World!",
      "End": true
    }
  }
}

に対して、Pythonだと

>>> hwp = Pass("HelloWorld", Result="Hello World!", End=True)
>>> hwsm = StateMachine(StartAt=hwp.Name,
                        Comment="A Hello World example of the Amazon States Language using a Pass state",
                        States=[hwp])

のように定義してあげると

>>> print(hwsm.to_json(indent=2))

{
  "StartAt": "HelloWorld",
  "Comment": "A Hello World example of the Amazon States Language using a Pass state",
  "States": {
    "HelloWorld": {
      "End": true,
      "Result": "Hello World!",
      "Type": "Pass"
    }
  }
}

という風にドキュメントの例と同様のJSONができます。これくらい小さいJSONだとあまり違いを感じないですが、コンポーネントが多くなってくるとPythonで定義してJSONを吐き出したほうが、圧倒的に楽だと思います。

おわりに

今回PythonでStep FunctionsのStateMachineのJSONを作るプログラムを実装してみましたが、Amazon States Languageがコンパクトな仕様だったので簡単にできました。
ただ、Choiceの中のChoicesのコンポーネントはAmazon States Languageに定義されていないので、Pythonでもdictのままだったり、クラスをdictにするときに呼び出している関数で再帰を使っているので、コンポーネントが多くなるとStackOverflowになりそうとかありますが、とりあえず使えそうなので、使いながらより良くしていきたいです。

追記

実際に少しStep Functionsで使ってみると、まだまだ不便なところが多くて、例えばStatesのNextを直に定義するのではなく、apache-airflowでDAGを定義するときにつかうset_upstreamやset_downstreamのようなメソッドがあると良いかもと思いました。