Add remote state package

.github/workflows/inngest_remote_state.yml

@@ -0,0 +1,120 @@
+name: "inngest_remote_state"
+
+on:
+  push:
+    branches:
+      - "main"
+    tags:
+      - "*"
+  pull_request:
+    paths:
+      - "!examples/**"
+      - "**/*.py"
+      - ".github/**"
+      - "Makefile"
+      - "constraints.txt"
+      - "pyproject.toml"
+      - "pytest.ini"
+      - "ruff.toml"
+
+jobs:
+  itest:
+    runs-on: "ubuntu-latest"
+    strategy:
+      matrix:
+        python-version: ["3.9", "3.12"]
+    timeout-minutes: 2
+    steps:
+      - uses: "actions/checkout@v2"
+      - name: "Set up Python ${{ matrix.python-version }}"
+        uses: "actions/setup-python@v2"
+        with:
+          python-version: "${{ matrix.python-version }}"
+      - name: "Install"
+        run: "make install"
+      - name: "Integration test"
+        run: "make itest"
+        working-directory: "./pkg/inngest_remote_state"
+
+  lint:
+    runs-on: "ubuntu-latest"
+    strategy:
+      matrix:
+        python-version: ["3.9", "3.12"]
+    steps:
+      - uses: "actions/checkout@v2"
+      - name: "Set up Python ${{ matrix.python-version }}"
+        uses: "actions/setup-python@v2"
+        with:
+          python-version: "${{ matrix.python-version }}"
+      - name: "Install"
+        run: "make install"
+      - name: "Lint"
+        run: "make lint"
+        working-directory: "./pkg/inngest_remote_state"
+
+  publish-pypi:
+    runs-on: "ubuntu-latest"
+    needs:
+      - "itest"
+      - "lint"
+      - "type-check"
+      - "utest"
+
+    # Only publish tagged versions.
+    # TODO: Add a check to ensure that the git tag matches the version.
+    if: "startsWith(github.ref, 'refs/tags/inngest_remote_state@')"
+    permissions:
+      id-token: write
+    strategy:
+      matrix:
+        python-version: ["3.9"]
+    steps:
+      - uses: "actions/checkout@v2"
+      - name: "Set up Python ${{ matrix.python-version }}"
+        uses: "actions/setup-python@v2"
+        with:
+          python-version: "${{ matrix.python-version }}"
+      - name: "Install"
+        run: "make install"
+      - name: "Build"
+        run: "make build"
+        working-directory: "./pkg/inngest_remote_state"
+      - name: "Upload package to PyPI"
+        uses: "pypa/gh-action-pypi-publish@release/v1"
+        with:
+          packages-dir: "./pkg/inngest_remote_state/dist"
+
+  type-check:
+    runs-on: "ubuntu-latest"
+    strategy:
+      matrix:
+        python-version: ["3.9", "3.12"]
+    steps:
+      - uses: "actions/checkout@v2"
+      - name: "Set up Python ${{ matrix.python-version }}"
+        uses: "actions/setup-python@v2"
+        with:
+          python-version: "${{ matrix.python-version }}"
+      - name: "Install"
+        run: "make install"
+      - name: "Type check"
+        run: "make type-check"
+        working-directory: "./pkg/inngest_remote_state"
+
+  utest:
+    runs-on: "ubuntu-latest"
+    strategy:
+      matrix:
+        python-version: ["3.9", "3.12"]
+    steps:
+      - uses: "actions/checkout@v2"
+      - name: "Set up Python ${{ matrix.python-version }}"
+        uses: "actions/setup-python@v2"
+        with:
+          python-version: "${{ matrix.python-version }}"
+      - name: "Install"
+        run: "make install"
+      - name: "Unit test"
+        run: "make utest"
+        working-directory: "./pkg/inngest_remote_state"

.vscode/settings.json

@@ -11,6 +11,7 @@
   "python.analysis.extraPaths": [
     "pkg/inngest",
     "pkg/inngest_encryption",
+    "pkg/inngest_remote_state",
     "pkg/test_core"
   ],
   "python.analysis.typeCheckingMode": "basic",

