Attach invocation

Author: igalshilman

python/restate/context.py

@@ -228,6 +228,12 @@ def cancel(self, invocation_id: str):
         Cancels the invocation with the given id.
         """
 
+    @abc.abstractmethod
+    def attach_invocation(self, invocation_id: str, serde: Serde[T] = JsonSerde()) -> T:
+        """
+        Attaches the invocation with the given id.
+        """
+
 
 class ObjectContext(Context, KeyValueStore):
     """

python/restate/server_context.py

@@ -436,3 +436,18 @@ def cancel(self, invocation_id: str):
         if invocation_id is None:
             raise ValueError("invocation_id cannot be None")
         self.vm.sys_cancel(invocation_id)
+
+    def attach_invocation(self, invocation_id: str, serde: Serde[T] = JsonSerde()) -> T:
+        if invocation_id is None:
+            raise ValueError("invocation_id cannot be None")
+        assert serde is not None
+        handle = self.vm.attach_invocation(invocation_id)
+        coro = self.create_poll_or_cancel_coroutine(handle)
+
+        async def await_point():
+            """Wait for this handle to be resolved."""
+            res = await coro
+            assert res is not None
+            return serde.deserialize(res)
+
+        return await_point()

src/lib.rs

@@ -658,6 +658,17 @@ impl PyVM {
             .map_err(Into::into)
     }
 
+    fn attach_invocation(
+        mut self_: PyRefMut<'_, Self>,
+        invocation_id: String,
+    ) -> Result<PyNotificationHandle, PyVMError> {
+        self_
+            .vm
+            .sys_attach_invocation(restate_sdk_shared_core::AttachInvocationTarget::InvocationId(invocation_id))
+            .map(Into::into)
+            .map_err(Into::into)
+    }
+
     fn sys_end(mut self_: PyRefMut<'_, Self>) -> Result<(), PyVMError> {
         self_.vm.sys_end().map(Into::into).map_err(Into::into)
     }