support metadata in sink

Signed-off-by: Sreekanth <[email protected]>

Author: BulkBeing

packages/pynumaflow/pynumaflow/sinker/_dtypes.py

@@ -6,6 +6,8 @@
 from collections.abc import Sequence, Iterator
 from warnings import warn
 
+from pynumaflow._metadata import SystemMetadata, UserMetadata
+
 R = TypeVar("R", bound="Response")
 Rs = TypeVar("Rs", bound="Responses")
 
@@ -120,14 +122,25 @@ class Datum:
     ...    )
     """
 
-    __slots__ = ("_keys", "_id", "_value", "_event_time", "_watermark", "_headers")
+    __slots__ = (
+        "_keys",
+        "_id",
+        "_value",
+        "_event_time",
+        "_watermark",
+        "_headers",
+        "_user_metadata",
+        "_system_metadata",
+    )
 
     _keys: list[str]
     _id: str
     _value: bytes
     _event_time: datetime
     _watermark: datetime
     _headers: dict[str, str]
+    _user_metadata: UserMetadata
+    _system_metadata: SystemMetadata
 
     def __init__(
         self,
@@ -137,6 +150,8 @@ def __init__(
         event_time: datetime,
         watermark: datetime,
         headers: Optional[dict[str, str]] = None,
+        user_metadata: Optional[UserMetadata] = None,
+        system_metadata: Optional[SystemMetadata] = None,
     ):
         self._keys = keys
         self._id = sink_msg_id or ""
@@ -148,6 +163,8 @@ def __init__(
             raise TypeError(f"Wrong data type: {type(watermark)} for Datum.watermark")
         self._watermark = watermark
         self._headers = headers or {}
+        self._user_metadata = user_metadata or UserMetadata()
+        self._system_metadata = system_metadata or SystemMetadata()
 
     def __str__(self):
         value_string = self._value.decode("utf-8")
@@ -192,6 +209,16 @@ def headers(self) -> dict[str, str]:
         """Returns the headers of the event."""
         return self._headers.copy()
 
+    @property
+    def user_metadata(self) -> UserMetadata:
+        """Returns the user metadata of the event."""
+        return self._user_metadata
+
+    @property
+    def system_metadata(self) -> SystemMetadata:
+        """Returns the system metadata of the event."""
+        return self._system_metadata
+
 
 class Sinker(metaclass=ABCMeta):
     """

packages/pynumaflow/pynumaflow/sinker/servicer/async_servicer.py

@@ -1,5 +1,5 @@
 import asyncio
-from collections.abc import AsyncIterable
+from collections.abc import AsyncIterator
 
 from google.protobuf import empty_pb2 as _empty_pb2
 from pynumaflow.shared.asynciter import NonBlockingIterator
@@ -33,7 +33,7 @@ def __init__(
 
     async def SinkFn(
         self,
-        request_iterator: AsyncIterable[sink_pb2.SinkRequest],
+        request_iterator: AsyncIterator[sink_pb2.SinkRequest],
         context: NumaflowServicerContext,
     ) -> sink_pb2.SinkResponse:
         """
@@ -89,7 +89,7 @@ async def SinkFn(
             return
 
     async def __invoke_sink(
-        self, request_queue: AsyncIterable[Datum], context: NumaflowServicerContext
+        self, request_queue: AsyncIterator[Datum], context: NumaflowServicerContext
     ):
         try:
             # invoke the user function with the request queue

packages/pynumaflow/pynumaflow/sinker/servicer/sync_servicer.py

@@ -1,4 +1,4 @@
-from collections.abc import Iterable
+from collections.abc import Iterator
 
 
 from pynumaflow._constants import _LOGGER, STREAM_EOF
@@ -26,8 +26,8 @@ def __init__(self, handler: SinkSyncCallable):
         self.handler: SinkSyncCallable = handler
 
     def SinkFn(
-        self, request_iterator: Iterable[sink_pb2.SinkRequest], context: NumaflowServicerContext
-    ) -> Iterable[sink_pb2.SinkResponse]:
+        self, request_iterator: Iterator[sink_pb2.SinkRequest], context: NumaflowServicerContext
+    ) -> Iterator[sink_pb2.SinkResponse]:
         """
         Applies a sink function to datum elements.
         """

packages/pynumaflow/pynumaflow/sinker/servicer/utils.py

@@ -1,3 +1,4 @@
+from pynumaflow._metadata import _user_and_system_metadata_from_proto
 from pynumaflow.proto.sinker import sink_pb2
 from pynumaflow.sinker._dtypes import Response, Datum, Responses
 
@@ -47,13 +48,16 @@ def datum_from_sink_req(d: sink_pb2.SinkRequest) -> Datum:
     Returns:
     Datum: A Datum object populated with the data from the input SinkRequest object.
     """
+    user_metadata, system_metadata = _user_and_system_metadata_from_proto(d.request.metadata)
     datum = Datum(
         keys=list(d.request.keys),
         sink_msg_id=d.request.id,
         value=d.request.value,
         event_time=d.request.event_time.ToDatetime(),
         watermark=d.request.watermark.ToDatetime(),
         headers=dict(d.request.headers),
+        user_metadata=user_metadata,
+        system_metadata=system_metadata,
     )
     return datum
 

packages/pynumaflow/tests/sink/test_async_sink.py

@@ -7,7 +7,7 @@
 
 import grpc
 from google.protobuf import empty_pb2 as _empty_pb2
-from grpc.aio._server import Server
+from grpc.aio import Server
 
 from pynumaflow import setup_logging
 from pynumaflow._constants import (
@@ -16,6 +16,7 @@
     FALLBACK_SINK_SOCK_PATH,
     FALLBACK_SINK_SERVER_INFO_FILE_PATH,
 )
+from pynumaflow.proto.common import metadata_pb2
 from pynumaflow.sinker import (
     Datum,
 )
@@ -41,6 +42,10 @@ async def udsink_handler(datums: AsyncIterable[Datum]) -> Responses:
         elif msg.value.decode("utf-8") == "test_mock_fallback_message":
             responses.append(Response.as_fallback(msg.id))
         else:
+            if msg.user_metadata.groups() != ["custom_info"]:
+                raise ValueError("user metadata groups do not match")
+            if msg.system_metadata.groups() != ["numaflow_version_info"]:
+                raise ValueError("system metadata groups do not match")
             responses.append(Response.as_success(msg.id))
     return responses
 
@@ -55,7 +60,21 @@ def start_sink_streaming_request(_id: str, req_type) -> (Datum, tuple):
         value = mock_fallback_message()
 
     request = sink_pb2.SinkRequest.Request(
-        value=value, event_time=event_time_timestamp, watermark=watermark_timestamp, id=_id
+        value=value,
+        event_time=event_time_timestamp,
+        watermark=watermark_timestamp,
+        id=_id,
+        metadata=metadata_pb2.Metadata(
+            previous_vertex="test-source",
+            user_metadata={
+                "custom_info": metadata_pb2.KeyValueGroup(key_value={"version": b"1.0.0"}),
+            },
+            sys_metadata={
+                "numaflow_version_info": metadata_pb2.KeyValueGroup(
+                    key_value={"version": b"1.0.0"}
+                ),
+            },
+        ),
     )
     return sink_pb2.SinkRequest(request=request)
 
@@ -64,10 +83,8 @@ def request_generator(count, req_type="success", session=1, handshake=True):
     if handshake:
         yield sink_pb2.SinkRequest(handshake=sink_pb2.Handshake(sot=True))
 
-    for j in range(session):
-        for i in range(count):
-            yield start_sink_streaming_request(str(i), req_type)
-
+    for _ in range(session):
+        yield from (start_sink_streaming_request(str(i), req_type) for i in range(count))
         yield sink_pb2.SinkRequest(status=sink_pb2.TransmissionStatus(eot=True))
 
 

packages/pynumaflow/tests/sink/test_server.py

@@ -15,6 +15,7 @@
     FALLBACK_SINK_SOCK_PATH,
     FALLBACK_SINK_SERVER_INFO_FILE_PATH,
 )