Makefile

@@ -11,26 +11,35 @@ format-check: check-venv
 	@ruff format --check .
 
 install: check-venv
-	@pip install -e '.[extra]' -e ./pkg/inngest -e ./pkg/inngest_encryption -e ./pkg/test_core -c constraints.txt
+	@pip install \
+		-e '.[extra]' -c constraints.txt \
+		-e ./pkg/inngest \
+		-e ./pkg/inngest_encryption \
+		-e ./pkg/inngest_remote_state \
+		-e ./pkg/test_core
 
 itest: check-venv
 	@cd pkg/inngest && make itest
 	@cd pkg/inngest_encryption && make itest
+	@cd pkg/inngest_remote_state && make itest
 
 pre-commit: format-check lint type-check utest
 
 lint: check-venv
 	@cd examples && make lint
 	@cd pkg/inngest && make lint
 	@cd pkg/inngest_encryption && make lint
+	@cd pkg/inngest_remote_state && make lint
 	@cd pkg/test_core && make lint
 
 type-check: check-venv
 	@cd examples && make type-check
 	@cd pkg/inngest && make type-check
 	@cd pkg/inngest_encryption && make type-check
+	@cd pkg/inngest_remote_state && make type-check
 	@cd pkg/test_core && make type-check
 
 utest: check-venv
 	@cd pkg/inngest && make utest
 	@cd pkg/inngest_encryption && make utest
+	@cd pkg/inngest_remote_state && make utest

mypy.ini

@@ -1,7 +1,7 @@
 [mypy]
 enable_error_code = possibly-undefined, redundant-expr, truthy-bool
 incremental = false
-mypy_path = ./pkg/inngest, ./pkg/inngest_encryption, ./pkg/test_core
+mypy_path = ./pkg/inngest, ./pkg/inngest_encryption, ./pkg/inngest_remote_state, ./pkg/test_core
 strict = true
 warn_unreachable = true
 

pkg/inngest_remote_state/Makefile

@@ -0,0 +1,32 @@
+export MYPYPATH=../inngest
+
+.PHONY: build
+build:
+	@if [ -d "dist" ]; then rm -rf dist; fi
+	@python -m build
+
+.PHONY: check-venv
+check-venv:
+	@if [ -z "$${CI}" ] && [ -z "$${VIRTUAL_ENV}" ]; then \
+		echo "virtual environment is not activated"; \
+		exit 1; \
+	fi
+
+.PHONY: itest
+itest: check-venv
+	@cd ../../tests/test_inngest_remote_state && pytest -n 4 -v .
+
+.PHONY: lint
+lint: check-venv
+	@ruff check .
+
+release:
+	@grep "version = \"$${VERSION}\"" pyproject.toml && git tag inngest_remote_state@$${VERSION} && git push origin inngest_remote_state@$${VERSION} || echo "pyproject.toml version does not match"
+
+.PHONY: type-check
+type-check: check-venv
+	@mypy --config-file=../../mypy.ini .
+
+.PHONY: utest
+utest: check-venv
+	@echo "inngest_remote_state unit tests not implemented"

pkg/inngest_remote_state/README.md

@@ -0,0 +1,38 @@
+# Inngest Python SDK: Remote state
+
+This package provides the tools for storing step output in a custom store (e.g. AWS S3). This can drastically reduce bandwidth to/from the Inngest server, since step output is stored within your infrastructure rather than Inngest's.
+
+## Usage
+
+Setting remote state middleware on the client will turn on remote state for steps in all functions:
+
+```py
+import inngest
+from inngest_remote_state import RemoteStateMiddleware
+from inngest_remote_state.s3 import S3Driver
+
+inngest.Inngest(
+    app_id="my-app",
+    middleware=[
+        RemoteStateMiddleware.factory(
+            S3Driver(
+                bucket="inngest-remote-state",
+                client=boto3.client("s3"),
+            )
+        )
+    ],
+)
+```
+
+The entire `step.run` output is stored in the remote store:
+
+```py
+def _my_step() -> dict[str, object]:
+    # Stored in the remote store.
+    return {"msg": "hello"}
+
+output = await step.run("my-step", _my_step)
+
+# Available within this function (it's automatically loaded by the middleware).
+print(output)
+```

