From 09425fe8d4589ee0deb0ad16ea89aa4d79990325 Mon Sep 17 00:00:00 2001 From: Sreekanth Date: Sat, 10 Oct 2026 08:11:11 +0530 Subject: [PATCH 1/3] pynumaflow-lite: Pythonic APIs for Sourcetransformer Signed-off-by: Sreekanth --- .../manifests/sourcetransform/.dockerignore | 1 + .../manifests/sourcetransform/Dockerfile | 77 +-- .../manifests/sourcetransform/pyproject.toml | 8 +- .../sourcetransform_event_filter.py | 69 +-- .../manifests/sourcetransform/uv.lock | 8 + .../pynumaflow_lite/__init__.py | 2 + .../_sourcetransform_server.py | 125 +++++ .../_sourcetransformer_dtypes.py | 12 +- .../pynumaflow_lite/sourcetransformer.pyi | 119 ++--- .../src/sourcetransform/mod.rs | 487 ++++++++---------- .../src/sourcetransform/server.rs | 113 ++-- .../tests/bin/sourcetransform.rs | 11 + .../examples/sourcetransform_event_filter.py | 121 ++--- .../sourcetransform_event_filter_func.py | 46 ++ .../tests/test_sourcetransform.py | 1 + .../tests/test_sourcetransform_unit.py | 119 +++++ 16 files changed, 750 insertions(+), 569 deletions(-) create mode 100644 packages/pynumaflow-lite/manifests/sourcetransform/.dockerignore create mode 100644 packages/pynumaflow-lite/manifests/sourcetransform/uv.lock create mode 100644 packages/pynumaflow-lite/pynumaflow_lite/_sourcetransform_server.py create mode 100644 packages/pynumaflow-lite/tests/examples/sourcetransform_event_filter_func.py create mode 100644 packages/pynumaflow-lite/tests/test_sourcetransform_unit.py diff --git a/packages/pynumaflow-lite/manifests/sourcetransform/.dockerignore b/packages/pynumaflow-lite/manifests/sourcetransform/.dockerignore new file mode 100644 index 00000000..21d0b898 --- /dev/null +++ b/packages/pynumaflow-lite/manifests/sourcetransform/.dockerignore @@ -0,0 +1 @@ +.venv/ diff --git a/packages/pynumaflow-lite/manifests/sourcetransform/Dockerfile b/packages/pynumaflow-lite/manifests/sourcetransform/Dockerfile index fd4d0ac5..d49578e1 100644 --- a/packages/pynumaflow-lite/manifests/sourcetransform/Dockerfile +++ b/packages/pynumaflow-lite/manifests/sourcetransform/Dockerfile @@ -1,38 +1,43 @@ -FROM python:3.11-slim-bullseye AS builder - -ENV PYTHONFAULTHANDLER=1 \ - PYTHONUNBUFFERED=1 \ - PYTHONHASHSEED=random \ - PIP_NO_CACHE_DIR=on \ - PIP_DISABLE_PIP_VERSION_CHECK=on \ - PIP_DEFAULT_TIMEOUT=100 \ - POETRY_HOME="/opt/poetry" \ - POETRY_VIRTUALENVS_IN_PROJECT=true \ - POETRY_NO_INTERACTION=1 \ - PYSETUP_PATH="/opt/pysetup" - - ENV PATH="$POETRY_HOME/bin:$PATH" - -RUN apt-get update \ - && apt-get install --no-install-recommends -y \ - curl \ - wget \ - # deps for building python deps - build-essential \ - && apt-get install -y git \ - && apt-get clean && rm -rf /var/lib/apt/lists/* \ - && curl -sSL https://install.python-poetry.org | python3 - - -FROM builder AS udf - -WORKDIR $PYSETUP_PATH -COPY ./ ./ - -RUN pip install $PYSETUP_PATH/pynumaflow_lite-0.1.0-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl - -RUN poetry lock -RUN poetry install --no-cache --no-root && \ - rm -rf ~/.cache/pypoetry/ +FROM python:3.11-slim-trixie AS builder -CMD ["python", "sourcetransform_event_filter.py"] +COPY --from=ghcr.io/astral-sh/uv:0.12.13 /uv /uvx /bin/ + +ENV UV_COMPILE_BYTECODE=1 \ + UV_LINK_MODE=copy \ + UV_PYTHON_DOWNLOADS=never + +WORKDIR /app + +RUN --mount=type=cache,target=/root/.cache/uv \ + --mount=type=bind,source=uv.lock,target=uv.lock \ + --mount=type=bind,source=pyproject.toml,target=pyproject.toml \ + uv sync --locked --no-install-project --no-dev + +COPY . . + +RUN --mount=type=cache,target=/root/.cache/uv \ + uv sync --locked --no-dev + +# Install the local pynumaflow-lite wheel that matches the build platform. +RUN uv pip install --no-index --find-links . pynumaflow-lite + +FROM python:3.11-slim-trixie +# Setup a non-root user +RUN groupadd --system --gid 999 nonroot \ + && useradd --system --gid 999 --uid 999 --create-home nonroot + +COPY --from=builder --chown=nonroot:nonroot /app /app + +ENV PATH="/app/.venv/bin:$PATH" + +# Keeps Python from buffering stdout and stderr to avoid situations where +# the application crashes without emitting any logs due to buffering. +ENV PYTHONUNBUFFERED=1 + +# Use the non-root user to run our application +USER nonroot + +WORKDIR /app + +CMD ["python", "sourcetransform_event_filter.py"] diff --git a/packages/pynumaflow-lite/manifests/sourcetransform/pyproject.toml b/packages/pynumaflow-lite/manifests/sourcetransform/pyproject.toml index e6b8b866..12e2a438 100644 --- a/packages/pynumaflow-lite/manifests/sourcetransform/pyproject.toml +++ b/packages/pynumaflow-lite/manifests/sourcetransform/pyproject.toml @@ -6,12 +6,6 @@ authors = [ { name = "Vigith Maurice", email = "vigith@gmail.com" } ] readme = "README.md" -requires-python = ">=3.11" +requires-python = "==3.11.*" dependencies = [ ] - - -[build-system] -requires = ["poetry-core>=2.0.0,<3.0.0"] -build-backend = "poetry.core.masonry.api" - diff --git a/packages/pynumaflow-lite/manifests/sourcetransform/sourcetransform_event_filter.py b/packages/pynumaflow-lite/manifests/sourcetransform/sourcetransform_event_filter.py index ea79646f..332e15c3 100644 --- a/packages/pynumaflow-lite/manifests/sourcetransform/sourcetransform_event_filter.py +++ b/packages/pynumaflow-lite/manifests/sourcetransform/sourcetransform_event_filter.py @@ -1,16 +1,13 @@ import asyncio -import signal -from collections.abc import Callable from datetime import datetime, timezone -from pynumaflow_lite import sourcetransformer +from pynumaflow_lite.sourcetransformer import Datum, Message, SourceTransformAsyncServer, SourceTransformer -# Define epoch timestamps for filtering january_first_2022 = datetime(2022, 1, 1, tzinfo=timezone.utc) january_first_2023 = datetime(2023, 1, 1, tzinfo=timezone.utc) -class EventFilter(sourcetransformer.SourceTransformer): +class EventFilter(SourceTransformer): """ A source transformer that filters and routes messages based on event time. @@ -19,65 +16,23 @@ class EventFilter(sourcetransformer.SourceTransformer): - Messages after 2022 are tagged with "after_year_2022" """ - async def handler(self, keys: list[str], datum: sourcetransformer.Datum) -> sourcetransformer.Messages: - val = datum.value + async def handler(self, datum: Datum) -> list[Message]: event_time = datum.event_time - messages = sourcetransformer.Messages() - if event_time < january_first_2022: print(f"Got event time: {event_time}, it is before 2022, so dropping") - messages.append(sourcetransformer.Message.message_to_drop(event_time)) - elif event_time < january_first_2023: - print(f"Got event time: {event_time}, it is within year 2022, so forwarding to within_year_2022") - messages.append( - sourcetransformer.Message( - value=val, - event_time=january_first_2022, - keys=keys, - tags=["within_year_2022"], - ) - ) - else: - print(f"Got event time: {event_time}, it is after year 2022, so forwarding to after_year_2022") - messages.append( - sourcetransformer.Message( - value=val, - event_time=january_first_2023, - keys=keys, - tags=["after_year_2022"], - ) - ) - - return messages + return [Message.to_drop(event_time)] + if event_time < january_first_2023: + print(f"Got event time: {event_time}, it is within year 2022, so forwarding to within_year_2022") + return [Message(datum.value, january_first_2022, keys=datum.keys, tags=["within_year_2022"])] -async def start( - f: Callable[[list[str], sourcetransformer.Datum], sourcetransformer.Messages], -): - server = sourcetransformer.SourceTransformAsyncServer() + print(f"Got event time: {event_time}, it is after year 2022, so forwarding to after_year_2022") + return [Message(datum.value, january_first_2023, keys=datum.keys, tags=["after_year_2022"])] - # Register loop-level signal handlers so we control shutdown and avoid asyncio.run - # converting it into KeyboardInterrupt/CancelledError traces. - loop = asyncio.get_running_loop() - loop.set_debug(True) - print("Registering signal handlers", loop) - try: - loop.add_signal_handler(signal.SIGINT, lambda: server.stop()) - loop.add_signal_handler(signal.SIGTERM, lambda: server.stop()) - except (NotImplementedError, RuntimeError): - print("Failed to register signal handlers") - # add_signal_handler may not be available on some platforms/contexts; fallback below. - pass - try: - await server.start(f) - print("Shutting down gracefully...") - except asyncio.CancelledError: - # Fallback in case the task was cancelled by the runner - server.stop() - return +async def main() -> None: + await SourceTransformAsyncServer(EventFilter()).serve() if __name__ == "__main__": - async_handler = EventFilter() - asyncio.run(start(async_handler)) + asyncio.run(main()) diff --git a/packages/pynumaflow-lite/manifests/sourcetransform/uv.lock b/packages/pynumaflow-lite/manifests/sourcetransform/uv.lock new file mode 100644 index 00000000..d38ecf5f --- /dev/null +++ b/packages/pynumaflow-lite/manifests/sourcetransform/uv.lock @@ -0,0 +1,8 @@ +version = 1 +revision = 3 +requires-python = "==3.11.*" + +[[package]] +name = "sourcetransform-event-filter" +version = "0.1.0" +source = { virtual = "." } diff --git a/packages/pynumaflow-lite/pynumaflow_lite/__init__.py b/packages/pynumaflow-lite/pynumaflow_lite/__init__.py index 8b23548b..5f8f0c14 100644 --- a/packages/pynumaflow-lite/pynumaflow_lite/__init__.py +++ b/packages/pynumaflow-lite/pynumaflow_lite/__init__.py @@ -18,6 +18,7 @@ from ._sink_dtypes import Sinker from ._sink_server import SinkAsyncServer from ._source_dtypes import Sourcer +from ._sourcetransform_server import SourceTransformAsyncServer from ._sourcetransformer_dtypes import SourceTransformer from .pynumaflow_lite import * # noqa: F403 # Rust extension; exports resolved at runtime @@ -51,6 +52,7 @@ sinker.SinkAsyncServer = SinkAsyncServer sourcer.Sourcer = Sourcer sourcetransformer.SourceTransformer = SourceTransformer +sourcetransformer.SourceTransformAsyncServer = SourceTransformAsyncServer sideinputer.SideInput = SideInput # Public API diff --git a/packages/pynumaflow-lite/pynumaflow_lite/_sourcetransform_server.py b/packages/pynumaflow-lite/pynumaflow_lite/_sourcetransform_server.py new file mode 100644 index 00000000..fb4c01b8 --- /dev/null +++ b/packages/pynumaflow-lite/pynumaflow_lite/_sourcetransform_server.py @@ -0,0 +1,125 @@ +from __future__ import annotations + +import asyncio +import contextlib +import signal +from collections.abc import Awaitable, Callable +from types import TracebackType +from typing import TypeAlias + +from .pynumaflow_lite import sourcetransformer as _sourcetransformer + +Datum: TypeAlias = _sourcetransformer.Datum +Message: TypeAlias = _sourcetransformer.Message + +_SHUTDOWN_SIGNALS = (signal.SIGINT, signal.SIGTERM) + + +class SourceTransformAsyncServer: + def __init__( + self, + handler: Callable[[Datum], Awaitable[list[Message]]], + *, + sock_file: str | None = None, + server_info_file: str | None = None, + install_signal_handlers: bool = True, + ) -> None: + self._core = _sourcetransformer._SourceTransformAsyncServer(sock_file, server_info_file) + self._handler = handler + self._install_signal_handlers = install_signal_handlers + self._task: asyncio.Task[None] | None = None + self._serving = False + self._installed_signals: list[signal.Signals] = [] + + async def serve(self) -> None: + """Run the sourcetransform server until it stops. + + This is the entrypoint for an application that already runs an event + loop. It returns when a shutdown signal arrives or when `stop()` runs. + """ + await self._serve(install_signal_handlers=self._install_signal_handlers) + + async def _serve(self, *, install_signal_handlers: bool) -> None: + if self._serving: + raise RuntimeError("sourcetransform server is already serving") + self._serving = True + try: + if install_signal_handlers: + self._add_signal_handlers() + await self._core.start(self._handler) + finally: + self._remove_signal_handlers() + self._serving = False + + 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 wait_for_termination(self) -> None: + """Wait until the background server task ends. + + Use this inside an `async with` block. It raises the handler error if + the server task failed. + """ + if self._task is None: + raise RuntimeError("sourcetransform server is not serving") + await asyncio.shield(self._task) + + def _add_signal_handlers(self) -> None: + loop = asyncio.get_running_loop() + for sig in _SHUTDOWN_SIGNALS: + try: + loop.add_signal_handler(sig, self.stop) + except (NotImplementedError, RuntimeError, OSError): + continue + self._installed_signals.append(sig) + + def _remove_signal_handlers(self) -> None: + if not self._installed_signals: + return + try: + loop = asyncio.get_running_loop() + except RuntimeError: + self._installed_signals.clear() + return + for sig in self._installed_signals: + with contextlib.suppress(NotImplementedError, OSError): + loop.remove_signal_handler(sig) + self._installed_signals.clear() + + async def __aenter__(self) -> SourceTransformAsyncServer: + """Start the server in a background task and wait until it is ready. + + This form is for tests and for code that must run other work next to + the server. It never installs signal handlers. + """ + if self._task is not None and not self._task.done(): + raise RuntimeError("sourcetransform server is already serving") + + self._task = asyncio.create_task(self._serve(install_signal_handlers=False)) + try: + await self.wait_ready() + except BaseException: + self.stop() + task, self._task = self._task, None + if task is not None: + # Surface the server error, if there is one. It explains the + # failure better than the `wait_ready` error does. + await 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 diff --git a/packages/pynumaflow-lite/pynumaflow_lite/_sourcetransformer_dtypes.py b/packages/pynumaflow-lite/pynumaflow_lite/_sourcetransformer_dtypes.py index 000110dd..287a493b 100644 --- a/packages/pynumaflow-lite/pynumaflow_lite/_sourcetransformer_dtypes.py +++ b/packages/pynumaflow-lite/pynumaflow_lite/_sourcetransformer_dtypes.py @@ -1,6 +1,6 @@ from abc import ABCMeta, abstractmethod -from pynumaflow_lite.sourcetransformer import Datum, Messages +from pynumaflow_lite.sourcetransformer import Datum, Message class SourceTransformer(metaclass=ABCMeta): @@ -20,16 +20,16 @@ class instance is sent as a callable. return self.handler(*args, **kwargs) @abstractmethod - async def handler(self, keys: list[str], datum: Datum) -> Messages: + async def handler(self, datum: Datum) -> list[Message]: """ Implement this handler function which implements the SourceTransformer interface. Args: - keys: The keys associated with the message. - datum: The input datum containing value, event_time, watermark, and headers. + datum: The input datum containing keys, value, event_time, watermark, + headers, and metadata. Returns: - Messages: A collection of transformed messages with potentially modified - event times and tags for conditional forwarding. + The transformed messages, with potentially modified event times and + tags for conditional forwarding. """ pass diff --git a/packages/pynumaflow-lite/pynumaflow_lite/sourcetransformer.pyi b/packages/pynumaflow-lite/pynumaflow_lite/sourcetransformer.pyi index 72c9d93e..0dacc357 100644 --- a/packages/pynumaflow-lite/pynumaflow_lite/sourcetransformer.pyi +++ b/packages/pynumaflow-lite/pynumaflow_lite/sourcetransformer.pyi @@ -1,69 +1,11 @@ from __future__ import annotations import datetime as _dt -from collections.abc import Awaitable +from collections.abc import Awaitable, Callable +from types import TracebackType -# Re-export the Python ABC for user convenience and typing from ._sourcetransformer_dtypes import SourceTransformer as SourceTransformer -class SystemMetadata: - """System-generated metadata groups per message (read-only).""" - - 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-write).""" - - 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 create_group(self, group: str) -> None: - """Creates a new group in the user metadata.""" - ... - - def add_kv(self, group: str, key: str, value: bytes) -> None: - """Adds a key-value pair to the user metadata.""" - ... - - def remove_key(self, group: str, key: str) -> None: - """Removes a key from a group in the user metadata.""" - ... - - def remove_group(self, group: str) -> None: - """Removes a group from the user metadata.""" - ... - - def __repr__(self) -> str: ... - -class Messages: - def __init__(self) -> None: ... - def append(self, message: Message) -> None: ... - def __repr__(self) -> str: ... - def __str__(self) -> str: ... - class NackOptions: """Per-message redelivery options for a nack.""" @@ -86,7 +28,7 @@ class Message: value: bytes event_time: _dt.datetime tags: list[str] | None - user_metadata: UserMetadata | None + user_metadata: dict[str, dict[str, bytes]] | None nack_options: NackOptions | None def __init__( @@ -95,44 +37,79 @@ class Message: event_time: _dt.datetime, keys: list[str] | None = ..., tags: list[str] | None = ..., - user_metadata: UserMetadata | None = ..., + user_metadata: dict[str, dict[str, bytes]] | None = ..., ) -> None: ... @staticmethod - def message_to_drop(event_time: _dt.datetime) -> Message: ... + def to_drop(event_time: _dt.datetime) -> Message: ... @staticmethod def to_nack(event_time: _dt.datetime, nack_options: NackOptions | None = ...) -> Message: ... @staticmethod def to_fail(event_time: _dt.datetime) -> Message: ... + def __repr__(self) -> str: ... + def __str__(self) -> str: ... + def __eq__(self, other: object) -> bool: ... class Datum: - # Read-only attributes provided by the extension keys: list[str] value: bytes watermark: _dt.datetime event_time: _dt.datetime 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 = ..., + 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 _SourceTransformAsyncServer: + def __init__( + self, + sock_file: str | None = ..., + server_info_file: str | None = ..., + ) -> None: ... + def start( + self, + handler: Callable[[Datum], Awaitable[list[Message]]], + ) -> Awaitable[None]: ... + def wait_ready(self, timeout: float = ...) -> Awaitable[None]: ... + def stop(self) -> None: ... + class SourceTransformAsyncServer: def __init__( self, + handler: Callable[[Datum], Awaitable[list[Message]]], + *, sock_file: str | None = ..., - info_file: str | None = ..., + server_info_file: str | None = ..., + install_signal_handlers: bool = ..., ) -> None: ... - def start(self, py_func: SourceTransformer) -> Awaitable[None]: ... + async def serve(self) -> None: ... def stop(self) -> None: ... + async def wait_ready(self, timeout: float = ...) -> None: ... + async def wait_for_termination(self) -> None: ... + async def __aenter__(self) -> SourceTransformAsyncServer: ... + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + tb: TracebackType | None, + ) -> None: ... __all__ = [ "Datum", "Message", - "Messages", "NackOptions", "SourceTransformAsyncServer", "SourceTransformer", - "SystemMetadata", - "UserMetadata", ] diff --git a/packages/pynumaflow-lite/src/sourcetransform/mod.rs b/packages/pynumaflow-lite/src/sourcetransform/mod.rs index aa09a305..bb6bd3b9 100644 --- a/packages/pynumaflow-lite/src/sourcetransform/mod.rs +++ b/packages/pynumaflow-lite/src/sourcetransform/mod.rs @@ -1,231 +1,73 @@ use std::collections::HashMap; +use std::path::PathBuf; +use std::time::Duration; +use std::sync::Mutex; +use chrono::{DateTime, Utc}; +use tokio::net::UnixStream; +use tokio::time::Instant; +use tonic::transport::Uri; +use tower::service_fn; +use numaflow::proto::source_transformer::source_transform_client::SourceTransformClient; use numaflow::sourcetransform; +use pyo3::prelude::*; -use chrono::{DateTime, Utc}; +use crate::nack::NackOptions; +use crate::pyrs::{bytes_literal, py_repr}; /// SourceTransform interface managed by Python. It means Python code will start the server /// and can pass in the Python function. pub mod server; -use pyo3::prelude::*; -use std::sync::Mutex; - -use crate::nack::NackOptions; - -/// SystemMetadata wraps system-generated metadata groups per message. -/// It is read-only to UDFs. -#[pyclass(module = "pynumaflow_lite.sourcetransformer", from_py_object)] -#[derive(Clone, Default, Debug)] -pub struct SystemMetadata { - data: HashMap>>, -} - -#[pymethods] -impl SystemMetadata { - #[new] - #[pyo3(signature = () -> "SystemMetadata")] - fn new() -> Self { - Self::default() - } - /// Returns the groups of the system metadata. - /// If there are no groups, it returns an empty list. - #[pyo3(signature = () -> "list[str]")] - fn groups(&self) -> Vec { - self.data.keys().cloned().collect() - } - - /// Returns the keys of the system metadata for the given group. - /// If there are no keys or the group is not present, it returns an empty list. - #[pyo3(signature = (group: "str") -> "list[str]")] - fn keys(&self, group: &str) -> Vec { - self.data - .get(group) - .map(|kv| kv.keys().cloned().collect()) - .unwrap_or_default() - } - - /// Returns the value of the system metadata for the given group and key. - /// If there is no value or the group or key is not present, it returns an empty bytes. - #[pyo3(signature = (group: "str", key: "str") -> "bytes")] - fn value(&self, group: &str, key: &str) -> Vec { - self.data - .get(group) - .and_then(|kv| kv.get(key)) - .cloned() - .unwrap_or_default() - } - - fn __repr__(&self) -> String { - format!("SystemMetadata(groups={:?})", self.groups()) - } -} - -impl From for SystemMetadata { - fn from(value: sourcetransform::SystemMetadata) -> Self { - let mut data = HashMap::new(); - for group in value.groups() { - let mut kv = HashMap::new(); - for key in value.keys(&group) { - kv.insert(key.clone(), value.value(&group, &key)); - } - data.insert(group, kv); +fn system_metadata_to_hash_map( + value: sourcetransform::SystemMetadata, +) -> HashMap>> { + let mut data = HashMap::new(); + for group in value.groups() { + let mut kv = HashMap::new(); + for key in value.keys(&group) { + kv.insert(key.clone(), value.value(&group, &key)); } - Self { data } + data.insert(group, kv); } + data } -/// UserMetadata wraps user-defined metadata groups per message. -/// Users can read and write to this metadata. -#[pyclass(module = "pynumaflow_lite.sourcetransformer", from_py_object)] -#[derive(Clone, Default, Debug)] -pub struct UserMetadata { - data: HashMap>>, -} - -#[pymethods] -impl UserMetadata { - #[new] - #[pyo3(signature = () -> "UserMetadata")] - fn new() -> Self { - Self::default() - } - - /// Returns the groups of the user metadata. - /// If there are no groups, it returns an empty list. - #[pyo3(signature = () -> "list[str]")] - fn groups(&self) -> Vec { - self.data.keys().cloned().collect() - } - - /// Returns the keys of the user metadata for the given group. - /// If there are no keys or the group is not present, it returns an empty list. - #[pyo3(signature = (group: "str") -> "list[str]")] - fn keys(&self, group: &str) -> Vec { - self.data - .get(group) - .map(|kv| kv.keys().cloned().collect()) - .unwrap_or_default() - } - - /// Returns the value of the user metadata for the given group and key. - /// If there is no value or the group or key is not present, it returns an empty bytes. - #[pyo3(signature = (group: "str", key: "str") -> "bytes")] - fn value(&self, group: &str, key: &str) -> Vec { - self.data - .get(group) - .and_then(|kv| kv.get(key)) - .cloned() - .unwrap_or_default() - } - - /// Creates a new group in the user metadata. - /// If the group already exists, this is a no-op. - #[pyo3(signature = (group: "str"))] - fn create_group(&mut self, group: String) { - self.data.entry(group).or_default(); - } - - /// Adds a key-value pair to the user metadata. - /// If the group is not present, it creates a new group. - #[pyo3(signature = (group: "str", key: "str", value: "bytes"))] - fn add_kv(&mut self, group: String, key: String, value: Vec) { - self.data.entry(group).or_default().insert(key, value); - } - - /// Removes a key from a group in the user metadata. - /// If the key or group is not present, it's a no-op. - #[pyo3(signature = (group: "str", key: "str"))] - fn remove_key(&mut self, group: &str, key: &str) { - if let Some(kv) = self.data.get_mut(group) { - kv.remove(key); - } - } - - /// Removes a group from the user metadata. - /// If the group is not present, it's a no-op. - #[pyo3(signature = (group: "str"))] - fn remove_group(&mut self, group: &str) { - self.data.remove(group); - } - - fn __repr__(&self) -> String { - format!("UserMetadata(groups={:?})", self.groups()) - } -} - -impl From for UserMetadata { - fn from(value: sourcetransform::UserMetadata) -> Self { - let mut data = HashMap::new(); - for group in value.groups() { - let mut kv = HashMap::new(); - for key in value.keys(&group) { - kv.insert(key.clone(), value.value(&group, &key)); - } - data.insert(group, kv); +fn user_metadata_to_hash_map( + value: sourcetransform::UserMetadata, +) -> HashMap>> { + let mut data = HashMap::new(); + for group in value.groups() { + let mut kv = HashMap::new(); + for key in value.keys(&group) { + kv.insert(key.clone(), value.value(&group, &key)); } - Self { data } - } -} - -impl From for sourcetransform::UserMetadata { - fn from(value: UserMetadata) -> Self { - let mut user_metadata = sourcetransform::UserMetadata::new(); - for (group, kv_map) in value.data { - for (key, val) in kv_map { - user_metadata.add_kv(group.clone(), key, val); - } - } - user_metadata - } -} - -/// A collection of [Message]s. -#[pyclass(module = "pynumaflow_lite.sourcetransformer", from_py_object)] -#[derive(Clone, Debug)] -pub struct Messages { - pub(crate) messages: Vec, -} - -#[pymethods] -impl Messages { - #[new] - #[pyo3(signature = () -> "Messages")] - fn new() -> Self { - Self { messages: vec![] } - } - - /// Append a [Message] to the collection. - #[pyo3(signature = (message: "Message"))] - fn append(&mut self, message: Message) { - self.messages.push(message); - } - - fn __repr__(&self) -> String { - format!("Messages({:?})", self.messages) - } - - fn __str__(&self) -> String { - format!("Messages({:?})", self.messages) + data.insert(group, kv); } + data } /// A message to be sent to the next vertex with event time transformation. -#[pyclass(module = "pynumaflow_lite.sourcetransformer", from_py_object)] -#[derive(Clone, Default, Debug)] +#[pyclass(module = "pynumaflow_lite.sourcetransformer", from_py_object, eq)] +#[derive(Clone, Debug, PartialEq)] pub struct Message { /// Keys are a collection of strings which will be passed on to the next vertex as is. It can /// be an empty collection. + #[pyo3(get)] pub keys: Option>, /// Value is the value passed to the next vertex. + #[pyo3(get)] pub value: Vec, /// Time for the given event. This will be used for tracking watermarks. + #[pyo3(get)] pub event_time: DateTime, - /// Tags are used for conditional forwarding. + /// Tags are used for [conditional forwarding](https://numaflow.numaproj.io/user-guide/reference/conditional-forwarding/). + #[pyo3(get)] pub tags: Option>, /// User metadata for the message. - pub user_metadata: Option, + #[pyo3(get)] + pub user_metadata: Option>>>, /// Options sent back to the source when nacking the message. #[pyo3(get)] pub nack_options: Option, @@ -233,16 +75,16 @@ pub struct Message { #[pymethods] impl Message { - /// Create a new [Message] with the given value, event_time, keys, tags, and user_metadata. + /// Create a new Message with the given value and event_time. Keys, tags, and + /// user_metadata are optional. #[new] - #[pyo3(signature = (value: "bytes", event_time: "datetime.datetime", keys: "list[str] | None"=None, tags: "list[str] | None"=None, user_metadata: "UserMetadata | None"=None) -> "Message" - )] + #[pyo3(signature = (value: "bytes", event_time: "datetime.datetime", keys: "list[str] | None"=None, tags: "list[str] | None"=None, user_metadata: "dict[str, dict[str, bytes]] | None"=None) -> "Message")] fn new( value: Vec, event_time: DateTime, keys: Option>, tags: Option>, - user_metadata: Option, + user_metadata: Option>>>, ) -> Self { Self { keys, @@ -254,12 +96,12 @@ impl Message { } } - /// Drop a [Message], do not forward to the next vertex. - /// Event time is required because even though a message is dropped, - /// it is still considered as being processed, hence the watermark should be updated. - #[pyo3(signature = (event_time: "datetime.datetime"))] + /// A Message marked to be dropped, i.e. not forwarded to the next vertex. + /// Event time is required because a dropped message is still processed, so the + /// watermark must be updated. #[staticmethod] - fn message_to_drop(event_time: DateTime) -> Self { + #[pyo3(signature = (event_time: "datetime.datetime") -> "Message")] + fn to_drop(event_time: DateTime) -> Self { Self { keys: None, value: vec![], @@ -272,8 +114,8 @@ impl Message { /// A Message marked to be negatively acknowledged (retried), with optional nack options. /// Event time is required so the watermark can still be updated. - #[pyo3(signature = (event_time: "datetime.datetime", nack_options: "NackOptions | None"=None) -> "Message")] #[staticmethod] + #[pyo3(signature = (event_time: "datetime.datetime", nack_options: "NackOptions | None"=None) -> "Message")] fn to_nack(event_time: DateTime, nack_options: Option) -> Self { Self { keys: None, @@ -287,8 +129,8 @@ impl Message { /// A Message marked to be failed. /// Event time is required so the watermark can still be updated. - #[pyo3(signature = (event_time: "datetime.datetime") -> "Message")] #[staticmethod] + #[pyo3(signature = (event_time: "datetime.datetime") -> "Message")] fn to_fail(event_time: DateTime) -> Self { Self { keys: None, @@ -299,6 +141,21 @@ impl Message { nack_options: None, } } + + fn __repr__(&self, py: Python<'_>) -> PyResult { + Ok(format!( + "Message(value={}, event_time={}, keys={}, tags={}, user_metadata={})", + bytes_literal(&self.value), + self.event_time, + py_repr(py, &self.keys)?, + py_repr(py, &self.tags)?, + py_repr(py, &self.user_metadata)?, + )) + } + + fn __str__(&self, py: Python<'_>) -> PyResult { + self.__repr__(py) + } } impl From for sourcetransform::Message { @@ -308,7 +165,13 @@ impl From for sourcetransform::Message { .with_tags(value.tags.unwrap_or_default()); if let Some(user_metadata) = value.user_metadata { - msg = msg.with_user_metadata(user_metadata.into()); + let mut umd = sourcetransform::UserMetadata::new(); + for (group, kv) in user_metadata { + for (key, val) in kv { + umd.add_kv(group.clone(), key, val); + } + } + msg = msg.with_user_metadata(umd); } msg.nack_options = value.nack_options.map(Into::into); @@ -317,7 +180,8 @@ impl From for sourcetransform::Message { } } -/// The incoming [SourceTransformRequest] accessible in Python function. +/// The incoming Datum passed to the source transform handler. It carries the event's keys, +/// value, event_time, watermark, headers, and the user/system metadata. #[pyclass(module = "pynumaflow_lite.sourcetransformer")] pub struct Datum { /// Set of keys in the (key, value) terminology of map/reduce paradigm. @@ -326,10 +190,11 @@ pub struct Datum { /// The value in the (key, value) terminology of map/reduce paradigm. #[pyo3(get)] pub value: Vec, - /// Watermark represented by time is a guarantee that we will not see an element older than this time. + /// [watermark](https://numaflow.numaproj.io/core-concepts/watermarks/) represented by time is a + /// guarantee that we will not see an element older than this time. #[pyo3(get)] pub watermark: DateTime, - /// Time of the element as seen at source or aligned after a reduce operation. + /// Time of the element as seen at source. #[pyo3(get)] pub event_time: DateTime, /// Headers for the message. @@ -337,112 +202,181 @@ pub struct Datum { pub headers: HashMap, /// User metadata for the message. #[pyo3(get)] - pub user_metadata: UserMetadata, + pub user_metadata: HashMap>>, /// System metadata for the message. #[pyo3(get)] - pub system_metadata: SystemMetadata, + pub system_metadata: HashMap>>, } +#[pymethods] impl Datum { + #[new] + #[pyo3(signature = ( + *, + keys: "list[str] | None"=None, + value: "bytes | None"=None, + event_time: "datetime.datetime | None"=None, + watermark: "datetime.datetime | None"=None, + headers: "dict[str, str] | None"=None, + user_metadata: "dict[str, dict[str, bytes]] | None"=None, + system_metadata: "dict[str, dict[str, bytes]] | None"=None, + ) -> "Datum")] fn new( - keys: Vec, - value: Vec, - watermark: DateTime, - event_time: DateTime, - headers: HashMap, - user_metadata: UserMetadata, - system_metadata: SystemMetadata, + keys: Option>, + value: Option>, + event_time: Option>, + watermark: Option>, + headers: Option>, + user_metadata: Option>>>, + system_metadata: Option>>>, ) -> Self { Self { - keys, - value, - watermark, - event_time, - headers, - user_metadata, - system_metadata, + keys: keys.unwrap_or_default(), + value: value.unwrap_or_default(), + watermark: watermark.unwrap_or(DateTime::::UNIX_EPOCH), + event_time: event_time.unwrap_or(DateTime::::UNIX_EPOCH), + headers: headers.unwrap_or_default(), + user_metadata: user_metadata.unwrap_or_default(), + system_metadata: system_metadata.unwrap_or_default(), } } - fn __repr__(&self) -> String { - format!( - "Datum(keys={:?}, value={:?}, watermark={}, event_time={}, headers={:?}, user_metadata={:?}, system_metadata={:?})", - self.keys, - self.value, + fn __repr__(&self, py: Python<'_>) -> PyResult { + Ok(format!( + "Datum(keys={}, value={}, watermark={}, event_time={}, headers={}, user_metadata={}, system_metadata={})", + py_repr(py, &self.keys)?, + bytes_literal(&self.value), self.watermark, self.event_time, - self.headers, - self.user_metadata, - self.system_metadata - ) + py_repr(py, &self.headers)?, + py_repr(py, &self.user_metadata)?, + py_repr(py, &self.system_metadata)?, + )) } - fn __str__(&self) -> String { - format!( - "Datum(keys={:?}, value={:?}, watermark={}, event_time={}, headers={:?}, user_metadata={:?}, system_metadata={:?})", - self.keys, - String::from_utf8_lossy(&self.value), - self.watermark, - self.event_time, - self.headers, - self.user_metadata, - self.system_metadata - ) + fn __str__(&self, py: Python<'_>) -> PyResult { + self.__repr__(py) } } impl From for Datum { fn from(value: sourcetransform::SourceTransformRequest) -> Self { - Datum::new( - value.keys, - value.value, - value.watermark, - value.eventtime, - value.headers, - value.user_metadata.into(), - value.system_metadata.into(), - ) + Self { + keys: value.keys, + value: value.value, + watermark: value.watermark, + event_time: value.eventtime, + headers: value.headers, + user_metadata: user_metadata_to_hash_map(value.user_metadata), + system_metadata: system_metadata_to_hash_map(value.system_metadata), + } + } +} + +async fn source_transform_client( + sock_file: String, +) -> PyResult> { + let endpoint = tonic::transport::Endpoint::try_from("http://[::]:50051") + .map_err(|e| pyo3::PyErr::new::(e.to_string()))?; + + let channel = endpoint + .connect_with_connector(service_fn(move |_: Uri| { + let sock = PathBuf::from(sock_file.clone()); + async move { + Ok::<_, std::io::Error>(hyper_util::rt::TokioIo::new( + UnixStream::connect(sock).await?, + )) + } + })) + .await + .map_err(|e| pyo3::PyErr::new::(e.to_string()))?; + + Ok(SourceTransformClient::new(channel)) +} + +async fn wait_for_ready(sock_file: String, timeout: Duration) -> PyResult<()> { + let deadline = Instant::now() + timeout; + + loop { + if let Ok(mut client) = source_transform_client(sock_file.clone()).await + && let Ok(response) = client.is_ready(()).await + && response.into_inner().ready + { + return Ok(()); + } + + let now = Instant::now(); + if now >= deadline { + return Err(pyo3::PyErr::new::( + "timed out waiting for sourcetransform server readiness", + )); + } + + tokio::time::sleep(std::cmp::min( + Duration::from_millis(100), + deadline.saturating_duration_since(now), + )) + .await; } } /// Async SourceTransform Server that can be started from Python code which will run the Python UDF function. -#[pyclass(module = "pynumaflow_lite.sourcetransformer")] +#[pyclass( + name = "_SourceTransformAsyncServer", + module = "pynumaflow_lite.sourcetransformer" +)] pub struct SourceTransformAsyncServer { sock_file: String, - info_file: String, + server_info_file: String, shutdown_tx: Mutex>>, } #[pymethods] impl SourceTransformAsyncServer { #[new] - #[pyo3(signature = (sock_file: "str | None"=sourcetransform::SOCK_ADDR.to_string(), info_file: "str | None"=sourcetransform::SERVER_INFO_FILE.to_string()) -> "SourceTransformAsyncServer" - )] - fn new(sock_file: String, info_file: String) -> Self { + #[pyo3(signature = ( + sock_file: "str | None"=None, + server_info_file: "str | None"=None, + ) -> "_SourceTransformAsyncServer")] + fn new(sock_file: Option, server_info_file: Option) -> Self { Self { - sock_file, - info_file, + sock_file: sock_file.unwrap_or_else(|| sourcetransform::SOCK_ADDR.to_string()), + server_info_file: server_info_file + .unwrap_or_else(|| sourcetransform::SERVER_INFO_FILE.to_string()), shutdown_tx: Mutex::new(None), } } /// Start the server with the given Python function. - #[pyo3(signature = (py_func: "callable") -> "None")] - pub fn start<'a>(&self, py: Python<'a>, py_func: Py) -> PyResult> { + #[pyo3(signature = (handler: "callable") -> "None")] + pub fn start<'a>(&self, py: Python<'a>, handler: Py) -> PyResult> { let sock_file = self.sock_file.clone(); - let info_file = self.info_file.clone(); + let server_info_file = self.server_info_file.clone(); let (tx, rx) = tokio::sync::oneshot::channel::<()>(); { let mut guard = self.shutdown_tx.lock().unwrap(); *guard = Some(tx); } - pyo3_async_runtimes::tokio::future_into_py(py, async move { - crate::sourcetransform::server::start(py_func, sock_file, info_file, rx) - .await - .expect("server failed to start"); - Ok(()) - }) + pyo3_async_runtimes::tokio::future_into_py( + py, + crate::sourcetransform::server::start(handler, sock_file, server_info_file, rx), + ) + } + + /// Wait until the Numaflow IsReady probe succeeds over the sourcetransform UDS. + #[pyo3(signature = (timeout: "float"=30.0) -> "None")] + pub fn wait_ready<'a>(&self, py: Python<'a>, timeout: f64) -> PyResult> { + if !timeout.is_finite() || timeout < 0.0 { + return Err(pyo3::PyErr::new::( + "timeout must be a non-negative finite float", + )); + } + + let sock_file = self.sock_file.clone(); + let timeout = Duration::from_secs_f64(timeout); + + pyo3_async_runtimes::tokio::future_into_py(py, wait_for_ready(sock_file, timeout)) } /// Trigger server shutdown from Python (idempotent). @@ -457,9 +391,6 @@ impl SourceTransformAsyncServer { /// Helper to populate a PyModule with sourcetransform types/functions. pub(crate) fn populate_py_module(m: &Bound) -> PyResult<()> { - m.add_class::()?; - m.add_class::()?; - m.add_class::()?; m.add_class::()?; m.add_class::()?; m.add_class::()?; diff --git a/packages/pynumaflow-lite/src/sourcetransform/server.rs b/packages/pynumaflow-lite/src/sourcetransform/server.rs index 31ee6b7e..2c3325f8 100644 --- a/packages/pynumaflow-lite/src/sourcetransform/server.rs +++ b/packages/pynumaflow-lite/src/sourcetransform/server.rs @@ -1,13 +1,31 @@ -use crate::sourcetransform::{Datum, Messages}; +use std::sync::{Arc, Mutex}; + use numaflow::shared::ServerExtras; use numaflow::sourcetransform; - +use pyo3::exceptions::PyTypeError; use pyo3::prelude::*; -use std::sync::Arc; + +use crate::pyrs::{combine_errors, format_error}; pub(crate) struct PySourceTransformRunner { pub(crate) event_loop: Arc>, pub(crate) py_func: Arc>, + pub(crate) errors: Arc>>, +} + +impl PySourceTransformRunner { + fn fail(&self, error: PyErr) -> ! { + // numaflow calls transform() concurrently, so each error belongs to a + // different message. Keep all of them. start() raises them together, which + // lets Python format every traceback instead of Rust printing them by hand. + let message = Python::attach(|py| format_error(py, &error)); + self.errors.lock().unwrap().push(error); + + // numaflow catches this panic, sends a gRPC error for this message, and + // starts the server shutdown. An empty result would instead look like a + // message that the handler dropped on purpose. + panic!("{message}"); + } } #[tonic::async_trait] @@ -16,28 +34,52 @@ impl sourcetransform::SourceTransformer for PySourceTransformRunner { &self, input: sourcetransform::SourceTransformRequest, ) -> Vec { - let fut = Python::attach(|py| { - let keys = input.keys.clone(); - let input: Datum = input.into(); - let py_func = self.py_func.clone(); - + let fut = match Python::attach(|py| -> PyResult<_> { let locals = pyo3_async_runtimes::TaskLocals::new(self.event_loop.bind(py).clone()); - - let coro = py_func.call1(py, (keys, input)).unwrap().into_bound(py); - - pyo3_async_runtimes::into_future_with_locals(&locals, coro).unwrap() - }); - - let result = fut.await.unwrap(); - - let result = Python::attach(|py| { - let x: Messages = result.extract(py).unwrap(); - x - }); - - println!("{:?}", result); - - result.messages.into_iter().map(|m| m.into()).collect() + let datum: crate::sourcetransform::Datum = input.into(); + let coro = self.py_func.call1(py, (datum,))?.into_bound(py); + let is_awaitable: bool = py + .import("inspect")? + .call_method1("isawaitable", (&coro,))? + .extract()?; + if !is_awaitable { + return Err(PyErr::new::( + "sourcetransform handler must be an async function (coroutine)", + )); + } + pyo3_async_runtimes::into_future_with_locals(&locals, coro).map_err(|_| { + PyErr::new::( + "sourcetransform handler must be an async function (coroutine)", + ) + }) + }) { + Ok(fut) => fut, + Err(error) => self.fail(error), + }; + + let result = match fut.await { + Ok(result) => result, + Err(error) => self.fail(error), + }; + + let messages: Vec = match Python::attach(|py| { + result.extract(py).map_err(|_| { + let type_name = result + .bind(py) + .get_type() + .name() + .map(|name| name.to_string_lossy().into_owned()) + .unwrap_or_else(|_| "".to_string()); + PyErr::new::(format!( + "sourcetransform handler must return list[Message], got {type_name}" + )) + }) + }) { + Ok(messages) => messages, + Err(error) => self.fail(error), + }; + + messages.into_iter().map(|m| m.into()).collect() } } @@ -49,14 +91,24 @@ pub(super) async fn start( shutdown_rx: tokio::sync::oneshot::Receiver<()>, ) -> Result<(), pyo3::PyErr> { let (tx, rx) = tokio::sync::oneshot::channel(); - let py_asyncio_loop_handle = tokio::task::spawn_blocking(move || crate::pyrs::run_asyncio(tx)); + let py_asyncio_loop_handle = tokio::task::spawn_blocking({ + println!( + "Starting SourceTransform UDF. socket={}, server_info={}", + sock_file, info_file + ); + move || crate::pyrs::run_asyncio(tx) + }); let event_loop = rx.await.unwrap(); - let (sig_handle, combined_rx) = crate::pyrs::setup_sig_handler(shutdown_rx); + let errors = Arc::new(Mutex::new(Vec::new())); + // Shutdown has two sources, and neither one needs a channel here. The Python + // side signals stop() through shutdown_rx. An uncaught Python error panics in + // fail(), and numaflow then shuts the server down on its own. let py_sourcetransform_runner = PySourceTransformRunner { py_func: Arc::new(py_func), event_loop: event_loop.clone(), + errors: Arc::clone(&errors), }; let server = numaflow::sourcetransform::Server::new(py_sourcetransform_runner) @@ -64,7 +116,7 @@ pub(super) async fn start( .with_server_info_file(info_file); let result = server - .start_with_shutdown(combined_rx) + .start_with_shutdown(shutdown_rx) .await .map_err(|e| pyo3::PyErr::new::(e.to_string())); @@ -80,10 +132,9 @@ pub(super) async fn start( // Wait for the blocking asyncio thread to finish. let _ = py_asyncio_loop_handle.await; - // if not finished, abort it - if !sig_handle.is_finished() { - println!("Aborting signal handler"); - sig_handle.abort(); + let errors = std::mem::take(&mut *errors.lock().unwrap()); + if !errors.is_empty() { + return Err(Python::attach(|py| combine_errors(py, errors))); } result diff --git a/packages/pynumaflow-lite/tests/bin/sourcetransform.rs b/packages/pynumaflow-lite/tests/bin/sourcetransform.rs index 46473945..0cedbc6e 100644 --- a/packages/pynumaflow-lite/tests/bin/sourcetransform.rs +++ b/packages/pynumaflow-lite/tests/bin/sourcetransform.rs @@ -81,6 +81,17 @@ async fn main() -> Result<(), Box> { assert_eq!(msg.value, "hello".as_bytes()); // Verify event_time is set assert!(msg.event_time.is_some()); + assert!(msg.tags.contains(&"after_year_2022".to_string())); + // Verify the user metadata that the handler created + let filter_info = &msg + .metadata + .as_ref() + .expect("metadata is set") + .user_metadata["filter_info"]; + assert_eq!( + filter_info.key_value["filter_result"], + b"after_year_2022".to_vec() + ); // Request 2 - message to be dropped (event time in 2021) let request_2 = proto::source_transformer::SourceTransformRequest { diff --git a/packages/pynumaflow-lite/tests/examples/sourcetransform_event_filter.py b/packages/pynumaflow-lite/tests/examples/sourcetransform_event_filter.py index fb238f38..c2e28873 100644 --- a/packages/pynumaflow-lite/tests/examples/sourcetransform_event_filter.py +++ b/packages/pynumaflow-lite/tests/examples/sourcetransform_event_filter.py @@ -1,16 +1,13 @@ import asyncio -import signal -from collections.abc import Callable from datetime import datetime, timezone -from pynumaflow_lite import sourcetransformer +from pynumaflow_lite.sourcetransformer import Datum, Message, SourceTransformAsyncServer, SourceTransformer -# Define epoch timestamps for filtering january_first_2022 = datetime(2022, 1, 1, tzinfo=timezone.utc) january_first_2023 = datetime(2023, 1, 1, tzinfo=timezone.utc) -class EventFilter(sourcetransformer.SourceTransformer): +class EventFilter(SourceTransformer): """ A source transformer that filters and routes messages based on event time. @@ -21,97 +18,55 @@ class EventFilter(sourcetransformer.SourceTransformer): Also demonstrates reading and creating metadata. """ - async def handler(self, keys: list[str], datum: sourcetransformer.Datum) -> sourcetransformer.Messages: - val = datum.value - event_time = datum.event_time - messages = sourcetransformer.Messages() - + async def handler(self, datum: Datum) -> list[Message]: # Read system metadata (read-only) - print(f"System metadata groups: {datum.system_metadata.groups()}") - for group in datum.system_metadata.groups(): - for key in datum.system_metadata.keys(group): - value = datum.system_metadata.value(group, key) + print(f"System metadata groups: {list(datum.system_metadata)}") + for group, key_values in datum.system_metadata.items(): + for key, value in key_values.items(): print(f" System[{group}][{key}] = {value}") - # Read user metadata (read-only from input) - print(f"User metadata groups: {datum.user_metadata.groups()}") - for group in datum.user_metadata.groups(): - for key in datum.user_metadata.keys(group): - value = datum.user_metadata.value(group, key) + # Read user metadata + print(f"User metadata groups: {list(datum.user_metadata)}") + for group, key_values in datum.user_metadata.items(): + for key, value in key_values.items(): print(f" User[{group}][{key}] = {value}") + event_time = datum.event_time if event_time < january_first_2022: print(f"Got event time: {event_time}, it is before 2022, so dropping") - messages.append(sourcetransformer.Message.message_to_drop(event_time)) - elif event_time < january_first_2023: - print(f"Got event time: {event_time}, it is within year 2022, so forwarding to within_year_2022") - - # Create user metadata for the outgoing message - user_metadata = sourcetransformer.UserMetadata() - user_metadata.create_group("filter_info") - user_metadata.add_kv("filter_info", "filter_result", b"within_year_2022") - user_metadata.add_kv("filter_info", "original_event_time", str(event_time).encode()) + return [Message.to_drop(event_time)] - messages.append( - sourcetransformer.Message( - value=val, - event_time=january_first_2022, - keys=keys, - tags=["within_year_2022"], - user_metadata=user_metadata, - ) - ) + if event_time < january_first_2023: + print(f"Got event time: {event_time}, it is within year 2022, so forwarding to within_year_2022") + tag, new_event_time = "within_year_2022", january_first_2022 else: print(f"Got event time: {event_time}, it is after year 2022, so forwarding to after_year_2022") - - # Create user metadata for the outgoing message - user_metadata = sourcetransformer.UserMetadata() - user_metadata.create_group("filter_info") - user_metadata.add_kv("filter_info", "filter_result", b"after_year_2022") - user_metadata.add_kv("filter_info", "original_event_time", str(event_time).encode()) - - messages.append( - sourcetransformer.Message( - value=val, - event_time=january_first_2023, - keys=keys, - tags=["after_year_2022"], - user_metadata=user_metadata, - ) + tag, new_event_time = "after_year_2022", january_first_2023 + + user_metadata = { + "filter_info": { + "filter_result": tag.encode(), + "original_event_time": str(event_time).encode(), + } + } + return [ + Message( + datum.value, + new_event_time, + keys=datum.keys, + tags=[tag], + user_metadata=user_metadata, ) + ] - return messages - - -async def start( - f: Callable[[list[str], sourcetransformer.Datum], sourcetransformer.Messages], -): - sock_file = "/tmp/var/run/numaflow/sourcetransform.sock" - server_info_file = "/tmp/var/run/numaflow/sourcetransformer-server-info" - server = sourcetransformer.SourceTransformAsyncServer(sock_file, server_info_file) - - # Register loop-level signal handlers so we control shutdown and avoid asyncio.run - # converting it into KeyboardInterrupt/CancelledError traces. - loop = asyncio.get_running_loop() - loop.set_debug(True) - print("Registering signal handlers", loop) - try: - loop.add_signal_handler(signal.SIGINT, lambda: server.stop()) - loop.add_signal_handler(signal.SIGTERM, lambda: server.stop()) - except (NotImplementedError, RuntimeError): - print("Failed to register signal handlers") - # add_signal_handler may not be available on some platforms/contexts; fallback below. - pass - try: - await server.start(f) - print("Shutting down gracefully...") - except asyncio.CancelledError: - # Fallback in case the task was cancelled by the runner - server.stop() - return +async def main() -> None: + await SourceTransformAsyncServer( + EventFilter(), + sock_file="/tmp/var/run/numaflow/sourcetransform.sock", + server_info_file="/tmp/var/run/numaflow/sourcetransformer-server-info", + ).serve() if __name__ == "__main__": - async_handler = EventFilter() - asyncio.run(start(async_handler)) + asyncio.run(main()) diff --git a/packages/pynumaflow-lite/tests/examples/sourcetransform_event_filter_func.py b/packages/pynumaflow-lite/tests/examples/sourcetransform_event_filter_func.py new file mode 100644 index 00000000..ac2a337b --- /dev/null +++ b/packages/pynumaflow-lite/tests/examples/sourcetransform_event_filter_func.py @@ -0,0 +1,46 @@ +import asyncio +from datetime import datetime, timezone + +from pynumaflow_lite.sourcetransformer import Datum, Message, SourceTransformAsyncServer + +january_first_2022 = datetime(2022, 1, 1, tzinfo=timezone.utc) +january_first_2023 = datetime(2023, 1, 1, tzinfo=timezone.utc) + + +async def event_filter(datum: Datum) -> list[Message]: + event_time = datum.event_time + if event_time < january_first_2022: + return [Message.to_drop(event_time)] + + if event_time < january_first_2023: + tag, new_event_time = "within_year_2022", january_first_2022 + else: + tag, new_event_time = "after_year_2022", january_first_2023 + + user_metadata = { + "filter_info": { + "filter_result": tag.encode(), + "original_event_time": str(event_time).encode(), + } + } + return [ + Message( + datum.value, + new_event_time, + keys=datum.keys, + tags=[tag], + user_metadata=user_metadata, + ) + ] + + +async def main() -> None: + await SourceTransformAsyncServer( + event_filter, + sock_file="/tmp/var/run/numaflow/sourcetransform.sock", + server_info_file="/tmp/var/run/numaflow/sourcetransformer-server-info", + ).serve() + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/packages/pynumaflow-lite/tests/test_sourcetransform.py b/packages/pynumaflow-lite/tests/test_sourcetransform.py index 01274f15..cfd04b46 100644 --- a/packages/pynumaflow-lite/tests/test_sourcetransform.py +++ b/packages/pynumaflow-lite/tests/test_sourcetransform.py @@ -9,6 +9,7 @@ SCRIPTS = [ "sourcetransform_event_filter.py", + "sourcetransform_event_filter_func.py", ] diff --git a/packages/pynumaflow-lite/tests/test_sourcetransform_unit.py b/packages/pynumaflow-lite/tests/test_sourcetransform_unit.py new file mode 100644 index 00000000..03e39d88 --- /dev/null +++ b/packages/pynumaflow-lite/tests/test_sourcetransform_unit.py @@ -0,0 +1,119 @@ +import asyncio +import os +import subprocess +import uuid +from contextlib import suppress +from datetime import datetime, timezone +from pathlib import Path + +import pytest + +from pynumaflow_lite import sourcetransformer + +CARGO_ROOT = Path(__file__).resolve().parent.parent + +EVENT_TIME = datetime(2023, 6, 1, tzinfo=timezone.utc) + + +async def _unit_handler(datum: sourcetransformer.Datum) -> list[sourcetransformer.Message]: + assert datum.user_metadata["custom"]["version"] == b"1.0.0" + return [sourcetransformer.Message(datum.value, datum.event_time, keys=datum.keys, tags=["seen"])] + + +def test_sourcetransform_handler_can_be_unit_tested_without_server(): + datum = sourcetransformer.Datum( + keys=["key"], + value=b"value", + event_time=EVENT_TIME, + user_metadata={"custom": {"version": b"1.0.0"}}, + ) + + messages = asyncio.run(_unit_handler(datum)) + + assert messages == [sourcetransformer.Message(b"value", EVENT_TIME, keys=["key"], tags=["seen"])] + + +def test_message_attributes_and_repr(): + message = sourcetransformer.Message( + b"value", + EVENT_TIME, + keys=["key"], + user_metadata={"group": {"k": b"v"}}, + ) + + assert message.value == b"value" + assert message.event_time == EVENT_TIME + assert message.keys == ["key"] + assert message.tags is None + assert message.user_metadata == {"group": {"k": b"v"}} + assert repr(message).startswith("Message(value=b'value', ") + assert sourcetransformer.Message.to_drop(EVENT_TIME).event_time == EVENT_TIME + + +async def _run_sourcetransform_client(sock_path: Path) -> subprocess.CompletedProcess[str]: + return await asyncio.to_thread( + subprocess.run, + ["cargo", "run", "--quiet", "--bin", "test_sourcetransform", "--", str(sock_path)], + cwd=CARGO_ROOT, + capture_output=True, + text=True, + timeout=60, + check=False, + ) + + +async def _exercise_server( + tmp_path: Path, + handler, + clients: list[subprocess.CompletedProcess[str]] | None = None, +) -> None: + path_id = f"{tmp_path.name[:12]}-{uuid.uuid4().hex[:12]}" + sock_path = Path(f"/tmp/pnl-{path_id}.sock") + server_info_path = Path(f"/tmp/pnl-{path_id}.info") + server = sourcetransformer.SourceTransformAsyncServer( + handler, + sock_file=str(sock_path), + server_info_file=str(server_info_path), + ) + + try: + async with server: + client = await _run_sourcetransform_client(sock_path) + if clients is not None: + clients.append(client) + finally: + for path in (sock_path, server_info_path): + with suppress(FileNotFoundError): + os.unlink(path) + + +def test_sourcetransform_server_propagates_handler_exception(tmp_path: Path): + async def handler(datum: sourcetransformer.Datum) -> list[sourcetransformer.Message]: + raise RuntimeError("transform exploded") + + clients: list[subprocess.CompletedProcess[str]] = [] + with pytest.raises(RuntimeError, match="transform exploded"): + asyncio.run(_exercise_server(tmp_path, handler, clients)) + + # The client must get a gRPC error, not an empty result. Do not check for the + # handler text: numaflow-rs keeps only the first panic of the process, and an + # earlier test in the same pytest process can own it. + [client] = clients + assert client.returncode != 0 + assert "UDF_EXECUTION_ERROR" in client.stderr + + +def test_sourcetransform_server_rejects_non_list_response(tmp_path: Path): + async def handler(datum: sourcetransformer.Datum) -> str: + return "not messages" + + with pytest.raises(TypeError, match=r"sourcetransform handler must return list\[Message\]"): + asyncio.run(_exercise_server(tmp_path, handler)) + + +def test_sourcetransform_server_rejects_sync_handler(tmp_path: Path): + def handler(datum: sourcetransformer.Datum) -> list[sourcetransformer.Message]: + return [] + + with pytest.raises(TypeError, match="sourcetransform handler must be an async function"): + asyncio.run(_exercise_server(tmp_path, handler)) From 2bc555276ca2fbcc181b8c169c89aa64723f891c Mon Sep 17 00:00:00 2001 From: Sreekanth Date: Sat, 10 Oct 2026 08:28:05 +0530 Subject: [PATCH 2/3] Python's repr like output for when printing Signed-off-by: Sreekanth --- packages/pynumaflow-lite/src/batchmap/mod.rs | 4 +- packages/pynumaflow-lite/src/map/mod.rs | 61 +++++++------------ packages/pynumaflow-lite/src/mapstream/mod.rs | 26 +++----- packages/pynumaflow-lite/src/nack.rs | 21 +++---- packages/pynumaflow-lite/src/sink/mod.rs | 4 +- .../src/sourcetransform/mod.rs | 5 +- 6 files changed, 46 insertions(+), 75 deletions(-) diff --git a/packages/pynumaflow-lite/src/batchmap/mod.rs b/packages/pynumaflow-lite/src/batchmap/mod.rs index 3c73b253..11c2c13f 100644 --- a/packages/pynumaflow-lite/src/batchmap/mod.rs +++ b/packages/pynumaflow-lite/src/batchmap/mod.rs @@ -92,9 +92,7 @@ impl Message { bytes_literal(&self.value), py_repr(py, &self.keys)?, py_repr(py, &self.tags)?, - self.nack_options - .as_ref() - .map_or_else(|| "None".to_string(), NackOptions::__repr__), + py_repr(py, self.nack_options.clone())?, )) } } diff --git a/packages/pynumaflow-lite/src/map/mod.rs b/packages/pynumaflow-lite/src/map/mod.rs index 8bcdfdea..6a2e4fc3 100644 --- a/packages/pynumaflow-lite/src/map/mod.rs +++ b/packages/pynumaflow-lite/src/map/mod.rs @@ -20,21 +20,7 @@ use pyo3::prelude::*; use std::sync::Mutex; use crate::nack::NackOptions; -use crate::pyrs::bytes_literal; - -fn metadata_literal(metadata: &HashMap>>) -> String { - let groups: Vec = metadata - .iter() - .map(|(group, kv)| { - let entries: Vec = kv - .iter() - .map(|(key, value)| format!("{:?}: {}", key, bytes_literal(value))) - .collect(); - format!("{:?}: {{{}}}", group, entries.join(", ")) - }) - .collect(); - format!("{{{}}}", groups.join(", ")) -} +use crate::pyrs::{bytes_literal, py_repr}; fn system_metadata_to_hash_map( value: map::SystemMetadata, @@ -145,24 +131,19 @@ impl Message { } } - fn __repr__(&self) -> String { - format!( - "Message(value={}, keys={}, tags={}, user_metadata={})", + fn __repr__(&self, py: Python<'_>) -> PyResult { + Ok(format!( + "Message(value={}, keys={}, tags={}, user_metadata={}, nack_options={})", bytes_literal(&self.value), - self.keys - .as_ref() - .map_or_else(|| "None".to_string(), |keys| format!("{keys:?}")), - self.tags - .as_ref() - .map_or_else(|| "None".to_string(), |tags| format!("{tags:?}")), - self.user_metadata - .as_ref() - .map_or_else(|| "None".to_string(), metadata_literal), - ) + py_repr(py, &self.keys)?, + py_repr(py, &self.tags)?, + py_repr(py, &self.user_metadata)?, + py_repr(py, self.nack_options.clone())?, + )) } - fn __str__(&self) -> String { - self.__repr__() + fn __str__(&self, py: Python<'_>) -> PyResult { + self.__repr__(py) } } @@ -247,21 +228,21 @@ impl Datum { } } - fn __repr__(&self) -> String { - format!( - "Datum(keys={:?}, value={}, watermark={}, event_time={}, headers={:?}, user_metadata={}, system_metadata={})", - self.keys, + fn __repr__(&self, py: Python<'_>) -> PyResult { + Ok(format!( + "Datum(keys={}, value={}, watermark={}, event_time={}, headers={}, user_metadata={}, system_metadata={})", + py_repr(py, &self.keys)?, bytes_literal(&self.value), self.watermark, self.event_time, - self.headers, - metadata_literal(&self.user_metadata), - metadata_literal(&self.system_metadata) - ) + py_repr(py, &self.headers)?, + py_repr(py, &self.user_metadata)?, + py_repr(py, &self.system_metadata)?, + )) } - fn __str__(&self) -> String { - self.__repr__() + fn __str__(&self, py: Python<'_>) -> PyResult { + self.__repr__(py) } } diff --git a/packages/pynumaflow-lite/src/mapstream/mod.rs b/packages/pynumaflow-lite/src/mapstream/mod.rs index c35c5560..c5bda782 100644 --- a/packages/pynumaflow-lite/src/mapstream/mod.rs +++ b/packages/pynumaflow-lite/src/mapstream/mod.rs @@ -10,7 +10,7 @@ pub mod server; use crate::map::wait_for_ready; use crate::nack::NackOptions; -use crate::pyrs::bytes_literal; +use crate::pyrs::{bytes_literal, py_repr}; /// Streaming Datum mirrors MapStreamRequest for Python #[pyclass(module = "pynumaflow_lite.mapstreamer", from_py_object)] @@ -154,24 +154,18 @@ impl Message { } } - fn __repr__(&self) -> String { - format!( - "Message(value={}, keys={}, tags={}, user_metadata={})", + fn __repr__(&self, py: Python<'_>) -> PyResult { + Ok(format!( + "Message(value={}, keys={}, tags={}, nack_options={})", bytes_literal(&self.value), - self.keys - .as_ref() - .map_or_else(|| "None".to_string(), |keys| format!("{keys:?}")), - self.tags - .as_ref() - .map_or_else(|| "None".to_string(), |tags| format!("{tags:?}")), - self.nack_options - .as_ref() - .map_or_else(|| "None".to_string(), NackOptions::__repr__), - ) + py_repr(py, &self.keys)?, + py_repr(py, &self.tags)?, + py_repr(py, self.nack_options.clone())?, + )) } - fn __str__(&self) -> String { - self.__repr__() + fn __str__(&self, py: Python<'_>) -> PyResult { + self.__repr__(py) } } diff --git a/packages/pynumaflow-lite/src/nack.rs b/packages/pynumaflow-lite/src/nack.rs index 3539875e..3c9377f9 100644 --- a/packages/pynumaflow-lite/src/nack.rs +++ b/packages/pynumaflow-lite/src/nack.rs @@ -2,6 +2,8 @@ use std::collections::HashMap; use pyo3::prelude::*; +use crate::pyrs::py_repr; + /// Per-message redelivery options for a nack. /// /// Shared across all UDF types (mapper, batchmapper, mapstreamer, @@ -41,17 +43,14 @@ impl NackOptions { } } - pub(crate) fn __repr__(&self) -> String { - let opt = |v: &Option| v.as_ref().map_or_else(|| "None".to_string(), |s| s.clone()); - format!( - "NackOptions(delay={}, max_deliveries={}, reason={}, nack_map={:?})", - self.delay - .map_or_else(|| "None".to_string(), |v| v.to_string()), - self.max_deliveries - .map_or_else(|| "None".to_string(), |v| v.to_string()), - opt(&self.reason), - self.nack_map - ) + pub(crate) fn __repr__(&self, py: Python<'_>) -> PyResult { + Ok(format!( + "NackOptions(delay={}, max_deliveries={}, reason={}, nack_map={})", + py_repr(py, self.delay)?, + py_repr(py, self.max_deliveries)?, + py_repr(py, &self.reason)?, + py_repr(py, &self.nack_map)?, + )) } } diff --git a/packages/pynumaflow-lite/src/sink/mod.rs b/packages/pynumaflow-lite/src/sink/mod.rs index bb927eac..e63e1a13 100644 --- a/packages/pynumaflow-lite/src/sink/mod.rs +++ b/packages/pynumaflow-lite/src/sink/mod.rs @@ -235,9 +235,7 @@ impl Response { ), ResponseType::Nack => format!( "Response.nack(id={id}, nack_options={})", - self.nack_options - .as_ref() - .map_or_else(|| "None".to_string(), NackOptions::__repr__) + py_repr(py, self.nack_options.clone())? ), }) } diff --git a/packages/pynumaflow-lite/src/sourcetransform/mod.rs b/packages/pynumaflow-lite/src/sourcetransform/mod.rs index bb6bd3b9..c8de8967 100644 --- a/packages/pynumaflow-lite/src/sourcetransform/mod.rs +++ b/packages/pynumaflow-lite/src/sourcetransform/mod.rs @@ -144,12 +144,13 @@ impl Message { fn __repr__(&self, py: Python<'_>) -> PyResult { Ok(format!( - "Message(value={}, event_time={}, keys={}, tags={}, user_metadata={})", + "Message(value={}, event_time={}, keys={}, tags={}, user_metadata={}, nack_options={})", bytes_literal(&self.value), - self.event_time, + py_repr(py, self.event_time)?, py_repr(py, &self.keys)?, py_repr(py, &self.tags)?, py_repr(py, &self.user_metadata)?, + py_repr(py, self.nack_options.clone())?, )) } From 2d261f808fff69b8327e0d9979f72658ae83efdb Mon Sep 17 00:00:00 2001 From: Sreekanth Date: Sat, 10 Oct 2026 08:35:43 +0530 Subject: [PATCH 3/3] improve tests, formatting Signed-off-by: Sreekanth --- .../src/sourcetransform/mod.rs | 9 +-- .../tests/bin/sourcetransform.rs | 77 +++++++++++++++++-- .../examples/sourcetransform_event_filter.py | 9 ++- .../sourcetransform_event_filter_func.py | 9 ++- 4 files changed, 89 insertions(+), 15 deletions(-) diff --git a/packages/pynumaflow-lite/src/sourcetransform/mod.rs b/packages/pynumaflow-lite/src/sourcetransform/mod.rs index c8de8967..c2ab11a6 100644 --- a/packages/pynumaflow-lite/src/sourcetransform/mod.rs +++ b/packages/pynumaflow-lite/src/sourcetransform/mod.rs @@ -1,16 +1,16 @@ use std::collections::HashMap; use std::path::PathBuf; -use std::time::Duration; use std::sync::Mutex; +use std::time::Duration; use chrono::{DateTime, Utc}; +use numaflow::proto::source_transformer::source_transform_client::SourceTransformClient; +use numaflow::sourcetransform; +use pyo3::prelude::*; use tokio::net::UnixStream; use tokio::time::Instant; use tonic::transport::Uri; use tower::service_fn; -use numaflow::proto::source_transformer::source_transform_client::SourceTransformClient; -use numaflow::sourcetransform; -use pyo3::prelude::*; use crate::nack::NackOptions; use crate::pyrs::{bytes_literal, py_repr}; @@ -19,7 +19,6 @@ use crate::pyrs::{bytes_literal, py_repr}; /// and can pass in the Python function. pub mod server; - fn system_metadata_to_hash_map( value: sourcetransform::SystemMetadata, ) -> HashMap>> { diff --git a/packages/pynumaflow-lite/tests/bin/sourcetransform.rs b/packages/pynumaflow-lite/tests/bin/sourcetransform.rs index 0cedbc6e..0ae59f90 100644 --- a/packages/pynumaflow-lite/tests/bin/sourcetransform.rs +++ b/packages/pynumaflow-lite/tests/bin/sourcetransform.rs @@ -1,3 +1,4 @@ +use std::collections::HashMap; use std::env; use std::path::PathBuf; @@ -52,7 +53,7 @@ async fn main() -> Result<(), Box> { let handshake_response = handshake_response.unwrap(); assert!(handshake_response.handshake.is_some()); - // Request 1 - normal message (event time in 2023) + // Request 1 - normal message (event time in 2023), with user and system metadata let request_1 = proto::source_transformer::SourceTransformRequest { request: Some( proto::source_transformer::source_transform_request::Request { @@ -65,7 +66,17 @@ async fn main() -> Result<(), Box> { nanos: 0, }), headers: Default::default(), - metadata: None, + metadata: Some(proto::metadata::Metadata { + previous_vertex: "in".to_string(), + sys_metadata: HashMap::from([( + "sys_group".to_string(), + kv_group(&[("sys_key", b"sys_value")]), + )]), + user_metadata: HashMap::from([( + "user_group".to_string(), + kv_group(&[("user_key", b"user_value")]), + )]), + }), }, ), handshake: None, @@ -79,19 +90,35 @@ async fn main() -> Result<(), Box> { let msg = &r.results[0]; assert_eq!(msg.keys.first(), Some(&"first".to_owned())); assert_eq!(msg.value, "hello".as_bytes()); - // Verify event_time is set - assert!(msg.event_time.is_some()); + // Verify the handler set the event time to 2023-01-01 00:00:00 UTC + assert_eq!( + msg.event_time, + Some(prost_types::Timestamp { + seconds: 1672531200, + nanos: 0, + }) + ); assert!(msg.tags.contains(&"after_year_2022".to_string())); - // Verify the user metadata that the handler created - let filter_info = &msg + let user_metadata = &msg .metadata .as_ref() .expect("metadata is set") - .user_metadata["filter_info"]; + .user_metadata; + // Verify the user metadata that the handler created assert_eq!( - filter_info.key_value["filter_result"], + user_metadata["filter_info"].key_value["filter_result"], b"after_year_2022".to_vec() ); + // Verify the handler received the incoming user metadata in Datum.user_metadata + assert_eq!( + user_metadata["user_group"].key_value["user_key"], + b"user_value".to_vec() + ); + // Verify the handler received the incoming system metadata in Datum.system_metadata + assert_eq!( + user_metadata["system_echo"].key_value["sys_group.sys_key"], + b"sys_value".to_vec() + ); // Request 2 - message to be dropped (event time in 2021) let request_2 = proto::source_transformer::SourceTransformRequest { @@ -146,10 +173,44 @@ async fn main() -> Result<(), Box> { assert_eq!(r.id, "3"); let msg = &r.results[0]; assert_eq!(msg.value, "year_2022_message".as_bytes()); + // Verify the handler set the event time to 2022-01-01 00:00:00 UTC + assert_eq!( + msg.event_time, + Some(prost_types::Timestamp { + seconds: 1640995200, + nanos: 0, + }) + ); assert!(msg.tags.contains(&"within_year_2022".to_string())); + // Verify the user metadata that the handler created + let user_metadata = &msg + .metadata + .as_ref() + .expect("metadata is set") + .user_metadata; + assert_eq!( + user_metadata["filter_info"].key_value["filter_result"], + b"within_year_2022".to_vec() + ); + // No metadata was sent, so the handler forwards no user groups and echoes no system metadata + assert!(!user_metadata.contains_key("user_group")); + assert!( + user_metadata + .get("system_echo") + .is_none_or(|g| g.key_value.is_empty()) + ); // close request stream drop(tx); Ok(()) } + +fn kv_group(pairs: &[(&str, &[u8])]) -> proto::metadata::KeyValueGroup { + proto::metadata::KeyValueGroup { + key_value: pairs + .iter() + .map(|(k, v)| (k.to_string(), v.to_vec())) + .collect(), + } +} diff --git a/packages/pynumaflow-lite/tests/examples/sourcetransform_event_filter.py b/packages/pynumaflow-lite/tests/examples/sourcetransform_event_filter.py index c2e28873..c19bb503 100644 --- a/packages/pynumaflow-lite/tests/examples/sourcetransform_event_filter.py +++ b/packages/pynumaflow-lite/tests/examples/sourcetransform_event_filter.py @@ -43,11 +43,18 @@ async def handler(self, datum: Datum) -> list[Message]: print(f"Got event time: {event_time}, it is after year 2022, so forwarding to after_year_2022") tag, new_event_time = "after_year_2022", january_first_2023 + # Forward the incoming user metadata, echo the system metadata, and add the filter result user_metadata = { + **datum.user_metadata, + "system_echo": { + f"{group}.{key}": value + for group, key_values in datum.system_metadata.items() + for key, value in key_values.items() + }, "filter_info": { "filter_result": tag.encode(), "original_event_time": str(event_time).encode(), - } + }, } return [ Message( diff --git a/packages/pynumaflow-lite/tests/examples/sourcetransform_event_filter_func.py b/packages/pynumaflow-lite/tests/examples/sourcetransform_event_filter_func.py index ac2a337b..9c96d3c2 100644 --- a/packages/pynumaflow-lite/tests/examples/sourcetransform_event_filter_func.py +++ b/packages/pynumaflow-lite/tests/examples/sourcetransform_event_filter_func.py @@ -17,11 +17,18 @@ async def event_filter(datum: Datum) -> list[Message]: else: tag, new_event_time = "after_year_2022", january_first_2023 + # Forward the incoming user metadata, echo the system metadata, and add the filter result user_metadata = { + **datum.user_metadata, + "system_echo": { + f"{group}.{key}": value + for group, key_values in datum.system_metadata.items() + for key, value in key_values.items() + }, "filter_info": { "filter_result": tag.encode(), "original_event_time": str(event_time).encode(), - } + }, } return [ Message(