pgdispatch 0.1.0__py3-none-any.whl

pgdispatch/models.py

@@ -0,0 +1,208 @@
+"""Immutable records returned by the public API."""
+
+from __future__ import annotations
+
+import json
+import uuid
+from collections.abc import Mapping
+from dataclasses import dataclass, field
+from datetime import datetime
+from enum import Enum
+from typing import Any
+
+__all__ = [
+    "Attempt",
+    "CircuitState",
+    "Delivery",
+    "DeliveryStatus",
+    "Endpoint",
+    "Event",
+    "PublishResult",
+]
+
+
+class DeliveryStatus(str, Enum):
+    """Lifecycle of a delivery."""
+
+    PENDING = "pending"
+    """Waiting for its next attempt."""
+    IN_FLIGHT = "in_flight"
+    """Claimed by a worker that is sending it right now."""
+    SUCCEEDED = "succeeded"
+    """The endpoint answered with a 2xx status."""
+    FAILED = "failed"
+    """All attempts were exhausted; can be retried manually."""
+    CANCELLED = "cancelled"
+    """Abandoned, e.g. because the endpoint answered ``410 Gone``."""
+
+    @property
+    def is_final(self) -> bool:
+        return self in (DeliveryStatus.SUCCEEDED, DeliveryStatus.FAILED, DeliveryStatus.CANCELLED)
+
+
+class CircuitState(str, Enum):
+    """State of an endpoint's circuit breaker."""
+
+    CLOSED = "closed"
+    """Healthy: deliveries flow normally."""
+    OPEN = "open"
+    """Failing: no deliveries are sent until the cooldown elapses."""
+    HALF_OPEN = "half_open"
+    """Cooldown elapsed: a single probe delivery decides whether to close again."""
+
+
+@dataclass(frozen=True, slots=True)
+class Endpoint:
+    """A receiver URL with its own signing secret and settings."""
+
+    id: uuid.UUID
+    url: str
+    owner_id: str | None
+    description: str | None
+    event_types: tuple[str, ...]
+    headers: Mapping[str, str]
+    rate_limit_rps: float | None
+    is_enabled: bool
+    disabled_reason: str | None
+    circuit_state: CircuitState
+    consecutive_failures: int
+    failing_since: datetime | None
+    created_at: datetime
+    updated_at: datetime
+    secret: str | None = field(default=None, repr=False)
+    """The signing secret. Only populated by ``create``, ``rotate_secret`` and ``get_secret``."""
+
+    @classmethod
+    def from_row(cls, row: Mapping[str, Any], *, secret: str | None = None) -> Endpoint:
+        return cls(
+            id=to_uuid(row["id"]),
+            url=row["url"],
+            owner_id=row["owner_id"],
+            description=row["description"],
+            event_types=tuple(row["event_types"]),
+            headers=load_json(row["headers"]),
+            rate_limit_rps=row["rate_limit_rps"],
+            is_enabled=row["is_enabled"],
+            disabled_reason=row["disabled_reason"],
+            circuit_state=CircuitState(row["circuit_state"]),
+            consecutive_failures=row["consecutive_failures"],
+            failing_since=row["failing_since"],
+            created_at=row["created_at"],
+            updated_at=row["updated_at"],
+            secret=secret,
+        )
+
+
+@dataclass(frozen=True, slots=True)
+class Event:
+    """An immutable fact that happened in your system."""
+
+    id: uuid.UUID
+    event_type: str
+    owner_id: str | None
+    payload: Any
+    idempotency_key: str | None
+    created_at: datetime
+
+    @classmethod
+    def from_row(cls, row: Mapping[str, Any]) -> Event:
+        return cls(
+            id=to_uuid(row["id"]),
+            event_type=row["event_type"],
+            owner_id=row["owner_id"],
+            payload=json.loads(row["payload"]),
+            idempotency_key=row["idempotency_key"],
+            created_at=row["created_at"],
+        )
+
+
+@dataclass(frozen=True, slots=True)
+class Delivery:
+    """The task of delivering one event to one endpoint."""
+
+    id: int
+    event_id: uuid.UUID
+    endpoint_id: uuid.UUID
+    event_type: str
+    ordering_key: str | None
+    status: DeliveryStatus
+    attempt_count: int
+    next_attempt_at: datetime
+    last_status_code: int | None
+    last_error: str | None
+    created_at: datetime
+    finished_at: datetime | None
+
+    @classmethod
+    def from_row(cls, row: Mapping[str, Any]) -> Delivery:
+        return cls(
+            id=row["id"],
+            event_id=to_uuid(row["event_id"]),
+            endpoint_id=to_uuid(row["endpoint_id"]),
+            event_type=row["event_type"],
+            ordering_key=row["ordering_key"],
+            status=DeliveryStatus(row["status"]),
+            attempt_count=row["attempt_count"],
+            next_attempt_at=row["next_attempt_at"],
+            last_status_code=row["last_status_code"],
+            last_error=row["last_error"],
+            created_at=row["created_at"],
+            finished_at=row["finished_at"],
+        )
+
+
+@dataclass(frozen=True, slots=True)
+class Attempt:
+    """A single HTTP request made for a delivery, with the receiver's answer."""
+
+    id: int
+    delivery_id: int
+    attempt_number: int
+    request_headers: Mapping[str, str]
+    response_status: int | None
+    response_body: str | None
+    duration_ms: int | None
+    error: str | None
+    created_at: datetime
+
+    @property
+    def succeeded(self) -> bool:
+        return self.response_status is not None and 200 <= self.response_status < 300
+
+    @classmethod
+    def from_row(cls, row: Mapping[str, Any]) -> Attempt:
+        return cls(
+            id=row["id"],
+            delivery_id=row["delivery_id"],
+            attempt_number=row["attempt_number"],
+            request_headers=load_json(row["request_headers"]) or {},
+            response_status=row["response_status"],
+            response_body=row["response_body"],
+            duration_ms=row["duration_ms"],
+            error=row["error"],
+            created_at=row["created_at"],
+        )
+
+
+@dataclass(frozen=True, slots=True)
+class PublishResult:
+    """Outcome of :meth:`Webhooks.publish <pgdispatch.Webhooks.publish>`."""
+
+    event_id: uuid.UUID
+    deliveries: int
+    """Number of deliveries created, one per subscribed endpoint."""
+    duplicate: bool = False
+    """``True`` if an event with the same idempotency key already existed; nothing was created."""
+
+
+def to_uuid(value: Any) -> uuid.UUID:
+    # asyncpg returns its own UUID subclass; normalise for equality and hashing.
+    return value if type(value) is uuid.UUID else uuid.UUID(str(value))
+
+
+def load_json(value: Any) -> Any:
+    # Columns are selected as text so every driver behaves the same, but accept
+    # already-decoded values too.
+    if isinstance(value, (str, bytes)):
+        return json.loads(value)
+    return value

