Implement streaming (#84)

Author: Marenz

RELEASE_NOTES.md

@@ -16,3 +16,4 @@ This release includes a new feature for pagination support in the dispatch list
 - `Client.__init__`:
  - Has a new parameter `connect` which is a boolean that determines if the client should connect to the server on initialization.
  - Automatically sets up the channel for encrypted TLS communication.
+- A new method `stream()` to receive dispatch events in real-time.

src/frequenz/client/dispatch/__main__.py

@@ -105,6 +105,15 @@ async def list_(ctx: click.Context, /, **filters: Any) -> None:
     click.echo(f"{num_dispatches} dispatches total.")
 
 
+@cli.command("stream")
+@click.pass_context
+@click.argument("microgrid-id", required=True, type=int)
+async def stream(ctx: click.Context, microgrid_id: int) -> None:
+    """Stream dispatches."""
+    async for message in ctx.obj["client"].stream(microgrid_id=microgrid_id):
+        click.echo(pformat(message, compact=True))
+
+
 def parse_recurrence(kwargs: dict[str, Any]) -> RecurrenceRule | None:
     """Parse recurrence rule from kwargs."""
     interval = kwargs.pop("interval", 0)
@@ -387,7 +396,16 @@ async def interactive_mode(url: str, key: str) -> None:
     hist_file = os.path.expanduser("~/.dispatch_cli_history.txt")
     session: PromptSession[str] = PromptSession(history=FileHistory(filename=hist_file))
 
-    user_commands = ["list", "create", "update", "get", "delete", "exit", "help"]
+    user_commands = [
+        "list",
+        "stream",
+        "create",
+        "update",
+        "get",
+        "delete",
+        "exit",
+        "help",
+    ]
 
     async def display_help() -> None:
         await cli.main(args=["--help"], standalone_mode=False)

src/frequenz/client/dispatch/_client.py

@@ -18,6 +18,8 @@
     GetMicrogridDispatchResponse,
     ListMicrogridDispatchesRequest,
     ListMicrogridDispatchesResponse,
+    StreamMicrogridDispatchesRequest,
+    StreamMicrogridDispatchesResponse,
 )
 from frequenz.api.dispatch.v1.dispatch_pb2 import (
     TimeIntervalFilter as PBTimeIntervalFilter,
@@ -27,14 +29,17 @@
     UpdateMicrogridDispatchResponse,
 )
 
+from frequenz import channels
 from frequenz.client.base.channel import ChannelOptions, SslOptions
 from frequenz.client.base.client import BaseApiClient
 from frequenz.client.base.conversion import to_timestamp
+from frequenz.client.base.streaming import GrpcStreamBroadcaster
 
 from ._internal_types import DispatchCreateRequest
 from .types import (
     ComponentSelector,
     Dispatch,
+    DispatchEvent,
     RecurrenceRule,
     component_selector_to_protobuf,
 )
@@ -46,6 +51,11 @@
 class Client(BaseApiClient[dispatch_pb2_grpc.MicrogridDispatchServiceStub]):
     """Dispatch API client."""
 
+    streams: dict[
+        int, GrpcStreamBroadcaster[StreamMicrogridDispatchesResponse, DispatchEvent]
+    ] = {}
+    """A dictionary of streamers, keyed by microgrid_id."""
+
     def __init__(
         self,
         *,
@@ -170,6 +180,50 @@ def to_interval(
             else:
                 break
 
+    def stream(self, microgrid_id: int) -> channels.Receiver[DispatchEvent]:
+        """Receive a stream of dispatch events.
+
+        This function returns a receiver channel that can be used to receive
+        dispatch events.
+        An event is one of [CREATE, UPDATE, DELETE].
+
+        Example usage:
+
+        ```
+        client = Client(key="key", server_url="grpc://fz-0004.frequenz.io")
+        async for message in client.stream(microgrid_id=1):
+            print(message.event, message.dispatch)
+        ```
+
+        Args:
+            microgrid_id: The microgrid_id to receive dispatches for.
+
+        Returns:
+            A receiver channel to receive the stream of dispatch events.
+        """
+        return self._get_stream(microgrid_id).new_receiver()
+
+    def _get_stream(
+        self, microgrid_id: int
+    ) -> GrpcStreamBroadcaster[StreamMicrogridDispatchesResponse, DispatchEvent]:
+        """Get an instance to the streaming helper."""
+        broadcaster = self.streams.get(microgrid_id)
+        if broadcaster is None:
+            request = StreamMicrogridDispatchesRequest(microgrid_id=microgrid_id)
+            broadcaster = GrpcStreamBroadcaster(
+                stream_name="StreamMicrogridDispatches",
+                stream_method=lambda: cast(
+                    AsyncIterator[StreamMicrogridDispatchesResponse],
+                    self.stub.StreamMicrogridDispatches(
+                        request, metadata=self._metadata
+                    ),
+                ),
+                transform=DispatchEvent.from_protobuf,
+            )
+            self.streams[microgrid_id] = broadcaster
+
+        return broadcaster
+
     async def create(
         self,
         microgrid_id: int,

src/frequenz/client/dispatch/types.py

@@ -15,6 +15,7 @@
 )
 from frequenz.api.dispatch.v1.dispatch_pb2 import Dispatch as PBDispatch
 from frequenz.api.dispatch.v1.dispatch_pb2 import RecurrenceRule as PBRecurrenceRule
+from frequenz.api.dispatch.v1.dispatch_pb2 import StreamMicrogridDispatchesResponse
 from google.protobuf.json_format import MessageToDict
 
 from frequenz.client.base.conversion import to_datetime, to_timestamp
@@ -348,3 +349,40 @@ def to_protobuf(self) -> PBDispatch:
         pb_dispatch.data.recurrence.CopyFrom(self.recurrence.to_protobuf())
 
         return pb_dispatch
+
+
+class Event(IntEnum):
+    """Enum representing the type of event that occurred during a dispatch operation."""
+
+    UNSPECIFIED = StreamMicrogridDispatchesResponse.Event.EVENT_UNSPECIFIED
+    CREATED = StreamMicrogridDispatchesResponse.Event.EVENT_CREATED
+    UPDATED = StreamMicrogridDispatchesResponse.Event.EVENT_UPDATED
+    DELETED = StreamMicrogridDispatchesResponse.Event.EVENT_DELETED
+
+
+@dataclass(kw_only=True, frozen=True)
+class DispatchEvent:
+    """Represents an event that occurred during a dispatch operation."""
+
+    dispatch: Dispatch
+    """The dispatch associated with the event."""
+
+    event: Event
+    """The type of event that occurred."""
+
+    @classmethod
+    def from_protobuf(
+        cls, pb_object: StreamMicrogridDispatchesResponse
+    ) -> "DispatchEvent":
+        """Convert a protobuf dispatch event to a dispatch event.
+
+        Args:
+            pb_object: The protobuf dispatch event to convert.
+
+        Returns:
+            The converted dispatch event.
+        """
+        return DispatchEvent(
+            dispatch=Dispatch.from_protobuf(pb_object.dispatch),
+            event=Event(pb_object.event),
+        )