estuary-cdk 0.3.0.dev31__py3-none-any.whl

estuary_cdk/__init__.py

@@ -0,0 +1,176 @@
+from dataclasses import dataclass
+from logging import Logger
+from pydantic import BaseModel
+from typing import Any, TypeVar, Callable, AsyncGenerator, Generic
+import abc
+from types import FrameType
+import asyncio
+import signal
+import sys
+import traceback
+
+from .logger import init_logger, FlowLogger
+from .flow import ValidationError
+from .utils import get_running_tasks_info
+
+# Request type served by this connector.
+Request = TypeVar("Request", bound=BaseModel)
+
+
+# stdin_jsonl parses newline-delimited JSON instances of `cls` from stdin and yields them.
+async def stdin_jsonl(cls: type[Request]) -> AsyncGenerator[Request, None]:
+    loop = asyncio.get_event_loop()
+    reader = asyncio.StreamReader(limit=1 << 27)  # 128 MiB.
+    protocol = asyncio.StreamReaderProtocol(reader)
+    await loop.connect_read_pipe(lambda: protocol, sys.stdin)
+
+    while line := await reader.readline():
+        request = cls.model_validate_json(line)
+        yield request
+
+
+# A Mixin is a supporting class which implements a utility on behalf of a connector.
+# Mixins may implement a protocol like Python's Asynchronous Context Manager pattern:
+# they're entered prior to handling any requests, and exited after all requests have
+# been fully processed.
+class Mixin:
+    # Parameters are positional-only so overrides may name them freely, and the
+    # return is Any because serve() discards it while implementations commonly
+    # return self.
+    async def _mixin_enter(self, log: Logger, /) -> Any: ...
+    async def _mixin_exit(self, log: Logger, /) -> Any: ...
+
+
+@dataclass
+class Stopped(Exception):
+    """
+    Stopped is raised by a connector's handle() function to indicate that the
+    connector has stopped and this process should exit.
+
+    `error`, if set, is a fatal error condition that caused the connector to exit.
+    """
+
+    error: str | None
+
+
+# BaseConnector from which all Flow Connectors inherit.
+class BaseConnector(Generic[Request], abc.ABC):
+
+    # request_class() returns the concrete class instance for Request served
+    # by this connector. It's required to implement to enable parsing of
+    # Requests prior to their being dispatched by the connector.
+    #
+    # Python classes are type-erased, so it's not possible to determine the
+    # concrete class at runtime using only type annotations.
+    # It is an instance method, not a classmethod, because connectors
+    # idiomatically override it as `def request_class(self)`; serve() calls it
+    # through the instance, so either form works at runtime.
+    @abc.abstractmethod
+    def request_class(self) -> type[Request]:
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    async def handle(self, log: FlowLogger, request: Request) -> None:
+        raise NotImplementedError()
+
+    # Serve this connector by invoking `handle()` for all incoming instances of
+    # Request as a concurrent asyncio.Task using the provided `responder`.
+    # When incoming `requests` are exhausted, all pending tasks are awaited
+    # before returning.
+    # All handler exceptions are caught and logged, and serve() returns an
+    # exit code which indicates whether an error occurred.
+    async def serve(
+        self,
+        log: FlowLogger | None = None,
+        requests: Callable[
+            [type[Request]], AsyncGenerator[Request, None]
+        ] = stdin_jsonl,
+    ):
+        if not log:
+            log = init_logger()
+
+        assert isinstance(log, FlowLogger)  # Narrow type to non-None.
+
+        loop = asyncio.get_running_loop()
+        this_task = asyncio.current_task(loop)
+        original_sigquit = signal.getsignal(signal.SIGQUIT)
+
+        def dump_all_tasks(signum: int, frame: FrameType | None):
+            task_infos = get_running_tasks_info(exclude_tasks={this_task} if this_task else None)
+
+            for info in task_infos:
+                log.info("Attempting to inject SIGQUIT exception into task.", {
+                    "task_name": info.task_name,
+                    "coro_name": info.coro_name,
+                })
+
+                # Reach inside the task coroutine to inject an exception, which
+                # will unwind the task stack and lets us print a precise stack trace.
+                coro = info.task.get_coro()
+                if coro:
+                    try:
+                        coro.throw(RuntimeError("injected SIGQUIT exception"))
+                    except Exception as exc:
+                        msg, args = type(exc).__name__, exc.args
+                        if len(args) != 0:
+                            msg = f"{msg}: {args[0]}"
+                        log.exception(msg, args)
+
+                # We manually injected an exception into the coroutine,
+                # so the asyncio event loop will attempt to await it again
+                # and get a new "cannot reuse already awaited coroutine".
+                # Cancel the task now and add a callback to clear the exception.
+                info.task.cancel()
+                info.task.add_done_callback(lambda task: task.exception())
+
+            # Wake event loop.
+            loop.call_soon_threadsafe(lambda: None)
+
+        signal.signal(signal.SIGQUIT, dump_all_tasks)
+
+        # Call _mixin_enter() on all mixed-in base classes.
+        for base in self.__class__.__bases__:
+            if enter := getattr(base, "_mixin_enter", None):
+                await enter(self, log)
+
+        failed = False
+        try:
+            async with asyncio.TaskGroup() as group:
+                async for request in requests(self.request_class()):
+                    group.create_task(self.handle(log, request))
+
+        except ExceptionGroup as exc_group:
+            for exc in exc_group.exceptions:
+                if isinstance(exc, ValidationError):
+                    if len(exc.errors) == 1:
+                        log.error(exc.errors[0])
+                    else:
+                        log.error(
+                            "Multiple validation errors:\n - "
+                            + "\n - ".join(exc.errors)
+                        )
+                    failed = True
+                elif isinstance(exc, Stopped):
+                    if exc.error:
+                        log.error(f"{exc.error}")
+                        failed = True
+                else:
+                    log.error("".join(traceback.format_exception(exc)))
+                    failed = True
+
+        finally:
+            # Call _mixin_exit() on all mixed-in base classes, in reverse order.
+            for base in reversed(self.__class__.__bases__):
+                if exit := getattr(base, "_mixin_exit", None):
+                    await exit(self, log)
+
+            await self.exit(log)
+
+            # Restore the original signal handler
+            signal.signal(signal.SIGQUIT, original_sigquit)
+
+        if failed:
+            raise SystemExit(1)
+
+    async def exit(self, log: FlowLogger):
+        return  # no-op default

