feat: udsink gRPC stream (#34)

Author: jy4096

examples/sink/simplesink/README.md

@@ -3,7 +3,7 @@
 1. Build the docker image, and push
    ```shell
    make image
-   # Privilege requried
+   # Privilege required
    docker push quay.io/numaio/simplesink-example:python
    ```
 

pynumaflow/sink/_udfunction_pb2.pyi

@@ -51,12 +51,6 @@ class Datum(_message.Message):
         watermark: _Optional[_timestamp_pb2.Timestamp] = ...,
     ) -> None: ...
 
-class DatumList(_message.Message):
-    __slots__ = ["elements"]
-    ELEMENTS_FIELD_NUMBER: _ClassVar[int]
-    elements: List[Datum]
-    def __init__(self, elements: _Optional[List[Datum]]) -> None: ...
-
 class Response(_message.Message):
     __slots__ = ["id", "success", "err_msg"]
     ID_FIELD_NUMBER: _ClassVar[int]

pynumaflow/sink/generated/udsink_pb2.py

@@ -15,9 +15,9 @@ def __init__(self, channel):
         Args:
             channel: A grpc.Channel.
         """
-        self.SinkFn = channel.unary_unary(
+        self.SinkFn = channel.stream_unary(
             "/sink.v1.UserDefinedSink/SinkFn",
-            request_serializer=udsink__pb2.DatumList.SerializeToString,
+            request_serializer=udsink__pb2.Datum.SerializeToString,
             response_deserializer=udsink__pb2.ResponseList.FromString,
         )
         self.IsReady = channel.unary_unary(
@@ -30,7 +30,7 @@ def __init__(self, channel):
 class UserDefinedSinkServicer(object):
     """Missing associated documentation comment in .proto file."""
 
-    def SinkFn(self, request, context):
+    def SinkFn(self, request_iterator, context):
         """SinkFn writes the Datum to a user defined sink."""
         context.set_code(grpc.StatusCode.UNIMPLEMENTED)
         context.set_details("Method not implemented!")
@@ -45,9 +45,9 @@ def IsReady(self, request, context):
 
 def add_UserDefinedSinkServicer_to_server(servicer, server):
     rpc_method_handlers = {
-        "SinkFn": grpc.unary_unary_rpc_method_handler(
+        "SinkFn": grpc.stream_unary_rpc_method_handler(
             servicer.SinkFn,
-            request_deserializer=udsink__pb2.DatumList.FromString,
+            request_deserializer=udsink__pb2.Datum.FromString,
             response_serializer=udsink__pb2.ResponseList.SerializeToString,
         ),
         "IsReady": grpc.unary_unary_rpc_method_handler(
@@ -68,7 +68,7 @@ class UserDefinedSink(object):
 
     @staticmethod
     def SinkFn(
-        request,
+        request_iterator,
         target,
         options=(),
         channel_credentials=None,
@@ -79,11 +79,11 @@ def SinkFn(
         timeout=None,
         metadata=None,
     ):
-        return grpc.experimental.unary_unary(
-            request,
+        return grpc.experimental.stream_unary(
+            request_iterator,
             target,
             "/sink.v1.UserDefinedSink/SinkFn",
-            udsink__pb2.DatumList.SerializeToString,
+            udsink__pb2.Datum.SerializeToString,
             udsink__pb2.ResponseList.FromString,
             options,
             channel_credentials,

pynumaflow/sink/generated/udsink_pb2_grpc.py

@@ -3,7 +3,7 @@
 import multiprocessing
 import os
 from concurrent.futures import ThreadPoolExecutor
-from typing import Callable, List
+from typing import Callable, Iterator
 
 import grpc
 from google.protobuf import empty_pb2 as _empty_pb2
@@ -21,7 +21,7 @@
 
 _LOGGER = logging.getLogger(__name__)
 
-UDSinkCallable = Callable[[List[Datum]], Responses]
+UDSinkCallable = Callable[[Iterator[Datum]], Responses]
 _PROCESS_COUNT = multiprocessing.cpu_count()
 MAX_THREADS = int(os.getenv("MAX_THREADS", 0)) or (_PROCESS_COUNT * 4)
 
@@ -41,7 +41,7 @@ class UserDefinedSinkServicer(udsink_pb2_grpc.UserDefinedSinkServicer):
     Example invocation:
     >>> from typing import List
     >>> from pynumaflow.sink import Datum, Responses, Response, UserDefinedSinkServicer
-    >>> def udsink_handler(datums: List[Datum]) -> Responses:
+    >>> def udsink_handler(datums: Iterator[Datum]) -> Responses:
     ...   responses = Responses()
     ...   for msg in datums:
     ...     responses.append(Response.as_success(msg.id))
@@ -69,20 +69,20 @@ def __init__(
         ]
 
     def SinkFn(
-        self, request: udsink_pb2.DatumList, context: NumaflowServicerContext
+        self, request_iterator: Iterator[Datum], context: NumaflowServicerContext
     ) -> udsink_pb2.ResponseList:
         """
         Applies a sink function to a list of datum elements.
         The pascal case function name comes from the generated udsink_pb2_grpc.py file.
         """
         # if there is an exception, we will mark all the responses as a failure
         try:
-            rspns = self.__sink_handler(request.elements)
+            rspns = self.__sink_handler(request_iterator)
         except Exception as err:
             err_msg = "UDSinkError: %r" % err
             _LOGGER.critical(err_msg, exc_info=True)
             rspns = Responses()
-            for _datum in request.elements:
+            for _datum in request_iterator:
                 rspns.append(Response.as_failure(_datum.id, err_msg))
 
         responses = []

pynumaflow/sink/server.py

@@ -1,6 +1,6 @@
 import unittest
 from datetime import datetime, timezone
-from typing import List
+from typing import Iterator
 
 from google.protobuf import empty_pb2 as _empty_pb2
 from google.protobuf import timestamp_pb2 as _timestamp_pb2
@@ -11,7 +11,7 @@
 from pynumaflow.sink.generated import udsink_pb2
 
 
-def udsink_handler(datums: List[Datum]) -> Responses:
+def udsink_handler(datums: Iterator[Datum]) -> Responses:
     responses = Responses()
     for msg in datums:
         if "err" in msg.value.decode("utf-8"):
@@ -21,7 +21,7 @@ def udsink_handler(datums: List[Datum]) -> Responses:
     return responses
 
 
-def err_udsink_handler(_: List[Datum]) -> Responses:
+def err_udsink_handler(_: Iterator[Datum]) -> Responses:
     raise RuntimeError("Something is fishy!")
 
 
@@ -91,17 +91,18 @@ def test_udsink_err(self):
             ),
         ]
 
-        request = udsink_pb2.DatumList(elements=test_datums)
-
-        method = self.test_server.invoke_unary_unary(
+        method = self.test_server.invoke_stream_unary(
             method_descriptor=(
                 udsink_pb2.DESCRIPTOR.services_by_name["UserDefinedSink"].methods_by_name["SinkFn"]
             ),
             invocation_metadata={},
-            request=request,
             timeout=1,
         )
 
+        method.send_request(test_datums[0])
+        method.send_request(test_datums[1])
+        method.requests_closed()
+
         response, metadata, code, details = method.termination()
         self.assertEqual(2, len(response.responses))
         self.assertEqual("test_id_0", response.responses[0].id)
@@ -133,17 +134,18 @@ def test_forward_message(self):
             ),
         ]
 
-        request = udsink_pb2.DatumList(elements=test_datums)
-
-        method = self.test_server.invoke_unary_unary(
+        method = self.test_server.invoke_stream_unary(
             method_descriptor=(
                 udsink_pb2.DESCRIPTOR.services_by_name["UserDefinedSink"].methods_by_name["SinkFn"]
             ),
             invocation_metadata={},
-            request=request,
             timeout=1,
         )
 
+        method.send_request(test_datums[0])
+        method.send_request(test_datums[1])
+        method.requests_closed()
+
         response, metadata, code, details = method.termination()
         self.assertEqual(2, len(response.responses))
         self.assertEqual("test_id_0", response.responses[0].id)