pkg/inngest_remote_state/inngest_remote_state/__init__.py

@@ -0,0 +1,5 @@
+"""Public entrypoint for the Inngest SDK encryption package."""
+
+from ._internal import RemoteStateMiddleware, StateDriver
+
+__all__ = ["RemoteStateMiddleware", "StateDriver"]

pkg/inngest_remote_state/inngest_remote_state/_internal/__init__.py

@@ -0,0 +1,10 @@
+"""
+Remote state middleware for Inngest. This middleware allows you to store state
+where you want, rather than in Inngest's infrastructure. This is useful for:
+- Reducing bandwidth to/from the Inngest server.
+- Avoiding step output size limits.
+"""
+
+from .middleware import RemoteStateMiddleware, StateDriver
+
+__all__ = ["RemoteStateMiddleware", "StateDriver"]

pkg/inngest_remote_state/inngest_remote_state/_internal/middleware.py

@@ -0,0 +1,129 @@
+from __future__ import annotations
+
+import typing
+
+import inngest
+
+
+class StateDriver(typing.Protocol):
+    """
+    Protocol for the state driver.
+    """
+
+    def load_steps(self, steps: inngest.StepMemos) -> None:
+        """
+        Retrieve the value associated with the key.
+
+        Args:
+        ----
+            steps: Steps whose output may need to be loaded from the remote store.
+        """
+
+        ...
+
+    def save_step(
+        self,
+        run_id: str,
+        value: object,
+    ) -> dict[str, object]:
+        """
+        Store the value and return a key to retrieve it later.
+
+        Args:
+        ----
+            run_id: Run ID.
+            value: Output for an ended step.
+        """
+
+        ...
+
+
+class RemoteStateMiddleware(inngest.MiddlewareSync):
+    """
+    Middleware that reads/writes step output in a custom store (e.g. AWS S3).
+    This can drastically reduce bandwidth to/from the Inngest server, since step
+    output is stored within your infrastructure rather than Inngest's.
+    """
+
+    _run_id: typing.Optional[str] = None
+
+    def __init__(
+        self,
+        client: inngest.Inngest,
+        raw_request: object,
+        driver: StateDriver,
+    ) -> None:
+        """
+        Args:
+        ----
+            client: Inngest client.
+            raw_request: Framework/platform specific request object.
+            driver: State driver.
+        """
+
+        super().__init__(client, raw_request)
+
+        self._driver = driver
+
+    @classmethod
+    def factory(
+        cls,
+        driver: StateDriver,
+    ) -> typing.Callable[[inngest.Inngest, object], RemoteStateMiddleware]:
+        """
+        Create a remote state middleware that can be passed to an Inngest client
+        or function.
+
+        Args:
+        ----
+            driver: State driver.
+        """
+
+        def _factory(
+            client: inngest.Inngest,
+            raw_request: object,
+        ) -> RemoteStateMiddleware:
+            return cls(
+                client,
+                raw_request,
+                driver,
+            )
+
+        return _factory
+
+    def transform_input(
+        self,
+        ctx: inngest.Context,
+        function: inngest.Function,
+        steps: inngest.StepMemos,
+    ) -> None:
+        """
+        Inject remote state.
+        """
+
+        self._driver.load_steps(steps)
+        self._run_id = ctx.run_id
+
+    def transform_output(self, result: inngest.TransformOutputResult) -> None:
+        """
+        Store step output externally and replace with a marker and key.
+        """
+
+        if result.step is None:
+            return None
+
+        # Only support step.run, but that may change in the future.
+        if result.step.op.value != "StepRun":
+            return None
+
+        if result.has_output() is False:
+            return None
+
+        if self._run_id is None:
+            # Unreachable
+            raise Exception("missing run ID")
+
+        result.output = self._driver.save_step(
+            self._run_id,
+            result.output,
+        )