planar 0.5.0__py3-none-any.whl

planar/sse/hub.py

@@ -0,0 +1,216 @@
+import argparse
+import atexit
+import logging
+import logging.config
+import re
+import shutil
+import time
+from asyncio import CancelledError, Queue, create_task, sleep
+from contextlib import asynccontextmanager
+from dataclasses import dataclass, field
+from fnmatch import translate
+from typing import Any, Optional
+from uuid import UUID, uuid4
+from weakref import WeakSet
+
+import uvicorn
+from fastapi import FastAPI, Request
+from starlette.responses import StreamingResponse
+
+from planar.config import PlanarConfig
+from planar.logging import get_logger
+from planar.sse.constants import SSE_ENDPOINT
+from planar.sse.model import Event
+
+logger = get_logger(__name__)
+
+
+@dataclass(kw_only=True)
+class EventWithTime:
+    # Since event storage is volatile, "time" also uniquely identifies the event
+    time: int = field(default_factory=time.monotonic_ns)
+    event: Event
+
+
+@dataclass(kw_only=True)
+class Client:
+    event_index: int = 0
+    patterns: list[str] = field(default_factory=list)
+    compiled_patterns: list[re.Pattern] = field(default_factory=list)
+    queue: Queue[EventWithTime] = field(default_factory=Queue)
+    id: UUID = field(default_factory=uuid4)
+
+    def __post_init__(self):
+        # Compile glob patterns to regex patterns
+        self.compiled_patterns = []
+        for pattern in self.patterns:
+            translated = translate(pattern)
+            logger.debug(
+                "compiling glob pattern to regex", pattern=pattern, regex=translated
+            )
+            self.compiled_patterns.append(re.compile(translated))
+
+    def forward(self, events: list[EventWithTime]):
+        while self.event_index < len(events):
+            event_time = events[self.event_index]
+            event = event_time.event
+            self.event_index += 1
+            if not self.compiled_patterns:
+                self.queue.put_nowait(event_time)
+                continue
+            for i, compiled_pattern in enumerate(self.compiled_patterns):
+                logger.debug(
+                    "matching event against pattern",
+                    event_name=event.name,
+                    pattern=self.patterns[i],
+                )
+                if compiled_pattern.fullmatch(event.name):
+                    logger.debug(
+                        "matched event against pattern",
+                        event_name=event.name,
+                        pattern=self.patterns[i],
+                    )
+                    self.queue.put_nowait(event_time)
+                    break
+
+    def __eq__(self, other):
+        return self.id == other.id
+
+    def __hash__(self):
+        return hash(self.id)
+
+
+events: list[EventWithTime] = []
+clients: WeakSet[Client] = WeakSet()
+
+
+# Periodically delete events older than 30 seconds.
+async def cleanup_old_events():
+    global events
+    while True:
+        await sleep(1)
+        cutoff = time.monotonic_ns() - 30 * 1_000_000_000
+        prune_index = None
+        # find the index of the first event that is older than 30 seconds,
+        # starting from the latest events
+        prune_index = -1
+        for i in range(len(events) - 1, -1, -1):
+            event = events[i]
+            if event.time < cutoff:
+                prune_index = i
+                break
+        remove_count = prune_index + 1
+        if remove_count > 0:
+            for client in clients:
+                client.event_index = max(0, client.event_index - remove_count)
+            events = events[remove_count:]
+            logger.debug("removed old events", count=remove_count)
+
+
+@asynccontextmanager
+async def lifespan(_: FastAPI):
+    prune_task = create_task(cleanup_old_events())
+    yield
+    prune_task.cancel()
+    try:
+        await prune_task
+    except CancelledError:
+        pass
+
+
+app = FastAPI(title="Planar SSE Hub", lifespan=lifespan)
+
+
+@app.post("/push")
+def push(event: Event):
+    events.append(EventWithTime(event=event))
+    for client in clients:
+        client.forward(events)
+
+
+@app.get(SSE_ENDPOINT)
+async def sse(
+    request: Request, subscribe: Optional[str] = None, new_events_only: bool = False
+):
+    # Get Last-Event-ID from headers
+    last_event_id = request.headers.get("Last-Event-ID")
+
+    client = Client(patterns=subscribe.split(",") if subscribe else [])
+    clients.add(client)
+    logger.debug("client connected", client_id=client.id)
+
+    # Handle Last-Event-ID so that reconnects are handled correctly
+    if last_event_id and not new_events_only:
+        try:
+            last_id = int(last_event_id)
+            # Find the index of the event with the matching ID or the first one after it
+            for i, event_time in enumerate(events):
+                if event_time.time > last_id:
+                    client.event_index = i
+                    break
+            logger.debug(
+                "client resuming from event index",
+                event_index=client.event_index,
+                last_event_id=last_event_id,
+            )
+        except ValueError:
+            logger.warning("invalid last-event-id format", last_event_id=last_event_id)
+
+    if not new_events_only:
+        client.forward(events)
+
+    async def event_stream():
+        try:
+            while True:
+                event_time = await client.queue.get()
+                # I don't think ordering matters, but in the docs ids always
+                # come last, so repeat that here
+                yield f"data: {event_time.event.model_dump_json()}\nid: {event_time.time}\n\n"
+        except CancelledError:
+            logger.debug("client disconnected", client_id=client.id)
+
+    return StreamingResponse(
+        event_stream(),
+        media_type="text/event-stream",
+        headers={
+            "Cache-Control": "no-cache",
+            "X-Accel-Buffering": "no",
+        },
+    )
+
+
+def builtin_run(server_url: str, socket_dir: str, planar_config_dict: dict[str, Any]):
+    # we only use this to configure logging
+    config = PlanarConfig.model_validate(planar_config_dict)
+    config.configure_logging()
+
+    def cleanup():
+        shutil.rmtree(socket_dir)
+
+    # ensure the socket directory is removed on exit
+    atexit.register(cleanup)
+
+    # explicitly set 1 worker since this is going to aggregate all events
+    # we have to pass log_config=None to prevent uvicorn from trying to
+    # configure logging by itself
+    uvicorn.run(app, uds=server_url, workers=1, log_config=None)
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--host", default="0.0.0.0")
+    parser.add_argument("--port", default=8888, type=int)
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+    logger.setLevel(logging.DEBUG)
+    logging.basicConfig(
+        level=logging.DEBUG, format="[%(levelname)s] %(asctime)s - %(message)s"
+    )
+    uvicorn.run(app, host=args.host, port=args.port, workers=1)
+
+
+if __name__ == "__main__":
+    main()

