Repository navigation
Expand file tree
/
Copy path_interceptor.py
More file actions
583 lines (466 loc) · 19.6 KB
/
Copy path_interceptor.py
File metadata and controls
583 lines (466 loc) · 19.6 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
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
"""Worker interceptor."""
from __future__ import annotations
import concurrent.futures
from collections.abc import Awaitable, Callable, Mapping, MutableMapping, Sequence
from dataclasses import dataclass, field
from datetime import timedelta
from typing import (
Any,
Generic,
NoReturn,
)
import nexusrpc
from nexusrpc import InputT, OutputT
import temporalio.activity
import temporalio.api.common.v1
import temporalio.common
import temporalio.nexus
import temporalio.nexus._util
import temporalio.workflow
from temporalio.nexus.system.workflow_service._system_nexus_interceptor import (
_SystemNexusWorkflowOutboundInterceptorBase,
)
from temporalio.workflow import ContinueAsNewVersioningBehavior, VersioningIntent
class Interceptor:
"""Interceptor for workers.
This should be extended by any worker interceptors.
"""
def intercept_activity(
self, next: ActivityInboundInterceptor
) -> ActivityInboundInterceptor:
"""Method called for intercepting an activity.
Args:
next: The underlying inbound interceptor this interceptor should
delegate to.
Returns:
The new interceptor that will be used to for the activity.
"""
return next
def workflow_interceptor_class(
self,
input: WorkflowInterceptorClassInput, # type:ignore[reportUnusedParameter]
) -> type[WorkflowInboundInterceptor] | None:
"""Class that will be instantiated and used to intercept workflows.
This method is called on workflow start. The class must have the same
init as :py:meth:`WorkflowInboundInterceptor.__init__`. The input can be
altered to do things like add additional extern functions.
Args:
input: Input to this method that contains mutable properties that
can be altered by this interceptor.
Returns:
The class to construct to intercept each workflow.
"""
return None
def intercept_nexus_operation(
self, next: NexusOperationInboundInterceptor
) -> NexusOperationInboundInterceptor:
"""Method called for intercepting a Nexus operation.
Args:
next: The underlying inbound this interceptor
should delegate to.
Returns:
The new interceptor that should be used for the Nexus operation.
"""
return next
@dataclass(frozen=True)
class WorkflowInterceptorClassInput:
"""Input for :py:meth:`Interceptor.workflow_interceptor_class`."""
unsafe_extern_functions: MutableMapping[str, Callable]
"""Set of external functions that can be called from the sandbox.
.. warning::
Exposing external functions to the workflow sandbox is dangerous and
should be avoided. Use at your own risk.
.. warning::
This API is experimental and subject to removal.
"""
@dataclass
class ExecuteActivityInput:
"""Input for :py:meth:`ActivityInboundInterceptor.execute_activity`."""
fn: Callable[..., Any]
args: Sequence[Any]
executor: concurrent.futures.Executor | None
headers: Mapping[str, temporalio.api.common.v1.Payload]
class ActivityInboundInterceptor:
"""Inbound interceptor to wrap outbound creation and activity execution.
This should be extended by any activity inbound interceptors.
"""
def __init__(self, next: ActivityInboundInterceptor) -> None:
"""Create the inbound interceptor.
Args:
next: The next interceptor in the chain. The default implementation
of all calls is to delegate to the next interceptor.
"""
self.next = next
def init(self, outbound: ActivityOutboundInterceptor) -> None:
"""Initialize with an outbound interceptor.
To add a custom outbound interceptor, wrap the given interceptor before
sending to the next ``init`` call.
"""
self.next.init(outbound)
async def execute_activity(self, input: ExecuteActivityInput) -> Any:
"""Called to invoke the activity."""
return await self.next.execute_activity(input)
class ActivityOutboundInterceptor:
"""Outbound interceptor to wrap calls made from within activities.
This should be extended by any activity outbound interceptors.
"""
def __init__(self, next: ActivityOutboundInterceptor) -> None:
"""Create the outbound interceptor.
Args:
next: The next interceptor in the chain. The default implementation
of all calls is to delegate to the next interceptor.
"""
self.next = next
def info(self) -> temporalio.activity.Info:
"""Called for every :py:func:`temporalio.activity.info` call."""
return self.next.info()
def heartbeat(self, *details: Any) -> None:
"""Called for every :py:func:`temporalio.activity.heartbeat` call."""
self.next.heartbeat(*details)
@dataclass
class ContinueAsNewInput:
"""Input for :py:meth:`WorkflowOutboundInterceptor.continue_as_new`."""
workflow: str | None
args: Sequence[Any]
task_queue: str | None
run_timeout: timedelta | None
task_timeout: timedelta | None
backoff_start_interval: timedelta | None
retry_policy: temporalio.common.RetryPolicy | None
memo: Mapping[str, Any] | None
search_attributes: None | (
temporalio.common.SearchAttributes | temporalio.common.TypedSearchAttributes
)
headers: Mapping[str, temporalio.api.common.v1.Payload]
versioning_intent: VersioningIntent | None
initial_versioning_behavior: ContinueAsNewVersioningBehavior | None
event_groups: Sequence[temporalio.workflow.EventGroup] | None = None
""".. warning::
Event Groups is an experimental API and may change without notice.
"""
# The types may be absent
arg_types: list[type] | None = None
@dataclass
class ExecuteWorkflowInput:
"""Input for :py:meth:`WorkflowInboundInterceptor.execute_workflow`."""
type: type
# Note, this is an unbound method
run_fn: Callable[..., Awaitable[Any]]
args: Sequence[Any]
headers: Mapping[str, temporalio.api.common.v1.Payload]
@dataclass
class HandleSignalInput:
"""Input for :py:meth:`WorkflowInboundInterceptor.handle_signal`."""
signal: str
args: Sequence[Any]
headers: Mapping[str, temporalio.api.common.v1.Payload]
@dataclass
class HandleQueryInput:
"""Input for :py:meth:`WorkflowInboundInterceptor.handle_query`."""
id: str
query: str
args: Sequence[Any]
headers: Mapping[str, temporalio.api.common.v1.Payload]
@dataclass
class HandleUpdateInput:
"""Input for :py:meth:`WorkflowInboundInterceptor.handle_update_validator`
and :py:meth:`WorkflowInboundInterceptor.handle_update_handler`.
"""
id: str
update: str
args: Sequence[Any]
headers: Mapping[str, temporalio.api.common.v1.Payload]
@dataclass
class SignalChildWorkflowInput:
"""Input for :py:meth:`WorkflowOutboundInterceptor.signal_child_workflow`."""
signal: str
args: Sequence[Any]
child_workflow_id: str
headers: Mapping[str, temporalio.api.common.v1.Payload]
event_groups: Sequence[temporalio.workflow.EventGroup] | None = None
""".. warning::
Event Groups is an experimental API and may change without notice.
"""
@dataclass
class SignalExternalWorkflowInput:
"""Input for :py:meth:`WorkflowOutboundInterceptor.signal_external_workflow`."""
signal: str
args: Sequence[Any]
namespace: str
workflow_id: str
workflow_run_id: str | None
headers: Mapping[str, temporalio.api.common.v1.Payload]
event_groups: Sequence[temporalio.workflow.EventGroup] | None = None
""".. warning::
Event Groups is an experimental API and may change without notice.
"""
@dataclass
class StartActivityInput:
"""Input for :py:meth:`WorkflowOutboundInterceptor.start_activity`."""
activity: str
args: Sequence[Any]
activity_id: str | None
task_queue: str | None
schedule_to_close_timeout: timedelta | None
schedule_to_start_timeout: timedelta | None
start_to_close_timeout: timedelta | None
heartbeat_timeout: timedelta | None
retry_policy: temporalio.common.RetryPolicy | None
cancellation_type: temporalio.workflow.ActivityCancellationType
headers: Mapping[str, temporalio.api.common.v1.Payload]
disable_eager_execution: bool
versioning_intent: VersioningIntent | None
summary: str | None
event_groups: Sequence[temporalio.workflow.EventGroup] | None = field(
default=None, kw_only=True
)
""".. warning::
Event Groups is an experimental API and may change without notice.
"""
priority: temporalio.common.Priority
# The types may be absent
arg_types: list[type] | None
ret_type: type | None
@dataclass
class StartChildWorkflowInput:
"""Input for :py:meth:`WorkflowOutboundInterceptor.start_child_workflow`."""
workflow: str
args: Sequence[Any]
id: str
task_queue: str | None
cancellation_type: temporalio.workflow.ChildWorkflowCancellationType
parent_close_policy: temporalio.workflow.ParentClosePolicy
execution_timeout: timedelta | None
run_timeout: timedelta | None
task_timeout: timedelta | None
id_reuse_policy: temporalio.common.WorkflowIDReusePolicy
retry_policy: temporalio.common.RetryPolicy | None
cron_schedule: str
memo: Mapping[str, Any] | None
search_attributes: None | (
temporalio.common.SearchAttributes | temporalio.common.TypedSearchAttributes
)
headers: Mapping[str, temporalio.api.common.v1.Payload]
versioning_intent: VersioningIntent | None
static_summary: str | None
static_details: str | None
versioning_override: temporalio.common.VersioningOverride | None = field(
default=None, kw_only=True
)
""".. warning::
Child workflow versioning overrides are experimental and may change without notice.
"""
event_groups: Sequence[temporalio.workflow.EventGroup] | None = field(
default=None, kw_only=True
)
""".. warning::
Event Groups is an experimental API and may change without notice.
"""
priority: temporalio.common.Priority
# The types may be absent
arg_types: list[type] | None
ret_type: type | None
@dataclass
class StartNexusOperationInput(Generic[InputT, OutputT]):
"""Input for :py:meth:`WorkflowOutboundInterceptor.start_nexus_operation`."""
endpoint: str
service: str
operation: nexusrpc.Operation[InputT, OutputT] | str | Callable[..., Any]
input: InputT
schedule_to_close_timeout: timedelta | None
schedule_to_start_timeout: timedelta | None
start_to_close_timeout: timedelta | None
cancellation_type: temporalio.workflow.NexusOperationCancellationType
headers: Mapping[str, str] | None
summary: str | None
event_groups: Sequence[temporalio.workflow.EventGroup] | None = None
""".. warning::
Event Groups is an experimental API and may change without notice.
"""
output_type: type[OutputT] | None = None
def __post_init__(self) -> None:
"""Initialize operation-specific attributes after dataclass creation."""
if isinstance(self.operation, nexusrpc.Operation):
self.output_type = self.operation.output_type
elif callable(self.operation):
_, op = temporalio.nexus._util.get_operation_factory(self.operation)
if isinstance(op, nexusrpc.Operation):
self.output_type = op.output_type
else:
raise ValueError(
f"Operation callable is not a Nexus operation: {self.operation}"
)
elif isinstance(self.operation, str):
pass
else:
raise ValueError(f"Operation is not a Nexus operation: {self.operation}")
@property
def operation_name(self) -> str:
"""Get the name of the Nexus operation."""
if isinstance(self.operation, nexusrpc.Operation):
return self.operation.name
elif isinstance(self.operation, str):
return self.operation
elif callable(self.operation):
_, op = temporalio.nexus._util.get_operation_factory(self.operation)
if isinstance(op, nexusrpc.Operation):
return op.name
else:
raise ValueError(
f"Operation callable is not a Nexus operation: {self.operation}"
)
else:
raise ValueError(f"Operation is not a Nexus operation: {self.operation}")
@dataclass
class StartLocalActivityInput:
"""Input for :py:meth:`WorkflowOutboundInterceptor.start_local_activity`."""
activity: str
args: Sequence[Any]
activity_id: str | None
schedule_to_close_timeout: timedelta | None
schedule_to_start_timeout: timedelta | None
start_to_close_timeout: timedelta | None
retry_policy: temporalio.common.RetryPolicy | None
local_retry_threshold: timedelta | None
cancellation_type: temporalio.workflow.ActivityCancellationType
headers: Mapping[str, temporalio.api.common.v1.Payload]
summary: str | None
event_groups: Sequence[temporalio.workflow.EventGroup] | None = None
""".. warning::
Event Groups is an experimental API and may change without notice.
"""
# The types may be absent
arg_types: list[type] | None = None
ret_type: type | None = None
class WorkflowInboundInterceptor:
"""Inbound interceptor to wrap outbound creation, workflow execution, and
signal/query handling.
This should be extended by any workflow inbound interceptors.
"""
def __init__(self, next: WorkflowInboundInterceptor) -> None:
"""Create the inbound interceptor.
Args:
next: The next interceptor in the chain. The default implementation
of all calls is to delegate to the next interceptor.
"""
self.next = next
def init(self, outbound: WorkflowOutboundInterceptor) -> None:
"""Initialize with an outbound interceptor.
To add a custom outbound interceptor, wrap the given interceptor before
sending to the next ``init`` call.
"""
self.next.init(outbound)
async def execute_workflow(self, input: ExecuteWorkflowInput) -> Any:
"""Called to run the workflow."""
return await self.next.execute_workflow(input)
async def handle_signal(self, input: HandleSignalInput) -> None:
"""Called to handle a signal."""
return await self.next.handle_signal(input)
async def handle_query(self, input: HandleQueryInput) -> Any:
"""Called to handle a query."""
return await self.next.handle_query(input)
def handle_update_validator(self, input: HandleUpdateInput) -> None:
"""Called to handle an update's validation stage."""
self.next.handle_update_validator(input)
async def handle_update_handler(self, input: HandleUpdateInput) -> Any:
"""Called to handle an update's handler."""
return await self.next.handle_update_handler(input)
class WorkflowOutboundInterceptor(_SystemNexusWorkflowOutboundInterceptorBase):
"""Outbound interceptor to wrap calls made from within workflows.
This should be extended by any workflow outbound interceptors.
"""
next: WorkflowOutboundInterceptor
def __init__(self, next: WorkflowOutboundInterceptor) -> None:
"""Create the outbound interceptor.
Args:
next: The next interceptor in the chain. The default implementation
of all calls is to delegate to the next interceptor.
"""
self.next = next
def _next_system_nexus_interceptor(
self,
) -> _SystemNexusWorkflowOutboundInterceptorBase:
return self.next
def continue_as_new(self, input: ContinueAsNewInput) -> NoReturn:
"""Called for every :py:func:`temporalio.workflow.continue_as_new` call."""
self.next.continue_as_new(input)
def info(self) -> temporalio.workflow.Info:
"""Called for every :py:func:`temporalio.workflow.info` call."""
return self.next.info()
async def signal_child_workflow(self, input: SignalChildWorkflowInput) -> None:
"""Called for every
:py:meth:`temporalio.workflow.ChildWorkflowHandle.signal` call.
"""
return await self.next.signal_child_workflow(input)
async def signal_external_workflow(
self, input: SignalExternalWorkflowInput
) -> None:
"""Called for every
:py:meth:`temporalio.workflow.ExternalWorkflowHandle.signal` call.
"""
return await self.next.signal_external_workflow(input)
def start_activity(
self, input: StartActivityInput
) -> temporalio.workflow.ActivityHandle[Any]:
"""Called for every :py:func:`temporalio.workflow.start_activity` and
:py:func:`temporalio.workflow.execute_activity` call.
"""
return self.next.start_activity(input)
async def start_child_workflow(
self, input: StartChildWorkflowInput
) -> temporalio.workflow.ChildWorkflowHandle[Any, Any]:
"""Called for every :py:func:`temporalio.workflow.start_child_workflow`
and :py:func:`temporalio.workflow.execute_child_workflow` call.
"""
return await self.next.start_child_workflow(input)
def start_local_activity(
self, input: StartLocalActivityInput
) -> temporalio.workflow.ActivityHandle[Any]:
"""Called for every :py:func:`temporalio.workflow.start_local_activity`
and :py:func:`temporalio.workflow.execute_local_activity` call.
"""
return self.next.start_local_activity(input)
async def start_nexus_operation(
self, input: StartNexusOperationInput[InputT, OutputT]
) -> temporalio.workflow.NexusOperationHandle[OutputT]:
"""Called for every :py:func:`temporalio.workflow.NexusClient.start_operation` call."""
return await self.next.start_nexus_operation(input)
async def start_system_nexus_operation(
self, input: StartNexusOperationInput[InputT, OutputT]
) -> temporalio.workflow.NexusOperationHandle[OutputT]:
"""Intercept a Temporal System Nexus operation started by a workflow."""
return await self.next.start_system_nexus_operation(input)
@dataclass
class ExecuteNexusOperationStartInput:
"""Input for :pyt:meth:`NexusOperationInboundInterceptor.start_operation"""
ctx: nexusrpc.handler.StartOperationContext
input: Any
@dataclass
class ExecuteNexusOperationCancelInput:
"""Input for :pyt:meth:`NexusOperationInboundInterceptor.cancel_operation"""
ctx: nexusrpc.handler.CancelOperationContext
token: str
class NexusOperationInboundInterceptor:
"""Inbound interceptor to wrap Nexus operation starting and cancelling.
This should be extended by any Nexus operation inbound interceptors.
"""
def __init__(self, next: NexusOperationInboundInterceptor) -> None:
"""Create the inbound interceptor.
Args:
next: The next interceptor in the chain. The default implementation
of all calls is to delegate to the next interceptor.
"""
self.next = next
async def execute_nexus_operation_start(
self, input: ExecuteNexusOperationStartInput
) -> (
nexusrpc.handler.StartOperationResultSync[Any]
| nexusrpc.handler.StartOperationResultAsync
):
"""Called to start a Nexus operation"""
return await self.next.execute_nexus_operation_start(input)
async def execute_nexus_operation_cancel(
self, input: ExecuteNexusOperationCancelInput
) -> None:
"""Called to cancel an in progress Nexus operation"""
return await self.next.execute_nexus_operation_cancel(input)