Journal codec (#229)

Author: slinkydeveloper

python/restate/__init__.py

@@ -38,6 +38,7 @@
     SendHandle,
     RunOptions,
 )
+from .entry_codec import JournalValueCodec, JournalValueCodecProvider
 from .exceptions import TerminalError, SdkInternalBaseException, is_internal_exception
 from .asyncio import as_completed, gather, wait_completed, select
 
@@ -58,6 +59,7 @@ def create_test_harness(
         restate_image: str = "docker.io/restatedev/restate:latest",
         always_replay: bool = False,
         disable_retries: bool = False,
+        journal_value_codec: typing.Optional[JournalValueCodec] = None,
     ) -> typing.AsyncGenerator[HarnessEnvironment, None]:
         """a dummy harness constructor to raise ImportError. Install restate-sdk[harness] to use this feature"""
         raise ImportError("Install restate-sdk[harness] to use this feature")
@@ -81,7 +83,9 @@ def test_harness(
 
     @asynccontextmanager
     async def create_client(
-        ingress: str, headers: typing.Optional[dict] = None
+        ingress: str,
+        headers: typing.Optional[dict] = None,
+        journal_value_codec: typing.Optional[JournalValueCodec] = None,
     ) -> typing.AsyncGenerator[RestateClient, None]:
         """a dummy client constructor to raise ImportError. Install restate-sdk[client] to use this feature"""
         raise ImportError("Install restate-sdk[client] to use this feature")
@@ -105,6 +109,8 @@ async def create_client(
     "ScopedContext",
     "RunOptions",
     "TerminalError",
+    "JournalValueCodec",
+    "JournalValueCodecProvider",
     "app",
     "create_test_harness",
     "test_harness",

python/restate/client.py

@@ -20,6 +20,7 @@
 from .client_types import RestateClient, RestateClientSendHandle, RestateScopedClient, HttpError
 
 from .context import HandlerType
+from .entry_codec import JournalValueCodec
 from .serde import BytesSerde, JsonSerde, Serde
 from .handler import handler_from_callable
 
@@ -32,9 +33,15 @@ class Client(RestateClient):
     A basic client for connecting to the Restate service.
     """
 
-    def __init__(self, client: httpx.AsyncClient, headers: typing.Optional[dict] = None):
+    def __init__(
+        self,
+        client: httpx.AsyncClient,
+        headers: typing.Optional[dict] = None,
+        journal_codec: typing.Optional[JournalValueCodec] = None,
+    ):
         self.headers = headers or {}
         self.client = client
+        self.journal_codec = journal_codec
 
     def scope(self, scope: str) -> RestateScopedClient:
         return ScopedClient(self, scope)
@@ -103,6 +110,8 @@ async def do_raw_call(
     ) -> O:
         """Make an RPC call to the given handler"""
         parameter = input_serde.serialize(input_param)
+        if self.journal_codec is not None:
+            parameter = self.journal_codec.encode(parameter)
         if headers is not None:
             headers_kvs = list(headers.items())
         else:
@@ -124,6 +133,9 @@ async def do_raw_call(
             scope=scope,
             limit_key=limit_key,
         )
+        # A send returns the invocation-id envelope, not a codec'd payload, so skip decode for it.
+        if not send and self.journal_codec is not None:
+            res = await self.journal_codec.decode(res)
         return output_serde.deserialize(res)  # type: ignore
 
     async def post(
@@ -471,10 +483,19 @@ async def workflow_send(
 
 @asynccontextmanager
 async def create_client(
-    ingress: str, headers: typing.Optional[dict] = None
+    ingress: str,
+    headers: typing.Optional[dict] = None,
+    journal_value_codec: typing.Optional[JournalValueCodec] = None,
 ) -> typing.AsyncGenerator[RestateClient, None]:
     """
     Create a new Restate client.
+
+    Args:
+        ingress: The base URL of the Restate ingress.
+        headers: Optional default headers to send with every request.
+        journal_value_codec: Optional journal value codec. When set, request bodies are encoded and
+            call responses decoded through it. It must match the codec configured on the endpoint.
+            NOTE: This is experimental and may change in future releases.
     """
     async with httpx.AsyncClient(base_url=ingress, headers=headers, http2=True) as http_client:
-        yield Client(http_client, headers)
+        yield Client(http_client, headers, journal_codec=journal_value_codec)

python/restate/endpoint.py

@@ -17,6 +17,7 @@
 from restate.service import Service
 from restate.object import VirtualObject
 from restate.workflow import Workflow
+from restate.entry_codec import JournalValueCodec, JournalValueCodecProvider
 
 
 # disable too few methods in a class
@@ -31,6 +32,7 @@ class Endpoint:
     services: typing.Dict[str, typing.Union[Service, VirtualObject, Workflow]]
     protocol: typing.Optional[typing.Literal["bidi", "request_response"]]
     identity_keys: typing.List[str]
+    journal_value_codec: typing.Optional[typing.Union[JournalValueCodec, JournalValueCodecProvider]]
 
     def __init__(self):
         """
@@ -44,6 +46,10 @@ def __init__(self):
 
         self.identity_keys = []
 
+        # An optional journal value codec (or async provider building one). When set, the SDK
+        # transforms serialized journal values through it. None means no codec is configured.
+        self.journal_value_codec = None
+
     def bind(self, *services: typing.Union[Service, VirtualObject, Workflow]):
         """
         Bind a service to the endpoint
@@ -79,6 +85,23 @@ def identity_key(self, identity_key: str):
         """Add an identity key to this endpoint."""
         self.identity_keys.append(identity_key)
 
+    def set_journal_value_codec(self, codec_or_provider: typing.Union[JournalValueCodec, JournalValueCodecProvider]):
+        """
+        Set the journal value codec for this endpoint.
+
+        NOTE: This is experimental and may change in future releases.
+
+        Args:
+            codec_or_provider: Either a :class:`JournalValueCodec` instance, or an async provider
+                (a zero-arg callable returning an awaitable of a codec) that is invoked once and
+                cached for the lifetime of the endpoint.
+
+        Returns:
+            The updated Endpoint instance.
+        """
+        self.journal_value_codec = codec_or_provider
+        return self
+
     def app(self):
         """
         Returns the ASGI application for this endpoint.
@@ -101,6 +124,7 @@ def app(
     services: typing.Iterable[typing.Union[Service, VirtualObject, Workflow]],
     protocol: typing.Optional[typing.Literal["bidi", "request_response"]] = None,
     identity_keys: typing.Optional[typing.List[str]] = None,
+    journal_value_codec: typing.Optional[typing.Union[JournalValueCodec, JournalValueCodecProvider]] = None,
 ):
     """A restate ASGI application that hosts the given services."""
     endpoint = Endpoint()
@@ -113,4 +137,6 @@ def app(
     if identity_keys:
         for key in identity_keys:
             endpoint.identity_key(key)
+    if journal_value_codec is not None:
+        endpoint.set_journal_value_codec(journal_value_codec)
     return endpoint.app()

python/restate/entry_codec.py

@@ -0,0 +1,71 @@
+#
+#  Copyright (c) 2023-2026 - Restate Software, Inc., Restate GmbH
+#
+#  This file is part of the Restate SDK for Python,
+#  which is released under the MIT license.
+#
+#  You can find a copy of the license in file LICENSE in the root
+#  directory of this repository or package, or at
+#  https://github.com/restatedev/sdk-typescript/blob/main/LICENSE
+#
+"""This module contains the journal value codec interface."""
+
+import abc
+import typing
+
+# disable too few public methods
+# pylint: disable=R0903
+
+
+class JournalValueCodec(abc.ABC):
+    """
+    Journal values codec.
+
+    This allows to transform journal values after being serialized, before writing them to the
+    wire, and vice-versa. It sits *between* the ``Serde`` layer and the journal/wire.
+    The canonical use case is to encrypt or compress everything the SDK persists to the journal.
+
+    Values that are passed through the codec:
+
+    * Handlers input and success output
+    * ``ctx.run`` success results
+    * Awakeables/Promise success results
+    * State values
+    * Call/send request parameters and call responses
+
+    Failures are never passed through the codec.
+
+    NOTE: This is preview and may change in future releases.
+    """
+
+    @abc.abstractmethod
+    def encode(self, buf: bytes) -> bytes:
+        """
+        Encodes the given buffer. This will be applied *after* serialization.
+
+        Args:
+            buf: The buffer to encode. Empty byte buffers should be appropriately handled as well.
+
+        Returns:
+            The encoded buffer.
+        """
+
+    @abc.abstractmethod
+    async def decode(self, buf: bytes) -> bytes:
+        """
+        Decodes the given buffer. This will be applied *before* deserialization.
+
+        Args:
+            buf: The buffer to decode.
+
+        Returns:
+            The decoded buffer.
+        """
+
+
+JournalValueCodecProvider = typing.Callable[[], typing.Awaitable[JournalValueCodec]]
+"""A provider that asynchronously builds a :class:`JournalValueCodec`.
+
+It is invoked once, and the resulting codec is reused for the lifetime of the endpoint.
+This is useful to perform async setup at startup, e.g. loading an encryption key.
+"""

python/restate/handler.py

@@ -34,6 +34,7 @@
 from restate.retry_policy import InvocationRetryPolicy
 
 from restate.context import HandlerType
+from restate.entry_codec import JournalValueCodec
 from restate.exceptions import TerminalError
 from restate.serde import DefaultSerde, PydanticJsonSerde, MsgspecJsonSerde, Serde, is_pydantic, Msgspec
 from restate.types import extract_core_type
@@ -350,11 +351,18 @@ def handler_from_callable(wrapper: HandlerType[I, O]) -> Handler[I, O]:
         raise ValueError("Handler not found")  # pylint: disable=raise-missing-from
 
 
-async def invoke_handler(handler: Handler[I, O], ctx: Any, in_buffer: bytes) -> bytes:
+async def invoke_handler(
+    handler: Handler[I, O], ctx: Any, in_buffer: bytes, journal_codec: Optional[JournalValueCodec] = None
+) -> bytes:
     """
     Invoke the handler with the given context and input.
     """
     if handler.arity == 2:
+        if journal_codec is not None:
+            try:
+                in_buffer = await journal_codec.decode(in_buffer)
+            except Exception as e:
+                raise TerminalError(message="Failed to decode input using journal value codec", status_code=400) from e
         try:
             in_arg = handler.handler_io.input_serde.deserialize(in_buffer)
         except Exception as e:
@@ -363,4 +371,6 @@ async def invoke_handler(handler: Handler[I, O], ctx: Any, in_buffer: bytes) ->
     else:
         out_arg = await handler.fn(ctx)  # type: ignore [call-arg]
     out_buffer = handler.handler_io.output_serde.serialize(out_arg)
+    if journal_codec is not None:
+        out_buffer = journal_codec.encode(bytes(out_buffer))
     return bytes(out_buffer)

python/restate/harness.py

@@ -21,6 +21,7 @@
 from hypercorn.config import Config
 from hypercorn.asyncio import serve
 from restate.client import create_client
+from restate.entry_codec import JournalValueCodec
 from restate.server_types import RestateAppT
 from restate.types import HarnessEnvironment
 from testcontainers.core.container import DockerContainer  # type: ignore
@@ -338,6 +339,7 @@ async def create_test_harness(
     restate_image: str = "docker.io/restatedev/restate:latest",
     always_replay: bool = False,
     disable_retries: bool = False,
+    journal_value_codec: typing.Optional[JournalValueCodec] = None,
 ) -> typing.AsyncGenerator[HarnessEnvironment, None]:
     """
     Creates a test harness for running Restate services together with restate-server.
@@ -360,6 +362,8 @@ async def create_test_harness(
                           on a suspension point. This is useful to hunt non-deterministic bugs
                           that might prevent your code to replay correctly (default is False).
     :param disable_retries: When True, retries are disabled (default is False).
+    :param journal_value_codec: Optional journal value codec to configure on the ingress client, so
+                          it matches a codec configured on the endpoint under test (default is None).
     """
     with (
         create_restate_container(
@@ -377,7 +381,7 @@ async def create_test_harness(
             msg = f"unable to register the services at {bind_address} - {res.status_code} {res.text}"
             raise AssertionError(msg)
 
-        async with create_client(runtime.ingress_url()) as client:
+        async with create_client(runtime.ingress_url(), journal_value_codec=journal_value_codec) as client:
             yield HarnessEnvironment(
                 ingress_url=runtime.ingress_url(), admin_api_url=runtime.admin_url(), client=client
             )

python/restate/server.py

@@ -11,12 +11,14 @@
 """This module contains the ASGI server for the restate framework."""
 
 import asyncio
+import functools
 import logging
 import signal
-from typing import Dict, Set, TypedDict, Literal
+from typing import Dict, Optional, Set, TypedDict, Literal
 
 from restate.discovery import compute_discovery_json
 from restate.endpoint import Endpoint
+from restate.entry_codec import JournalValueCodec
 from restate.server_context import ServerInvocationContext, DisconnectedException
 from restate.server_types import Receive, ReceiveChannel, RestateAppT, Scope, Send, binary_to_header, header_to_binary  # pylint: disable=line-too-long
 from restate.vm import VMWrapper
@@ -125,7 +127,12 @@ async def send_health_check(send: Send):
 
 
 async def process_invocation_to_completion(
-    vm: VMWrapper, handler, attempt_headers: Dict[str, str], receive: ReceiveChannel, send: Send
+    vm: VMWrapper,
+    handler,
+    attempt_headers: Dict[str, str],
+    receive: ReceiveChannel,
+    send: Send,
+    journal_codec: Optional[JournalValueCodec] = None,
 ):
     """Invoke the user code."""
     status, res_headers = vm.get_response_head()
@@ -155,7 +162,13 @@ async def process_invocation_to_completion(
     # ========================================
     invocation = vm.sys_input()
     context = ServerInvocationContext(
-        vm=vm, handler=handler, invocation=invocation, attempt_headers=attempt_headers, send=send, receive=receive
+        vm=vm,
+        handler=handler,
+        invocation=invocation,
+        attempt_headers=attempt_headers,
+        send=send,
+        receive=receive,
+        journal_codec=journal_codec,
     )
     try:
         await context.enter()
@@ -217,6 +230,27 @@ def asgi_app(endpoint: Endpoint) -> RestateAppT:
     active_channels: Set[ReceiveChannel] = set()
     sigterm_installed = False
 
+    # Journal value codec resolution.
+    # A codec (or an async provider building one) may be configured on the endpoint. Whether one is
+    # configured is known synchronously and drives the VM's payload-checks flag. An async provider
+    # is resolved once, lazily on the first invocation, and the resulting codec is reused.
+    codec_config = endpoint.journal_value_codec
+    disable_payload_checks = codec_config is not None
+
+    # Memoize the provider's resolution. functools.cache can't wrap the async provider directly
+    # (it would cache a coroutine, which can't be awaited twice), so we cache a sync factory that
+    # returns the Task instead: the provider runs once and every caller awaits the same Task.
+    @functools.cache
+    def _codec_task() -> "asyncio.Task[JournalValueCodec]":
+        assert callable(codec_config)
+        return asyncio.ensure_future(codec_config())
+
+    async def get_codec() -> Optional[JournalValueCodec]:
+        if codec_config is None or isinstance(codec_config, JournalValueCodec):
+            # Nothing configured, or an already-built instance: nothing to await.
+            return codec_config
+        return await _codec_task()
+
     def _on_sigterm() -> None:
         """Notify all active receive channels of graceful shutdown."""
         for ch in active_channels:
@@ -284,8 +318,14 @@ async def app(scope: Scope, receive: Receive, send: Send):
             receive_channel = ReceiveChannel(receive)
             active_channels.add(receive_channel)
             try:
+                journal_codec = await get_codec()
                 await process_invocation_to_completion(
-                    VMWrapper(request_headers), handler, dict(request_headers), receive_channel, send
+                    VMWrapper(request_headers, disable_payload_checks=disable_payload_checks),
+                    handler,
+                    dict(request_headers),
+                    receive_channel,
+                    send,
+                    journal_codec=journal_codec,
                 )
             finally:
                 active_channels.discard(receive_channel)