planar/sse/model.py

@@ -0,0 +1,8 @@
+from typing import Any
+
+from pydantic import BaseModel
+
+
+class Event(BaseModel):
+    name: str
+    payload: dict[str, Any]

planar/sse/proxy.py

@@ -0,0 +1,257 @@
+import os
+import tempfile
+import uuid
+from asyncio import (
+    CancelledError,
+    Queue,
+    QueueFull,
+    Task,
+    create_task,
+    current_task,
+    sleep,
+)
+from contextlib import asynccontextmanager
+from multiprocessing import Process
+from typing import Any
+from weakref import WeakSet
+
+import httpx
+from fastapi import APIRouter, Request
+from starlette.responses import StreamingResponse
+
+from planar.config import PlanarConfig
+from planar.logging import get_logger
+from planar.sse.constants import SSE_ENDPOINT
+from planar.sse.hub import builtin_run
+from planar.sse.model import Event
+
+logger = get_logger(__name__)
+
+
+def parse_hub_url(
+    hub_url: str,
+) -> tuple[httpx.AsyncHTTPTransport | None, str]:
+    transport = None
+    if hub_url[0] == "/":
+        transport = httpx.AsyncHTTPTransport(uds=hub_url)
+        # Use a dummy server URL for formatting the target URL.
+        # The transport will ignore it and connect to the domain socket
+        hub_url = "http://planar-sse"
+    return transport, hub_url
+
+
+def get_builtin_hub_socket_path():
+    # This env var trick has two purposes:
+    #
+    #   - Generate a unique socket path
+    #   - Distribute the same socket path to all workers.
+    #
+    # Initially the env var will not be set, so we generate the uuid
+    # and set it. When uvicorn spawns workers, they will inherit the
+    # parent environment and use it instead of generating their own.
+    builtin_hub_uuid = os.getenv("PLANAR_BUILTIN_SSE_UUID", "")
+    if builtin_hub_uuid == "":
+        builtin_hub_uuid = str(uuid.uuid4())
+        os.environ["PLANAR_BUILTIN_SSE_UUID"] = builtin_hub_uuid
+    return f"{tempfile.gettempdir()}/planar-sse-{builtin_hub_uuid}/socket"
+
+
+class SSEProxy:
+    def __init__(
+        self,
+        config: PlanarConfig,
+    ):
+        sse_hub = config.sse_hub
+        self.config = config
+        self.enable_builtin_hub = False
+        self.hub_url = ""
+        self.transport: httpx.AsyncHTTPTransport | None = None
+        self.stream_tasks: WeakSet[Task] = WeakSet()
+
+        if isinstance(sse_hub, str):
+            # Connect to a separate SSE hub listening on a TCP address
+            self.hub_url = sse_hub
+        elif sse_hub is True:
+            # Use builtin hub spawned as a subprocess listening on an UNIX socket.
+            self.enable_builtin_hub = True
+            self.hub_url = get_builtin_hub_socket_path()
+
+        self.builtin_process: Process | None = None
+        self.router = APIRouter()
+        self.queue: Queue[Event] = Queue(maxsize=1000)
+        self.forward_task: Task | None = None
+
+        if self.hub_url:
+            self.setup_proxy_endpoint()
+
+    def push(self, name: str, payload: dict[str, Any]):
+        try:
+            self.queue.put_nowait(Event(name=name, payload=payload))
+        except QueueFull:
+            # not processing events fast enough, so just ignore it
+            logger.warning("sse proxy queue is full, dropping event", event_name=name)
+
+    def start_builtin_hub(self):
+        assert self.hub_url[0] == "/"
+        socket_dir = os.path.dirname(self.hub_url)
+        logger.debug("attempting to create socket directory", socket_dir=socket_dir)
+        try:
+            # try to create the socket directory, only one of the workers will
+            # succeed
+            os.mkdir(socket_dir)
+            logger.debug("socket directory created", socket_dir=socket_dir)
+        except FileExistsError:
+            # another worker created the directory first, ignore
+            logger.debug("socket directory already exists", socket_dir=socket_dir)
+            return
+
+        # start the builtin SSE hub in a separate process
+        logger.info(
+            "starting builtin sse hub process for socket", socket_url=self.hub_url
+        )
+        self.builtin_process = Process(
+            target=builtin_run,
+            args=(self.hub_url, socket_dir, self.config.model_dump()),
+        )
+        self.builtin_process.start()
+
+    def start(self):
+        logger.debug(
+            "sseproxy start called",
+            hub_url=self.hub_url,
+            enable_builtin=self.enable_builtin_hub,
+        )
+        if not self.hub_url:
+            raise ValueError("hub_url is not set")
+
+        if self.enable_builtin_hub:
+            self.start_builtin_hub()
+
+        self.transport, self.hub_url = parse_hub_url(self.hub_url)
+        forward_url = f"{self.hub_url}/push"
+
+        async def forward():
+            logger.debug("sse event forwarding task started", url=forward_url)
+            async with httpx.AsyncClient(transport=self.transport) as client:
+                while True:
+                    event = await self.queue.get()
+                    logger.debug(
+                        "got event from queue to forward", event_name=event.name
+                    )
+                    while True:
+                        try:
+                            await client.post(
+                                forward_url,
+                                content=event.model_dump_json(),
+                                headers={"Content-Type": "application/json"},
+                            )
+                            logger.info(
+                                "successfully forwarded event",
+                                event_name=event.name,
+                                url=forward_url,
+                            )
+                            break
+                        except Exception:
+                            logger.exception(
+                                "exception while forwarding sse event to hub, will retry"
+                            )
+                            await sleep(5)
+
+        self.forward_task = create_task(forward())
+        logger.info("sse event forwarding task created")
+
+    async def stop(self):
+        logger.debug("sseproxy stop called")
+        if self.forward_task:
+            logger.info("cancelling sse event forwarding task")
+            self.forward_task.cancel()
+            try:
+                await self.forward_task
+            except CancelledError:
+                logger.debug("sse event forwarding task cancelled")
+                pass
+
+        for stream_task in self.stream_tasks:
+            logger.debug("cancelling sse stream task", task=stream_task.get_name())
+            stream_task.cancel()
+
+        if self.builtin_process:
+            # when using multiple workers, only one worker will have started the
+            # builtin hub. That's why we check for self.builtin_process instead of
+            # self.enable_builtin_hub. Also helps with type checking.
+            logger.info("terminating builtin sse hub process")
+            self.builtin_process.terminate()
+
+    def setup_proxy_endpoint(self):
+        @self.router.get("/")
+        async def proxy(request: Request):
+            logger.debug(
+                "sse proxy endpoint called",
+                query=request.url.query,
+                headers=dict(request.headers),
+            )
+            async with self.connect(request.url.query, dict(request.headers)) as (
+                status,
+                headers,
+                stream,
+            ):
+                return StreamingResponse(
+                    stream(),
+                    status_code=status,
+                    headers=headers,
+                )
+
+        return proxy  # dummy return to prevent unused warning
+
+    @asynccontextmanager
+    async def connect(self, query: str = "", headers: dict[str, str] = {}):
+        logger.debug("sseproxy connect called", query=query, headers=headers)
+        if not self.hub_url or not self.transport:
+            raise ValueError("hub_url is not set")
+
+        hub_url = self.hub_url
+        transport = self.transport
+
+        client = httpx.AsyncClient(
+            transport=transport, base_url=hub_url, timeout=httpx.Timeout(None)
+        )
+
+        while True:
+            try:
+                # Construct the target URL
+                url = httpx.URL(
+                    path=SSE_ENDPOINT,
+                    query=query.encode(),
+                )
+
+                # Build the outgoing request
+                proxy_request = client.build_request(
+                    method="GET",
+                    url=url,
+                    headers=headers,
+                )
+
+                # Send the request and stream the response
+                response = await client.send(proxy_request, stream=True)
+                logger.debug("connected to sse hub", hub_url=hub_url)
+
+                async def stream(lines: bool = False):
+                    stream_task = current_task()
+                    assert stream_task
+                    self.stream_tasks.add(stream_task)
+                    try:
+                        async for chunk in (
+                            response.aiter_lines() if lines else response.aiter_bytes()
+                        ):
+                            yield chunk
+                    except CancelledError:
+                        logger.debug("sse stream task cancelled")
+                        raise
+                    finally:
+                        await client.aclose()
+
+                yield response.status_code, dict(response.headers), stream
+                break
+            except Exception:
+                logger.exception("exception while connecting to sse hub, will retry")
+                await sleep(5)

