Make sinker more Pythonic

Signed-off-by: Sreekanth <prsreekanth920@gmail.com>

Author: BulkBeing

.github/workflows/pynumaflow-lite.yml

@@ -59,7 +59,7 @@ jobs:
           rustup default stable
 
       - name: Install dependencies
-        run: uv sync --group dev
+        run: uv sync --locked --group dev
 
       - name: Check Rust formatting
         run: cargo fmt --all --check
@@ -112,7 +112,7 @@ jobs:
             echo "::group::Python ${python_version}"
             uv python install "${python_version}"
             export UV_PYTHON="${python_version}"
-            uv sync --group dev --python "${python_version}"
+            uv sync --locked --group dev --python "${python_version}"
             uv run --python "${python_version}" maturin develop --locked
 
             uv run --python "${python_version}" pytest -v

packages/pynumaflow-lite/manifests/sink/sink_log.py

@@ -1,10 +1,8 @@
-import asyncio
 import logging
-import signal
-from collections.abc import AsyncIterable, AsyncIterator, Awaitable, Callable
+from collections.abc import AsyncIterable
 
 from pynumaflow_lite import sinker
-from pynumaflow_lite._sink_dtypes import Sinker
+from pynumaflow_lite.sinker import Sinker
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
@@ -16,37 +14,15 @@ class SimpleLogSink(Sinker):
     Simple log sink that logs each message and returns success responses.
     """
 
-    async def handler(self, datums: AsyncIterable[sinker.Datum]) -> sinker.Responses:
-        responses = sinker.Responses()
+    async def handler(self, datums: AsyncIterable[sinker.Datum]) -> list[sinker.Response]:
+        responses = []
         async for msg in datums:
             _LOGGER.info("User Defined Sink: %s", msg.value.decode("utf-8"))
-            responses.append(sinker.Response.as_success(msg.id))
+            responses.append(sinker.Response.success(msg.id))
             # if we are not able to write to sink and if we have a fallback sink configured
-            # we can use Response.as_fallback(msg.id) to write the message to fallback sink
+            # we can use Response.fallback(msg.id) to write the message to fallback sink
         return responses
 
 
-async def start(
-    f: Callable[[AsyncIterator[sinker.Datum]], Awaitable[sinker.Responses]],
-):
-    server = sinker.SinkAsyncServer()
-
-    # Register loop-level signal handlers so we control shutdown and avoid asyncio.run
-    loop = asyncio.get_running_loop()
-    try:
-        loop.add_signal_handler(signal.SIGINT, lambda: server.stop())
-        loop.add_signal_handler(signal.SIGTERM, lambda: server.stop())
-    except (NotImplementedError, RuntimeError):
-        pass
-
-    try:
-        await server.start(f)
-        print("Shutting down gracefully...")
-    except asyncio.CancelledError:
-        server.stop()
-        return
-
-
 if __name__ == "__main__":
-    async_handler = SimpleLogSink()
-    asyncio.run(start(async_handler))
+    sinker.SinkAsyncServer(SimpleLogSink()).run()

packages/pynumaflow-lite/pynumaflow_lite/__init__.py

@@ -72,6 +72,7 @@
 from ._session_reduce_dtypes import SessionReducer
 from ._sideinput_dtypes import SideInput
 from ._sink_dtypes import Sinker
+from ._sink_server import SinkAsyncServer
 from ._source_dtypes import Sourcer
 from ._sourcetransformer_dtypes import SourceTransformer
 
@@ -98,6 +99,7 @@
 
 if sinker is not None:
     sinker.Sinker = Sinker
+    sinker.SinkAsyncServer = SinkAsyncServer
 
 if sourcer is not None:
     sourcer.Sourcer = Sourcer

packages/pynumaflow-lite/pynumaflow_lite/_sink_dtypes.py

@@ -1,7 +1,7 @@
 from abc import ABCMeta, abstractmethod
 from collections.abc import AsyncIterable
 
-from pynumaflow_lite.sinker import Datum, Responses
+from pynumaflow_lite.sinker import Datum, Response
 
 
 class Sinker(metaclass=ABCMeta):
@@ -13,7 +13,7 @@ def __call__(self, *args, **kwargs):
         return self.handler(*args, **kwargs)
 
     @abstractmethod
-    async def handler(self, datums: AsyncIterable[Datum]) -> Responses:
+    async def handler(self, datums: AsyncIterable[Datum]) -> list[Response]:
         """
         Implement this handler function for sink.
         Process the stream of datums and return responses.

