encore-dev 2.0.0a7__py3-none-any.whl

encore/storage/redis/pubsub.py

@@ -0,0 +1,276 @@
+"""Redis channel subscriptions and publishing.
+
+These channels are separate from Encore's pubsub.Topic API. Publish bytes, UTF-8
+text, integers, floats, or RedisValue implementations. Select a payload decoder
+with db.pubsub.subscribe[T](channel, *channels) or
+psubscribe[T](pattern, *patterns). Channel names remain bytes. Use an async
+context manager to close each subscription when iteration ends early or raises
+an exception.
+"""
+
+from __future__ import annotations
+
+import asyncio
+from dataclasses import dataclass
+from types import TracebackType
+from typing import Generic, cast, overload
+
+from ...internal.runtime import _native
+from ...tracing._context import active_ctx_bytes
+from ._base import CommandInterface, TypedCommand
+from ._values import T, Value, decode, encode, validate_decoder, U
+
+
+@dataclass(frozen=True)
+class Message(Generic[T]):
+    """A message received from a Redis channel.
+
+    Attributes:
+        channel: The channel name as bytes, without the Encore key prefix.
+        payload: The message body decoded using the subscription's type
+            argument.
+    """
+
+    channel: bytes
+    payload: T
+
+
+class Subscription(Generic[T]):
+    """An async iterator owning a dedicated Redis connection.
+
+    Use ``async with await db.pubsub.subscribe("events") as messages`` to
+    close the connection when iteration ends early or the body raises.
+    """
+
+    @overload
+    def __init__(
+        self: Subscription[bytes],
+        handle: _native.RedisSubscription,
+        *,
+        value_type: type[bytes] = bytes,
+    ) -> None: ...
+
+    @overload
+    def __init__(
+        self, handle: _native.RedisSubscription, *, value_type: type[T]
+    ) -> None: ...
+
+    def __init__(
+        self, handle: _native.RedisSubscription, *, value_type: type[Value] = bytes
+    ) -> None:
+        """Wrap an active native subscription.
+
+        Args:
+            handle: The native subscription. Application code should obtain
+                subscriptions with subscribe() or psubscribe().
+            value_type: Decode payloads as this type; defaults to bytes.
+        """
+        validate_decoder(value_type)
+        self._handle = handle
+        self._value_type = cast("type[T]", value_type)
+
+    async def receive(self) -> Message[T] | None:
+        """Wait for the next message.
+
+        Custom decoding errors propagate to the caller. The message is consumed
+        before decoding; a decoding error does not put it back on the
+        subscription. Closing the subscription interrupts a pending receive.
+
+        Returns:
+            The received message, or None if the subscription is closed or its
+            message stream ends.
+
+        Raises:
+            UnicodeDecodeError: The payload is not UTF-8 when T is str.
+            ValueError: The payload cannot be parsed when T is int or float.
+            EncoreError: Receiving the message fails.
+        """
+        message = await self._handle.receive()
+        return (
+            None
+            if message is None
+            else Message(message[0], decode(message[1], self._value_type))
+        )
+
+    def __aiter__(self) -> Subscription[T]:
+        """Return this subscription as an asynchronous message iterator."""
+        return self
+
+    async def __anext__(self) -> Message[T]:
+        """Wait for the next message during asynchronous iteration.
+
+        Custom decoding errors propagate to the caller. The message is consumed
+        before decoding; a decoding error does not put it back on the
+        subscription. Closing the subscription interrupts a pending receive.
+
+        Returns:
+            The received message.
+
+        Raises:
+            StopAsyncIteration: The subscription is closed or its stream ends.
+            UnicodeDecodeError: The payload is not UTF-8 when T is str.
+            ValueError: The payload cannot be parsed when T is int or float.
+            EncoreError: Receiving the message fails.
+        """
+        message = await self.receive()
+        if message is None:
+            raise StopAsyncIteration
+        return message
+
+    async def close(self) -> None:
+        """Close the connection and interrupt pending receives.
+
+        Calling this method again after closure has no effect.
+        """
+        await self._handle.close()
+
+    async def __aenter__(self) -> Subscription[T]:
+        """Return this subscription for use in an async context manager."""
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> None:
+        """Close the subscription when leaving its async context.
+
+        Args:
+            exc_type: The exception type raised in the context, if any.
+            exc: The exception raised in the context, if any.
+            traceback: The exception traceback, if any.
+        """
+        await asyncio.shield(self.close())
+
+
+class _SubscribeCommand(TypedCommand[T]):
+    """A callable SUBSCRIBE command with selectable value decoding."""
+
+    def __getitem__(self, value_type: type[U]) -> _SubscribeCommand[U]:
+        """Select the value type for this command.
+
+        Use subscribe[T](...) to decode message payloads as T. Selecting a type
+        does not open a connection or change the decoder of an existing command.
+
+        Args:
+            value_type: bytes, str, int, float, or a class implementing
+                RedisValue.
+
+        Returns:
+            An independent command with the same arguments and the selected
+            decoder. Keys and other command metadata retain their existing
+            types.
+        """
+        return _SubscribeCommand(self._get_handle, value_type)
+
+    async def __call__(self, channel: str, *channels: str) -> Subscription[T]:
+        """Subscribe to the named channels.
+
+        Args:
+            channel: The first channel name, without an Encore prefix.
+            *channels: Additional channel names, without Encore prefixes.
+
+        Returns:
+            A subscription owning a dedicated connection. Use it with ``async
+            with await db.pubsub.subscribe("events") as messages`` or call
+            close() when finished.
+
+        Raises:
+            EncoreError: The subscription cannot be established.
+        """
+        channels = (channel, *channels)
+        return Subscription(
+            await self._get_handle().pubsub_subscribe(
+                channels, source=active_ctx_bytes()
+            ),
+            value_type=self._value_type,
+        )
+
+
+class _PsubscribeCommand(TypedCommand[T]):
+    """A callable PSUBSCRIBE command with selectable value decoding."""
+
+    def __getitem__(self, value_type: type[U]) -> _PsubscribeCommand[U]:
+        """Select the value type for this command.
+
+        Use psubscribe[T](...) to decode message payloads as T. Selecting a type
+        does not open a connection or change the decoder of an existing command.
+
+        Args:
+            value_type: bytes, str, int, float, or a class implementing
+                RedisValue.
+
+        Returns:
+            An independent command with the same arguments and the selected
+            decoder. Keys and other command metadata retain their existing
+            types.
+        """
+        return _PsubscribeCommand(self._get_handle, value_type)
+
+    async def __call__(self, pattern: str, *patterns: str) -> Subscription[T]:
+        """Subscribe to channels matching Redis glob patterns.
+
+        Args:
+            pattern: The first channel pattern, without an Encore prefix.
+            *patterns: Additional channel patterns, without Encore prefixes.
+
+        Returns:
+            A subscription owning a dedicated connection. Each message contains
+            the matched channel name. Use an async context manager or call
+            close() when finished.
+
+        Raises:
+            EncoreError: The subscription cannot be established.
+        """
+        patterns = (pattern, *patterns)
+        return Subscription(
+            await self._get_handle().pubsub_psubscribe(
+                patterns, source=active_ctx_bytes()
+            ),
+            value_type=self._value_type,
+        )
+
+
+class ReadOnlyInterface(CommandInterface):
+    """Redis channel subscriptions available through a read-only handle."""
+
+    @property
+    def subscribe(self) -> _SubscribeCommand[bytes]:
+        """The callable SUBSCRIBE command.
+
+        Use subscribe[T](...) to decode message payloads as T. Plain calls keep
+        payloads as bytes; channel names are always bytes.
+        """
+        return _SubscribeCommand(self._get_handle, bytes)
+
+    @property
+    def psubscribe(self) -> _PsubscribeCommand[bytes]:
+        """The callable PSUBSCRIBE command.
+
+        Use psubscribe[T](...) to decode message payloads as T. Plain calls keep
+        payloads as bytes; channel names are always bytes.
+        """
+        return _PsubscribeCommand(self._get_handle, bytes)
+
+
+class Interface(ReadOnlyInterface):
+    """Redis publishing and subscriptions exposed by a writable handle."""
+
+    async def publish(self, channel: str, message: Value) -> int:
+        """Publish a message to a Redis channel.
+
+        Args:
+            channel: The channel name without an Encore prefix.
+            message: The value to encode as the message body.
+
+        Returns:
+            The number of receiving subscriptions reported by Redis. This does
+            not confirm that subscribers processed the message.
+
+        Raises:
+            EncoreError: Redis rejects the command or the request fails.
+        """
+        return await self._get_handle().pubsub_publish(
+            channel, encode(message), source=active_ctx_bytes()
+        )