datacurrent 1.0.0__py3-none-any.whl

datacurrent/transports/sse.py

@@ -0,0 +1,374 @@
+"""Server-Sent Events (SSE).
+
+The parser follows the WHATWG event-stream rules: ``data`` lines are joined with
+newlines, blank lines dispatch, comment lines are ignored, ``event`` type,
+``id`` and ``retry`` fields are honoured.  The client supports ``Last-Event-ID``
+resumption, heartbeat comments and configurable reconnection limits.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+from collections.abc import AsyncIterable, AsyncIterator, Mapping
+from dataclasses import dataclass
+from typing import Any
+
+from ..config import DEFAULT_TIMEOUT
+from ..exceptions import (
+    ConfigurationError,
+    ConnectionFailedError,
+    SSEProtocolError,
+    StreamClosedError,
+)
+from ..observability.logging import redact_url
+from ._deps import require_module
+
+__all__ = [
+    "SSEEvent",
+    "SSELineDecoder",
+    "SSEParser",
+    "SSEStream",
+    "format_heartbeat",
+    "format_sse",
+    "parse_sse",
+]
+
+
+@dataclass(frozen=True, slots=True)
+class SSEEvent:
+    """A dispatched Server-Sent Event."""
+
+    data: str = ""
+    event: str = "message"
+    id: str | None = None
+    retry: int | None = None
+
+    def json(self) -> Any:
+        """Parse ``data`` as JSON."""
+        return json.loads(self.data)
+
+
+class SSELineDecoder:
+    """Split arbitrary text chunks into lines following the SSE line rules."""
+
+    __slots__ = ("_buffer",)
+
+    def __init__(self) -> None:
+        self._buffer = ""
+
+    def feed(self, chunk: str) -> list[str]:
+        if chunk:
+            self._buffer += chunk
+        lines: list[str] = []
+        while True:
+            index = -1
+            for candidate in ("\n", "\r"):
+                position = self._buffer.find(candidate)
+                if position != -1 and (index == -1 or position < index):
+                    index = position
+            if index == -1:
+                break
+            character = self._buffer[index]
+            if character == "\r" and index == len(self._buffer) - 1:
+                # A CRLF pair may be split across chunks; wait for the next chunk.
+                break
+            line = self._buffer[:index]
+            skip = 1
+            if (
+                character == "\r"
+                and index + 1 < len(self._buffer)
+                and self._buffer[index + 1] == "\n"
+            ):
+                skip = 2
+            self._buffer = self._buffer[index + skip :]
+            lines.append(line)
+        return lines
+
+    def flush(self) -> list[str]:
+        """Return a trailing line that had no newline terminator."""
+        if not self._buffer:
+            return []
+        line = self._buffer
+        self._buffer = ""
+        return [line]
+
+
+class SSEParser:
+    """Incremental SSE event parser.
+
+    Feed complete lines with :meth:`feed_line`; a non-``None`` return value is a
+    dispatched event.
+    """
+
+    __slots__ = ("_data", "_event", "_has_data", "_last_id", "_retry")
+
+    def __init__(self) -> None:
+        self._data: list[str] = []
+        self._event = "message"
+        self._has_data = False
+        self._last_id: str | None = None
+        self._retry: int | None = None
+
+    @property
+    def last_event_id(self) -> str | None:
+        return self._last_id
+
+    def feed_line(self, line: str) -> SSEEvent | None:
+        line = line.rstrip("\r\n")
+        if line == "":
+            return self._dispatch()
+        if line.startswith(":"):
+            return None
+        field, separator, value = line.partition(":")
+        if not separator:
+            value = ""
+        elif value.startswith(" "):
+            value = value[1:]
+
+        if field == "data":
+            self._data.append(value)
+            self._has_data = True
+        elif field == "event":
+            self._event = value or "message"
+        elif field == "id":
+            if "\x00" not in value:
+                self._last_id = value
+        elif field == "retry":
+            try:
+                parsed = int(value)
+            except ValueError:
+                return None
+            if parsed >= 0:
+                self._retry = parsed
+        return None
+
+    def _dispatch(self) -> SSEEvent | None:
+        if not self._has_data:
+            self._event = "message"
+            return None
+        event = SSEEvent(
+            data="\n".join(self._data),
+            event=self._event or "message",
+            id=self._last_id,
+            retry=self._retry,
+        )
+        self._data = []
+        self._has_data = False
+        self._event = "message"
+        return event
+
+
+async def parse_sse(lines: AsyncIterable[str | bytes]) -> AsyncIterator[SSEEvent]:
+    """Parse an async iterable of text or bytes chunks into events."""
+    decoder = SSELineDecoder()
+    parser = SSEParser()
+    async for chunk in lines:
+        text = (
+            chunk.decode("utf-8", errors="replace")
+            if isinstance(chunk, (bytes, bytearray))
+            else str(chunk)
+        )
+        for line in decoder.feed(text):
+            event = parser.feed_line(line)
+            if event is not None:
+                yield event
+    for line in decoder.flush():
+        event = parser.feed_line(line)
+        if event is not None:
+            yield event
+    trailing = parser.feed_line("")
+    if trailing is not None:
+        yield trailing
+
+
+def format_sse(event: SSEEvent) -> str:
+    """Encode an :class:`SSEEvent` for a server response."""
+    if not isinstance(event, SSEEvent):
+        raise TypeError(f"format_sse() expects an SSEEvent, got {type(event).__name__}")
+    lines: list[str] = []
+    if event.id is not None:
+        lines.append(f"id: {event.id}")
+    if event.event and event.event != "message":
+        lines.append(f"event: {event.event}")
+    if event.retry is not None:
+        lines.append(f"retry: {int(event.retry)}")
+    for line in event.data.split("\n"):
+        lines.append(f"data: {line}")
+    return "\n".join(lines) + "\n\n"
+
+
+def format_heartbeat(comment: str = "keep-alive") -> str:
+    """Encode an SSE comment line, used as a heartbeat."""
+    return f": {comment}\n\n"
+
+
+class SSEStream:
+    """Async client for an SSE endpoint.
+
+    Args:
+        url: ``http`` or ``https`` endpoint.
+        headers: Extra request headers (credentials are never logged).
+        params: Query parameters.
+        method: HTTP method, ``GET`` by default.
+        timeout: Per-request timeout in seconds.
+        client: Optional ``httpx.AsyncClient``; one is created when omitted.
+        max_reconnects: Additional connections after the first one ends or
+            fails.  ``0`` (default) never reconnects.
+        reconnect_delay / max_reconnect_delay: Exponential backoff bounds.
+        last_event_id: Resume from this event id on the first request.
+        check_content_type: Require ``text/event-stream`` responses.
+
+    The client yields :class:`SSEEvent` objects and tracks the last event id so a
+    reconnect sends ``Last-Event-ID`` automatically.
+    """
+
+    def __init__(
+        self,
+        url: str,
+        *,
+        headers: Mapping[str, str] | None = None,
+        params: Mapping[str, Any] | None = None,
+        method: str = "GET",
+        timeout: float = DEFAULT_TIMEOUT,
+        client: Any = None,
+        max_reconnects: int = 0,
+        reconnect_delay: float = 1.0,
+        max_reconnect_delay: float = 30.0,
+        last_event_id: str | None = None,
+        check_content_type: bool = True,
+    ) -> None:
+        if not isinstance(url, str) or not url.lower().startswith(("http://", "https://")):
+            raise ConfigurationError("SSEStream requires an http:// or https:// URL")
+        if isinstance(max_reconnects, bool) or not isinstance(max_reconnects, int):
+            raise ConfigurationError("max_reconnects must be an integer")
+        if max_reconnects < 0:
+            raise ConfigurationError("max_reconnects must be >= 0")
+        if reconnect_delay < 0 or max_reconnect_delay < 0:
+            raise ConfigurationError("reconnect delays must be >= 0")
+        self._url = url
+        self._headers = {"Accept": "text/event-stream", **dict(headers or {})}
+        self._params = dict(params) if params else None
+        self._method = method.upper()
+        self._timeout = float(timeout)
+        self._client = client
+        self._owned_client: Any = None
+        self._max_reconnects = max_reconnects
+        self._reconnect_delay = float(reconnect_delay)
+        self._max_reconnect_delay = float(max_reconnect_delay)
+        self._last_event_id = last_event_id
+        self._retry_hint: float | None = None
+        self._check_content_type = check_content_type
+        self._closed = False
+
+    # ------------------------------------------------------------------ #
+    # Construction / lifecycle
+    # ------------------------------------------------------------------ #
+    @classmethod
+    def connect(cls, url: str, **kwargs: Any) -> SSEStream:
+        """Create a client; use it with ``async for`` or ``async with``."""
+        return cls(url, **kwargs)
+
+    @property
+    def url(self) -> str:
+        return self._url
+
+    @property
+    def last_event_id(self) -> str | None:
+        return self._last_event_id
+
+    @property
+    def closed(self) -> bool:
+        return self._closed
+
+    async def __aenter__(self) -> SSEStream:
+        return self
+
+    async def __aexit__(self, exc_type: Any, exc: Any, tb: Any) -> None:
+        await self.aclose()
+
+    async def aclose(self) -> None:
+        """Close the client and drop the owned HTTP client, if any."""
+        self._closed = True
+        client, self._owned_client = self._owned_client, None
+        if client is not None:
+            await client.aclose()
+
+    def __aiter__(self) -> AsyncIterator[SSEEvent]:
+        return self._events()
+
+    # ------------------------------------------------------------------ #
+    # Streaming
+    # ------------------------------------------------------------------ #
+    def _next_delay(self, reconnects: int) -> float:
+        if self._retry_hint is not None:
+            return self._retry_hint
+        base = self._reconnect_delay * (2 ** min(reconnects, 10))
+        return min(self._max_reconnect_delay, base)
+
+    async def _events(self) -> AsyncIterator[SSEEvent]:
+        httpx = require_module("httpx", purpose="SSEStream")
+        if self._closed:
+            raise StreamClosedError("SSEStream is closed")
+        if self._client is not None:
+            client = self._client
+        else:
+            client = httpx.AsyncClient(follow_redirects=True)
+            self._owned_client = client
+
+        reconnects = 0
+        try:
+            while True:
+                headers = dict(self._headers)
+                if self._last_event_id is not None:
+                    headers["Last-Event-ID"] = self._last_event_id
+                try:
+                    async with client.stream(
+                        self._method,
+                        self._url,
+                        headers=headers,
+                        params=self._params,
+                        timeout=self._timeout,
+                    ) as response:
+                        if response.status_code >= 400:
+                            response.raise_for_status()
+                        if self._check_content_type:
+                            content_type = response.headers.get("content-type", "")
+                            if "text/event-stream" not in content_type.lower():
+                                raise SSEProtocolError(
+                                    f"expected a 'text/event-stream' response, got {content_type!r}"
+                                )
+                        # aiter_bytes() gives decoded body chunks; parse_sse splits
+                        # lines itself so events spanning chunks are handled.
+                        async for event in parse_sse(response.aiter_bytes()):
+                            if event.id is not None:
+                                self._last_event_id = event.id
+                            if event.retry is not None:
+                                self._retry_hint = event.retry / 1000.0
+                            yield event
+                except asyncio.CancelledError:
+                    raise
+                except SSEProtocolError:
+                    raise
+                except httpx.HTTPStatusError:
+                    raise
+                except (httpx.HTTPError, OSError) as error:
+                    if reconnects >= self._max_reconnects:
+                        raise ConnectionFailedError(
+                            f"SSE connection to {redact_url(self._url)} failed: {error}"
+                        ) from error
+                    delay = self._next_delay(reconnects)
+                    reconnects += 1
+                    await asyncio.sleep(delay)
+                    continue
+
+                if reconnects >= self._max_reconnects:
+                    return
+                delay = self._next_delay(reconnects)
+                reconnects += 1
+                if delay:
+                    await asyncio.sleep(delay)
+        finally:
+            if self._owned_client is client:
+                self._owned_client = None
+                await client.aclose()

datacurrent/transports/websocket.py

@@ -0,0 +1,277 @@
+"""WebSocket streaming over the maintained ``websockets`` library."""
+
+from __future__ import annotations
+
+import asyncio
+import inspect
+import json
+from collections.abc import AsyncIterator, Mapping
+from typing import Any
+
+from ..config import DEFAULT_MAX_MESSAGE_SIZE
+from ..exceptions import (
+    ConfigurationError,
+    ConnectionFailedError,
+    MessageTooLargeError,
+    StreamClosedError,
+    WebSocketClosedError,
+)
+from ..observability.logging import get_logger, redact_headers, redact_url
+from ._deps import require_module
+
+__all__ = ["WebSocketStream"]
+
+
+def _close_details(error: Any) -> tuple[int | None, str | None]:
+    """Read close code/reason without the deprecated attributes."""
+    frame = getattr(error, "rcvd", None) or getattr(error, "sent", None)
+    code = getattr(frame, "code", None)
+    reason = getattr(frame, "reason", None)
+    if code is None:
+        code = getattr(error, "code", None)
+    if reason is None:
+        reason = getattr(error, "reason", None)
+    return code, reason
+
+
+_log = get_logger("websocket")
+
+_VALID_MODES = {"auto", "text", "json"}
+
+
+class WebSocketStream:
+    """A WebSocket client that behaves like an async sequence.
+
+    Args:
+        url: ``ws://`` or ``wss://`` endpoint.
+        headers: Handshake headers; values are redacted in logs and errors.
+        mode: ``auto`` yields ``str``/``bytes`` as received, ``text`` decodes
+            binary frames to text, ``json`` parses text frames as JSON.
+        max_size: Maximum accepted message size in bytes.
+        ping_interval / ping_timeout: Protocol heartbeats in seconds.
+        open_timeout / close_timeout: Handshake and close deadlines.
+        reconnect_attempts: Automatic reconnects after an abnormal close.
+        reconnect_delay / max_reconnect_delay: Exponential backoff bounds.
+        subprotocols: Optional WebSocket subprotocols.
+        ssl: Optional TLS context passed to ``websockets``.
+        origin: Optional ``Origin`` header.
+
+    TLS certificates are verified by default: ``wss://`` uses the system trust
+    store through ``websockets``.  Only an explicit ``ssl=`` context can change
+    that, and doing so is the caller's decision.
+    """
+
+    def __init__(
+        self,
+        url: str,
+        *,
+        headers: Mapping[str, str] | None = None,
+        mode: str = "auto",
+        max_size: int = DEFAULT_MAX_MESSAGE_SIZE,
+        ping_interval: float | None = 20.0,
+        ping_timeout: float | None = 20.0,
+        open_timeout: float = 10.0,
+        close_timeout: float = 5.0,
+        reconnect_attempts: int = 0,
+        reconnect_delay: float = 0.5,
+        max_reconnect_delay: float = 30.0,
+        subprotocols: list[str] | None = None,
+        ssl: Any = None,
+        origin: str | None = None,
+    ) -> None:
+        if not isinstance(url, str) or not url.lower().startswith(("ws://", "wss://")):
+            raise ConfigurationError("WebSocketStream requires a ws:// or wss:// URL")
+        if mode not in _VALID_MODES:
+            raise ConfigurationError(f"mode must be one of {sorted(_VALID_MODES)}, got {mode!r}")
+        if isinstance(reconnect_attempts, bool) or not isinstance(reconnect_attempts, int):
+            raise ConfigurationError("reconnect_attempts must be an integer")
+        if reconnect_attempts < 0:
+            raise ConfigurationError("reconnect_attempts must be >= 0")
+        if max_size < 1:
+            raise ConfigurationError("max_size must be >= 1")
+        if ping_interval is not None and ping_interval < 0:
+            raise ConfigurationError("ping_interval must be >= 0 or None")
+        self._url = url
+        self._headers = dict(headers or {})
+        self._mode = mode
+        self._max_size = max_size
+        self._ping_interval = ping_interval
+        self._ping_timeout = ping_timeout
+        self._open_timeout = open_timeout
+        self._close_timeout = close_timeout
+        self._reconnect_attempts = reconnect_attempts
+        self._reconnect_delay = reconnect_delay
+        self._max_reconnect_delay = max_reconnect_delay
+        self._subprotocols = subprotocols
+        self._ssl = ssl
+        self._origin = origin
+        self._ws: Any = None
+        self._closed = False
+
+    # ------------------------------------------------------------------ #
+    # Lifecycle
+    # ------------------------------------------------------------------ #
+    @classmethod
+    def connect(cls, url: str, **kwargs: Any) -> WebSocketStream:
+        """Create a client; use ``async with WebSocketStream.connect(...)``."""
+        return cls(url, **kwargs)
+
+    @property
+    def url(self) -> str:
+        return self._url
+
+    @property
+    def connected(self) -> bool:
+        return self._ws is not None
+
+    @property
+    def protocol(self) -> str | None:
+        return None if self._ws is None else getattr(self._ws, "subprotocol", None)
+
+    async def __aenter__(self) -> WebSocketStream:
+        await self._connect()
+        return self
+
+    async def __aexit__(self, exc_type: Any, exc: Any, tb: Any) -> None:
+        await self.aclose()
+
+    async def _connect(self) -> None:
+        websockets = require_module("websockets", purpose="WebSocketStream")
+        parameters = inspect.signature(websockets.connect).parameters
+        header_key = "additional_headers" if "additional_headers" in parameters else "extra_headers"
+        kwargs: dict[str, Any] = {
+            header_key: self._headers,
+            "max_size": self._max_size,
+            "ping_interval": self._ping_interval,
+            "ping_timeout": self._ping_timeout,
+            "open_timeout": self._open_timeout,
+            "close_timeout": self._close_timeout,
+        }
+        if self._subprotocols:
+            kwargs["subprotocols"] = self._subprotocols
+        if self._ssl is not None:
+            kwargs["ssl"] = self._ssl
+        if self._origin is not None:
+            kwargs["origin"] = self._origin
+        try:
+            self._ws = await websockets.connect(self._url, **kwargs)
+        except asyncio.CancelledError:
+            raise
+        except Exception as error:
+            raise ConnectionFailedError(
+                f"could not connect to {redact_url(self._url)}: {error}"
+            ) from error
+        _log.debug(
+            "websocket connected url=%s headers=%s",
+            redact_url(self._url),
+            redact_headers(self._headers),
+        )
+
+    async def aclose(self) -> None:
+        """Close the connection without raising on an already closed socket."""
+        self._closed = True
+        ws, self._ws = self._ws, None
+        if ws is None:
+            return
+        try:
+            await ws.close()
+        except asyncio.CancelledError:
+            raise
+        except Exception:  # pragma: no cover - close races are harmless
+            return
+
+    # ------------------------------------------------------------------ #
+    # Sending
+    # ------------------------------------------------------------------ #
+    def _require_ws(self) -> Any:
+        if self._ws is None:
+            raise StreamClosedError(
+                "websocket is not connected; use 'async with WebSocketStream.connect(...)'"
+            )
+        return self._ws
+
+    async def send(self, data: str | bytes) -> None:
+        """Send a text or binary frame."""
+        await self._require_ws().send(data)
+
+    async def send_json(self, value: Any) -> None:
+        """Serialize ``value`` and send it as a text frame."""
+        await self._require_ws().send(json.dumps(value))
+
+    async def ping(self) -> None:
+        """Send a ping and wait for the pong."""
+        waiter = await self._require_ws().ping()
+        await waiter
+
+    # ------------------------------------------------------------------ #
+    # Receiving
+    # ------------------------------------------------------------------ #
+    def _decode(self, message: str | bytes) -> Any:
+        size = len(message.encode("utf-8") if isinstance(message, str) else message)
+        if size > self._max_size:
+            raise MessageTooLargeError(size, self._max_size)
+        if self._mode == "json":
+            if isinstance(message, bytes):
+                message = message.decode("utf-8")
+            return json.loads(message)
+        if self._mode == "text" and isinstance(message, bytes):
+            return message.decode("utf-8")
+        return message
+
+    async def recv(self) -> Any:
+        """Receive one message, reconnecting when configured."""
+        websockets = require_module("websockets", purpose="WebSocketStream")
+        if self._ws is None:
+            await self._connect()
+        assert self._ws is not None
+        try:
+            message = await self._ws.recv()
+        except websockets.ConnectionClosed as error:
+            code, reason = _close_details(error)
+            raise WebSocketClosedError(code, reason) from error
+        return self._decode(message)
+
+    def __aiter__(self) -> AsyncIterator[Any]:
+        return self._messages()
+
+    async def _messages(self) -> AsyncIterator[Any]:
+        websockets = require_module("websockets", purpose="WebSocketStream")
+        if self._closed:
+            raise StreamClosedError("WebSocketStream is closed")
+        if self._ws is None:
+            await self._connect()
+        attempts = 0
+        while True:
+            ws = self._ws
+            if ws is None:
+                raise StreamClosedError("websocket is not connected")
+            try:
+                async for message in ws:
+                    yield self._decode(message)
+                return
+            except asyncio.CancelledError:
+                raise
+            except websockets.ConnectionClosed as error:
+                code, reason = _close_details(error)
+                if self._closed:
+                    raise WebSocketClosedError(code, reason) from error
+                if code == 1000:
+                    return
+                if attempts >= self._reconnect_attempts:
+                    raise WebSocketClosedError(code, reason) from error
+                attempts += 1
+                await self.aclose()
+                self._closed = False
+                delay = min(
+                    self._max_reconnect_delay,
+                    self._reconnect_delay * (2 ** (attempts - 1)),
+                )
+                if delay:
+                    await asyncio.sleep(delay)
+                await self._connect()
+
+    def __repr__(self) -> str:  # pragma: no cover - trivial
+        return (
+            f"WebSocketStream(url={self._url!r}, mode={self._mode!r}, "
+            f"connected={self._ws is not None})"
+        )