packages/pynumaflow-lite/pynumaflow_lite/_sink_server.py

@@ -0,0 +1,78 @@
+from __future__ import annotations
+
+import asyncio
+import signal
+from types import TracebackType
+from typing import Any
+
+from .pynumaflow_lite import sinker as _sinker
+
+
+class SinkAsyncServer:
+    def __init__(
+        self,
+        handler: Any,
+        *,
+        sock_file: str | None = None,
+        server_info_file: str | None = None,
+    ) -> None:
+        self._core = _sinker._SinkAsyncServer(sock_file, server_info_file)
+        self._handler = handler
+        self._task: asyncio.Task[None] | None = None
+
+    async def serve(self) -> None:
+        await self._core.start(self._handler)
+
+    def stop(self) -> None:
+        self._core.stop()
+
+    async def wait_ready(self, timeout: float = 30.0) -> None:
+        await self._core.wait_ready(timeout)
+
+    async def __aenter__(self) -> SinkAsyncServer:
+        if self._task is not None and not self._task.done():
+            raise RuntimeError("sink server is already serving")
+
+        self._task = asyncio.create_task(self.serve())
+        try:
+            await self.wait_ready()
+        except asyncio.CancelledError:
+            self.stop()
+            if self._task is not None:
+                await self._task
+            raise
+        except Exception:
+            self.stop()
+            if self._task is not None:
+                await self._task
+            raise
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc: BaseException | None,
+        tb: TracebackType | None,
+    ) -> None:
+        self.stop()
+        if self._task is not None:
+            try:
+                await self._task
+            finally:
+                self._task = None
+
+    async def _main(self) -> None:
+        loop = asyncio.get_running_loop()
+        try:
+            loop.add_signal_handler(signal.SIGINT, self.stop)
+            loop.add_signal_handler(signal.SIGTERM, self.stop)
+        except (NotImplementedError, RuntimeError):
+            pass
+
+        await self.serve()
+
+    def run(self) -> None:
+        try:
+            asyncio.run(self._main())
+        except KeyboardInterrupt:
+            self.stop()

packages/pynumaflow-lite/pynumaflow_lite/sinker.pyi

@@ -2,113 +2,103 @@ from __future__ import annotations
 
 import datetime as _dt
 from collections.abc import AsyncIterator, Awaitable, Callable
-
-class SystemMetadata:
-    """System-generated metadata groups per message (read-only for sink)."""
-
-    def __init__(self) -> None: ...
-    def groups(self) -> list[str]:
-        """Returns the groups of the system metadata."""
-        ...
-
-    def keys(self, group: str) -> list[str]:
-        """Returns the keys of the system metadata for the given group."""
-        ...
-
-    def value(self, group: str, key: str) -> bytes:
-        """Returns the value of the system metadata for the given group and key."""
-        ...
-
-    def __repr__(self) -> str: ...
-
-class UserMetadata:
-    """User-defined metadata groups per message (read-only for sink)."""
-
-    def __init__(self) -> None: ...
-    def groups(self) -> list[str]:
-        """Returns the groups of the user metadata."""
-        ...
-
-    def keys(self, group: str) -> list[str]:
-        """Returns the keys of the user metadata for the given group."""
-        ...
-
-    def value(self, group: str, key: str) -> bytes:
-        """Returns the value of the user metadata for the given group and key."""
-        ...
-
-    def __repr__(self) -> str: ...
-
-class KeyValueGroup:
-    key_value: dict[str, bytes]
-
-    def __init__(self, key_value: dict[str, bytes] | None = ...) -> None: ...
-    @staticmethod
-    def from_dict(key_value: dict[str, bytes]) -> KeyValueGroup: ...
+from types import TracebackType
 
 class Message:
     keys: list[str] | None
     value: bytes