estuary_cdk/_background.py

@@ -0,0 +1,20 @@
+"""Fire-and-forget task spawning that keeps a strong reference.
+
+asyncio only holds a weak reference to a running task, so a task whose only
+reference is the `create_task()` return value can be garbage-collected
+mid-flight. Callers here want to start work without awaiting it, so the
+reference has to live somewhere until the task finishes.
+"""
+
+import asyncio
+from typing import Any, Coroutine
+
+_tasks: set[asyncio.Task[Any]] = set()
+
+
+def spawn(coro: Coroutine[Any, Any, Any], *, name: str | None = None) -> asyncio.Task[Any]:
+    """Start `coro` as a task which is not awaited, and keep it alive."""
+    task = asyncio.create_task(coro, name=name)
+    _tasks.add(task)
+    _ = task.add_done_callback(_tasks.discard)
+    return task

estuary_cdk/async_streams.py

@@ -0,0 +1,112 @@
+import asyncio
+from collections.abc import AsyncGenerator
+from typing import AsyncIterator, TypeVar, cast
+
+T = TypeVar("T")
+
+
+async def merge(*iterators: AsyncGenerator[T, None]) -> AsyncIterator[T]:
+    """
+    Merge multiple async iterators into a single stream in completion order.
+
+    Items are yielded as they become available from any iterator, without
+    preserving the original ordering. All iterators run concurrently.
+
+    If any iterator raises an exception, all other iterators are cancelled
+    and the exception is propagated after cleanup. Similarly, if the consumer
+    stops iterating early (e.g., using break), all iterators are cancelled.
+
+    Args:
+        *iterators: Variable number of async iterators to merge. Can be
+            async generators or any object implementing AsyncIterator protocol.
+
+    Yields:
+        Items from any of the input iterators as they become available.
+
+    Raises:
+        Exception: Any exception raised by the input iterators is propagated
+            to the caller after all iterators are properly cleaned up.
+
+    Example:
+        >>> async def gen1():
+        ...     for i in [1, 2, 3]:
+        ...         yield i
+        ...
+        >>> async def gen2():
+        ...     for i in [4, 5, 6]:
+        ...         yield i
+        ...
+        >>> async for item in merge(gen1(), gen2()):
+        ...     print(item)  # Prints 1-6 in arbitrary order
+
+    Note:
+        The bounded queue (size 16) provides backpressure to prevent fast
+        iterators from consuming excessive memory while waiting for slow
+        consumers. This makes it safe to use with high-throughput streams.
+    """
+    if not iterators:
+        return
+
+    _DONE = object()
+    queue: asyncio.Queue[T | BaseException | object] = asyncio.Queue(16)
+    active_count = len(iterators)
+    cancelling = False
+
+    async def _drain(iterator: AsyncIterator[T]) -> None:
+        nonlocal cancelling
+        try:
+            async for item in iterator:
+                await queue.put(item)
+        except Exception as e:
+            if not cancelling:
+                cancelling = True
+                await queue.put(e)
+        finally:
+            await queue.put(_DONE)
+
+    tasks = [asyncio.create_task(_drain(aiter)) for aiter in iterators]
+
+    async def _reap() -> None:
+        """Discard queued items so blocked producers can finish unwinding."""
+        while True:
+            await queue.get()
+
+    def _cancel_all_tasks() -> None:
+        nonlocal cancelling
+        cancelling = True
+        for task in tasks:
+            if not task.done():
+                task.cancel()
+
+    try:
+        while active_count > 0:
+            item = await queue.get()
+
+            if isinstance(item, BaseException):
+                raise item
+
+            if item is _DONE:
+                active_count -= 1
+            else:
+                # Neither the _DONE sentinel nor an exception, so an item of T.
+                yield cast(T, item)
+
+    finally:
+        _cancel_all_tasks()
+
+        # A `_drain` task cancelled while blocked on a full queue still runs its
+        # own `finally`, which awaits `queue.put(_DONE)`. By now nothing is
+        # consuming the queue, so that put would never complete and the gather
+        # below would wait forever. Keep draining until the tasks have exited.
+        reaper = asyncio.create_task(_reap())
+        try:
+            await asyncio.gather(*tasks, return_exceptions=True)
+        finally:
+            reaper.cancel()
+            await asyncio.gather(reaper, return_exceptions=True)
+
+        for iterator in iterators:
+            try:
+                await iterator.aclose()
+            except Exception:
+                pass

