From 08791a02010f6089ed86cd91ebc88406abcea9dd Mon Sep 17 00:00:00 2001 From: slinkydeveloper Date: Wed, 4 Nov 2026 15:42:35 +0100 Subject: [PATCH] Support eager state whitelist Port of restatedev/sdk-typescript#798. Add the `state_preload` option on VirtualObject/Workflow constructors and handler decorators, accepting True (preload all), False (preload nothing) or a list of state keys to preload eagerly. The list form is advertised through the new `eagerStateKeysWhitelist` manifest field, which requires endpoint manifest v5 (restate-server >= 1.8). `enable_lazy_state` is kept but deprecated; setting both raises ValueError. Co-Authored-By: Claude Fable 5.1 --- README.md | 3 + python/restate/discovery.py | 52 ++++++++++++- python/restate/handler.py | 27 ++++++- python/restate/object.py | 27 +++++-- python/restate/server.py | 8 +- python/restate/service.py | 1 + python/restate/workflow.py | 51 +++++++++---- tests/discovery.py | 146 +++++++++++++++++++++++++++++++++++- 8 files changed, 286 insertions(+), 29 deletions(-) diff --git a/README.md b/README.md index 73dd2cb3..8ae73092 100644 --- a/README.md +++ b/README.md @@ -64,11 +64,14 @@ The compatibility with Restate is described in the following table. sdk-python v | 1.3 - 1.4 | ✅ (1) | ✅ (1)(2) | | 1.5 - 1.6 | ✅ | ✅ (2) | | 1.7 | ✅ | ✅ | +| 1.8 | ✅ | ✅ (3) | (1) **Note** The new Service/Object/Workflow constructor field and the decorator field `invocation_retry_policy` works only from Restate 1.4 onward. Check the in-code documentation for more details. (2) **Note** Scopes and limit keys work only from Restate 1.7 onward. Check the in-code documentation for more details. +(3) **Note** The list form of `state_preload` (eager state keys whitelist) works only from Restate 1.8 onward. Check the in-code documentation for more details. + ## Contributing We’re excited if you join the Restate community and start contributing! diff --git a/python/restate/discovery.py b/python/restate/discovery.py index 0bd22dc0..38d26e2b 100644 --- a/python/restate/discovery.py +++ b/python/restate/discovery.py @@ -27,7 +27,7 @@ import json import typing from enum import Enum -from typing import Dict, Optional, Any, List, get_args, get_origin +from typing import Dict, Optional, Any, List, Tuple, Union, get_args, get_origin from restate.endpoint import Endpoint as RestateEndpoint @@ -86,6 +86,7 @@ def __init__( idempotencyRetention: Optional[int] = None, workflowCompletionRetention: Optional[int] = None, enableLazyState: Optional[bool] = None, + eagerStateKeysWhitelist: Optional[List[str]] = None, ingressPrivate: Optional[bool] = None, retryPolicyInitialInterval: Optional[int] = None, retryPolicyMaxInterval: Optional[int] = None, @@ -105,6 +106,7 @@ def __init__( self.idempotencyRetention = idempotencyRetention self.workflowCompletionRetention = workflowCompletionRetention self.enableLazyState = enableLazyState + self.eagerStateKeysWhitelist = eagerStateKeysWhitelist self.ingressPrivate = ingressPrivate self.retryPolicyInitialInterval = retryPolicyInitialInterval self.retryPolicyMaxInterval = retryPolicyMaxInterval @@ -127,6 +129,7 @@ def __init__( journalRetention: Optional[int] = None, idempotencyRetention: Optional[int] = None, enableLazyState: Optional[bool] = None, + eagerStateKeysWhitelist: Optional[List[str]] = None, ingressPrivate: Optional[bool] = None, retryPolicyInitialInterval: Optional[int] = None, retryPolicyMaxInterval: Optional[int] = None, @@ -144,6 +147,7 @@ def __init__( self.journalRetention = journalRetention self.idempotencyRetention = idempotencyRetention self.enableLazyState = enableLazyState + self.eagerStateKeysWhitelist = eagerStateKeysWhitelist self.ingressPrivate = ingressPrivate self.retryPolicyInitialInterval = retryPolicyInitialInterval self.retryPolicyMaxInterval = retryPolicyMaxInterval @@ -225,6 +229,27 @@ def json_schema_from_type_hint(type_hint: Optional[TypeHint[Any]]) -> Any: return type_hint_to_json_schema(type_hint.annotation) +def resolve_state_preload( + state_preload: Optional[Union[bool, List[str]]], enable_lazy_state: Optional[bool] +) -> Tuple[Optional[bool], Optional[List[str]]]: + """ + Resolve the user facing state preload options into the manifest pair (enableLazyState, eagerStateKeysWhitelist). + """ + if state_preload is None: + # Legacy passthrough (may itself be None). + return enable_lazy_state, None + + if isinstance(state_preload, bool): + # True => eager (preload all) => not lazy + # False => lazy (preload none) => lazy + return not state_preload, None + + # state_preload is a whitelist of keys to preload eagerly (selective preloading). + # An empty whitelist is equivalent to fully lazy state. + # Copy, so later mutations of the user's list don't affect discovery. + return True, list(state_preload) if len(state_preload) > 0 else None + + # pylint: disable=R0912,R0915 def compute_discovery_json( endpoint: RestateEndpoint, version: int, discovered_as: typing.Literal["bidi", "request_response"] @@ -236,6 +261,15 @@ def compute_discovery_json( ep = compute_discovery(endpoint, discovered_as) # Validate that new discovery fields aren't used with older protocol versions + if version <= 4: + # Note: enableLazyState itself is available since v3; only the eager state whitelist requires v5. + for service in ep.services: + if service.eagerStateKeysWhitelist is not None: + raise ValueError("eagerStateKeysWhitelist is only supported in discovery protocol version 5") + for handler in service.handlers: + if handler.eagerStateKeysWhitelist is not None: + raise ValueError("eagerStateKeysWhitelist is only supported in discovery protocol version 5") + if version <= 3: # Strip lambdaCompression for older discovery versions ep.lambdaCompression = None @@ -336,6 +370,9 @@ def compute_discovery(endpoint: RestateEndpoint, discovered_as: typing.Literal[" contentType=handler.handler_io.content_type, jsonSchema=json_schema_from_type_hint(handler.handler_io.output_type), ) + handler_enable_lazy_state, handler_eager_state_keys_whitelist = resolve_state_preload( + handler.state_preload, handler.enable_lazy_state + ) # add the handler service_handlers.append( Handler( @@ -360,7 +397,8 @@ def compute_discovery(endpoint: RestateEndpoint, discovered_as: typing.Literal[" workflowCompletionRetention=int(handler.workflow_retention.total_seconds() * 1000) if handler.workflow_retention is not None else None, - enableLazyState=handler.enable_lazy_state, + enableLazyState=handler_enable_lazy_state, + eagerStateKeysWhitelist=handler_eager_state_keys_whitelist, ingressPrivate=handler.ingress_private, retryPolicyInitialInterval=int( handler.invocation_retry_policy.initial_interval.total_seconds() * 1000 @@ -386,6 +424,13 @@ def compute_discovery(endpoint: RestateEndpoint, discovered_as: typing.Literal[" ) ) # add the service + if isinstance(service, (Workflow, VirtualObject)): + service_enable_lazy_state, service_eager_state_keys_whitelist = resolve_state_preload( + service.state_preload, service.enable_lazy_state + ) + else: + # Stateless services must not advertise state options. + service_enable_lazy_state, service_eager_state_keys_whitelist = None, None description = service.service_tag.description metadata = service.service_tag.metadata services.append( @@ -407,7 +452,8 @@ def compute_discovery(endpoint: RestateEndpoint, discovered_as: typing.Literal[" idempotencyRetention=int(service.idempotency_retention.total_seconds() * 1000) if service.idempotency_retention is not None else None, - enableLazyState=service.enable_lazy_state if isinstance(service, (Workflow, VirtualObject)) else None, + enableLazyState=service_enable_lazy_state, + eagerStateKeysWhitelist=service_eager_state_keys_whitelist, ingressPrivate=service.ingress_private, retryPolicyInitialInterval=int(service.invocation_retry_policy.initial_interval.total_seconds() * 1000) if service.invocation_retry_policy and service.invocation_retry_policy.initial_interval diff --git a/python/restate/handler.py b/python/restate/handler.py index 306e12c6..3253b7e6 100644 --- a/python/restate/handler.py +++ b/python/restate/handler.py @@ -29,6 +29,7 @@ Literal, Optional, TypeVar, + Union, ) from restate.retry_policy import InvocationRetryPolicy @@ -216,6 +217,19 @@ def update_handler_io_with_type_hints(handler_io: HandlerIO[I, O], signature: Si update_handler_io_with_return_type_hints(handler_io, signature) +def check_state_preload_options( + state_preload: Optional[Union[bool, List[str]]], enable_lazy_state: Optional[bool] +) -> None: + """ + Validate that the new `state_preload` option and the deprecated `enable_lazy_state` option are not both set. + """ + if state_preload is not None and enable_lazy_state is not None: + raise ValueError( + "Cannot set both the 'state_preload' option and the deprecated 'enable_lazy_state' option at the same time. " + "Use only 'state_preload'." + ) + + # pylint: disable=R0902 @dataclass class Handler(Generic[I, O]): @@ -236,7 +250,9 @@ class Handler(Generic[I, O]): journal_retention: Journal retention duration. idempotency_retention: Idempotency retention duration. workflow_retention: Workflow completion retention duration. - enable_lazy_state: If true, lazy state is enabled. + enable_lazy_state: Deprecated, use `state_preload` instead. + state_preload: Configures which state keys are preloaded ("eager state") at the start of each invocation. + True preloads all state, False preloads nothing, a list preloads only the listed keys. ingress_private: If true, the handler cannot be invoked from the HTTP nor Kafka ingress. invocation_retry_policy: Optional retry policy configuration applied to this handler. explicit_recovery: If true, the handler runs with explicit recovery, see restate._preview.explicit_recovery. @@ -256,6 +272,7 @@ class Handler(Generic[I, O]): idempotency_retention: Optional[timedelta] = None workflow_retention: Optional[timedelta] = None enable_lazy_state: Optional[bool] = None + state_preload: Optional[Union[bool, List[str]]] = None ingress_private: Optional[bool] = None invocation_retry_policy: Optional[InvocationRetryPolicy] = None context_managers: Optional[List[Callable[[], AsyncContextManager[None]]]] = None @@ -280,6 +297,7 @@ def make_handler( idempotency_retention: Optional[timedelta] = None, workflow_retention: Optional[timedelta] = None, enable_lazy_state: Optional[bool] = None, + state_preload: Optional[Union[bool, List[str]]] = None, ingress_private: Optional[bool] = None, invocation_retry_policy: Optional[InvocationRetryPolicy] = None, context_managers: Optional[List[Callable[[], AsyncContextManager[None]]]] = None, @@ -301,7 +319,9 @@ def make_handler( journal_retention: Journal retention duration. idempotency_retention: Idempotency retention duration. workflow_retention: Workflow completion retention duration. - enable_lazy_state: If true, lazy state is enabled. + enable_lazy_state: Deprecated, use `state_preload` instead. + state_preload: Configures which state keys are preloaded ("eager state") at the start of each invocation. + True preloads all state, False preloads nothing, a list preloads only the listed keys. ingress_private: If true, the handler cannot be invoked from the HTTP nor Kafka ingress. invocation_retry_policy: Retry policy used by Restate when invoking this handler. NOTE: only supported with restate-server >= 1.5. @@ -316,6 +336,8 @@ def make_handler( if len(signature.parameters) == 0: raise ValueError("Handler must have at least one parameter") + check_state_preload_options(state_preload, enable_lazy_state) + arity = len(signature.parameters) update_handler_io_with_type_hints(handler_io, signature) # mutates handler_io @@ -334,6 +356,7 @@ def make_handler( idempotency_retention=idempotency_retention, workflow_retention=workflow_retention, enable_lazy_state=enable_lazy_state, + state_preload=state_preload, ingress_private=ingress_private, invocation_retry_policy=invocation_retry_policy, context_managers=context_managers, diff --git a/python/restate/object.py b/python/restate/object.py index cff1953a..ffde1d4f 100644 --- a/python/restate/object.py +++ b/python/restate/object.py @@ -21,7 +21,7 @@ from restate.serde import Serde, DefaultSerde from restate.retry_policy import InvocationRetryPolicy -from restate.handler import Handler, HandlerIO, ServiceTag, make_handler +from restate.handler import Handler, HandlerIO, ServiceTag, check_state_preload_options, make_handler I = typing.TypeVar("I") O = typing.TypeVar("O") @@ -63,9 +63,13 @@ class VirtualObject: idempotency_retention (timedelta, optional): The retention duration of idempotent requests to this service. NOTE: You can set this field only if you register this service against restate-server >= 1.4, otherwise the service discovery will fail. - enable_lazy_state (bool, optional): When set to True, lazy state will be enabled for all invocations to this service. - NOTE: You can set this field only if you register this service against restate-server >= 1.4, - otherwise the service discovery will fail. + enable_lazy_state (bool, optional): Deprecated, use `state_preload` instead. + state_preload (bool | List[str], optional): Configures which state keys are preloaded ("eager state") + at the start of each invocation to this virtual object. Can be overridden per handler. + - `True`: preload all state (this is the default behavior). + - `False`: preload nothing (lazy state). + - `["key1", "key2"]`: preload only the listed keys eagerly. + NOTE: The list form is supported only by restate-server >= 1.8, otherwise the service discovery will fail. ingress_private (bool, optional): When set to True this service, with all its handlers, cannot be invoked from the restate-server HTTP and Kafka ingress, but only from other services. NOTE: You can set this field only if you register this service against restate-server >= 1.4, @@ -85,6 +89,7 @@ def __init__( journal_retention: typing.Optional[timedelta] = None, idempotency_retention: typing.Optional[timedelta] = None, enable_lazy_state: typing.Optional[bool] = None, + state_preload: typing.Optional[typing.Union[bool, typing.List[str]]] = None, ingress_private: typing.Optional[bool] = None, invocation_retry_policy: typing.Optional[InvocationRetryPolicy] = None, invocation_context_managers: typing.Optional[ @@ -97,7 +102,9 @@ def __init__( self.abort_timeout = abort_timeout self.journal_retention = journal_retention self.idempotency_retention = idempotency_retention + check_state_preload_options(state_preload, enable_lazy_state) self.enable_lazy_state = enable_lazy_state + self.state_preload = state_preload self.ingress_private = ingress_private self.invocation_retry_policy = invocation_retry_policy self.context_managers = invocation_context_managers @@ -125,6 +132,7 @@ def handler( journal_retention: typing.Optional[timedelta] = None, idempotency_retention: typing.Optional[timedelta] = None, enable_lazy_state: typing.Optional[bool] = None, + state_preload: typing.Optional[typing.Union[bool, typing.List[str]]] = None, ingress_private: typing.Optional[bool] = None, invocation_retry_policy: typing.Optional[InvocationRetryPolicy] = None, invocation_context_managers: typing.Optional[ @@ -163,9 +171,13 @@ def handler( idempotency_retention: The retention duration of idempotent requests to this service. NOTE: You can set this field only if you register this service against restate-server >= 1.4, otherwise the service discovery will fail. - enable_lazy_state: When set to True, lazy state will be enabled for all invocations to this handler. - NOTE: You can set this field only if you register this service against restate-server >= 1.4, - otherwise the service discovery will fail. + enable_lazy_state: Deprecated, use `state_preload` instead. + state_preload: Configures which state keys are preloaded ("eager state") at the start of each + invocation to this handler. Overrides the value set on the service. + - `True`: preload all state (this is the default behavior). + - `False`: preload nothing (lazy state). + - `["key1", "key2"]`: preload only the listed keys eagerly. + NOTE: The list form is supported only by restate-server >= 1.8, otherwise the service discovery will fail. ingress_private: When set to True this handler cannot be invoked from the restate-server HTTP and Kafka ingress, but only from other services. NOTE: You can set this field only if you register this service against restate-server >= 1.4, @@ -212,6 +224,7 @@ def wrapped(*args, **kwargs): idempotency_retention, None, enable_lazy_state, + state_preload, ingress_private, invocation_retry_policy, combined_context_managers, diff --git a/python/restate/server.py b/python/restate/server.py index b3c77bf4..723858a7 100644 --- a/python/restate/server.py +++ b/python/restate/server.py @@ -105,14 +105,16 @@ async def send_discovery(scope: Scope, send: Send, endpoint: Endpoint): # Negotiate discovery protocol version version = 2 if accept_header: - if "application/vnd.restate.endpointmanifest.v4+json" in accept_header: + if "application/vnd.restate.endpointmanifest.v5+json" in accept_header: + version = 5 + elif "application/vnd.restate.endpointmanifest.v4+json" in accept_header: version = 4 elif "application/vnd.restate.endpointmanifest.v3+json" in accept_header: version = 3 elif "application/vnd.restate.endpointmanifest.v2+json" in accept_header: version = 2 else: - await send_status_with_error_text(send, 415, f"Unsupported discovery version ${accept_header}") + await send_status_with_error_text(send, 415, f"Unsupported discovery version {accept_header}") return try: @@ -128,7 +130,7 @@ async def send_discovery(scope: Scope, send: Send, endpoint: Endpoint): } ) except ValueError as e: - await send_status_with_error_text(send, 500, f"Error when computing discovery ${e}") + await send_status_with_error_text(send, 500, f"Error when computing discovery {e}") return diff --git a/python/restate/service.py b/python/restate/service.py index 94899da3..ea695430 100644 --- a/python/restate/service.py +++ b/python/restate/service.py @@ -203,6 +203,7 @@ def wrapped(*args, **kwargs): idempotency_retention, None, None, + None, ingress_private, invocation_retry_policy, combined_context_managers, diff --git a/python/restate/workflow.py b/python/restate/workflow.py index f03b5a80..e1d0056b 100644 --- a/python/restate/workflow.py +++ b/python/restate/workflow.py @@ -22,7 +22,7 @@ from restate.retry_policy import InvocationRetryPolicy from restate.serde import DefaultSerde, Serde -from restate.handler import Handler, HandlerIO, ServiceTag, make_handler +from restate.handler import Handler, HandlerIO, ServiceTag, check_state_preload_options, make_handler I = typing.TypeVar("I") O = typing.TypeVar("O") @@ -68,9 +68,13 @@ class Workflow: idempotency_retention (timedelta, optional): The retention duration of idempotent requests to this service. NOTE: You can set this field only if you register this service against restate-server >= 1.4, otherwise the service discovery will fail. - enable_lazy_state (bool, optional): When set to True, lazy state will be enabled for all invocations to this service. - NOTE: You can set this field only if you register this service against restate-server >= 1.4, - otherwise the service discovery will fail. + enable_lazy_state (bool, optional): Deprecated, use `state_preload` instead. + state_preload (bool | List[str], optional): Configures which state keys are preloaded ("eager state") + at the start of each invocation to this workflow. Can be overridden per handler. + - `True`: preload all state (this is the default behavior). + - `False`: preload nothing (lazy state). + - `["key1", "key2"]`: preload only the listed keys eagerly. + NOTE: The list form is supported only by restate-server >= 1.8, otherwise the service discovery will fail. ingress_private (bool, optional): When set to True this service, with all its handlers, cannot be invoked from the restate-server HTTP and Kafka ingress, but only from other services. NOTE: You can set this field only if you register this service against restate-server >= 1.4, @@ -91,6 +95,7 @@ def __init__( journal_retention: typing.Optional[timedelta] = None, idempotency_retention: typing.Optional[timedelta] = None, enable_lazy_state: typing.Optional[bool] = None, + state_preload: typing.Optional[typing.Union[bool, typing.List[str]]] = None, ingress_private: typing.Optional[bool] = None, invocation_retry_policy: typing.Optional[InvocationRetryPolicy] = None, invocation_context_managers: typing.Optional[ @@ -103,7 +108,9 @@ def __init__( self.abort_timeout = abort_timeout self.journal_retention = journal_retention self.idempotency_retention = idempotency_retention + check_state_preload_options(state_preload, enable_lazy_state) self.enable_lazy_state = enable_lazy_state + self.state_preload = state_preload self.ingress_private = ingress_private self.invocation_retry_policy = invocation_retry_policy self.context_managers = invocation_context_managers @@ -129,6 +136,7 @@ def main( journal_retention: typing.Optional[timedelta] = None, workflow_retention: typing.Optional[timedelta] = None, enable_lazy_state: typing.Optional[bool] = None, + state_preload: typing.Optional[typing.Union[bool, typing.List[str]]] = None, ingress_private: typing.Optional[bool] = None, invocation_retry_policy: typing.Optional[InvocationRetryPolicy] = None, invocation_context_managers: typing.Optional[ @@ -166,9 +174,13 @@ def main( workflow_retention: The retention duration for this workflow handler. NOTE: You can set this field only if you register this service against restate-server >= 1.4, otherwise the service discovery will fail. - enable_lazy_state: When set to True, lazy state will be enabled for all invocations to this handler. - NOTE: You can set this field only if you register this service against restate-server >= 1.4, - otherwise the service discovery will fail. + enable_lazy_state: Deprecated, use `state_preload` instead. + state_preload: Configures which state keys are preloaded ("eager state") at the start of each + invocation to this handler. Overrides the value set on the service. + - `True`: preload all state (this is the default behavior). + - `False`: preload nothing (lazy state). + - `["key1", "key2"]`: preload only the listed keys eagerly. + NOTE: The list form is supported only by restate-server >= 1.8, otherwise the service discovery will fail. ingress_private: When set to True this handler cannot be invoked from the restate-server HTTP and Kafka ingress, but only from other services. NOTE: You can set this field only if you register this service against restate-server >= 1.4, @@ -189,6 +201,7 @@ def main( idempotency_retention=None, workflow_retention=workflow_retention, enable_lazy_state=enable_lazy_state, + state_preload=state_preload, ingress_private=ingress_private, invocation_retry_policy=invocation_retry_policy, invocation_context_managers=invocation_context_managers, @@ -207,6 +220,7 @@ def handler( journal_retention: typing.Optional[timedelta] = None, idempotency_retention: typing.Optional[timedelta] = None, enable_lazy_state: typing.Optional[bool] = None, + state_preload: typing.Optional[typing.Union[bool, typing.List[str]]] = None, ingress_private: typing.Optional[bool] = None, invocation_retry_policy: typing.Optional[InvocationRetryPolicy] = None, invocation_context_managers: typing.Optional[ @@ -244,9 +258,13 @@ def handler( idempotency_retention: The retention duration of idempotent requests to this service. NOTE: You can set this field only if you register this service against restate-server >= 1.4, otherwise the service discovery will fail. - enable_lazy_state: When set to True, lazy state will be enabled for all invocations to this handler. - NOTE: You can set this field only if you register this service against restate-server >= 1.4, - otherwise the service discovery will fail. + enable_lazy_state: Deprecated, use `state_preload` instead. + state_preload: Configures which state keys are preloaded ("eager state") at the start of each + invocation to this handler. Overrides the value set on the service. + - `True`: preload all state (this is the default behavior). + - `False`: preload nothing (lazy state). + - `["key1", "key2"]`: preload only the listed keys eagerly. + NOTE: The list form is supported only by restate-server >= 1.8, otherwise the service discovery will fail. ingress_private: When set to True this handler cannot be invoked from the restate-server HTTP and Kafka ingress, but only from other services. NOTE: You can set this field only if you register this service against restate-server >= 1.4, @@ -268,6 +286,7 @@ def handler( idempotency_retention, None, enable_lazy_state, + state_preload, ingress_private, invocation_retry_policy, invocation_context_managers, @@ -289,6 +308,7 @@ def _add_handler( idempotency_retention: typing.Optional[timedelta] = None, workflow_retention: typing.Optional[timedelta] = None, enable_lazy_state: typing.Optional[bool] = None, + state_preload: typing.Optional[typing.Union[bool, typing.List[str]]] = None, ingress_private: typing.Optional[bool] = None, invocation_retry_policy: typing.Optional["InvocationRetryPolicy"] = None, invocation_context_managers: typing.Optional[ @@ -330,9 +350,13 @@ def _add_handler( workflow_retention: The retention duration for this workflow handler. NOTE: You can set this field only if you register this service against restate-server >= 1.4, otherwise the service discovery will fail. - enable_lazy_state: When set to True, lazy state will be enabled for all invocations to this handler. - NOTE: You can set this field only if you register this service against restate-server >= 1.4, - otherwise the service discovery will fail. + enable_lazy_state: Deprecated, use `state_preload` instead. + state_preload: Configures which state keys are preloaded ("eager state") at the start of each + invocation to this handler. Overrides the value set on the service. + - `True`: preload all state (this is the default behavior). + - `False`: preload nothing (lazy state). + - `["key1", "key2"]`: preload only the listed keys eagerly. + NOTE: The list form is supported only by restate-server >= 1.8, otherwise the service discovery will fail. ingress_private: When set to True this handler cannot be invoked from the restate-server HTTP and Kafka ingress, but only from other services. NOTE: You can set this field only if you register this service against restate-server >= 1.4, @@ -380,6 +404,7 @@ def wrapped(*args, **kwargs): idempotency_retention=idempotency_retention, workflow_retention=workflow_retention, enable_lazy_state=enable_lazy_state, + state_preload=state_preload, ingress_private=ingress_private, invocation_retry_policy=invocation_retry_policy, context_managers=combined_context_managers, diff --git a/tests/discovery.py b/tests/discovery.py index 8fe26fc9..4ef17979 100644 --- a/tests/discovery.py +++ b/tests/discovery.py @@ -8,10 +8,13 @@ # directory of this repository or package, or at # https://github.com/restatedev/sdk-typescript/blob/main/LICENSE # +import json from datetime import timedelta +import pytest + import restate -from restate.discovery import compute_discovery +from restate.discovery import compute_discovery, compute_discovery_json from restate.endpoint import Endpoint @@ -55,3 +58,144 @@ async def run(ctx: restate.WorkflowContext): workflow_handler_manifest = manifest.services[1].handlers[0] assert workflow_handler_manifest.workflowCompletionRetention == 0 + + +def _object_manifest(state_preload=None, enable_lazy_state=None, handler_kwargs=None): + obj = restate.VirtualObject("Obj", state_preload=state_preload, enable_lazy_state=enable_lazy_state) + + @obj.handler(**(handler_kwargs or {})) + async def handle(ctx: restate.ObjectContext): + pass + + return compute_discovery(Endpoint().bind(obj), "bidi").services[0] + + +def test_state_preload_whitelist_on_object_and_handler(): + keys = ["a", "b"] + service_manifest = _object_manifest(state_preload=keys, handler_kwargs={"state_preload": ["c"]}) + assert service_manifest.enableLazyState is True + assert service_manifest.eagerStateKeysWhitelist == ["a", "b"] + assert service_manifest.handlers[0].enableLazyState is True + assert service_manifest.handlers[0].eagerStateKeysWhitelist == ["c"] + # The whitelist must be a copy of the user list + keys.append("z") + assert service_manifest.eagerStateKeysWhitelist == ["a", "b"] + + +def test_state_preload_booleans(): + service_manifest = _object_manifest(state_preload=False, handler_kwargs={"state_preload": True}) + assert service_manifest.enableLazyState is True + assert service_manifest.eagerStateKeysWhitelist is None + assert service_manifest.handlers[0].enableLazyState is False + assert service_manifest.handlers[0].eagerStateKeysWhitelist is None + + +def test_state_preload_empty_list_is_fully_lazy(): + service_manifest = _object_manifest(state_preload=[]) + assert service_manifest.enableLazyState is True + assert service_manifest.eagerStateKeysWhitelist is None + + +def test_state_preload_unset_keeps_legacy_enable_lazy_state(): + service_manifest = _object_manifest(enable_lazy_state=True, handler_kwargs={"enable_lazy_state": False}) + assert service_manifest.enableLazyState is True + assert service_manifest.eagerStateKeysWhitelist is None + assert service_manifest.handlers[0].enableLazyState is False + assert service_manifest.handlers[0].eagerStateKeysWhitelist is None + + service_manifest = _object_manifest() + assert service_manifest.enableLazyState is None + assert service_manifest.handlers[0].enableLazyState is None + + +def test_state_preload_and_enable_lazy_state_are_mutually_exclusive(): + with pytest.raises(ValueError): + restate.VirtualObject("Obj", state_preload=False, enable_lazy_state=True) + with pytest.raises(ValueError): + restate.Workflow("Wf", state_preload=["a"], enable_lazy_state=True) + + obj = restate.VirtualObject("Obj") + with pytest.raises(ValueError): + + @obj.handler(state_preload=False, enable_lazy_state=True) + async def handle(ctx: restate.ObjectContext): + pass + + wf = restate.Workflow("Wf") + with pytest.raises(ValueError): + + @wf.main(state_preload=["a"], enable_lazy_state=False) + async def run(ctx: restate.WorkflowContext): + pass + + +def test_state_preload_on_workflow(): + wf = restate.Workflow("Wf", state_preload=["a"]) + + @wf.main(state_preload=["b"]) + async def run(ctx: restate.WorkflowContext): + pass + + @wf.handler(state_preload=False) + async def query(ctx: restate.WorkflowSharedContext): + pass + + service_manifest = compute_discovery(Endpoint().bind(wf), "bidi").services[0] + assert service_manifest.enableLazyState is True + assert service_manifest.eagerStateKeysWhitelist == ["a"] + handlers = {h.name: h for h in service_manifest.handlers} + assert handlers["run"].eagerStateKeysWhitelist == ["b"] + assert handlers["query"].enableLazyState is True + assert handlers["query"].eagerStateKeysWhitelist is None + + +def test_stateless_service_never_advertises_state_options(): + service = restate.Service("Svc") + + @service.handler() + async def handle(ctx: restate.Context): + pass + + service_manifest = compute_discovery(Endpoint().bind(service), "bidi").services[0] + assert service_manifest.enableLazyState is None + assert service_manifest.eagerStateKeysWhitelist is None + assert service_manifest.handlers[0].enableLazyState is None + assert service_manifest.handlers[0].eagerStateKeysWhitelist is None + + +def test_eager_state_keys_whitelist_requires_manifest_v5(): + obj = restate.VirtualObject("Obj", state_preload=["a"]) + + @obj.handler() + async def handle(ctx: restate.ObjectContext): + pass + + endpoint = Endpoint().bind(obj) + with pytest.raises(ValueError, match="eagerStateKeysWhitelist"): + compute_discovery_json(endpoint, 4, "bidi") + manifest = json.loads(compute_discovery_json(endpoint, 5, "bidi")) + assert manifest["services"][0]["enableLazyState"] is True + assert manifest["services"][0]["eagerStateKeysWhitelist"] == ["a"] + assert "eagerStateKeysWhitelist" not in manifest["services"][0]["handlers"][0] + + # Handler-level whitelist is rejected too + obj2 = restate.VirtualObject("Obj2") + + @obj2.handler(state_preload=["a"]) + async def handle2(ctx: restate.ObjectContext): + pass + + with pytest.raises(ValueError, match="eagerStateKeysWhitelist"): + compute_discovery_json(Endpoint().bind(obj2), 4, "bidi") + + +def test_state_preload_false_works_on_manifest_v4(): + obj = restate.VirtualObject("Obj", state_preload=False) + + @obj.handler() + async def handle(ctx: restate.ObjectContext): + pass + + manifest = json.loads(compute_discovery_json(Endpoint().bind(obj), 4, "bidi")) + assert manifest["services"][0]["enableLazyState"] is True + assert "eagerStateKeysWhitelist" not in manifest["services"][0]