Skip to content

Commit 9db2da2

Browse files
committed
Standardize codec, entity fixes
1 parent e84790b commit 9db2da2

9 files changed

Lines changed: 317 additions & 249 deletions

File tree

durabletask/client.py

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -845,7 +845,7 @@ def get_entity(self,
845845
return None
846846
if self._payload_store is not None:
847847
payload_helpers.deexternalize_payloads(res, self._payload_store)
848-
return EntityMetadata.from_entity_metadata(res.entity, include_state)
848+
return EntityMetadata.from_entity_metadata(res.entity, include_state, self._data_converter)
849849

850850
def get_all_entities(self,
851851
entity_query: EntityQuery | None = None) -> list[EntityMetadata]:
@@ -862,7 +862,7 @@ def get_all_entities(self,
862862
resp: pb.QueryEntitiesResponse = self._stub.QueryEntities(query_request)
863863
if self._payload_store is not None:
864864
payload_helpers.deexternalize_payloads(resp, self._payload_store)
865-
entities += [EntityMetadata.from_entity_metadata(entity, query_request.query.includeState) for entity in resp.entities]
865+
entities += [EntityMetadata.from_entity_metadata(entity, query_request.query.includeState, self._data_converter) for entity in resp.entities]
866866
if check_continuation_token(resp.continuationToken, _continuation_token, self._logger):
867867
_continuation_token = resp.continuationToken
868868
else:
@@ -1328,7 +1328,7 @@ async def get_entity(self,
13281328
return None
13291329
if self._payload_store is not None:
13301330
await payload_helpers.deexternalize_payloads_async(res, self._payload_store)
1331-
return EntityMetadata.from_entity_metadata(res.entity, include_state)
1331+
return EntityMetadata.from_entity_metadata(res.entity, include_state, self._data_converter)
13321332

13331333
async def get_all_entities(self,
13341334
entity_query: EntityQuery | None = None) -> list[EntityMetadata]:
@@ -1345,7 +1345,7 @@ async def get_all_entities(self,
13451345
resp: pb.QueryEntitiesResponse = await self._stub.QueryEntities(query_request)
13461346
if self._payload_store is not None:
13471347
await payload_helpers.deexternalize_payloads_async(resp, self._payload_store)
1348-
entities += [EntityMetadata.from_entity_metadata(entity, query_request.query.includeState) for entity in resp.entities]
1348+
entities += [EntityMetadata.from_entity_metadata(entity, query_request.query.includeState, self._data_converter) for entity in resp.entities]
13491349
if check_continuation_token(resp.continuationToken, _continuation_token, self._logger):
13501350
_continuation_token = resp.continuationToken
13511351
else:

durabletask/entities/entity_metadata.py

Lines changed: 28 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,11 @@
11
from datetime import datetime, timezone
2-
from typing import Any, TypeVar, overload
2+
from typing import TYPE_CHECKING, Any, TypeVar, overload
33
from durabletask.entities.entity_instance_id import EntityInstanceId
44

55
import durabletask.internal.orchestrator_service_pb2 as pb
6-
from durabletask.internal import shared
6+
7+
if TYPE_CHECKING:
8+
from durabletask.serialization import DataConverter
79

810
TState = TypeVar("TState")
911

@@ -30,7 +32,8 @@ def __init__(self,
3032
backlog_queue_size: int,
3133
locked_by: str,
3234
includes_state: bool,
33-
state: Any | None):
35+
state: Any | None,
36+
data_converter: "DataConverter | None" = None):
3437
"""Initializes a new instance of the EntityMetadata class.
3538
3639
Args:
@@ -42,13 +45,20 @@ def __init__(self,
4245
self._locked_by = locked_by
4346
self.includes_state = includes_state
4447
self._state = state
48+
if data_converter is None:
49+
from durabletask.serialization import JsonDataConverter
50+
data_converter = JsonDataConverter()
51+
self._data_converter = data_converter
4552

4653
@staticmethod
47-
def from_entity_response(entity_response: pb.GetEntityResponse, includes_state: bool):
48-
return EntityMetadata.from_entity_metadata(entity_response.entity, includes_state)
54+
def from_entity_response(entity_response: pb.GetEntityResponse, includes_state: bool,
55+
data_converter: "DataConverter | None" = None):
56+
return EntityMetadata.from_entity_metadata(
57+
entity_response.entity, includes_state, data_converter)
4958

5059
@staticmethod
51-
def from_entity_metadata(entity: pb.EntityMetadata, includes_state: bool):
60+
def from_entity_metadata(entity: pb.EntityMetadata, includes_state: bool,
61+
data_converter: "DataConverter | None" = None):
5262
try:
5363
entity_id = EntityInstanceId.parse(entity.instanceId)
5464
except ValueError:
@@ -62,7 +72,8 @@ def from_entity_metadata(entity: pb.EntityMetadata, includes_state: bool):
6272
backlog_queue_size=entity.backlogQueueSize,
6373
locked_by=entity.lockedBy.value,
6474
includes_state=includes_state,
65-
state=entity_state
75+
state=entity_state,
76+
data_converter=data_converter,
6677
)
6778

6879
@overload
@@ -74,11 +85,17 @@ def get_state(self, intended_type: None = None) -> Any:
7485
...
7586

7687
def get_state(self, intended_type: type[TState] | None = None) -> TState | Any | None:
77-
"""Get the current state of the entity, optionally converting it to a specified type."""
78-
if intended_type is None or self._state is None:
79-
return self._state
88+
"""Get the current state of the entity, optionally converting it to a specified type.
89+
90+
The state is stored as its raw serialized JSON payload and deserialized
91+
here. When ``intended_type`` is provided the payload is reconstructed as
92+
that type (dataclasses, ``from_json()``-capable types, etc.); otherwise
93+
the plain deserialized JSON value is returned.
94+
"""
95+
if self._state is None:
96+
return None
8097

81-
return shared.coerce_to_type(self._state, intended_type)
98+
return self._data_converter.deserialize(self._state, intended_type)
8299

83100
def get_locked_by(self) -> EntityInstanceId | None:
84101
"""Get the identifier of the worker that currently holds the lock on the entity.

durabletask/internal/entity_state_shim.py

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,17 +1,23 @@
1-
from typing import Any, TypeVar, overload
1+
from typing import TYPE_CHECKING, Any, TypeVar, overload
22

33
import durabletask.internal.orchestrator_service_pb2 as pb
4-
from durabletask.internal import shared
4+
5+
if TYPE_CHECKING:
6+
from durabletask.serialization import DataConverter
57

68
TState = TypeVar("TState")
79

810

911
class StateShim:
10-
def __init__(self, start_state: Any):
12+
def __init__(self, start_state: Any, data_converter: "DataConverter | None" = None):
1113
self._current_state: Any = start_state
1214
self._checkpoint_state: Any = start_state
1315
self._operation_actions: list[pb.OperationAction] = []
1416
self._actions_checkpoint_state: int = 0
17+
if data_converter is None:
18+
from durabletask.serialization import JsonDataConverter
19+
data_converter = JsonDataConverter()
20+
self._data_converter = data_converter
1521

1622
@overload
1723
def get_state(self, intended_type: type[TState], default: TState) -> TState:
@@ -32,7 +38,7 @@ def get_state(self, intended_type: type[TState] | None = None, default: TState |
3238
if intended_type is None:
3339
return self._current_state
3440

35-
return shared.coerce_to_type(self._current_state, intended_type)
41+
return self._data_converter.coerce(self._current_state, intended_type)
3642

3743
def set_state(self, state: Any) -> None:
3844
self._current_state = state

durabletask/internal/json_codec.py

Lines changed: 183 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,183 @@
1+
# Copyright (c) Microsoft Corporation.
2+
# Licensed under the MIT License.
3+
4+
"""Internal JSON codec for Durable Task payloads.
5+
6+
This module holds the low-level serialization *mechanism* -- the JSON string
7+
encode/decode primitives and the value-level type coercion used to reconstruct
8+
custom objects. Serialization *policy* (the public, pluggable strategy) lives in
9+
:mod:`durabletask.serialization`; the default ``JsonDataConverter`` is the only
10+
production consumer of ``to_json`` / ``from_json``, while ``coerce_to_type`` is
11+
also used directly by entity state accessors that already hold a parsed value.
12+
"""
13+
14+
from __future__ import annotations
15+
16+
import dataclasses
17+
import json
18+
import types
19+
import typing
20+
from collections.abc import Sequence
21+
from types import SimpleNamespace
22+
from typing import Any, cast
23+
24+
# Marker formerly added to JSON payloads to flag objects for automatic
25+
# deserialization into a SimpleNamespace. New code no longer emits this marker
26+
# (objects are serialized as plain JSON), but the decoder still recognizes it so
27+
# that orchestration histories produced by older SDK versions continue to replay.
28+
AUTO_SERIALIZED = "__durabletask_autoobject__"
29+
30+
31+
def to_json(obj: Any) -> str:
32+
"""Serialize a value to a JSON string.
33+
34+
Builtins serialize to plain JSON. Dataclasses, ``SimpleNamespace``
35+
instances, and objects exposing a ``to_json()`` method are serialized to
36+
plain JSON as well (without any type marker); custom objects can be
37+
reconstructed on the receiving side by passing ``expected_type`` to
38+
:func:`from_json`.
39+
"""
40+
try:
41+
return json.dumps(obj, default=_encode_custom_object)
42+
except TypeError as e:
43+
# Preserve the original error as the cause so serialization failures are
44+
# easier to diagnose, while naming the offending top-level type.
45+
raise TypeError(
46+
f"Failed to serialize object of type '{type(obj).__name__}' to JSON: {e}"
47+
) from e
48+
49+
50+
def from_json(json_str: str | bytes | bytearray, expected_type: type | None = None) -> Any:
51+
"""Deserialize a JSON string, optionally coercing the result to a type.
52+
53+
When ``expected_type`` is ``None`` (the default) the raw parsed JSON is
54+
returned. For backwards compatibility, payloads carrying the legacy
55+
:data:`AUTO_SERIALIZED` marker are reconstructed as ``SimpleNamespace``
56+
instances so that in-flight orchestrations produced by older SDK versions
57+
continue to replay.
58+
59+
When ``expected_type`` is provided, the legacy marker (if present) is
60+
stripped and the parsed value is coerced to ``expected_type`` -- dataclasses
61+
are constructed from their dict payloads, types exposing a ``from_json()``
62+
classmethod are reconstructed via that hook, and ``Optional``/``Union`` and
63+
``list`` type hints are honored recursively. The destination type is always
64+
supplied by the caller; it is never read from the payload.
65+
"""
66+
if expected_type is None:
67+
return json.loads(json_str, object_hook=_legacy_object_hook)
68+
raw = json.loads(json_str, object_hook=_strip_legacy_marker)
69+
return coerce_to_type(raw, expected_type)
70+
71+
72+
def _encode_custom_object(o: Any) -> Any:
73+
"""``default`` hook for :func:`json.dumps` that emits plain JSON.
74+
75+
Called only for values the JSON encoder cannot natively serialize. Note that
76+
namedtuples are handled natively by the encoder (serialized as JSON arrays)
77+
and never reach this hook.
78+
"""
79+
if dataclasses.is_dataclass(o) and not isinstance(o, type):
80+
return dataclasses.asdict(o)
81+
if isinstance(o, SimpleNamespace):
82+
return vars(o)
83+
# Custom objects may opt in via a ``to_json`` hook. It is resolved off the
84+
# type and called with the instance (``type(o).to_json(o)``) so that both
85+
# instance methods and ``@staticmethod`` hooks work -- matching the calling
86+
# convention used by ``azure-functions-durable``. The hook returns a
87+
# JSON-serializable value (a structure or a string), not a JSON document.
88+
to_json_hook = getattr(cast(Any, type(o)), "to_json", None)
89+
if callable(to_json_hook):
90+
return to_json_hook(o)
91+
# This will raise a TypeError describing the unsupported type.
92+
raise TypeError(f"Object of type '{type(o).__name__}' is not JSON serializable")
93+
94+
95+
def _legacy_object_hook(d: dict[str, Any]) -> Any:
96+
# If the object carries the legacy marker, deserialize it as a SimpleNamespace.
97+
if d.pop(AUTO_SERIALIZED, False):
98+
return SimpleNamespace(**d)
99+
return d
100+
101+
102+
def _strip_legacy_marker(d: dict[str, Any]) -> dict[str, Any]:
103+
# Discard the legacy marker so typed coercion sees a plain dict.
104+
d.pop(AUTO_SERIALIZED, None)
105+
return d
106+
107+
108+
def coerce_to_type(value: Any, expected_type: Any) -> Any:
109+
"""Coerce an already-parsed JSON value to ``expected_type``.
110+
111+
Handles ``None``/``Optional``/``Union`` and ``list`` type hints recursively,
112+
types exposing a ``from_json()`` classmethod, and dataclasses (including
113+
nested dataclass fields). The destination type is always caller-supplied and
114+
never derived from the payload, keeping deserialization secure.
115+
"""
116+
if expected_type is None or value is None:
117+
return value
118+
119+
origin = typing.get_origin(expected_type)
120+
if origin is not None:
121+
return _coerce_generic(value, expected_type, origin)
122+
123+
if not isinstance(expected_type, type):
124+
# Not a concrete, instantiable type (e.g. a typing special form we don't
125+
# special-case) -- return the value unchanged.
126+
return value
127+
128+
if isinstance(value, expected_type):
129+
return value
130+
131+
from_json_hook = getattr(expected_type, "from_json", None)
132+
if callable(from_json_hook):
133+
return from_json_hook(value)
134+
135+
if dataclasses.is_dataclass(expected_type) and isinstance(value, dict):
136+
return _build_dataclass(expected_type, cast(dict[str, Any], value))
137+
138+
type_ctor = cast(Any, expected_type)
139+
try:
140+
return type_ctor(value)
141+
except Exception as e:
142+
type_name = getattr(type_ctor, "__name__", None) or str(type_ctor)
143+
raise TypeError(
144+
f"Could not coerce value of type '{type(value).__name__}' to "
145+
f"'{type_name}'"
146+
) from e
147+
148+
149+
def _coerce_generic(value: Any, expected_type: Any, origin: Any) -> Any:
150+
args = typing.get_args(expected_type)
151+
if origin is typing.Union or origin is types.UnionType:
152+
# If the value already matches a member type, keep it as-is.
153+
non_none = [a for a in args if a is not type(None)]
154+
for arg in non_none:
155+
if isinstance(arg, type) and isinstance(value, arg):
156+
return value
157+
# ``Optional[T]`` (exactly one non-None member): coerce to that member.
158+
# For a genuine multi-member ``Union`` where the value matched none of
159+
# the members, leave it untouched rather than guessing the first arg --
160+
# forcing a coercion there can silently mis-construct the wrong type.
161+
if len(non_none) == 1:
162+
return coerce_to_type(value, non_none[0])
163+
return value
164+
if origin in (list, Sequence) and isinstance(value, list):
165+
elem_type = args[0] if args else None
166+
return [coerce_to_type(item, elem_type) for item in cast(list[Any], value)]
167+
# Other generics (dict, tuple, ...) are returned as parsed JSON.
168+
return value
169+
170+
171+
def _build_dataclass(cls: Any, data: dict[str, Any]) -> Any:
172+
"""Construct a dataclass from its dict payload, recursing into typed fields."""
173+
try:
174+
hints = typing.get_type_hints(cls)
175+
except Exception:
176+
hints = {}
177+
kwargs: dict[str, Any] = {}
178+
for field in dataclasses.fields(cls):
179+
if field.name not in data:
180+
continue
181+
field_type = hints.get(field.name)
182+
kwargs[field.name] = coerce_to_type(data[field.name], field_type)
183+
return cls(**kwargs)

0 commit comments

Comments
 (0)