estuary_cdk/buffer_ordered.py

@@ -0,0 +1,122 @@
+import asyncio
+from dataclasses import dataclass
+from typing import AsyncGenerator, Awaitable, Generic, TypeVar
+
+T = TypeVar("T")
+
+
+@dataclass
+class BufferWork(Generic[T]):
+    aw: Awaitable[T]
+    result: asyncio.Queue[T | None]
+
+    next: asyncio.Future["BufferWork[T] | None"]
+
+
+async def buffer_ordered(
+    aws: AsyncGenerator[Awaitable[T], None],
+    concurrency: int,
+) -> AsyncGenerator[T, None]:
+    """
+    Run the stream of awaitables 'aws' concurrently and return the results of
+    each in order. There may be up to 'concurrency' results from the awaitables
+    held in memory at a time, so a lower concurrency may improve memory usage at
+    the expense of throughput.
+
+    Args:
+        aws: A stream of awaitables to run concurrently. concurrency: The
+        maximum number of concurrent awaitables to run.
+
+    Returns:
+        A stream of results from the awaitables, with the order matching the
+        order yielded by 'aws'.
+    """
+
+    work: asyncio.Queue[BufferWork[T] | None] = asyncio.Queue(1)
+    next: asyncio.Future[BufferWork[T] | None] = asyncio.Future()
+
+    async def _producer():
+        current = next
+        try:
+            async for aw in aws:
+                this_next: asyncio.Future[BufferWork[T] | None] = asyncio.Future()
+                this_work = BufferWork(
+                    aw=aw,
+                    result=asyncio.Queue(1),
+                    next=this_next,
+                )
+                current.set_result(this_work)
+                current = this_next
+                await work.put(this_work)
+
+                # Wait until the last item has been removed from the queue by a
+                # worker before requesting anything else from the awaitables
+                # generator. This prevents an awaitable from being held in limbo
+                # in this loop without being awaited if the producer or worker
+                # raises an exception and exits early.
+                await work.join()
+
+            # Send stopping signals to the output loop and workers.
+            current.set_result(None)
+            for _ in range(concurrency):
+                await work.put(None)
+        except Exception:
+            # Signal the output loop to stop.
+            current.set_result(None)
+            raise
+
+    async def _worker():
+        while True:
+            this_work = await work.get()
+            work.task_done()  # Signal removal from the queue, per the note above in _producer.
+            if this_work is None:
+                break
+
+            try:
+                this_work.result.put_nowait(await this_work.aw)
+            except Exception:
+                # Signal the output loop to stop.
+                this_work.result.put_nowait(None)
+                raise
+            # Do not get another awaitable until this one has been fully handled
+            # by the output loop. This limits the number of pending work items
+            # to output, which is important their result may not necessarily be
+            # small and will be held in memory.
+            await this_work.result.join()
+
+    try:
+        async with asyncio.TaskGroup() as tg:
+            for coro in [_producer(), *[_worker() for _ in range(concurrency)]]:
+                tg.create_task(coro)
+
+            # Output loop.
+            while True:
+                finished_work = await next
+                if finished_work is None:
+                    break
+
+                this_result = await finished_work.result.get()
+                if this_result is None:
+                    break
+
+                yield this_result
+
+                # Signal the worker so it can start on another awaitable.
+                finished_work.result.task_done()
+                # Output the next result per the original ordering of input awaitables.
+                next = finished_work.next
+    except ExceptionGroup as eg:
+        # Raise the first error from the producer or any of the workers.
+        for e in eg.exceptions:
+            raise e
+    except Exception:
+        raise
+    finally:
+        # Await any queued awaitables, discarding further errors.
+        while not work.empty():
+            remaining_work = work.get_nowait()
+            if remaining_work is not None:
+                try:
+                    await remaining_work.aw
+                except Exception:
+                    pass

estuary_cdk/capture/__init__.py

@@ -0,0 +1,9 @@
+from .task import Request, Response, Task
+from .base_capture_connector import BaseCaptureConnector
+
+__all__ = [
+    "Request",
+    "Response",
+    "Task",
+    "BaseCaptureConnector",
+]

estuary_cdk/capture/_emit.py

@@ -0,0 +1,26 @@
+import asyncio
+import shutil
+from typing import IO, BinaryIO
+
+# Global lock for serializing all emissions to stdout and prevent interleaving output.
+_emit_lock = asyncio.Lock()
+
+
+async def emit_bytes(data: bytes, output: BinaryIO) -> None:
+    async with _emit_lock:
+        await asyncio.to_thread(_write_bytes, data, output)
+
+
+async def emit_from_buffer(buffer: IO[bytes], output: BinaryIO) -> None:
+    async with _emit_lock:
+        await asyncio.to_thread(_copy_buffer, buffer, output)
+
+
+def _write_bytes(data: bytes, output: BinaryIO) -> None:
+    output.write(data)
+    output.flush()
+
+
+def _copy_buffer(buffer: IO[bytes], output: BinaryIO) -> None:
+    shutil.copyfileobj(buffer, output)
+    output.flush()