example.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# pystepfunction\n",
    "\n",
    "- Create AWS Stepfunction asl.json files (state machine definition files) using python.\n",
    "- Pre-made tasks that are easy to use. \n",
    "- Test dataflow through the stepfunction using the same python code.\n",
    "- Visualise the stepfunction using pyvis for easier dubugging without the use of the AWS console.\n",
    "\n",
    "\n",
    "## Installation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "shellscript"
    }
   },
   "outputs": [],
   "source": [
    "pip install pystepfunction"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tasks\n",
    "pystepfunction has Task implementations for various AWS resources."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'Comment': 'This is an easy branch', 'StartAt': 'LambdaTaskName', 'States': {'LambdaTaskName': {'Type': 'Task', 'Resource': 'arn:aws:states:::lambda:invoke', 'End': False, 'Next': 'GlueTaskName', 'Parameters': {'FunctionName': 'my-lambda-arn'}}, 'GlueTaskName': {'Type': 'Task', 'Resource': 'arn:aws:states:::glue:startJobRun.sync', 'End': False, 'Next': 'SucceedTaskName', 'Parameters': {'JobName': 'my-glue-job-name'}}, 'SucceedTaskName': {'Type': 'Succeed', 'Resource': '', 'End': False}}}\n"
     ]
    }
   ],
   "source": [
    "from pystepfunction.tasks import LambdaTask, GlueTask, SucceedTask\n",
    "from pystepfunction.branch import Branch\n",
    "\n",
    "\n",
    "# create a simple chain of tasks\n",
    "lambda_task = (\n",
    "    LambdaTask(name=\"LambdaTaskName\", function_arn=\"my-lambda-arn\")\n",
    "    .and_then(GlueTask(name=\"GlueTaskName\", job_name=\"my-glue-job-name\"))\n",
    "    .and_then(SucceedTask(name=\"SucceedTaskName\"))\n",
    ")\n",
    "# or\n",
    "lambda_task = (\n",
    "    LambdaTask(name=\"LambdaTaskName\", function_arn=\"my-lambda-arn\") \n",
    "    >> GlueTask(name=\"GlueTaskName\", job_name=\"my-glue-job-name\")\n",
    "    >> SucceedTask(name=\"SucceedTaskName\")\n",
    ")\n",
    "\n",
    "# create a branch\n",
    "easy_branch = Branch(comment=\"This is an easy branch\", start_task=lambda_task)\n",
    "# view the asl as a dict\n",
    "asl = easy_branch.to_asl()\n",
    "print(asl)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Add error handling to a task\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'Comment': '',\n",
       " 'StartAt': 'LambdaTaskName',\n",
       " 'States': {'LambdaTaskName': {'Type': 'Task',\n",
       "   'Resource': 'arn:aws:states:::lambda:invoke',\n",
       "   'End': False,\n",
       "   'Next': 'Succeed',\n",
       "   'Parameters': {'FunctionName': 'my-lambda-arn'},\n",
       "   'Retry': [{'ErrorEquals': ['States.ALL'],\n",
       "     'IntervalSeconds': 1,\n",
       "     'MaxAttempts': 2,\n",
       "     'BackoffRate': 1.0}],\n",
       "   'Catch': [{'ErrorEquals': 'States.ALL', 'Next': 'Fail'}]},\n",
       "  'Succeed': {'Type': 'Succeed', 'Resource': '', 'End': False},\n",
       "  'Fail': {'Type': 'Fail', 'Cause': '', 'Error': ''}}}"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from pystepfunction.errors import ERROR_STATE_ALL\n",
    "from pystepfunction.tasks import LambdaTask, SucceedTask, FailTask, Retry\n",
    "from pystepfunction.branch import Branch\n",
    "\n",
    "succeed = SucceedTask(name=\"Succeed\")\n",
    "fail = FailTask(name=\"Fail\")\n",
    "task = (\n",
    "    LambdaTask(name=\"LambdaTaskName\", function_arn=\"my-lambda-arn\")\n",
    "    .with_retries([Retry(error_equals=[ERROR_STATE_ALL], interval_seconds=1, max_attempts=2)])\n",
    "    .with_catcher(catcher={ERROR_STATE_ALL: fail})\n",
    "    >> succeed\n",
    ")\n",
    "branch = Branch(start_task=task)\n",
    "branch.to_asl()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Manipulate input and output state of a Task"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'LambdaTaskName': {'Type': 'Task',\n",
       "  'Resource': 'arn:aws:states:::lambda:invoke',\n",
       "  'End': False,\n",
       "  'InputPath': '$.Inputs',\n",
       "  'Parameters': {'FunctionName': 'my-lambda-arn',\n",
       "   'Input1.$': '$.Input1',\n",
       "   'Input2.$': '$.Input2'},\n",
       "  'ResultPath': '$.TaskResult',\n",
       "  'ResultSelector': {'Output1.$': '$.Output1', 'Output2.$': '$.Output2'}}}"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from pystepfunction.state import StateMachine\n",
    "from pystepfunction.tasks import LambdaTask, TaskInputState, TaskOutputState\n",
    "# Inputs use paramaters and result paths\n",
    "# Both use jsonpath syntax\n",
    "\n",
    "parameters = {\"Input1.$\": \"$.Input1\", \"Input2.$\": \"$.Input2\"}\n",
    "input_state = TaskInputState(parameters=parameters, input_path=\"$.Inputs\")\n",
    "\n",
    "# Inputs use paramaters and result paths\n",
    "# Both use jsonpath syntax\n",
    "output_state = TaskOutputState(\n",
    "    result_path=\"$.TaskResult\",\n",
    "    result_selector={\"Output1.$\": \"$.Output1\", \"Output2.$\": \"$.Output2\"},\n",
    ")\n",
    "\n",
    "lambda_task = LambdaTask(name=\"LambdaTaskName\", function_arn=\"my-lambda-arn\").with_input(input_state).with_output(output_state)\n",
    "lambda_task.to_asl()\n",
    "\n",
    "# Test the dataflow of a task\n",
    "# define what is expected back from the lambda resource for testing purposes\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'Inputs': {'Input1': 'val1', 'Input2': 'val2'}, '__msg__': 'Input state'}]"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "resource_result = {\"Output1\": \"val1\", \"Output2\": \"val2\", \"Output3\": \"val3\"}\n",
    "lambda_task.with_resource_result(resource_result)\n",
    "\n",
    "# create a new state and set its value\n",
    "state = StateMachine()\n",
    "state.update_state(state={\"Inputs\": {\"Input1\": \"val1\", \"Input2\": \"val2\"}}, msg=\"Input state\")\n",
    "\n",
    "# view the value log\n",
    "state.state_log\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'state' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[1], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m state\u001b[39m.\u001b[39mapply_task(lambda_task)\n\u001b[1;32m      2\u001b[0m state\u001b[39m.\u001b[39mstate_log\n",
      "\u001b[0;31mNameError\u001b[0m: name 'state' is not defined"
     ]
    }
   ],
   "source": [
    "state.apply_task(lambda_task)\n",
    "state.state_log\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# visualise with pyvis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "easy_branch.html\n"
     ]
    }
   ],
   "source": [
    "from pystepfunction.viz import BranchViz\n",
    "\n",
    "BranchViz(easy_branch).show(\"easy_branch.html\") "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Visualise with networkx and matplotlib\n",
    "TODO"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Add error handling"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Create choice and parallel flows"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Test data flows"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.11"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}