planar/task_local.py

@@ -0,0 +1,37 @@
+import asyncio
+from typing import Generic, TypeVar
+from weakref import WeakKeyDictionary
+
+T = TypeVar("T")
+
+
+class TaskLocal(Generic[T]):
+    def __init__(self):
+        self._data = WeakKeyDictionary()
+
+    def set(self, context: T):
+        current_task = asyncio.current_task()
+        if not current_task:
+            raise RuntimeError("No current task")
+        self._data[current_task] = context
+
+    def get(self) -> T:
+        current_task = asyncio.current_task()
+        if not current_task:
+            raise RuntimeError("No current task")
+        context = self._data.get(current_task)
+        if context is None:
+            raise RuntimeError("No execution context")
+        return context
+
+    def clear(self):
+        current_task = asyncio.current_task()
+        if not current_task:
+            raise RuntimeError("No current task")
+        del self._data[current_task]
+
+    def is_set(self) -> bool:
+        current_task = asyncio.current_task()
+        if not current_task:
+            raise RuntimeError("No current task")
+        return self._data.get(current_task) is not None

planar/test_app.py

@@ -0,0 +1,51 @@
+from dotenv import load_dotenv
+from fastapi import APIRouter
+from pydantic import BaseModel
+
+from examples.simple_service.models import (
+    Invoice,
+)
+from planar import PlanarApp, sqlite_config
+
+load_dotenv()
+
+
+router = APIRouter()
+
+
+class InvoiceRequest(BaseModel):
+    message: str
+
+
+class InvoiceResponse(BaseModel):
+    status: str
+    echo: str
+
+
+app = PlanarApp(
+    config=sqlite_config("simple_service.db"),
+    title="Sample Invoice API",
+    description="API for CRUD'ing invoices",
+)
+
+
+def test_register_model_deduplication():
+    """Test that registering the same model multiple times only adds it once to the registry."""
+
+    # Ensure Invoice is registered (ObjectRegistry gets reset before each test)
+    app.register_entity(Invoice)
+    initial_model_count = len(app._object_registry.get_entities())
+
+    # Register the Invoice model again
+    app.register_entity(Invoice)
+
+    assert len(app._object_registry.get_entities()) == initial_model_count
+
+    # Register the same model a second time
+    app.register_entity(Invoice)
+
+    assert len(app._object_registry.get_entities()) == initial_model_count
+
+    # Verify that the model in the registry is the Invoice model
+    registered_models = app._object_registry.get_entities()
+    assert any(model.__name__ == "Invoice" for model in registered_models)