+from pynumaflow.proto.common import metadata_pb2
 from pynumaflow.proto.sinker import sink_pb2
 from pynumaflow.sinker import Responses, Datum, Response, SinkServer
 from tests.testing_utils import mock_terminate_on_stop
@@ -32,6 +33,10 @@ def udsink_handler(datums: Iterator[Datum]) -> Responses:
         elif "fallback" in msg.value.decode("utf-8"):
             results.append(Response.as_fallback(msg.id))
         else:
+            if msg.user_metadata.groups() != ["custom_info"]:
+                raise ValueError("user metadata groups do not match")
+            if msg.system_metadata.groups() != ["numaflow_version_info"]:
+                raise ValueError("system metadata groups do not match")
             results.append(Response.as_success(msg.id))
     return results
 
@@ -68,6 +73,16 @@ def mock_watermark():
 # We are mocking the terminate function from the psutil to not exit the program during testing
 @patch("psutil.Process.kill", mock_terminate_on_stop)
 class TestServer(unittest.TestCase):
+    metadata = metadata_pb2.Metadata(
+        previous_vertex="test-source",
+        user_metadata={
+            "custom_info": metadata_pb2.KeyValueGroup(key_value={"version": b"1.0.0"}),
+        },
+        sys_metadata={
+            "numaflow_version_info": metadata_pb2.KeyValueGroup(key_value={"version": b"1.0.0"}),
+        },
+    )
+
     def setUp(self) -> None:
         server = SinkServer(sinker_instance=udsink_handler)
         my_servicer = server.servicer
@@ -147,6 +162,7 @@ def test_udsink_err(self):
                     value=mock_message(),
                     event_time=event_time_timestamp,
                     watermark=watermark_timestamp,
+                    metadata=self.metadata,
                 )
             ),
             sink_pb2.SinkRequest(
@@ -155,6 +171,7 @@ def test_udsink_err(self):
                     value=mock_err_message(),
                     event_time=event_time_timestamp,
                     watermark=watermark_timestamp,
+                    metadata=self.metadata,
                 )
             ),
             sink_pb2.SinkRequest(status=sink_pb2.TransmissionStatus(eot=True)),
@@ -201,6 +218,7 @@ def test_forward_message(self):
                     value=mock_message(),
                     event_time=event_time_timestamp,
                     watermark=watermark_timestamp,
+                    metadata=self.metadata,
                 )
             ),
             sink_pb2.SinkRequest(
@@ -209,6 +227,7 @@ def test_forward_message(self):
                     value=mock_err_message(),
                     event_time=event_time_timestamp,
                     watermark=watermark_timestamp,
+                    metadata=self.metadata,
                 )
             ),
             sink_pb2.SinkRequest(status=sink_pb2.TransmissionStatus(eot=True)),