amazon-polly-streaming 1.0.0__py3-none-any.whl

amazon_polly_streaming/__init__.py

@@ -0,0 +1,20 @@
+"""Amazon Polly bidirectional streaming over HTTP/2 with SigV4."""
+
+from amazon_polly_streaming.client import PollyStreamingClient
+from amazon_polly_streaming.exceptions import (
+    ServiceException,
+    ServiceFailureException,
+    ServiceQuotaExceededException,
+    ThrottlingException,
+    ValidationException,
+)
+
+__all__ = [
+    "PollyStreamingClient",
+    "ServiceException",
+    "ServiceFailureException",
+    "ServiceQuotaExceededException",
+    "ThrottlingException",
+    "ValidationException",
+]
+__version__ = "1.0.0"

amazon_polly_streaming/_buffered_stream.py

@@ -0,0 +1,102 @@
+"""Writable, non-blocking byte stream used as HTTP/2 request body channel.
+
+`awscrt`'s HTTP/2 client reads the request body from a file-like object on its
+own I/O thread, calling `read(size)` repeatedly until the stream signals EOF.
+For bidirectional streaming we want to write events to the body channel
+incrementally from the application's asyncio thread, after the request has
+already been opened.
+
+`BufferableByteStream` exposes that pattern by behaving as a non-blocking
+file-like object: `read` raises `BlockingIOError` when no data is currently
+available, signalling `awscrt` to retry later; once `end_stream` is called and
+all buffered chunks have been read, `read` returns `b""` (EOF) and the HTTP/2
+stream is closed cleanly.
+
+Pattern adapted from the `amazon-transcribe-streaming-sdk` Python package
+(Apache 2.0), which uses the same approach with `awscrt`.
+"""
+
+from __future__ import annotations
+
+from io import BufferedIOBase
+
+
+class BufferableByteStream(BufferedIOBase):
+    """Non-blocking, append-only byte buffer for HTTP/2 request body channels."""
+
+    def __init__(self) -> None:
+        """Initialize an empty, open stream."""
+        self._chunks: list[bytes] = []
+        self._done: bool = False
+        self._closed: bool = False
+
+    def read(self, size: int | None = -1) -> bytes:
+        """Return up to `size` bytes; raise `BlockingIOError` if no data is buffered.
+
+        Args:
+            size: Maximum number of bytes to return. ``-1`` or ``None`` means
+                "the next chunk in full".
+
+        Returns:
+            Bytes consumed from the head of the buffer. ``b""`` once the
+            stream has been ended (`end_stream`) and fully drained.
+
+        Raises:
+            BlockingIOError: when the stream is still open but no data has
+                been written yet. The caller (`awscrt` in production) is
+                expected to retry later.
+        """
+        if not self._chunks:
+            if self._done or self._closed:
+                return b""
+            msg = "no data buffered yet"
+            raise BlockingIOError(msg)
+
+        chunk = self._chunks.pop(0)
+        if size is None or size < 0 or size >= len(chunk):
+            return chunk
+        leftover = chunk[size:]
+        self._chunks.insert(0, leftover)
+        return chunk[:size]
+
+    def read1(self, size: int = -1) -> bytes:
+        """Read at most `size` bytes; same semantics as `read`."""
+        return self.read(size)
+
+    def write(self, b: bytes) -> int:  # pyright: ignore[reportIncompatibleMethodOverride]
+        """Append `b` to the buffer; return the number of bytes accepted.
+
+        Args:
+            b: Bytes to append. Non-bytes input raises `TypeError`.
+
+        Returns:
+            ``len(b)``.
+
+        Raises:
+            TypeError: if `b` is not exactly `bytes`.
+            OSError: if the stream has already been ended or closed.
+        """
+        if not isinstance(b, bytes):  # pyright: ignore[reportUnnecessaryIsInstance]
+            msg = f"BufferableByteStream.write requires bytes, got {type(b).__name__}"
+            raise TypeError(msg)
+        if self._done or self._closed:
+            msg = "stream is closed"
+            raise OSError(msg)
+        if b:
+            self._chunks.append(b)
+        return len(b)
+
+    def end_stream(self) -> None:
+        """Mark the stream as ended; future writes raise, future reads drain then return `b""`."""
+        self._done = True
+
+    def close(self) -> None:
+        """Close the stream and discard any buffered data."""
+        self._chunks = []
+        self._done = True
+        self._closed = True
+
+    @property
+    def closed(self) -> bool:
+        """True if `close` has been called."""
+        return self._closed

