feat: support headers for message (numaproj#138)

Signed-off-by: Yashash H L <[email protected]>

Author: yhl25

examples/sink/log/example.py

@@ -20,7 +20,9 @@ def handler(self, datums: Iterator[Datum]) -> Responses:
 def udsink_handler(datums: Iterator[Datum]) -> Responses:
     responses = Responses()
     for msg in datums:
-        _LOGGER.info("User Defined Sink %s", msg.value.decode("utf-8"))
+        _LOGGER.info(
+            "User Defined Sink: Payload %s , Headers %s", msg.value.decode("utf-8"), msg.headers
+        )
         responses.append(Response.as_success(msg.id))
     return responses
 

examples/source/async-source/example.py

@@ -1,3 +1,5 @@
+import uuid
+
 from collections.abc import AsyncIterable
 from datetime import datetime
 
@@ -37,10 +39,12 @@ async def read_handler(self, datum: ReadRequest) -> AsyncIterable[Message]:
             return
 
         for x in range(datum.num_records):
+            headers = {"x-txn-id": str(uuid.uuid4())}
             yield Message(
                 payload=str(self.read_idx).encode(),
                 offset=Offset.offset_with_default_partition_id(str(self.read_idx).encode()),
                 event_time=datetime.now(),
+                headers=headers,
             )
             self.to_ack_set.add(str(self.read_idx))
             self.read_idx += 1

examples/source/simple-source/example.py

@@ -1,3 +1,5 @@
+import uuid
+
 from collections.abc import Iterable
 from datetime import datetime
 
@@ -36,11 +38,14 @@ def read_handler(self, datum: ReadRequest) -> Iterable[Message]:
         if self.to_ack_set:
             return
 
+        headers = {"x-txn-id": str(uuid.uuid4())}
+
         for x in range(datum.num_records):
             yield Message(
                 payload=str(self.read_idx).encode(),
                 offset=Offset.offset_with_default_partition_id(str(self.read_idx).encode()),
                 event_time=datetime.now(),
+                headers=headers,
             )
             self.to_ack_set.add(str(self.read_idx))
             self.read_idx += 1

pynumaflow/mapper/_dtypes.py

@@ -2,7 +2,7 @@
 from collections.abc import Iterator, Sequence, Awaitable
 from dataclasses import dataclass
 from datetime import datetime
-from typing import TypeVar, Callable, Union
+from typing import TypeVar, Callable, Union, Optional
 from warnings import warn
 
 from pynumaflow._constants import DROP
@@ -106,33 +106,39 @@ class Datum:
         value: the payload of the event.
         event_time: the event time of the event.
         watermark: the watermark of the event.
+        headers: the headers of the event.
+
     >>> # Example usage
     >>> from pynumaflow.mapper import Datum
     >>> from datetime import datetime, timezone
     >>> payload = bytes("test_mock_message", encoding="utf-8")
     >>> t1 = datetime.fromtimestamp(1662998400, timezone.utc)
     >>> t2 = datetime.fromtimestamp(1662998460, timezone.utc)
+    >>> msg_headers = {"key1": "value1", "key2": "value2"}
     >>> d = Datum(
     ...       keys=["test_key"],
     ...       value=payload,
     ...       event_time=t1,
     ...       watermark=t2,
+    ...       headers=msg_headers,
     ...    )
     """
 
-    __slots__ = ("_keys", "_value", "_event_time", "_watermark")
+    __slots__ = ("_keys", "_value", "_event_time", "_watermark", "_headers")
 
     _keys: list[str]
     _value: bytes
     _event_time: datetime
     _watermark: datetime
+    _headers: dict[str, str]
 
     def __init__(
         self,
         keys: list[str],
         value: bytes,
         event_time: datetime,
         watermark: datetime,
+        headers: Optional[dict[str, str]] = None,
     ):
         self._keys = keys or list()
         self._value = value or b""
@@ -142,7 +148,9 @@ def __init__(
         if not isinstance(watermark, datetime):
             raise TypeError(f"Wrong data type: {type(watermark)} for Datum.watermark")
         self._watermark = watermark
+        self._headers = headers or {}
 
+    @property
     def keys(self) -> list[str]:
         """Returns the keys of the event"""
         return self._keys
@@ -162,6 +170,11 @@ def watermark(self) -> datetime:
         """Returns the watermark of the event."""
         return self._watermark
 
+    @property
+    def headers(self) -> dict[str, str]:
+        """Returns the headers of the event."""
+        return self._headers.copy()
+
 
 class Mapper(metaclass=ABCMeta):
     """

pynumaflow/mapper/servicer/async_servicer.py

@@ -38,6 +38,7 @@ async def MapFn(
                     value=request.value,
                     event_time=request.event_time.ToDatetime(),
                     watermark=request.watermark.ToDatetime(),
+                    headers=request.headers,
                 ),
             )
         except Exception as e:

pynumaflow/mapper/servicer/utils.py

@@ -20,6 +20,7 @@ def _map_fn_util(
                 value=request.value,
                 event_time=request.event_time.ToDatetime(),
                 watermark=request.watermark.ToDatetime(),
+                headers=dict(request.headers),
             ),
         )
     except Exception as err:

pynumaflow/mapstreamer/_dtypes.py

@@ -2,7 +2,7 @@
 from collections.abc import Iterator, Sequence
 from dataclasses import dataclass
 from datetime import datetime
-from typing import TypeVar, Callable, Union
+from typing import TypeVar, Callable, Union, Optional
 from collections.abc import AsyncIterable
 from warnings import warn
 
@@ -107,33 +107,39 @@ class Datum:
         value: the payload of the event.
         event_time: the event time of the event.
         watermark: the watermark of the event.
+        headers: the headers of the event.
+
     >>> # Example usage
     >>> from pynumaflow.mapstreamer import Datum
     >>> from datetime import datetime, timezone
     >>> payload = bytes("test_mock_message", encoding="utf-8")
     >>> t1 = datetime.fromtimestamp(1662998400, timezone.utc)
     >>> t2 = datetime.fromtimestamp(1662998460, timezone.utc)
+    >>> msg_headers = {"key1": "value1", "key2": "value2"}
     >>> d = Datum(
     ...       keys=["test_key"],
     ...       value=payload,
     ...       event_time=t1,
     ...       watermark=t2,
+    ...       headers=msg_headers,
     ...    )
     """
 
-    __slots__ = ("_keys", "_value", "_event_time", "_watermark")
+    __slots__ = ("_keys", "_value", "_event_time", "_watermark", "_headers")
 
     _keys: list[str]
     _value: bytes
     _event_time: datetime
     _watermark: datetime
+    _headers: dict[str, str]
 
     def __init__(
         self,
         keys: list[str],
         value: bytes,
         event_time: datetime,
         watermark: datetime,
+        headers: Optional[dict[str, str]] = None,
     ):
         self._keys = keys or list()
         self._value = value or b""
@@ -143,6 +149,7 @@ def __init__(
         if not isinstance(watermark, datetime):
             raise TypeError(f"Wrong data type: {type(watermark)} for Datum.watermark")
         self._watermark = watermark
+        self._headers = headers or {}
 
     def keys(self) -> list[str]:
         """Returns the keys of the event"""
@@ -163,6 +170,11 @@ def watermark(self) -> datetime:
         """Returns the watermark of the event."""
         return self._watermark
 
+    @property
+    def headers(self) -> dict[str, str]:
+        """Returns the headers of the event."""
+        return self._headers.copy()
+
 
 class MapStreamer(metaclass=ABCMeta):
     """

pynumaflow/mapstreamer/servicer/async_servicer.py

@@ -40,6 +40,7 @@ async def MapStreamFn(
                 value=request.value,
                 event_time=request.event_time.ToDatetime(),
                 watermark=request.watermark.ToDatetime(),
+                headers=request.headers,
             ),
         ):
             yield mapstream_pb2.MapStreamResponse(result=res)

pynumaflow/proto/mapper/map.proto

@@ -21,6 +21,7 @@ message MapRequest {
   bytes value = 2;
   google.protobuf.Timestamp event_time = 3;
   google.protobuf.Timestamp watermark = 4;
+  map<string, string> headers = 5;
 }
 
 /**

pynumaflow/proto/mapper/map_pb2.py

@@ -22,6 +22,7 @@ message MapStreamRequest {
   bytes value = 2;
   google.protobuf.Timestamp event_time = 3;
   google.protobuf.Timestamp watermark = 4;
+  map<string, string> headers = 5;
 }
 
 /**

pynumaflow/proto/mapstreamer/mapstream.proto

@@ -37,6 +37,7 @@ message ReduceRequest {
     bytes value = 2;
     google.protobuf.Timestamp event_time = 3;
     google.protobuf.Timestamp watermark = 4;
+    map<string, string> headers = 5;
   }
 
   Payload payload = 1;