Merge pull request #153 from stealthrocket/run-request-dispatch-id

Extract new proto fields

Author: chriso

src/dispatch/proto.py

@@ -12,6 +12,7 @@
 import tblib  # type: ignore[import-untyped]
 from google.protobuf import descriptor_pool, duration_pb2, message_factory
 
+from dispatch.id import DispatchID
 from dispatch.sdk.v1 import call_pb2 as call_pb
 from dispatch.sdk.v1 import error_pb2 as error_pb
 from dispatch.sdk.v1 import exit_pb2 as exit_pb
@@ -51,6 +52,9 @@ class Input:
     """
 
     __slots__ = (
+        "dispatch_id",
+        "parent_dispatch_id",
+        "root_dispatch_id",
         "_has_input",
         "_input",
         "_coroutine_state",
@@ -59,6 +63,10 @@ class Input:
     )
 
     def __init__(self, req: function_pb.RunRequest):
+        self.dispatch_id = req.dispatch_id
+        self.parent_dispatch_id = req.parent_dispatch_id
+        self.root_dispatch_id = req.root_dispatch_id
+
         self._has_input = req.HasField("input")
         if self._has_input:
             if req.input.Is(google.protobuf.wrappers_pb2.BytesValue.DESCRIPTOR):
@@ -285,6 +293,7 @@ class CallResult:
     correlation_id: Optional[int] = None
     output: Optional[Any] = None
     error: Optional[Error] = None
+    dispatch_id: DispatchID = ""
 
     def _as_proto(self) -> call_pb.CallResult:
         output_any = None
@@ -295,7 +304,10 @@ def _as_proto(self) -> call_pb.CallResult:
             error_proto = self.error._as_proto()
 
         return call_pb.CallResult(
-            correlation_id=self.correlation_id, output=output_any, error=error_proto
+            correlation_id=self.correlation_id,
+            output=output_any,
+            error=error_proto,
+            dispatch_id=self.dispatch_id,
         )
 
     @classmethod
@@ -308,7 +320,10 @@ def _from_proto(cls, proto: call_pb.CallResult) -> CallResult:
             error = Error._from_proto(proto.error)
 
         return CallResult(
-            correlation_id=proto.correlation_id, output=output, error=error
+            correlation_id=proto.correlation_id,
+            output=output,
+            error=error,
+            dispatch_id=proto.dispatch_id,
         )
 
     @classmethod

src/dispatch/sdk/v1/call_pb2.py

@@ -35,16 +35,19 @@ class Call(_message.Message):
     ) -> None: ...
 
 class CallResult(_message.Message):
-    __slots__ = ("correlation_id", "output", "error")
+    __slots__ = ("correlation_id", "output", "error", "dispatch_id")
     CORRELATION_ID_FIELD_NUMBER: _ClassVar[int]
     OUTPUT_FIELD_NUMBER: _ClassVar[int]
     ERROR_FIELD_NUMBER: _ClassVar[int]
+    DISPATCH_ID_FIELD_NUMBER: _ClassVar[int]
     correlation_id: int
     output: _any_pb2.Any
     error: _error_pb2.Error
+    dispatch_id: str
     def __init__(
         self,
         correlation_id: _Optional[int] = ...,
         output: _Optional[_Union[_any_pb2.Any, _Mapping]] = ...,
         error: _Optional[_Union[_error_pb2.Error, _Mapping]] = ...,
+        dispatch_id: _Optional[str] = ...,
     ) -> None: ...

src/dispatch/sdk/v1/call_pb2.pyi

@@ -14,18 +14,34 @@ from dispatch.sdk.v1 import status_pb2 as _status_pb2
 DESCRIPTOR: _descriptor.FileDescriptor
 
 class RunRequest(_message.Message):
-    __slots__ = ("function", "input", "poll_result")
+    __slots__ = (
+        "function",
+        "input",
+        "poll_result",
+        "dispatch_id",
+        "parent_dispatch_id",
+        "root_dispatch_id",
+    )
     FUNCTION_FIELD_NUMBER: _ClassVar[int]
     INPUT_FIELD_NUMBER: _ClassVar[int]
     POLL_RESULT_FIELD_NUMBER: _ClassVar[int]
+    DISPATCH_ID_FIELD_NUMBER: _ClassVar[int]
+    PARENT_DISPATCH_ID_FIELD_NUMBER: _ClassVar[int]
+    ROOT_DISPATCH_ID_FIELD_NUMBER: _ClassVar[int]
     function: str
     input: _any_pb2.Any
     poll_result: _poll_pb2.PollResult
+    dispatch_id: str
+    parent_dispatch_id: str
+    root_dispatch_id: str
     def __init__(
         self,
         function: _Optional[str] = ...,
         input: _Optional[_Union[_any_pb2.Any, _Mapping]] = ...,
         poll_result: _Optional[_Union[_poll_pb2.PollResult, _Mapping]] = ...,
+        dispatch_id: _Optional[str] = ...,
+        parent_dispatch_id: _Optional[str] = ...,
+        root_dispatch_id: _Optional[str] = ...,
     ) -> None: ...
 
 class RunResponse(_message.Message):

src/dispatch/sdk/v1/dispatch_pb2.py

@@ -110,6 +110,8 @@ def Dispatch(self, request: dispatch_pb.DispatchRequest, context):
                 run_request = function_pb.RunRequest(
                     function=call.function,
                     input=call.input,
+                    dispatch_id=dispatch_id,
+                    root_dispatch_id=dispatch_id,
                 )
                 self.queue.append((dispatch_id, run_request, CallType.CALL))
 
@@ -207,6 +209,8 @@ def dispatch_calls(self):
                 assert dispatch_id not in self.pollers
                 poller = Poller(
                     id=dispatch_id,
+                    parent_id=request.parent_dispatch_id,
+                    root_id=request.root_dispatch_id,
                     function=request.function,
                     coroutine_state=response.poll.coroutine_state,
                     waiting={},
@@ -219,6 +223,9 @@ def dispatch_calls(self):
                     child_request = function_pb.RunRequest(
                         function=call.function,
                         input=call.input,
+                        dispatch_id=child_dispatch_id,
+                        parent_dispatch_id=request.dispatch_id,
+                        root_dispatch_id=request.root_dispatch_id,
                     )
 
                     _next_queue.append(
@@ -239,6 +246,9 @@ def dispatch_calls(self):
                     tail_call_request = function_pb.RunRequest(
                         function=tail_call.function,
                         input=tail_call.input,
+                        dispatch_id=request.dispatch_id,
+                        parent_dispatch_id=request.parent_dispatch_id,
+                        root_dispatch_id=request.root_dispatch_id,
                     )
                     _next_queue.append((dispatch_id, tail_call_request, CallType.CALL))
 
@@ -269,6 +279,9 @@ def dispatch_calls(self):
                             len(poller.results),
                         )
                         poll_results_request = function_pb.RunRequest(
+                            dispatch_id=poller.id,
+                            parent_dispatch_id=poller.parent_id,
+                            root_dispatch_id=poller.root_id,
                             function=poller.function,
                             poll_result=poll_pb.PollResult(
                                 coroutine_state=poller.coroutine_state,
@@ -349,6 +362,9 @@ def __exit__(self, exc_type, exc_val, exc_tb):
 @dataclass
 class Poller:
     id: DispatchID
+    parent_id: DispatchID
+    root_id: DispatchID
+
     function: str
 
     coroutine_state: bytes