amazon_polly_streaming/_connection_pool.py

@@ -0,0 +1,229 @@
+"""HTTP/2 connection pool for `amazon-polly-streaming` with multi-connection lease semantics.
+
+Eliminates the TLS handshake, ALPN h2 negotiation, and HTTP/2 SETTINGS exchange
+between subsequent calls to the same Polly endpoint, while supporting fan-out:
+multiple concurrent leases on the same ``(host, port)`` get distinct underlying
+``HttpClientConnection`` instances, up to ``max_size_per_key``.
+
+The Polly bidirectional streaming endpoint advertises one active stream per
+HTTP/2 connection, so a single shared connection cannot serve more than one
+synthesis call at a time. The pool therefore keeps a per-key list of
+connections: idle entries are reused by subsequent acquires (preserving the
+TLS/H2 cache benefit), and concurrent leases each get their own connection.
+When the per-key cap is reached, a further acquire waits on a Condition until
+a release frees a slot.
+
+The underlying AWS Common Runtime resources (event loop group, host resolver,
+client bootstrap, TLS context) are shared across all entries and allocated
+lazily on the first connect.
+
+Pattern inspired by the AWS-maintained ``amazon-transcribe-streaming-sdk``
+(``AwsCrtHttpSessionManager._connections`` in
+``amazon_transcribe/httpsession.py``).
+"""
+
+# pyright: reportUnknownMemberType=false, reportUnknownVariableType=false
+# pyright: reportUnknownArgumentType=false, reportUnknownParameterType=false
+# pyright: reportAttributeAccessIssue=false, reportArgumentType=false
+# pyright: reportUnnecessaryComparison=false
+# Rationale: awscrt ships no type stubs and its inline annotations are
+# incomplete; suppressions are scoped to this module.
+from __future__ import annotations
+
+import asyncio
+from contextlib import asynccontextmanager
+from typing import TYPE_CHECKING
+
+from awscrt import http, io
+from awscrt.http import HttpClientConnection
+
+if TYPE_CHECKING:
+    from collections.abc import AsyncGenerator
+
+_DEFAULT_MAX_SIZE_PER_KEY = 8
+
+
+async def _connect(
+    host: str,
+    port: int,
+    *,
+    bootstrap: io.ClientBootstrap,
+    tls_ctx: io.ClientTlsContext,
+) -> HttpClientConnection:
+    """Open an HTTP/2 connection to ``host:port`` using the provided awscrt resources.
+
+    The bootstrap (event loop group + DNS resolver) and TLS context are
+    intended to be shared across multiple connections. Only the
+    per-connection TLS connection options carrying SNI server name and ALPN
+    list are built here.
+
+    Args:
+        host: Hostname to connect to.
+        port: TCP port (typically ``443``).
+        bootstrap: Pre-built ``ClientBootstrap`` used as I/O backbone.
+        tls_ctx: Pre-built ``ClientTlsContext`` from which per-connection
+            options are spawned.
+
+    Returns:
+        An open ``HttpClientConnection`` with HTTP/2 negotiated.
+
+    Raises:
+        RuntimeError: if the connection cannot be opened or HTTP/2 cannot be
+            negotiated.
+    """
+    tls_conn_options = tls_ctx.new_connection_options()
+    tls_conn_options.set_server_name(host)
+    tls_conn_options.set_alpn_list(["h2"])
+    connect_future = HttpClientConnection.new(
+        host_name=host,
+        port=port,
+        bootstrap=bootstrap,
+        socket_options=io.SocketOptions(),
+        tls_connection_options=tls_conn_options,
+    )
+    connection = await asyncio.wrap_future(connect_future)
+    if not connection.is_open():
+        msg = f"Could not open connection to {host}:{port}"
+        raise RuntimeError(msg)
+    if connection.version is not http.HttpVersion.Http2:
+        connection.close()
+        msg = f"HTTP/2 could not be negotiated: got {connection.version!r}"
+        raise RuntimeError(msg)
+    return connection
+
+
+class _ConnectionPool:
+    """Bounded pool of HTTP/2 connections keyed on ``(host, port)``, with leases.
+
+    The pool keeps up to ``max_size_per_key`` connections per key. Each acquire
+    grants exclusive use of one connection (a "lease") for the duration of an
+    ``async with acquire_connection(...)`` block. Idle connections are reused
+    by subsequent acquires; concurrent acquires on the same key each get their
+    own connection (opening fresh ones up to the cap).
+
+    A single ``asyncio.Condition`` serializes pool state mutations and signals
+    waiters when a slot frees up. The lock is held only over in-memory state
+    updates; ``_connect`` runs outside the lock so different keys (and even
+    different acquires on the same key) connect in parallel.
+
+    The underlying awscrt resources (event loop group, host resolver,
+    client bootstrap, TLS context) are shared across every cached
+    connection and allocated lazily on the first connect.
+    """
+
+    def __init__(self, *, max_size_per_key: int = _DEFAULT_MAX_SIZE_PER_KEY) -> None:
+        self._max_size_per_key = max_size_per_key
+        self._idle: dict[tuple[str, int], list[HttpClientConnection]] = {}
+        self._in_use: dict[tuple[str, int], int] = {}
+        self._cond = asyncio.Condition()
+        self._closed = False
+        self._bootstrap: io.ClientBootstrap | None = None
+        self._tls_ctx: io.ClientTlsContext | None = None
+        # Hold refs to event loop group and resolver so they are not GC'd
+        # while the bootstrap is alive.
+        self._elg: io.EventLoopGroup | None = None
+        self._resolver: io.DefaultHostResolver | None = None
+
+    def _ensure_resources(self) -> tuple[io.ClientBootstrap, io.ClientTlsContext]:
+        """Lazily build the shared awscrt resources on first use."""
+        if self._bootstrap is None or self._tls_ctx is None:
+            elg = io.EventLoopGroup(1)
+            resolver = io.DefaultHostResolver(elg)
+            self._elg = elg
+            self._resolver = resolver
+            self._bootstrap = io.ClientBootstrap(elg, resolver)
+            self._tls_ctx = io.ClientTlsContext(io.TlsContextOptions())
+        return self._bootstrap, self._tls_ctx
+
+    @asynccontextmanager
+    async def acquire_connection(
+        self, *, host: str, port: int
+    ) -> AsyncGenerator[HttpClientConnection]:
+        """Acquire an exclusive lease on an HTTP/2 connection for ``(host, port)``.
+
+        On entry, returns either a cached idle connection or a freshly opened
+        one (up to ``max_size_per_key`` total per key). If the cap is reached,
+        the call waits until a concurrent lease releases its connection. On
+        exit, the connection returns to the idle list (or is closed if it has
+        gone stale or the pool was closed in the meantime).
+
+        Args:
+            host: Hostname.
+            port: TCP port.
+
+        Yields:
+            An open ``HttpClientConnection`` ready to accept a new HTTP/2
+            stream via ``connection.request(...)``.
+
+        Raises:
+            RuntimeError: if no connection slot is available and a fresh
+                connect cannot be opened.
+        """
+        key = (host, port)
+        connection = await self._acquire(key)
+        try:
+            yield connection
+        finally:
+            await self._release(key, connection)
+
+    async def _acquire(self, key: tuple[str, int]) -> HttpClientConnection:
+        """Reserve a slot for ``key`` and return an open connection."""
+        async with self._cond:
+            while True:
+                idle_list = self._idle.get(key, [])
+                while idle_list:
+                    candidate = idle_list.pop()
+                    if candidate.is_open():
+                        self._in_use[key] = self._in_use.get(key, 0) + 1
+                        return candidate
+                    candidate.close()
+                count = self._in_use.get(key, 0)
+                if count < self._max_size_per_key:
+                    self._in_use[key] = count + 1
+                    break
+                await self._cond.wait()
+        try:
+            bootstrap, tls_ctx = self._ensure_resources()
+            return await _connect(key[0], key[1], bootstrap=bootstrap, tls_ctx=tls_ctx)
+        except BaseException:
+            async with self._cond:
+                self._in_use[key] = max(0, self._in_use.get(key, 0) - 1)
+                self._cond.notify_all()
+            raise
+
+    async def _release(self, key: tuple[str, int], connection: HttpClientConnection) -> None:
+        """Return a leased connection to idle, or close it on stale/closed-pool."""
+        async with self._cond:
+            self._in_use[key] = max(0, self._in_use.get(key, 0) - 1)
+            if self._closed or not connection.is_open():
+                connection.close()
+            else:
+                self._idle.setdefault(key, []).append(connection)
+            self._cond.notify_all()
+
+    async def close_all(self) -> None:
+        """Close every idle connection and mark the pool as closed.
+
+        Connections currently leased survive until their lease releases; on
+        release they are closed (not returned to idle) because the pool is
+        flagged closed.
+        """
+        async with self._cond:
+            self._closed = True
+            idle_snapshot = list(self._idle.items())
+            self._idle.clear()
+            self._cond.notify_all()
+        for _key, conns in idle_snapshot:
+            for connection in conns:
+                connection.close()
+
+
+_default_pool: _ConnectionPool | None = None
+
+
+def get_default_pool() -> _ConnectionPool:
+    """Return the module-level default ``_ConnectionPool``, creating it on first call."""
+    global _default_pool  # noqa: PLW0603
+    if _default_pool is None:
+        _default_pool = _ConnectionPool()
+    return _default_pool

