-
Notifications
You must be signed in to change notification settings - Fork 208
Expand file tree
/
Copy path_data_converter.py
More file actions
422 lines (362 loc) · 16.1 KB
/
Copy path_data_converter.py
File metadata and controls
422 lines (362 loc) · 16.1 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
"""DataConverter: the top-level data conversion orchestrator."""
from __future__ import annotations
import dataclasses
import warnings
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from logging import getLogger
from typing import TYPE_CHECKING, Any, ClassVar
from typing_extensions import Self
import temporalio.api.common.v1
import temporalio.api.failure.v1
import temporalio.common
from temporalio.api.sdk.v1.external_storage_pb2 import ExternalStorageReference
from temporalio.converter._extstore import (
_REFERENCE_ENCODING,
ExternalStorage,
StorageDriverStoreContext,
)
from temporalio.converter._failure_converter import (
FailureConverter,
)
from temporalio.converter._payload_codec import (
PayloadCodec,
_apply_to_failure_payloads,
)
from temporalio.converter._payload_converter import (
PayloadConverter,
)
from temporalio.converter._payload_limits import (
PayloadLimitsConfig,
PayloadSizeWarning,
_PayloadSizeError,
_ServerPayloadErrorLimits,
)
from temporalio.converter._serialization_context import (
SerializationContext,
WithSerializationContext,
)
_REFERENCE_MESSAGE_TYPE = ExternalStorageReference.DESCRIPTOR.full_name.encode()
def _is_reference_payload(p: temporalio.api.common.v1.Payload) -> bool:
"""Return True if *p* is an external-storage reference payload."""
return p.metadata.get("encoding") == _REFERENCE_ENCODING or (
p.metadata.get("encoding") == b"json/protobuf"
and p.metadata.get("messageType") == _REFERENCE_MESSAGE_TYPE
)
# Import defaults from public API to avoid pydoctor cross-reference issues
if TYPE_CHECKING:
from temporalio.converter import DefaultFailureConverter, DefaultPayloadConverter
else:
# Import from private modules for runtime to avoid circular imports
from temporalio.converter._failure_converter import DefaultFailureConverter
from temporalio.converter._payload_converter import DefaultPayloadConverter
logger = getLogger("temporalio.converter")
@dataclass(frozen=True)
class DataConverter(WithSerializationContext):
"""Data converter for converting and encoding payloads to/from Python values.
This combines :py:class:`PayloadConverter` which converts values with
:py:class:`PayloadCodec` which encodes bytes.
"""
payload_converter_class: type[PayloadConverter] = DefaultPayloadConverter
"""Class to instantiate for payload conversion."""
payload_codec: PayloadCodec | None = None
"""Optional codec for encoding payload bytes."""
failure_converter_class: type[FailureConverter] = DefaultFailureConverter
"""Class to instantiate for failure conversion."""
payload_converter: PayloadConverter = dataclasses.field(init=False)
"""Payload converter created from the :py:attr:`payload_converter_class`."""
failure_converter: FailureConverter = dataclasses.field(init=False)
"""Failure converter created from the :py:attr:`failure_converter_class`."""
payload_limits: PayloadLimitsConfig = PayloadLimitsConfig()
"""Settings for payload size limits."""
external_storage: ExternalStorage | None = None
"""Options for external storage. If None, external storage is disabled.
.. warning::
This API is experimental.
"""
default: ClassVar[DataConverter]
"""Singleton default data converter."""
_payload_error_limits: _ServerPayloadErrorLimits | None = None
"""Server-reported limits for payloads."""
def __post_init__(self) -> None: # noqa: D105
object.__setattr__(self, "payload_converter", self.payload_converter_class())
object.__setattr__(self, "failure_converter", self.failure_converter_class())
async def encode(
self, values: Sequence[Any]
) -> list[temporalio.api.common.v1.Payload]:
"""Encode values into payloads.
First converts values to payloads then encodes payloads using codec.
Args:
values: Values to be converted and encoded.
Returns:
Converted and encoded payloads. Note, this does not have to be the
same number as values given, but must be at least one and cannot be
more than was given.
"""
payloads = self.payload_converter.to_payloads(values)
payloads = await self._encode_payload_sequence(payloads)
payloads = await self._external_store_payload_sequence(payloads)
self._validate_payload_limits(payloads)
return payloads
async def decode(
self,
payloads: Sequence[temporalio.api.common.v1.Payload],
type_hints: list[type] | None = None,
) -> list[Any]:
"""Decode payloads into values.
First decodes payloads using codec then converts payloads to values.
Args:
payloads: Payloads to be decoded and converted.
Returns:
Decoded and converted values.
"""
payloads = await self._external_retrieve_payload_sequence(payloads)
payloads = await self._decode_payload_sequence(payloads)
return self.payload_converter.from_payloads(payloads, type_hints)
async def encode_wrapper(
self, values: Sequence[Any]
) -> temporalio.api.common.v1.Payloads:
""":py:meth:`encode` for the
:py:class:`temporalio.api.common.v1.Payloads` wrapper.
"""
return temporalio.api.common.v1.Payloads(payloads=(await self.encode(values)))
async def decode_wrapper(
self,
payloads: temporalio.api.common.v1.Payloads | None,
type_hints: list[type] | None = None,
) -> list[Any]:
""":py:meth:`decode` for the
:py:class:`temporalio.api.common.v1.Payloads` wrapper.
"""
if not payloads or not payloads.payloads:
return []
return await self.decode(payloads.payloads, type_hints)
async def encode_failure(
self, exception: BaseException, failure: temporalio.api.failure.v1.Failure
) -> None:
"""Convert and encode failure."""
self.failure_converter.to_failure(exception, self.payload_converter, failure)
await _apply_to_failure_payloads(failure, self._transform_outbound_payloads)
async def decode_failure(
self, failure: temporalio.api.failure.v1.Failure
) -> BaseException:
"""Decode and convert failure."""
await _apply_to_failure_payloads(failure, self._transform_inbound_payloads)
return self.failure_converter.from_failure(failure, self.payload_converter)
def with_context(self, context: SerializationContext) -> Self:
"""Return an instance with context set on the component converters."""
payload_converter = self.payload_converter
payload_codec = self.payload_codec
failure_converter = self.failure_converter
external_storage = self.external_storage
if isinstance(payload_converter, WithSerializationContext):
payload_converter = payload_converter.with_context(context)
if isinstance(payload_codec, WithSerializationContext):
payload_codec = payload_codec.with_context(context)
if isinstance(failure_converter, WithSerializationContext):
failure_converter = failure_converter.with_context(context)
if isinstance(external_storage, WithSerializationContext):
external_storage = external_storage.with_context(context)
if all(
new is orig
for new, orig in [
(payload_converter, self.payload_converter),
(payload_codec, self.payload_codec),
(failure_converter, self.failure_converter),
(external_storage, self.external_storage),
]
):
return self
cloned = dataclasses.replace(self)
object.__setattr__(cloned, "payload_converter", payload_converter)
object.__setattr__(cloned, "payload_codec", payload_codec)
object.__setattr__(cloned, "failure_converter", failure_converter)
object.__setattr__(cloned, "external_storage", external_storage)
return cloned
def _with_store_context(
self, store_ctx: StorageDriverStoreContext
) -> DataConverter:
"""Return an instance with ``store_ctx`` bound into :attr:`external_storage`."""
if self.external_storage is None:
return self
return dataclasses.replace(
self,
external_storage=self.external_storage._with_store_context(store_ctx),
)
def _with_contexts(
self,
serialization_ctx: SerializationContext,
store_ctx: StorageDriverStoreContext,
) -> DataConverter:
"""Return an instance with both serialization and store contexts applied."""
return self.with_context(serialization_ctx)._with_store_context(store_ctx)
def _with_payload_error_limits(
self, limits: _ServerPayloadErrorLimits | None
) -> DataConverter:
return dataclasses.replace(self, _payload_error_limits=limits)
async def _decode_memo(
self,
source: temporalio.api.common.v1.Memo,
) -> Mapping[str, Any]:
mapping: dict[str, Any] = {}
for k, v in source.fields.items():
mapping[k] = (await self.decode([v]))[0]
return mapping
async def _decode_memo_field(
self,
source: temporalio.api.common.v1.Memo,
key: str,
default: Any,
type_hint: type | None,
) -> dict[str, Any]:
payload = source.fields.get(key)
if not payload:
if default is temporalio.common._arg_unset:
raise KeyError(f"Memo does not have a value for key {key}")
return default
return (await self.decode([payload], [type_hint] if type_hint else None))[0]
async def _encode_memo(
self, source: Mapping[str, Any]
) -> temporalio.api.common.v1.Memo:
memo = temporalio.api.common.v1.Memo()
await self._encode_memo_existing(source, memo)
return memo
async def _encode_memo_existing(
self, source: Mapping[str, Any], memo: temporalio.api.common.v1.Memo
):
for k, v in source.items():
payload = v
if not isinstance(v, temporalio.api.common.v1.Payload):
payload = (await self.encode([v]))[0]
memo.fields[k].CopyFrom(payload)
# Memos have their field payloads validated all together in one unit
DataConverter._validate_limits(
list(memo.fields.values()),
self._payload_error_limits.memo_size_error
if self._payload_error_limits
else None,
"[TMPRL1103] Attempted to upload memo with size that exceeded the error limit.",
self.payload_limits.memo_size_warning,
"[TMPRL1103] Attempted to upload memo with size that exceeded the warning limit.",
)
async def _transform_outbound_payload(
self, payload: temporalio.api.common.v1.Payload
) -> temporalio.api.common.v1.Payload:
if self.payload_codec:
payload = (await self.payload_codec.encode([payload]))[0]
if self.external_storage:
payload = await self.external_storage._store_payload(payload)
self._validate_payload_limits([payload])
return payload
async def _transform_outbound_payloads(
self, payloads: temporalio.api.common.v1.Payloads
):
if self.payload_codec:
await self.payload_codec.encode_wrapper(payloads)
if self.external_storage:
await self.external_storage._store_payloads(payloads)
self._validate_payload_limits(payloads.payloads)
async def _transform_inbound_payload(
self, payload: temporalio.api.common.v1.Payload
) -> temporalio.api.common.v1.Payload:
if self.external_storage:
payload = await self.external_storage._retrieve_payload(payload)
if self.payload_codec:
payload = (await self.payload_codec.decode([payload]))[0]
return payload
async def _transform_inbound_payloads(
self, payloads: temporalio.api.common.v1.Payloads
):
if self.external_storage:
await self.external_storage._retrieve_payloads(payloads)
else:
if any(_is_reference_payload(p) for p in payloads.payloads):
raise RuntimeError(
"[TMPRL1105] Detected externally stored payload(s) but external storage is not configured."
)
if self.payload_codec:
await self.payload_codec.decode_wrapper(payloads)
async def _encode_payload_sequence(
self, payloads: Sequence[temporalio.api.common.v1.Payload]
) -> list[temporalio.api.common.v1.Payload]:
"""Codec encode only."""
encoded_payloads = list(payloads)
if self.payload_codec:
encoded_payloads = await self.payload_codec.encode(encoded_payloads)
return encoded_payloads
async def _external_store_payload_sequence(
self, payloads: Sequence[temporalio.api.common.v1.Payload]
) -> list[temporalio.api.common.v1.Payload]:
"""External storage store, then validate payload limits."""
stored_payloads = list(payloads)
if self.external_storage:
stored_payloads = await self.external_storage._store_payload_sequence(
stored_payloads
)
return stored_payloads
async def _external_retrieve_payload_sequence(
self, payloads: Sequence[temporalio.api.common.v1.Payload]
) -> list[temporalio.api.common.v1.Payload]:
"""External storage retrieve only."""
retrieved_payloads = list(payloads)
if self.external_storage:
retrieved_payloads = await self.external_storage._retrieve_payload_sequence(
retrieved_payloads
)
else:
if any(_is_reference_payload(p) for p in retrieved_payloads):
raise RuntimeError(
"[TMPRL1105] Detected externally stored payload(s) but external storage is not configured."
)
return retrieved_payloads
async def _decode_payload_sequence(
self, payloads: Sequence[temporalio.api.common.v1.Payload]
) -> list[temporalio.api.common.v1.Payload]:
"""Codec decode only."""
decoded_payloads = list(payloads)
if self.payload_codec:
decoded_payloads = await self.payload_codec.decode(decoded_payloads)
return decoded_payloads
# Temporary shortcircuit detection while the _decode_* methods may no-op if
# a payload codec is not configured. Remove once those paths have more to them.
@property
def _decode_payload_has_effect(self) -> bool:
return self.payload_codec is not None or self.external_storage is not None
def _validate_payload_limits(
self,
payloads: Sequence[temporalio.api.common.v1.Payload],
):
DataConverter._validate_limits(
payloads,
self._payload_error_limits.payload_size_error
if self._payload_error_limits
else None,
"[TMPRL1103] Attempted to upload payloads with size that exceeded the error limit.",
self.payload_limits.payload_size_warning,
"[TMPRL1103] Attempted to upload payloads with size that exceeded the warning limit.",
)
@staticmethod
def _validate_limits(
payloads: Sequence[temporalio.api.common.v1.Payload],
error_limit: int | None,
error_message: str,
warning_limit: int,
warning_message: str,
):
total_size = sum(payload.ByteSize() for payload in payloads)
if error_limit and error_limit > 0 and total_size > error_limit:
raise _PayloadSizeError(
f"{error_message} Size: {total_size} bytes, Limit: {error_limit} bytes"
)
if warning_limit > 0 and total_size > warning_limit:
# TODO: Use a context aware logger to log extra information about workflow/activity/etc
warnings.warn(
f"{warning_message} Size: {total_size} bytes, Limit: {warning_limit} bytes",
PayloadSizeWarning,
)
def default() -> DataConverter:
"""Default data converter.
.. deprecated::
Use :py:meth:`DataConverter.default` instead.
"""
return DataConverter.default