pgdispatch/security/__init__.py

@@ -0,0 +1 @@
+"""Outbound request safety."""

pgdispatch/security/ssrf.py

@@ -0,0 +1,141 @@
+"""Server-side request forgery (SSRF) protection.
+
+Endpoint URLs are supplied by your customers, so without protection anyone
+could point the worker at ``http://169.254.169.254/`` (cloud metadata),
+``http://localhost:5432`` or an internal admin panel. pgdispatch defends in two
+places:
+
+1. :func:`validate_endpoint_url` runs when an endpoint is created or updated:
+   only ``https`` (``http`` with ``allow_http``), no credentials in the URL, and
+   no literal private addresses.
+2. :class:`SafeResolver` runs on every connection: the host name is resolved,
+   *all* resulting addresses must be public, and aiohttp connects to exactly
+   the addresses that were checked. Re-resolving the name later is impossible,
+   which closes the DNS-rebinding hole. TLS still verifies the certificate
+   against the original host name. Literal IP hosts never reach a resolver, so
+   :func:`check_url_host` vets them before each request.
+
+Redirects are never followed, so a public endpoint cannot bounce the worker to
+a private one either.
+"""
+
+from __future__ import annotations
+
+import ipaddress
+import socket
+from urllib.parse import urlsplit
+
+import aiohttp
+from aiohttp.abc import AbstractResolver, ResolveResult
+
+from pgdispatch.errors import UnsafeURLError
+
+__all__ = [
+    "BlockedAddressError",
+    "SafeResolver",
+    "check_url_host",
+    "is_public_address",
+    "validate_endpoint_url",
+]
+
+IPAddress = ipaddress.IPv4Address | ipaddress.IPv6Address
+
+MAX_URL_LENGTH = 2048
+_LOCAL_HOST_SUFFIXES = (".localhost", ".local", ".internal", ".localdomain")
+
+
+class BlockedAddressError(OSError):
+    """The endpoint host resolves to an address the worker must not connect to."""
+
+
+def is_public_address(address: IPAddress) -> bool:
+    """Tell whether ``address`` is a globally routable unicast address.
+
+    Rejects private, loopback, link-local, multicast, reserved, shared (CGNAT)
+    and documentation ranges, and IPv6 addresses that embed a non-public IPv4
+    address (IPv4-mapped, 6to4, Teredo).
+    """
+    if isinstance(address, ipaddress.IPv6Address):
+        embedded = address.ipv4_mapped or address.sixtofour
+        if embedded is not None:
+            return is_public_address(embedded)
+        if address.teredo is not None:
+            return all(is_public_address(part) for part in address.teredo)
+    return address.is_global and not address.is_multicast
+
+
+def validate_endpoint_url(
+    url: str, *, allow_http: bool = False, allow_private_ips: bool = False
+) -> str:
+    """Return ``url`` if it may be used as an endpoint, else raise :class:`UnsafeURLError`."""
+    if len(url) > MAX_URL_LENGTH:
+        msg = f"URL is longer than {MAX_URL_LENGTH} characters"
+        raise UnsafeURLError(msg)
+    try:
+        parts = urlsplit(url)
+        port = parts.port
+    except ValueError as exc:
+        msg = f"invalid URL: {exc}"
+        raise UnsafeURLError(msg) from exc
+
+    allowed_schemes = ("https", "http") if allow_http else ("https",)
+    if parts.scheme not in allowed_schemes:
+        hint = "" if allow_http else " (plain http needs allow_http=True, for development only)"
+        msg = f"URL scheme must be {' or '.join(allowed_schemes)}{hint}"
+        raise UnsafeURLError(msg)
+    if parts.username is not None or parts.password is not None:
+        msg = "URL must not contain credentials; use custom endpoint headers instead"
+        raise UnsafeURLError(msg)
+    host = parts.hostname
+    if not host:
+        msg = "URL must contain a host name"
+        raise UnsafeURLError(msg)
+    if port == 0:
+        msg = "URL port must not be 0"
+        raise UnsafeURLError(msg)
+
+    if not allow_private_ips:
+        address = _parse_ip(host)
+        if address is not None and not is_public_address(address):
+            msg = f"URL points to a non-public address ({address})"
+            raise UnsafeURLError(msg)
+        if address is None and (host == "localhost" or host.endswith(_LOCAL_HOST_SUFFIXES)):
+            msg = f"URL points to a local host name ({host})"
+            raise UnsafeURLError(msg)
+    return url
+
+
+class SafeResolver(AbstractResolver):
+    """An aiohttp resolver that refuses host names resolving to non-public addresses."""
+
+    def __init__(self, inner: AbstractResolver | None = None) -> None:
+        self._inner = inner or aiohttp.ThreadedResolver()
+
+    async def resolve(
+        self, host: str, port: int = 0, family: socket.AddressFamily = socket.AF_INET
+    ) -> list[ResolveResult]:
+        results = await self._inner.resolve(host, port, family)
+        for result in results:
+            address = ipaddress.ip_address(result["host"].split("%", 1)[0])
+            if not is_public_address(address):
+                msg = f"refusing to connect to {host}: it resolves to non-public address {address}"
+                raise BlockedAddressError(msg)
+        return results
+
+    async def close(self) -> None:
+        await self._inner.close()
+
+
+def check_url_host(url: str) -> None:
+    """Raise :class:`BlockedAddressError` if ``url`` has a literal non-public IP host."""
+    address = _parse_ip(urlsplit(url).hostname or "")
+    if address is not None and not is_public_address(address):
+        msg = f"refusing to connect to non-public address {address}"
+        raise BlockedAddressError(msg)
+
+
+def _parse_ip(host: str) -> IPAddress | None:
+    try:
+        return ipaddress.ip_address(host.strip("[]").split("%", 1)[0])
+    except ValueError:
+        return None