amazon_polly_streaming/_event_signer.py

@@ -0,0 +1,129 @@
+"""Rolling AWS event-stream chunk signer for HTTP/2 bidirectional streaming.
+
+Each event sent on the request body channel must be wrapped in a signed
+envelope. The wrapper carries two headers, `:date` and `:chunk-signature`,
+where `:chunk-signature` is computed as an HMAC-SHA256 of an event-stream-
+specific string-to-sign. The signature of the previous wrapped event (or, for
+the first event, the SigV4 signature of the initial HTTP request) participates
+in the string-to-sign, producing a rolling chain.
+
+The signing string follows the format documented at
+https://docs.aws.amazon.com/transcribe/latest/dg/streaming-setting-up.html
+under "AWS4-HMAC-SHA256-PAYLOAD". The same algorithm is used by Amazon Polly
+bidirectional streaming and by the open-source `amazon-transcribe-streaming-
+sdk` Python package, both based on the AWS Common Runtime event-stream signing
+contract.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import hmac
+import struct
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import datetime
+
+    from amazon_polly_streaming._eventstream import HeaderValue
+
+_TIMESTAMP_FMT = "%Y%m%dT%H%M%SZ"
+_HEADER_TYPE_TIMESTAMP = 8
+
+
+@dataclass(frozen=True)
+class EventSignerCredentials:
+    """Subset of AWS credentials needed to derive the signing key."""
+
+    access_key_id: str
+    secret_access_key: str
+    session_token: str | None
+
+
+class EventSigner:
+    """Compute rolling `:chunk-signature` headers for event-stream wrappers."""
+
+    def __init__(self, *, signing_name: str, region: str) -> None:
+        """Bind the signer to a service signing name and a region."""
+        self._signing_name = signing_name
+        self._region = region
+
+    def sign(
+        self,
+        *,
+        payload: bytes,
+        prior_signature: bytes,
+        credentials: EventSignerCredentials,
+        now: datetime.datetime,
+    ) -> dict[str, HeaderValue]:
+        """Return wrapper headers (`:date`, `:chunk-signature`) for `payload`.
+
+        Args:
+            payload: Inner event-stream message bytes (already encoded with its
+                own headers and payload).
+            prior_signature: 32-byte raw signature of the previous wrapped
+                event, or the SigV4 signature of the initial HTTP request for
+                the first event.
+            credentials: AWS credentials used to derive the signing key.
+            now: Timezone-aware UTC datetime used both as the `:date` header
+                value and as the timestamp embedded in the string-to-sign.
+
+        Returns:
+            A dict with two entries: `:date` (the input `now`) and
+            `:chunk-signature` (32 raw HMAC-SHA256 bytes).
+
+        Raises:
+            ValueError: if `now` is naive (no `tzinfo`).
+        """
+        if now.tzinfo is None:
+            msg = "EventSigner.sign requires a timezone-aware `now`"
+            raise ValueError(msg)
+
+        timestamp = now.strftime(_TIMESTAMP_FMT)
+        date_header_bytes = _encode_date_header_value(now)
+        string_to_sign = "\n".join(
+            [
+                "AWS4-HMAC-SHA256-PAYLOAD",
+                timestamp,
+                self._scope(timestamp),
+                prior_signature.hex(),
+                hashlib.sha256(date_header_bytes).hexdigest(),
+                hashlib.sha256(payload).hexdigest(),
+            ]
+        )
+        signing_key = self._derive_signing_key(
+            secret_access_key=credentials.secret_access_key,
+            timestamp=timestamp,
+        )
+        chunk_signature = hmac.new(
+            signing_key, string_to_sign.encode("utf-8"), hashlib.sha256
+        ).digest()
+        return {":date": now, ":chunk-signature": chunk_signature}
+
+    def _scope(self, timestamp: str) -> str:
+        return f"{timestamp[:8]}/{self._region}/{self._signing_name}/aws4_request"
+
+    def _derive_signing_key(self, *, secret_access_key: str, timestamp: str) -> bytes:
+        date = timestamp[:8].encode("utf-8")
+        k_date = _hmac(b"AWS4" + secret_access_key.encode("utf-8"), date)
+        k_region = _hmac(k_date, self._region.encode("utf-8"))
+        k_service = _hmac(k_region, self._signing_name.encode("utf-8"))
+        return _hmac(k_service, b"aws4_request")
+
+
+def _hmac(key: bytes, msg: bytes) -> bytes:
+    return hmac.new(key, msg, hashlib.sha256).digest()
+
+
+def _encode_date_header_value(when: datetime.datetime) -> bytes:
+    """Encode `:date` as it appears in the wrapper headers section, for hashing.
+
+    The string-to-sign uses the SHA256 of the encoded `:date` header bytes
+    (name-prefixed, type-tagged, big-endian int64 milliseconds since UTC
+    epoch). This mirrors the wire format produced by the event-stream encoder
+    when serializing the wrapper message.
+    """
+    name = b":date"
+    epoch_ms = int(when.timestamp() * 1000)
+    return bytes([len(name)]) + name + bytes([_HEADER_TYPE_TIMESTAMP]) + struct.pack(">q", epoch_ms)

amazon_polly_streaming/_eventstream.py

@@ -0,0 +1,155 @@
+"""AWS event-stream binary format parser and encoder.
+
+This module provides three interfaces:
+
+- `EventStreamParser`: stateful, sync-friendly. Accumulates bytes via `feed(...)`
+  and returns the list of complete messages decoded so far.
+- `parse_stream(...)`: async iterator wrapper. Consumes an async iterator of
+  byte chunks and yields each complete message.
+- `encode_message(...)` / `encode_messages(...)`: build the binary on-the-wire
+  representation of one or more event-stream messages.
+
+Parsing reuses `botocore.eventstream.EventStreamBuffer`. Encoding is
+implemented directly with `struct` and `binascii.crc32` because `awscrt`
+exposes the encoder only via its connection-oriented RPC client, which is
+heavier than what we need here.
+"""
+
+from __future__ import annotations
+
+import binascii
+import datetime
+import struct
+from typing import TYPE_CHECKING
+
+from botocore.eventstream import EventStreamBuffer
+
+if TYPE_CHECKING:
+    from collections.abc import AsyncIterator, Mapping, Sequence
+
+    from botocore.eventstream import EventStreamMessage
+
+# Header value type codes per the AWS event-stream binary spec.
+_HEADER_TYPE_BYTE_ARRAY = 6
+_HEADER_TYPE_STRING = 7
+_HEADER_TYPE_TIMESTAMP = 8
+# Each prelude is 12 bytes: total_length (u32) + headers_length (u32) + crc (u32).
+_PRELUDE_LENGTH = 12
+# The trailing message CRC is 4 bytes.
+_MESSAGE_CRC_LENGTH = 4
+
+HeaderValue = str | bytes | datetime.datetime
+
+
+class EventStreamParser:
+    """Sync stateful parser for AWS event-stream framing."""
+
+    def __init__(self) -> None:
+        """Initialize with an empty buffer."""
+        self._buffer = EventStreamBuffer()
+
+    def feed(self, data: bytes) -> list[EventStreamMessage]:
+        """Add bytes to the buffer; return any complete messages now decoded.
+
+        Args:
+            data: Raw bytes from the network. May be a partial message, a full
+                message, or multiple messages concatenated.
+
+        Returns:
+            List of complete `EventStreamMessage` objects decoded after this
+            chunk. Empty if the buffer does not yet contain a full message.
+        """
+        self._buffer.add_data(data)
+        return list(self._buffer)
+
+    def reset(self) -> None:
+        """Reset the internal buffer to an empty state."""
+        self._buffer = EventStreamBuffer()
+
+
+async def parse_stream(
+    chunks: AsyncIterator[bytes],
+) -> AsyncIterator[EventStreamMessage]:
+    """Consume an async iterator of bytes and yield each complete message.
+
+    Args:
+        chunks: Async iterator producing byte chunks from a streaming HTTP body.
+
+    Yields:
+        Each complete `EventStreamMessage` as soon as enough bytes have arrived
+        to decode it.
+    """
+    parser = EventStreamParser()
+    async for chunk in chunks:
+        for msg in parser.feed(chunk):
+            yield msg
+
+
+def _encode_header(name: str, value: HeaderValue) -> bytes:
+    """Encode a single header in event-stream wire format.
+
+    Dispatches on the value Python type:
+        - `str`   -> type 7 (UTF-8 string), u16 length prefix.
+        - `bytes` -> type 6 (byte array), u16 length prefix.
+        - `datetime.datetime` -> type 8 (timestamp), i64 ms since UTC epoch.
+
+    Layout: name_len (u8) | name | value_type (u8) | value-encoded-bytes.
+    """
+    name_bytes = name.encode("utf-8")
+    name_prefix = bytes([len(name_bytes)]) + name_bytes
+    if isinstance(value, str):
+        value_bytes = value.encode("utf-8")
+        return (
+            name_prefix
+            + bytes([_HEADER_TYPE_STRING])
+            + struct.pack(">H", len(value_bytes))
+            + value_bytes
+        )
+    if isinstance(value, datetime.datetime):
+        if value.tzinfo is None:
+            msg = f"Header {name!r} datetime value must be timezone-aware"
+            raise ValueError(msg)
+        epoch_ms = int(value.timestamp() * 1000)
+        return name_prefix + bytes([_HEADER_TYPE_TIMESTAMP]) + struct.pack(">q", epoch_ms)
+    return name_prefix + bytes([_HEADER_TYPE_BYTE_ARRAY]) + struct.pack(">H", len(value)) + value
+
+
+def encode_message(headers: Mapping[str, HeaderValue], payload: bytes) -> bytes:
+    """Encode one event-stream message.
+
+    Headers are emitted with a wire type matching their Python type (see
+    `_encode_header`): string (type 7), byte_array (type 6), timestamp
+    (type 8). The Polly inbound protocol uses string headers for inner events
+    (`:message-type`, `:event-type`, `:content-type`) and bytes/timestamp for
+    the rolling-signature wrapper (`:date`, `:chunk-signature`).
+
+    Args:
+        headers: Mapping of header name to typed value.
+        payload: Raw payload bytes (already JSON-encoded for ``TextEvent`` or
+            empty for ``CloseStreamEvent``).
+
+    Returns:
+        Binary message: prelude (with prelude CRC) + headers + payload +
+        trailing message CRC.
+    """
+    headers_bytes = b"".join(_encode_header(name, value) for name, value in headers.items())
+    headers_length = len(headers_bytes)
+    total_length = _PRELUDE_LENGTH + headers_length + len(payload) + _MESSAGE_CRC_LENGTH
+
+    prelude = struct.pack(">II", total_length, headers_length)
+    prelude_with_crc = prelude + struct.pack(">I", binascii.crc32(prelude))
+
+    body = prelude_with_crc + headers_bytes + payload
+    return body + struct.pack(">I", binascii.crc32(body))
+
+
+def encode_messages(messages: Sequence[tuple[Mapping[str, HeaderValue], bytes]]) -> bytes:
+    """Encode a sequence of (headers, payload) tuples into one byte string.
+
+    Args:
+        messages: Ordered list of messages to encode.
+
+    Returns:
+        Concatenated binary representation of all messages.
+    """
+    return b"".join(encode_message(headers, payload) for headers, payload in messages)