-    user_metadata: dict[str, KeyValueGroup] | None
+    user_metadata: dict[str, dict[str, bytes]] | None
 
     def __init__(
         self,
         value: bytes,
         keys: list[str] | None = ...,
-        user_metadata: dict[str, KeyValueGroup] | None = ...,
+        user_metadata: dict[str, dict[str, bytes]] | None = ...,
     ) -> None: ...
+    def __repr__(self) -> str: ...
+    def __eq__(self, other: object) -> bool: ...
 
 class Response:
     id: str
+    error: str | None
 
     @staticmethod
-    def as_success(id: str) -> Response: ...
+    def success(id: str) -> Response: ...
     @staticmethod
-    def as_failure(id: str, err_msg: str) -> Response: ...
+    def failure(id: str, error: str) -> Response: ...
     @staticmethod
-    def as_fallback(id: str) -> Response: ...
+    def fallback(id: str) -> Response: ...
     @staticmethod
-    def as_serve(id: str, payload: bytes) -> Response: ...
+    def serve(id: str, payload: bytes) -> Response: ...
     @staticmethod
-    def as_on_success(id: str, message: Message | None = ...) -> Response: ...
-
-class Responses:
-    def __init__(self) -> None: ...
-    def append(self, response: Response) -> None: ...
+    def on_success(id: str, message: Message | None = ...) -> Response: ...
+    def __repr__(self) -> str: ...
+    def __eq__(self, other: object) -> bool: ...
 
 class Datum:
     keys: list[str]
     value: bytes
     watermark: _dt.datetime
-    eventtime: _dt.datetime
+    event_time: _dt.datetime
     id: str
     headers: dict[str, str]
-    user_metadata: UserMetadata
-    system_metadata: SystemMetadata
+    user_metadata: dict[str, dict[str, bytes]]
+    system_metadata: dict[str, dict[str, bytes]]
 
+    def __init__(
+        self,
+        *,
+        keys: list[str] = ...,
+        value: bytes = ...,
+        id: str = ...,
+        event_time: _dt.datetime | None = ...,
+        watermark: _dt.datetime | None = ...,
+        headers: dict[str, str] = ...,
+        user_metadata: dict[str, dict[str, bytes]] = ...,
+        system_metadata: dict[str, dict[str, bytes]] = ...,
+    ) -> None: ...
     def __repr__(self) -> str: ...
     def __str__(self) -> str: ...
 
-class SinkAsyncServer:
+_SinkHandler = Callable[[AsyncIterator[Datum]], Awaitable[list[Response]]]
+
+class _SinkAsyncServer:
     def __init__(
         self,
         sock_file: str | None = ...,
-        info_file: str | None = ...,
+        server_info_file: str | None = ...,
     ) -> None: ...
-    def start(self, py_func: Callable[[AsyncIterator[Datum]], Awaitable[Responses]]) -> Awaitable[None]: ...
+    def start(self, handler: _SinkHandler) -> Awaitable[None]: ...
+    def wait_ready(self, timeout: float = ...) -> Awaitable[None]: ...
     def stop(self) -> None: ...
 
 class Sinker:
-    async def handler(self, datums: AsyncIterator[Datum]) -> Responses: ...
+    async def handler(self, datums: AsyncIterator[Datum]) -> list[Response]: ...
+
+class SinkAsyncServer:
+    def __init__(
+        self,
+        handler: _SinkHandler | Sinker,
+        *,
+        sock_file: str | None = ...,
+        server_info_file: str | None = ...,
+    ) -> None: ...
+    def run(self) -> None: ...
+    async def serve(self) -> None: ...
+    def stop(self) -> None: ...
+    async def wait_ready(self, timeout: float = ...) -> None: ...
+    async def __aenter__(self) -> SinkAsyncServer: ...
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc: BaseException | None,
+        tb: TracebackType | None,
+    ) -> None: ...
 
 __all__ = [
     "Datum",
-    "KeyValueGroup",
     "Message",
     "Response",
-    "Responses",
     "SinkAsyncServer",
     "Sinker",
-    "SystemMetadata",
-    "UserMetadata",
 ]