pgdispatch/sqlalchemy.py

@@ -0,0 +1,130 @@
+"""SQLAlchemy 2.0 (async) integration.
+
+::
+
+    from pgdispatch import Webhooks
+    from pgdispatch.sqlalchemy import SQLAlchemyBackend
+
+    webhooks = Webhooks(SQLAlchemyBackend(engine))
+
+    async with session.begin():
+        session.add(order)
+        await webhooks.publish(session, "order.created", {"id": order.id})
+
+Any async PostgreSQL driver works for publishing (``asyncpg``, ``psycopg``).
+The worker always uses asyncpg with connection settings derived from the
+engine URL; pass ``worker_dsn`` to override them.
+
+Alembic users add the pgdispatch schema to their own migration chain::
+
+    from alembic import op
+    from pgdispatch.sqlalchemy import run_migrations
+
+    def upgrade() -> None:
+        run_migrations(op.get_bind())
+"""
+
+from __future__ import annotations
+
+from collections.abc import AsyncIterator, Mapping
+from contextlib import asynccontextmanager
+from functools import lru_cache
+from typing import Any, cast
+
+import asyncpg
+
+try:
+    from sqlalchemy import Connection, CursorResult, TextClause, text
+    from sqlalchemy.engine import URL
+    from sqlalchemy.ext.asyncio import AsyncConnection, AsyncEngine, AsyncSession
+except ImportError as exc:  # pragma: no cover - depends on the environment
+    msg = 'pgdispatch.sqlalchemy requires SQLAlchemy: pip install "pgdispatch[sqlalchemy]"'
+    raise ImportError(msg) from exc
+
+from pgdispatch.backends.asyncpg import WORKER_POOL_SIZE
+from pgdispatch.backends.base import Row, run_sync
+from pgdispatch.migrations import migrate
+
+__all__ = ["SQLAlchemyBackend", "SQLAlchemyExecutor", "run_migrations"]
+
+
+@lru_cache(maxsize=512)
+def _compiled(sql: str) -> TextClause:
+    return text(sql)
+
+
+class SQLAlchemyExecutor:
+    """Runs statements on an ``AsyncSession`` or ``AsyncConnection``, inside its transaction."""
+
+    def __init__(self, connection: AsyncSession | AsyncConnection) -> None:
+        self._connection = connection
+
+    async def fetch(self, sql: str, params: Mapping[str, Any]) -> list[Row]:
+        result = cast(
+            "CursorResult[Any]", await self._connection.execute(_compiled(sql), dict(params))
+        )
+        if not result.returns_rows:
+            return []
+        return [dict(row) for row in result.mappings()]
+
+
+class _SyncExecutor:
+    def __init__(self, connection: Connection) -> None:
+        self._connection = connection
+
+    def fetch(self, sql: str, params: Mapping[str, Any]) -> list[Row]:
+        result = cast("CursorResult[Any]", self._connection.execute(_compiled(sql), dict(params)))
+        if not result.returns_rows:
+            return []
+        return [dict(row) for row in result.mappings()]
+
+
+class SQLAlchemyBackend:
+    """Use pgdispatch with an SQLAlchemy ``AsyncEngine``."""
+
+    def __init__(self, engine: AsyncEngine, *, worker_dsn: str | None = None) -> None:
+        self.engine = engine
+        self._worker_dsn = worker_dsn or asyncpg_dsn(engine.url)
+
+    def executor(self, connection: Any) -> SQLAlchemyExecutor:
+        if not isinstance(connection, (AsyncSession, AsyncConnection)):
+            msg = (
+                "SQLAlchemyBackend.publish() expects an AsyncSession or AsyncConnection, "
+                f"got {type(connection).__name__}"
+            )
+            raise TypeError(msg)
+        return SQLAlchemyExecutor(connection)
+
+    @asynccontextmanager
+    async def transaction(self) -> AsyncIterator[SQLAlchemyExecutor]:
+        async with self.engine.begin() as connection:
+            yield SQLAlchemyExecutor(connection)
+
+    async def create_worker_pool(self) -> asyncpg.Pool:
+        pool: asyncpg.Pool = await asyncpg.create_pool(
+            self._worker_dsn, min_size=1, max_size=WORKER_POOL_SIZE
+        )
+        return pool
+
+    async def close(self) -> None:
+        """The engine belongs to the application, so there is nothing to close."""
+
+
+def run_migrations(connection: Connection, *, schema: str = "pgdispatch") -> list[int]:
+    """Apply pgdispatch migrations on a synchronous connection.
+
+    In an Alembic revision pass ``op.get_bind()``.
+    """
+    return run_sync(migrate(), _SyncExecutor(connection), schema=schema)
+
+
+def asyncpg_dsn(url: URL) -> str:
+    """Turn an SQLAlchemy URL into a DSN that asyncpg understands.
+
+    Driver-specific query parameters are dropped; ``sslmode``/``ssl`` are kept.
+    """
+    query = {key: value for key, value in url.query.items() if key in ("ssl", "sslmode")}
+    if "ssl" in query and "sslmode" not in query:
+        query["sslmode"] = query.pop("ssl")
+    plain = url.set(drivername="postgresql", query=query)
+    return plain.render_as_string(hide_password=False)

pgdispatch/worker/__init__.py

@@ -0,0 +1,5 @@
+"""The background delivery worker."""
+
+from pgdispatch.worker.worker import Worker
+
+__all__ = ["Worker"]

pgdispatch/worker/fetcher.py

@@ -0,0 +1,183 @@
+"""Claiming due deliveries with ``FOR UPDATE SKIP LOCKED``.
+
+Claiming is a *lease*, not a long-held lock: a single short transaction marks
+deliveries ``in_flight`` with ``locked_until = now() + lease`` and commits
+before any HTTP request is made. If the worker dies, the lease simply expires
+and another worker picks the delivery up. ``attempt_count`` is incremented at
+claim time; it doubles as a fencing token when the result is recorded, so a
+worker whose lease expired cannot overwrite a newer attempt's outcome.
+"""
+
+from __future__ import annotations
+
+import json
+import uuid
+from collections.abc import Mapping
+from dataclasses import dataclass
+from datetime import datetime
+from typing import Any
+
+from pgdispatch.backends.base import Row
+from pgdispatch.models import Delivery, DeliveryStatus, to_uuid
+
+__all__ = ["CLAIM", "CLAIM_PROBES", "ClaimedDelivery"]
+
+# A delivery is due when it is pending (or its lease expired) and its time has
+# come. With an ordering key it must also be the oldest unfinished delivery
+# with that key for its endpoint.
+_DUE = """
+    d.status IN ('pending', 'in_flight')
+    AND (d.status = 'pending' OR d.locked_until < now())
+    AND d.next_attempt_at <= now()
+    AND (
+        d.ordering_key IS NULL
+        OR NOT EXISTS (
+            SELECT 1 FROM {schema}.deliveries AS prior
+            WHERE prior.endpoint_id = d.endpoint_id
+              AND prior.ordering_key = d.ordering_key
+              AND prior.id < d.id
+              AND prior.status IN ('pending', 'in_flight')
+        )
+    )
+"""
+
+_LEASE = """
+UPDATE {schema}.deliveries AS d
+SET status = 'in_flight',
+    locked_until = now() + CAST(:lease AS double precision) * interval '1 second',
+    attempt_count = d.attempt_count + 1,
+    updated_at = now()
+FROM picked, {schema}.events AS ev, {schema}.endpoints AS e
+WHERE d.id = picked.id AND ev.id = d.event_id AND e.id = d.endpoint_id
+RETURNING d.id, d.event_id, d.endpoint_id, d.ordering_key, d.attempt_count, d.schedule_offset,
+          d.next_attempt_at, d.created_at, ev.event_type, ev.payload::text AS payload,
+          ev.created_at AS event_created_at, e.url, e.secrets::text AS secrets,
+          e.headers::text AS headers, e.rate_limit_rps
+"""
+
+CLAIM = f"""
+WITH due AS (
+    SELECT d.id, d.endpoint_id, d.next_attempt_at
+    FROM {{schema}}.deliveries AS d
+    JOIN {{schema}}.endpoints AS e ON e.id = d.endpoint_id
+    WHERE {_DUE}
+      AND e.is_enabled
+      AND e.circuit_state = 'closed'
+      AND NOT (d.endpoint_id = ANY(CAST(:saturated AS uuid[])))
+    ORDER BY d.next_attempt_at
+    LIMIT CAST(:limit AS integer)
+    FOR UPDATE OF d SKIP LOCKED
+), ranked AS (
+    SELECT due.id, due.endpoint_id,
+           row_number() OVER (
+               PARTITION BY due.endpoint_id ORDER BY due.next_attempt_at, due.id
+           ) AS position
+    FROM due
+), picked AS (
+    -- Respect the per-endpoint concurrency limit across this batch and the
+    -- requests this worker already has in flight.
+    SELECT ranked.id
+    FROM ranked
+    LEFT JOIN unnest(CAST(:busy_ids AS uuid[]), CAST(:busy_counts AS integer[]))
+        AS busy(endpoint_id, in_flight) ON busy.endpoint_id = ranked.endpoint_id
+    WHERE ranked.position <= CAST(:endpoint_limit AS integer) - coalesce(busy.in_flight, 0)
+)
+{_LEASE}
+"""
+"""Claim due deliveries of endpoints whose circuit breaker is closed."""
+
+CLAIM_PROBES = f"""
+WITH probe AS (
+    UPDATE {{schema}}.endpoints AS e
+    SET circuit_state = 'half_open',
+        circuit_retry_at = now() + CAST(:lease AS double precision) * interval '1 second',
+        updated_at = now()
+    WHERE e.id IN (
+        SELECT c.id FROM {{schema}}.endpoints AS c
+        WHERE c.is_enabled
+          AND c.circuit_state <> 'closed'
+          AND c.circuit_retry_at <= now()
+          AND EXISTS (
+              SELECT 1 FROM {{schema}}.deliveries AS d
+              WHERE d.endpoint_id = c.id AND {_DUE}
+          )
+        ORDER BY c.circuit_retry_at
+        LIMIT CAST(:limit AS integer)
+        FOR UPDATE SKIP LOCKED
+    )
+    RETURNING e.id
+), picked AS (
+    SELECT candidate.id
+    FROM probe
+    CROSS JOIN LATERAL (
+        SELECT d.id FROM {{schema}}.deliveries AS d
+        WHERE d.endpoint_id = probe.id AND {_DUE}
+        ORDER BY d.next_attempt_at
+        LIMIT 1
+        FOR UPDATE SKIP LOCKED
+    ) AS candidate
+)
+{_LEASE}
+"""
+"""Move endpoints whose cooldown elapsed to half-open and claim one probe delivery each.
+
+The endpoint itself is leased too (``circuit_retry_at``): if the worker dies
+during the probe, another probe becomes possible once the lease expires.
+"""
+
+
+@dataclass(frozen=True, slots=True)
+class ClaimedDelivery:
+    """Everything a worker needs to send one delivery."""
+
+    id: int
+    event_id: uuid.UUID
+    endpoint_id: uuid.UUID
+    ordering_key: str | None
+    attempt_count: int
+    schedule_offset: int
+    next_attempt_at: datetime
+    created_at: datetime
+    event_type: str
+    payload: str
+    event_created_at: datetime
+    url: str
+    secrets: str
+    headers: Mapping[str, str]
+    rate_limit_rps: float | None
+
+    @property
+    def attempts_in_round(self) -> int:
+        """Attempts made since the retry schedule (re)started, including this one."""
+        return self.attempt_count - self.schedule_offset
+
+    @classmethod
+    def from_row(cls, row: Row) -> ClaimedDelivery:
+        values: dict[str, Any] = dict(row)
+        values["event_id"] = to_uuid(row["event_id"])
+        values["endpoint_id"] = to_uuid(row["endpoint_id"])
+        values["headers"] = json.loads(row["headers"])
+        return cls(**values)
+
+    def to_delivery(
+        self,
+        *,
+        status: DeliveryStatus,
+        status_code: int | None,
+        error: str | None,
+        finished_at: datetime | None,
+    ) -> Delivery:
+        return Delivery(
+            id=self.id,
+            event_id=self.event_id,
+            endpoint_id=self.endpoint_id,
+            event_type=self.event_type,
+            ordering_key=self.ordering_key,
+            status=status,
+            attempt_count=self.attempt_count,
+            next_attempt_at=self.next_attempt_at,
+            last_status_code=status_code,
+            last_error=error,
+            created_at=self.created_at,
+            finished_at=finished_at,
+        )