diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 7e7adc118..7646a6f71 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -201,6 +201,7 @@ jobs: tests/conformance/decisioning/test_pg_idempotency_reservation.py \ tests/test_decisioning_task_polling.py \ tests/test_decisioning_task_lifecycle.py \ + tests/test_decisioning_lazy_pg_stores.py \ tests/conformance/decisioning/test_pg_task_webhook_outbox.py \ tests/conformance/decisioning/test_pg_notification_outbox.py \ tests/test_notification_outbox.py \ diff --git a/docs/lazy-pg-stores.md b/docs/lazy-pg-stores.md new file mode 100644 index 000000000..f464e6335 --- /dev/null +++ b/docs/lazy-pg-stores.md @@ -0,0 +1,71 @@ +# Lazy PostgreSQL decisioning stores + +`adcp.decisioning.pg` exports `LazyProposalStore`, `LazyTaskRegistry`, and +`LazyTaskWebhookOutbox`. Each accepts a synchronous or asynchronous zero-argument +factory returning its original concrete PostgreSQL store. Task factories may +also return an original `(PgTaskRegistry, PgTaskWebhookOutbox)` pair. + +Factories are called on the serving event loop, on the first async operation or +explicit `await wrapper.resolve()`. Concurrent calls share one successful +resolution. Failed or canceled initialization is not cached, so later calls +retry. Factories own pool opening and cleanup on failed initialization; wrappers +never open, close, replace or reopen their caller-owned pools. Schema setup stays +explicit: do it in the factory or call the wrapper's `create_schema()`. + +```python +from psycopg_pool import AsyncConnectionPool +from adcp.decisioning.pg import ( + LazyTaskRegistry, LazyTaskWebhookOutbox, PgTaskRegistry, PgTaskWebhookOutbox, +) + +pool = AsyncConnectionPool("postgresql://localhost/seller", open=False) + +async def task_stores() -> tuple[PgTaskRegistry, PgTaskWebhookOutbox]: + await pool.open() + outbox = PgTaskWebhookOutbox( + pool=pool, sender=signed_sender, encryption_key=encryption_key, + delivery_retry_horizon_seconds=86400, + ) + registry = PgTaskRegistry(pool=pool, task_webhook_outbox=outbox) + await registry.create_schema() + await outbox.create_schema() + return registry, outbox + +registry = LazyTaskRegistry(task_stores) +worker_outbox = LazyTaskWebhookOutbox.from_registry(registry) +# await worker_outbox.run_worker() and registry.issue(...) share the same pair. +# Cancel/join workers before closing pool during application shutdown. +``` + +Construct the original registry and outbox together, passing the same pool to +both. Their existing constructor checks (including signing-scope wiring) remain +authoritative, and a returned pair is checked again for identity. The worker +facade from `from_registry()` resolves the same original outbox, rather than +running an independent pool factory. Do not pass an unresolved worker facade to +the concrete `PgTaskRegistry` constructor. + +The wrappers declare the concrete stores' durability before resolution and +reject a factory returning a lossy store. `LazyTaskRegistry` retains the concrete +SDK registry's optional listing protocol; no arbitrary delegate attributes or +optional APIs are forwarded. Its additive observer methods can be registered or +removed before and after resolution. They receive the first submitted event and +retain the concrete registry's commit/no-op/failure semantics. + +`resolved` inspects the cached store without running the factory. A task registry +exposes its original `task_webhook_outbox` and atomic-outbox marker after +resolution. Servers advertising SDK task-webhook signing must await resolution +before synchronous server construction: the boot validator needs the actual +sender, retry horizon and shared-pool proof. Resolved exact SDK wrappers are +accepted by that proof; unaudited subclasses remain rejected. Polling-only +servers can resolve on first use. Outbox synchronous registration/crypto helpers +and retry-horizon inspection also require explicit resolution; they raise a +clear error when called too early. + +After success, a wrapper is bound to the event loop that resolved it. A different +loop raises an error, even if the previous loop has closed; create a fresh wrapper +and pool for the new loop. A failed initialization with no cached store may retry +on a new loop after the old loop closes. Closing the borrowed pool causes normal +pool errors on later operations, rather than rerunning the factory. + +No database migration or change to the existing `LazyBackend` is required. The +new wrappers reuse the concrete stores' current SQL and schema assets. diff --git a/src/adcp/decisioning/pg/__init__.py b/src/adcp/decisioning/pg/__init__.py index 14c586fde..cf4c7f254 100644 --- a/src/adcp/decisioning/pg/__init__.py +++ b/src/adcp/decisioning/pg/__init__.py @@ -35,6 +35,14 @@ PG_AVAILABLE, PgBuyerAgentRegistry, ) +from adcp.decisioning.pg.lazy import ( + LazyProposalStore, + LazyProposalStoreFactory, + LazyTaskRegistry, + LazyTaskRegistryFactory, + LazyTaskWebhookOutbox, + LazyTaskWebhookOutboxFactory, +) from adcp.decisioning.pg.proposal_store import PgProposalStore from adcp.decisioning.pg.task_registry import ( PgTaskRegistry, @@ -46,6 +54,12 @@ __all__ = [ "DEFAULT_TABLE_NAME", "PG_AVAILABLE", + "LazyProposalStore", + "LazyProposalStoreFactory", + "LazyTaskRegistry", + "LazyTaskRegistryFactory", + "LazyTaskWebhookOutbox", + "LazyTaskWebhookOutboxFactory", "PgBuyerAgentRegistry", "PgProposalStore", "PgTaskRegistry", diff --git a/src/adcp/decisioning/pg/lazy.py b/src/adcp/decisioning/pg/lazy.py new file mode 100644 index 000000000..542a3b1c1 --- /dev/null +++ b/src/adcp/decisioning/pg/lazy.py @@ -0,0 +1,470 @@ +"""Resolve-once wrappers for caller-owned PostgreSQL decisioning stores. + +Factories open infrastructure on the serving loop, then construct the original +concrete stores so constructor validation remains authoritative. A successful +resolution is cached; failed initialization can be retried. These wrappers never +open, close, replace or reopen pools themselves. +""" + +from __future__ import annotations + +import asyncio +from collections.abc import Awaitable, Callable, Mapping +from datetime import datetime +from typing import Any, ClassVar, Generic, TypeAlias, TypeVar + +from adcp.decisioning.context import RequestContext +from adcp.decisioning.pg.proposal_store import DEFAULT_TABLE_NAME, PgProposalStore, _migration_sql +from adcp.decisioning.pg.task_registry import PgTaskRegistry +from adcp.decisioning.pg.task_webhook_outbox import PgTaskWebhookOutbox +from adcp.decisioning.proposal_store import ProposalRecord +from adcp.decisioning.recipe import Recipe +from adcp.decisioning.task_registry import ( + TaskTransition, + TaskWebhookAuthentication, + _TaskLifecycleObservers, +) + +_Store = TypeVar("_Store") +LazyProposalStoreFactory: TypeAlias = Callable[[], PgProposalStore | Awaitable[PgProposalStore]] +TaskRegistryStores: TypeAlias = PgTaskRegistry | tuple[PgTaskRegistry, PgTaskWebhookOutbox] +LazyTaskRegistryFactory: TypeAlias = Callable[ + [], TaskRegistryStores | Awaitable[TaskRegistryStores] +] +LazyTaskWebhookOutboxFactory: TypeAlias = Callable[ + [], PgTaskWebhookOutbox | Awaitable[PgTaskWebhookOutbox] +] + + +class _LazyStore(Generic[_Store]): + def __init__( + self, factory: Callable[[], _Store | Awaitable[_Store]], store_type: type[_Store] + ) -> None: + self._factory = factory + self._store_type = store_type + self._store: _Store | None = None + self._lock = asyncio.Lock() + self._loop: asyncio.AbstractEventLoop | None = None + + @property + def resolved(self) -> _Store | None: + """The cached concrete store, or None; inspection never invokes the factory.""" + return self._store + + def _bind_loop(self) -> None: + loop = asyncio.get_running_loop() + if self._loop is None: + self._loop = loop + elif self._loop is not loop: + # No successful resource was cached. A failed attempt on a now-dead + # loop may retry, with a fresh lock. The factory owns failed cleanup. + if self._store is None and self._loop.is_closed() and not self._lock.locked(): + self._loop = loop + self._lock = asyncio.Lock() + else: + raise RuntimeError( + "Lazy PostgreSQL stores cannot be reused on another event loop; " + "construct a new wrapper and pool for that loop" + ) + + async def resolve(self) -> _Store: + """Resolve once on the serving loop. Concurrent callers share success. + + Factory failures/cancellation are not cached. After successful resolution, + reusing this wrapper on another loop raises, even if the old loop has closed. + """ + self._bind_loop() + cached = self._store + if cached is not None: + return cached + async with self._lock: + cached = self._store + if cached is not None: + return cached + result = self._factory() + store = await result if isinstance(result, Awaitable) else result + if not isinstance(store, self._store_type): + raise TypeError(f"Factory must return {self._store_type.__name__}") + self._validate(store) + self._store = store + return store + + def _validate(self, store: _Store) -> None: + pass + + def _require_resolved(self) -> _Store: + store = self._store + if store is None: + raise RuntimeError("Await resolve() before using synchronous store helpers") + # Pure configuration/crypto helpers may run without a loop. When called + # on an active loop, reject accidental reuse just like async operations. + try: + asyncio.get_running_loop() + except RuntimeError: + return store + self._bind_loop() + return store + + +class LazyProposalStore(_LazyStore[PgProposalStore]): + """Deferred PgProposalStore; durable before and after resolution.""" + + is_durable: ClassVar[bool] = True + + def __init__(self, factory: LazyProposalStoreFactory) -> None: + super().__init__(factory, PgProposalStore) + + def _validate(self, store: PgProposalStore) -> None: + if store.is_durable is not True: + raise ValueError("LazyProposalStore requires a durable PostgreSQL store") + + @classmethod + def migration_sql(cls, table_name: str = DEFAULT_TABLE_NAME) -> dict[str, str]: + """Return the concrete store's migration SQL without opening infrastructure.""" + return _migration_sql(table_name) + + async def create_schema(self) -> None: + await (await self.resolve()).create_schema() + + async def put_draft( + self, + *, + proposal_id: str, + account_id: str, + recipes: Mapping[str, Recipe], + proposal_payload: Mapping[str, Any], + ) -> None: + await (await self.resolve()).put_draft( + proposal_id=proposal_id, + account_id=account_id, + recipes=recipes, + proposal_payload=proposal_payload, + ) + + async def get(self, proposal_id: str, *, expected_account_id: str) -> ProposalRecord | None: + return await (await self.resolve()).get( + proposal_id, expected_account_id=expected_account_id + ) + + async def commit( + self, + proposal_id: str, + *, + expires_at: datetime, + proposal_payload: Mapping[str, Any], + expected_account_id: str, + ) -> None: + await (await self.resolve()).commit( + proposal_id, + expires_at=expires_at, + proposal_payload=proposal_payload, + expected_account_id=expected_account_id, + ) + + async def try_reserve_consumption( + self, proposal_id: str, *, expected_account_id: str + ) -> ProposalRecord: + return await (await self.resolve()).try_reserve_consumption( + proposal_id, expected_account_id=expected_account_id + ) + + async def finalize_consumption( + self, proposal_id: str, *, media_buy_id: str, expected_account_id: str + ) -> None: + await (await self.resolve()).finalize_consumption( + proposal_id, media_buy_id=media_buy_id, expected_account_id=expected_account_id + ) + + async def release_consumption(self, proposal_id: str, *, expected_account_id: str) -> None: + await (await self.resolve()).release_consumption( + proposal_id, expected_account_id=expected_account_id + ) + + async def mark_consumed( + self, proposal_id: str, *, media_buy_id: str, expected_account_id: str + ) -> None: + await (await self.resolve()).mark_consumed( + proposal_id, media_buy_id=media_buy_id, expected_account_id=expected_account_id + ) + + async def discard(self, proposal_id: str, *, expected_account_id: str) -> None: + await (await self.resolve()).discard(proposal_id, expected_account_id=expected_account_id) + + async def get_by_media_buy_id( + self, media_buy_id: str, *, expected_account_id: str + ) -> ProposalRecord | None: + return await (await self.resolve()).get_by_media_buy_id( + media_buy_id, expected_account_id=expected_account_id + ) + + +class LazyTaskRegistry(_LazyStore[PgTaskRegistry], _TaskLifecycleObservers): + """Deferred PgTaskRegistry, including its listing and metrics contracts. + + A factory can return a registry, or an original (registry, outbox) pair. Build + them together using one caller-owned pool. Pair and signing-scope validation + are retained. Register observers before first use; they receive the first + submitted transition too. No arbitrary delegate attributes are forwarded. + + For a server advertising SDK task-webhook signing, await resolve() before + constructing the server: its synchronous boot validator needs the actual + outbox/sender/horizon configuration. Polling-only servers can resolve on the + first task operation. + """ + + is_durable: ClassVar[bool] = True + + def __init__(self, factory: LazyTaskRegistryFactory) -> None: + self._init_lifecycle_observers() + + async def resolve_registry() -> PgTaskRegistry: + result = factory() + stores = await result if isinstance(result, Awaitable) else result + if isinstance(stores, tuple): + registry, outbox = stores + if not isinstance(registry, PgTaskRegistry) or not isinstance( + outbox, PgTaskWebhookOutbox + ): + raise TypeError("Task factory must return a concrete registry/outbox pair") + if registry.task_webhook_outbox is not outbox or registry._pool is not outbox._pool: + raise ValueError( + "Task registry/outbox pair must share one pool and registration" + ) + return registry + return stores + + super().__init__(resolve_registry, PgTaskRegistry) + + def _validate(self, store: PgTaskRegistry) -> None: + if store.is_durable is not True: + raise ValueError("LazyTaskRegistry requires a durable PostgreSQL registry") + if not callable(store.list): + raise TypeError("LazyTaskRegistry requires PostgreSQL listing support") + outbox = store.task_webhook_outbox + if outbox is not None and outbox._pool is not store._pool: + raise ValueError("Task registry and outbox must share one pool") + store.add_lifecycle_observer(self._relay_transition) + + def _relay_transition( + self, + event: TaskTransition, + *, + task_id: str, + account_id: str, + task_type: str, + created_at: float, + updated_at: float, + ) -> None: + self._notify_lifecycle_observers( + event, + { + "task_id": task_id, + "account_id": account_id, + "task_type": task_type, + "created_at": created_at, + "updated_at": updated_at, + }, + ) + + @property + def task_webhook_outbox(self) -> PgTaskWebhookOutbox | None: + """Original coupled outbox, available after resolution; no eager opening.""" + store = self.resolved + return store.task_webhook_outbox if store is not None else None + + @property + def atomic_task_webhook_outbox(self) -> bool: + store = self.resolved + return store.atomic_task_webhook_outbox if store is not None else False + + async def create_schema(self) -> None: + await (await self.resolve()).create_schema() + + async def issue( + self, + *, + account_id: str, + task_type: str, + request_context: dict[str, Any] | None = None, + webhook_url: str | None = None, + webhook_operation_id: str | None = None, + webhook_token: str | None = None, + webhook_authentication: TaskWebhookAuthentication | None = None, + webhook_signing_scope_id: str | None = None, + **_extra: Any, + ) -> str: + return await (await self.resolve()).issue( + account_id=account_id, + task_type=task_type, + request_context=request_context, + webhook_url=webhook_url, + webhook_operation_id=webhook_operation_id, + webhook_token=webhook_token, + webhook_authentication=webhook_authentication, + webhook_signing_scope_id=webhook_signing_scope_id, + **_extra, + ) + + async def resolve_webhook_signing_scope(self, context: RequestContext[Any]) -> str | None: + return await (await self.resolve()).resolve_webhook_signing_scope(context) + + async def update_progress(self, task_id: str, progress: dict[str, Any]) -> None: + await (await self.resolve()).update_progress(task_id, progress) + + async def complete(self, task_id: str, result: dict[str, Any]) -> None: + await (await self.resolve()).complete(task_id, result) + + async def fail(self, task_id: str, error: dict[str, Any]) -> None: + await (await self.resolve()).fail(task_id, error) + + async def get( + self, task_id: str, *, expected_account_id: str | None = None + ) -> dict[str, Any] | None: + return await (await self.resolve()).get(task_id, expected_account_id=expected_account_id) + + async def list( + self, + *, + account_id: str, + filters: dict[str, Any] | None = None, + sort: dict[str, Any] | None = None, + pagination: dict[str, Any] | None = None, + ) -> dict[str, Any]: + return await (await self.resolve()).list( + account_id=account_id, filters=filters, sort=sort, pagination=pagination + ) + + async def discard(self, task_id: str) -> None: + await (await self.resolve()).discard(task_id) + + +class LazyTaskWebhookOutbox(_LazyStore[PgTaskWebhookOutbox]): + """Deferred concrete outbox. Synchronous crypto helpers require resolve(). + + Use from_registry() to share a registry factory's original outbox. Pass the + original objects, constructed inside that factory, to PgTaskRegistry; do not + substitute this worker facade into the concrete constructor's pool checks. + """ + + delivery_state_is_durable: ClassVar[bool] = True + supports_atomic_task_outbox: ClassVar[bool] = True + + def __init__(self, factory: LazyTaskWebhookOutboxFactory) -> None: + super().__init__(factory, PgTaskWebhookOutbox) + + def _validate(self, store: PgTaskWebhookOutbox) -> None: + if ( + store.delivery_state_is_durable is not True + or store.supports_atomic_task_outbox is not True + ): + raise ValueError("LazyTaskWebhookOutbox requires a durable atomic PostgreSQL outbox") + + @classmethod + def from_registry(cls, registry: LazyTaskRegistry) -> LazyTaskWebhookOutbox: + """Create a worker facade resolving the same original registry/outbox pair.""" + + async def resolve_outbox() -> PgTaskWebhookOutbox: + store = await registry.resolve() + if store.task_webhook_outbox is None: + raise ValueError("The registry factory did not configure a task webhook outbox") + return store.task_webhook_outbox + + return cls(resolve_outbox) + + @property + def delivery_retry_horizon_seconds(self) -> int: + return self._require_resolved().delivery_retry_horizon_seconds + + @property + def legacy_hmac_fallback(self) -> bool: + return self._require_resolved().legacy_hmac_fallback + + async def create_schema(self) -> None: + await (await self.resolve()).create_schema() + + async def enqueue_terminal( + self, + conn: Any, + *, + task_id: str, + account_id: str, + task_type: str, + status: str, + result: dict[str, Any], + url: str, + operation_id: str, + token: str | None, + authentication: TaskWebhookAuthentication | None = None, + signing_scope_id: str | None = None, + ) -> int: + return await (await self.resolve()).enqueue_terminal( + conn, + task_id=task_id, + account_id=account_id, + task_type=task_type, + status=status, + result=result, + url=url, + operation_id=operation_id, + token=token, + authentication=authentication, + signing_scope_id=signing_scope_id, + ) + + def validate_registration( + self, url: str, authentication: TaskWebhookAuthentication | None = None + ) -> None: + (self._require_resolved()).validate_registration(url, authentication) + + def protect_registration( + self, + *, + account_id: str, + task_id: str, + task_type: str, + url: str, + operation_id: str, + token: str | None, + authentication: TaskWebhookAuthentication | None = None, + signing_scope_id: str | None = None, + ) -> tuple[bytes, bytes]: + return (self._require_resolved()).protect_registration( + account_id=account_id, + task_id=task_id, + task_type=task_type, + url=url, + operation_id=operation_id, + token=token, + authentication=authentication, + signing_scope_id=signing_scope_id, + ) + + def open_registration( + self, + *, + account_id: str, + task_id: str, + task_type: str, + encrypted_registration: bytes, + nonce: bytes, + ) -> tuple[str, str, str | None]: + return (self._require_resolved()).open_registration( + account_id=account_id, + task_id=task_id, + task_type=task_type, + encrypted_registration=encrypted_registration, + nonce=nonce, + ) + + async def run_worker( + self, *, poll_interval: float = 1.0, purge_interval: float = 300.0 + ) -> None: + await (await self.resolve()).run_worker( + poll_interval=poll_interval, purge_interval=purge_interval + ) + + async def process_one(self) -> bool: + return await (await self.resolve()).process_one() + + async def purge_expired(self) -> None: + await (await self.resolve()).purge_expired() diff --git a/src/adcp/decisioning/webhook_emit.py b/src/adcp/decisioning/webhook_emit.py index b17159ad5..660a1bad9 100644 --- a/src/adcp/decisioning/webhook_emit.py +++ b/src/adcp/decisioning/webhook_emit.py @@ -44,10 +44,15 @@ def _sdk_task_outbox_pair_ready(registry: Any, task_outbox: Any) -> bool: if registry is None or task_outbox is None: return False try: + from adcp.decisioning.pg.lazy import LazyTaskRegistry from adcp.decisioning.pg.task_registry import PgTaskRegistry from adcp.decisioning.pg.task_webhook_outbox import PgTaskWebhookOutbox except ImportError: return False + if type(registry) is LazyTaskRegistry: + registry = registry.resolved + if registry is None: + return False return ( type(registry) is PgTaskRegistry and type(task_outbox) is PgTaskWebhookOutbox diff --git a/tests/test_decisioning_lazy_pg_stores.py b/tests/test_decisioning_lazy_pg_stores.py new file mode 100644 index 000000000..6c618119d --- /dev/null +++ b/tests/test_decisioning_lazy_pg_stores.py @@ -0,0 +1,509 @@ +"""Lazy PostgreSQL infrastructure resolves once without weakening store contracts.""" + +from __future__ import annotations + +import asyncio +import os +import secrets +from collections.abc import AsyncIterator +from typing import Any +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from adcp.decisioning import ListableTaskRegistry, TaskRegistry +from adcp.decisioning.pg import ( + LazyProposalStore, + LazyTaskRegistry, + LazyTaskWebhookOutbox, + PgProposalStore, + PgTaskRegistry, + PgTaskWebhookOutbox, +) +from adcp.decisioning.webhook_emit import _sdk_task_outbox_pair_ready +from adcp.webhook_sender import PreparedWebhook, WebhookDeliveryResult + + +@pytest.fixture +def fake_pg_driver(monkeypatch: pytest.MonkeyPatch) -> None: + """Fake pools never call psycopg; retain the real constructor validation. + + This explicit fixture enables only the optional-driver availability guards + for unit cases, following the existing task-outbox unit-test pattern. Real + database cases and the missing-extra regression do not use it. + """ + for module in ("proposal_store", "task_registry", "task_webhook_outbox"): + monkeypatch.setattr(f"adcp.decisioning.pg.{module}.PG_AVAILABLE", True) + + +def _sender() -> MagicMock: + import json + + sender = MagicMock() + sender._owns_client = True + sender._allow_private_destinations = False + sender._timeout = 10.0 + sender.signs_with_rfc9421 = True + sender._auth.alg = "ed25519" + + def prepare(**kwargs: Any) -> PreparedWebhook: + key = f"whk_{secrets.token_hex(16)}" + return PreparedWebhook( + url=kwargs["url"], + idempotency_key=key, + body=json.dumps({"idempotency_key": key, **kwargs}).encode(), + ) + + sender.prepare_mcp.side_effect = prepare + sender.send_prepared = AsyncMock( + side_effect=lambda prepared: WebhookDeliveryResult( + status_code=200, + idempotency_key=prepared.idempotency_key, + url=prepared.url, + response_headers={}, + response_body=b"{}", + sent_body=prepared.body, + ) + ) + return sender + + +def outbox(pool: Any) -> PgTaskWebhookOutbox: + return PgTaskWebhookOutbox( + pool=pool, sender=_sender(), encryption_key=b"e" * 32, delivery_retry_horizon_seconds=86400 + ) + + +@pytest.mark.parametrize("kind", ["proposal", "task", "outbox"]) +@pytest.mark.usefixtures("fake_pg_driver") +async def test_concurrent_resolution_and_retry(kind: str) -> None: + pool = MagicMock() + concrete: Any = { + "proposal": lambda: PgProposalStore(pool=pool), + "task": lambda: PgTaskRegistry(pool=pool), + "outbox": lambda: outbox(pool), + }[kind]() + classes = { + "proposal": LazyProposalStore, + "task": LazyTaskRegistry, + "outbox": LazyTaskWebhookOutbox, + } + calls = 0 + + async def factory() -> Any: + nonlocal calls + calls += 1 + await asyncio.sleep(0.01) + if calls == 1: + raise RuntimeError("Bootstrap unavailable") + return concrete + + lazy = classes[kind](factory) + assert lazy.resolved is None + assert not hasattr(lazy, "clear_all") + with pytest.raises(RuntimeError, match="Bootstrap"): + await lazy.resolve() + resolved = await asyncio.gather(*(lazy.resolve() for _ in range(20))) + assert calls == 2 + assert all(store is concrete for store in resolved) + assert concrete._pool is pool + + +@pytest.mark.parametrize("kind", ["proposal", "task", "outbox"]) +@pytest.mark.usefixtures("fake_pg_driver") +async def test_first_method_use_resolves_once_and_forwards_arguments(kind: str) -> None: + pool = MagicMock() + if kind == "proposal": + concrete = PgProposalStore(pool=pool) + method = AsyncMock(return_value=None) + concrete.get = method + lazy = LazyProposalStore(lambda: concrete) + assert lazy.is_durable is True + await lazy.get("proposal", expected_account_id="acct") + method.assert_awaited_once_with("proposal", expected_account_id="acct") + elif kind == "task": + concrete = PgTaskRegistry(pool=pool) + method = AsyncMock(return_value={"pagination": {"has_more": False}, "tasks": []}) + concrete.list = method + lazy = LazyTaskRegistry(lambda: concrete) + assert lazy.is_durable is True + assert isinstance(lazy, TaskRegistry) and isinstance(lazy, ListableTaskRegistry) + await lazy.list(account_id="acct", filters={"status": "working"}) + method.assert_awaited_once_with( + account_id="acct", filters={"status": "working"}, sort=None, pagination=None + ) + else: + concrete = outbox(pool) + method = AsyncMock(return_value=True) + concrete.process_one = method + lazy = LazyTaskWebhookOutbox(lambda: concrete) + assert lazy.delivery_state_is_durable and lazy.supports_atomic_task_outbox + with pytest.raises(RuntimeError, match="resolve"): + lazy.validate_registration("https://buyer.example/hooks") + assert await lazy.process_one() is True + lazy.validate_registration("https://buyer.example/hooks") + method.assert_awaited_once_with() + assert lazy.resolved is concrete + concrete.clear_all = AsyncMock() + assert not hasattr(lazy, "clear_all") + + +@pytest.mark.usefixtures("fake_pg_driver") +async def test_pair_shares_pool_resolution_and_observers() -> None: + pool = MagicMock() + calls = 0 + + async def factory() -> tuple[PgTaskRegistry, PgTaskWebhookOutbox]: + nonlocal calls + calls += 1 + await asyncio.sleep(0.01) + box = outbox(pool) + return PgTaskRegistry(pool=pool, task_webhook_outbox=box), box + + lazy = LazyTaskRegistry(factory) + facade = LazyTaskWebhookOutbox.from_registry(lazy) + events = [] + lazy.add_lifecycle_observer(lambda event, **metadata: events.append((event, metadata))) + assert lazy.task_webhook_outbox is None and lazy.atomic_task_webhook_outbox is False + registry, box = await asyncio.gather(lazy.resolve(), facade.resolve()) + assert calls == 1 + assert registry._pool is box._pool is pool + assert lazy.task_webhook_outbox is box + assert lazy.atomic_task_webhook_outbox is True + assert _sdk_task_outbox_pair_ready(lazy, box) + # The bridge is registered before any operation delegated by the facade. + registry._notify_lifecycle_observers( + "submitted", + { + "task_id": "task", + "account_id": "acct", + "task_type": "get_products", + "created_at": 1.0, + "updated_at": 1.0, + }, + ) + assert events[0][0] == "submitted" + + +@pytest.mark.usefixtures("fake_pg_driver") +async def test_original_pool_identity_and_signing_scope_checks_remain_authoritative() -> None: + pool = MagicMock() + other_pool = MagicMock() + attempts = 0 + + def factory() -> PgTaskRegistry: + nonlocal attempts + attempts += 1 + box = outbox(other_pool if attempts == 1 else pool) + return PgTaskRegistry(pool=pool, task_webhook_outbox=box) + + lazy = LazyTaskRegistry(factory) + with pytest.raises(ValueError, match="same connection pool"): + await lazy.resolve() + assert lazy.resolved is None + assert (await lazy.resolve()).task_webhook_outbox._pool is pool + assert attempts == 2 + + box = outbox(pool) + bad_pair = LazyTaskRegistry(lambda: (PgTaskRegistry(pool=pool), box)) + with pytest.raises(ValueError, match="pair must share"): + await bad_pair.resolve() + + def missing_resolver() -> PgTaskRegistry: + return PgTaskRegistry(pool=pool, webhook_signing_scope_resolver=lambda ctx: "tenant") + + with pytest.raises(ValueError, match="required exactly"): + await LazyTaskRegistry(missing_resolver).resolve() + + +@pytest.mark.usefixtures("fake_pg_driver") +async def test_cancelled_factory_retries_and_waiters_do_not_duplicate_success() -> None: + started, allow = asyncio.Event(), asyncio.Event() + pool = MagicMock() + attempts = 0 + + async def factory() -> PgTaskRegistry: + nonlocal attempts + attempts += 1 + started.set() + await allow.wait() + return PgTaskRegistry(pool=pool) + + lazy = LazyTaskRegistry(factory) + owner = asyncio.create_task(lazy.resolve()) + await started.wait() + waiter = asyncio.create_task(lazy.resolve()) + owner.cancel() + with pytest.raises(asyncio.CancelledError): + await owner + allow.set() + assert await waiter is await lazy.resolve() + assert attempts == 2 + + +@pytest.mark.usefixtures("fake_pg_driver") +def test_successful_store_refuses_cross_loop_reuse_without_reopening() -> None: + pool = MagicMock() + factory = MagicMock(return_value=PgTaskRegistry(pool=pool)) + lazy = LazyTaskRegistry(factory) + asyncio.run(lazy.resolve()) + with pytest.raises(RuntimeError, match="another event loop"): + asyncio.run(lazy.resolve()) + factory.assert_called_once_with() + + +@pytest.mark.usefixtures("fake_pg_driver") +def test_failed_initialization_can_retry_on_a_new_loop_after_old_loop_closes() -> None: + pool = MagicMock() + factory = MagicMock(side_effect=[RuntimeError("Bootstrap"), PgTaskRegistry(pool=pool)]) + lazy = LazyTaskRegistry(factory) + with pytest.raises(RuntimeError, match="Bootstrap"): + asyncio.run(lazy.resolve()) + assert asyncio.run(lazy.resolve())._pool is pool + assert factory.call_count == 2 + + +@pytest.mark.usefixtures("fake_pg_driver") +async def test_bad_factory_result_not_cached() -> None: + factory = MagicMock(side_effect=[object(), PgTaskRegistry(pool=MagicMock())]) + lazy = LazyTaskRegistry(factory) + with pytest.raises(TypeError, match="PgTaskRegistry"): + await lazy.resolve() + assert lazy.resolved is None + assert isinstance(await lazy.resolve(), PgTaskRegistry) + + +@pytest.fixture +async def real_stack() -> AsyncIterator[Any]: + url = os.environ.get("ADCP_PG_TEST_URL") + if not url: + pytest.skip("ADCP_PG_TEST_URL required") + from psycopg_pool import AsyncConnectionPool + + suffix = secrets.token_hex(6) + pool = AsyncConnectionPool(url, open=False, min_size=1, max_size=4) + sender = _sender() + calls = 0 + + async def factory() -> tuple[PgTaskRegistry, PgTaskWebhookOutbox]: + nonlocal calls + calls += 1 + await pool.open() + box = PgTaskWebhookOutbox( + pool=pool, + sender=sender, + encryption_key=b"e" * 32, + delivery_retry_horizon_seconds=86400, + table=f"test_lazy_box_{suffix}", + ) + registry = PgTaskRegistry( + pool=pool, task_webhook_outbox=box, _table=f"test_lazy_task_{suffix}" + ) + await registry.create_schema() + await box.create_schema() + return registry, box + + lazy = LazyTaskRegistry(factory) + facade = LazyTaskWebhookOutbox.from_registry(lazy) + try: + yield lazy, facade, pool, sender, lambda: calls + finally: + if not pool.closed: + async with pool.connection() as conn: + for table in ( + f"test_lazy_box_{suffix}", + f"test_lazy_task_{suffix}", + f"test_lazy_proposal_{suffix}", + ): + await conn.execute(f"DROP TABLE IF EXISTS {table}") + await pool.close() + + +async def test_real_lazy_pair_atomic_completion_listing_and_observer_commit( + real_stack: Any, +) -> None: + lazy, facade, pool, sender, calls = real_stack + events = [] + lazy.add_lifecycle_observer(lambda event, **metadata: events.append(event)) + task = await lazy.issue( + account_id="acct", + task_type="get_products", + webhook_url="https://buyer.example/hooks", + webhook_operation_id="buyer-operation", + ) + await lazy.update_progress(task, {"percentage": 25}) + await asyncio.gather(*(lazy.complete(task, {"products": []}) for _ in range(10))) + assert events == ["submitted", "working", "completed"] + page = await lazy.list(account_id="acct") + assert page["tasks"][0]["task_id"] == task + assert (await lazy.get(task, expected_account_id="acct"))["state"] == "completed" + assert await lazy.get(task, expected_account_id="other") is None + assert await facade.process_one() is True + assert sender.send_prepared.await_count == 1 + assert calls() == 1 + assert (await facade.resolve())._pool is (await lazy.resolve())._pool is pool + # A closed caller-owned pool is not implicitly reopened or replaced. + await pool.close() + from psycopg_pool import PoolClosed + + with pytest.raises(PoolClosed): + await lazy.get(task) + assert calls() == 1 + + +async def test_real_lazy_proposal_lifecycle(real_stack: Any) -> None: + from datetime import datetime, timedelta, timezone + + from adcp.decisioning.proposal_store import ProposalState + + lazy, facade, pool, sender, calls = real_stack + registry = await lazy.resolve() + table = registry._table.replace("test_lazy_task_", "test_lazy_proposal_") + store = LazyProposalStore(lambda: PgProposalStore(pool=pool, table_name=table)) + await store.create_schema() + await store.put_draft( + proposal_id="proposal", account_id="acct", recipes={}, proposal_payload={} + ) + assert await store.get("proposal", expected_account_id="other") is None + await store.commit( + "proposal", + expires_at=datetime.now(timezone.utc) + timedelta(hours=1), + proposal_payload={}, + expected_account_id="acct", + ) + reserved = await store.try_reserve_consumption("proposal", expected_account_id="acct") + assert reserved.state == ProposalState.CONSUMING + await store.release_consumption("proposal", expected_account_id="acct") + await store.try_reserve_consumption("proposal", expected_account_id="acct") + await store.finalize_consumption("proposal", media_buy_id="buy", expected_account_id="acct") + assert ( + await store.get_by_media_buy_id("buy", expected_account_id="acct") + ).state == ProposalState.CONSUMED + + +async def test_observer_registration_and_removal_survive_resolution(real_stack: Any) -> None: + lazy, facade, pool, sender, calls = real_stack + before, after = [], [] + + def observe_before(event: Any, **metadata: Any) -> None: + before.append(event) + + def observe_after(event: Any, **metadata: Any) -> None: + after.append(event) + + lazy.add_lifecycle_observer(observe_before) + task = await lazy.issue(account_id="acct", task_type="get_products") + assert lazy.remove_lifecycle_observer(observe_before) is True + lazy.add_lifecycle_observer(observe_after) + await lazy.update_progress(task, {"percentage": 50}) + await lazy.complete(task, {"products": []}) + assert before == ["submitted"] + assert after == ["working", "completed"] + lazy.remove_lifecycle_observer(observe_after) + await lazy.discard(task) + assert after == ["working", "completed"] + + +@pytest.mark.usefixtures("fake_pg_driver") +async def test_sync_outbox_crypto_and_worker_methods_delegate() -> None: + concrete = outbox(MagicMock()) + lazy = LazyTaskWebhookOutbox(lambda: concrete) + await lazy.resolve() + protected, nonce = lazy.protect_registration( + account_id="acct", + task_id="task", + task_type="get_products", + url="https://buyer.example/hooks", + operation_id="buyer", + token="token", + ) + assert lazy.open_registration( + account_id="acct", + task_id="task", + task_type="get_products", + encrypted_registration=protected, + nonce=nonce, + ) == ("https://buyer.example/hooks", "buyer", "token") + concrete.run_worker = AsyncMock() + concrete.purge_expired = AsyncMock() + await lazy.run_worker(poll_interval=0.25, purge_interval=10) + await lazy.purge_expired() + concrete.run_worker.assert_awaited_once_with(poll_interval=0.25, purge_interval=10) + concrete.purge_expired.assert_awaited_once_with() + with pytest.raises(ValueError): + lazy.validate_registration("http://buyer.example/hooks") + + +@pytest.mark.usefixtures("fake_pg_driver") +async def test_durability_is_not_silently_changed_by_factory() -> None: + class LossyRegistry(PgTaskRegistry): + is_durable = False + + lazy = LazyTaskRegistry(lambda: LossyRegistry(pool=MagicMock())) + assert lazy.is_durable is True + with pytest.raises(ValueError, match="durable"): + await lazy.resolve() + assert lazy.resolved is None + + +@pytest.mark.parametrize("flag", ["delivery_state_is_durable", "supports_atomic_task_outbox"]) +@pytest.mark.usefixtures("fake_pg_driver") +async def test_outbox_factory_cannot_weaken_declared_guarantees(flag: str) -> None: + box = outbox(MagicMock()) + setattr(box, flag, False) + lazy = LazyTaskWebhookOutbox(lambda: box) + with pytest.raises(ValueError, match="durable atomic"): + await lazy.resolve() + assert lazy.resolved is None + setattr(box, flag, True) + assert await lazy.resolve() is box + + +@pytest.mark.usefixtures("fake_pg_driver") +async def test_registry_factory_cannot_remove_declared_listing_support() -> None: + concrete = PgTaskRegistry(pool=MagicMock()) + original_list = concrete.list + concrete.list = None + lazy = LazyTaskRegistry(lambda: concrete) + assert isinstance(lazy, ListableTaskRegistry) + with pytest.raises(TypeError, match="listing support"): + await lazy.resolve() + assert lazy.resolved is None + concrete.list = original_list + assert await lazy.resolve() is concrete + + +@pytest.mark.parametrize("kind", ["proposal", "task", "outbox"]) +async def test_missing_pg_extra_remains_an_import_error_and_is_not_cached( + kind: str, monkeypatch: pytest.MonkeyPatch +) -> None: + modules = { + "proposal": "proposal_store", + "task": "task_registry", + "outbox": "task_webhook_outbox", + } + monkeypatch.setattr(f"adcp.decisioning.pg.{modules[kind]}.PG_AVAILABLE", False) + pool = MagicMock() + constructors = { + "proposal": lambda: PgProposalStore(pool=pool), + "task": lambda: PgTaskRegistry(pool=pool), + "outbox": lambda: outbox(pool), + } + wrappers = { + "proposal": LazyProposalStore, + "task": LazyTaskRegistry, + "outbox": LazyTaskWebhookOutbox, + } + calls = 0 + + def factory() -> Any: + nonlocal calls + calls += 1 + return constructors[kind]() + + lazy = wrappers[kind](factory) + for _ in range(2): + with pytest.raises(ImportError, match=r"adcp\[pg\]"): + await lazy.resolve() + assert lazy.resolved is None + assert calls == 2 diff --git a/tests/test_decisioning_task_polling.py b/tests/test_decisioning_task_polling.py index c0f30a28c..fffdf08e2 100644 --- a/tests/test_decisioning_task_polling.py +++ b/tests/test_decisioning_task_polling.py @@ -70,7 +70,7 @@ def get_media_buys(self, req: Any, ctx: Any) -> Any: raise NotImplementedError -@pytest.fixture(params=["memory", "postgres"]) +@pytest.fixture(params=["memory", "postgres", "lazy-postgres"]) async def registry(request: pytest.FixtureRequest) -> AsyncIterator[Any]: if request.param == "memory": yield InMemoryTaskRegistry() @@ -86,7 +86,12 @@ async def registry(request: pytest.FixtureRequest) -> AsyncIterator[Any]: async with AsyncConnectionPool(url, open=False) as pool: reg = PgTaskRegistry(pool=pool, _table=table) await reg.create_schema() - yield reg + if request.param == "lazy-postgres": + from adcp.decisioning.pg import LazyTaskRegistry + + yield LazyTaskRegistry(lambda: reg) + else: + yield reg async with pool.connection() as conn: await conn.execute(f"DROP TABLE {table}") diff --git a/tests/type_checks/lazy_pg_stores.py b/tests/type_checks/lazy_pg_stores.py new file mode 100644 index 000000000..1c31bcee2 --- /dev/null +++ b/tests/type_checks/lazy_pg_stores.py @@ -0,0 +1,66 @@ +"""Typed adopter factories preserve concrete store, listing and observer APIs.""" + +from typing import Any + +from psycopg_pool import AsyncConnectionPool + +from adcp.decisioning import ListableTaskRegistry, TaskRegistry, TaskTransition +from adcp.decisioning.pg import ( + LazyProposalStore, + LazyProposalStoreFactory, + LazyTaskRegistry, + LazyTaskRegistryFactory, + LazyTaskWebhookOutbox, + LazyTaskWebhookOutboxFactory, + PgProposalStore, + PgTaskRegistry, + PgTaskWebhookOutbox, +) +from adcp.decisioning.proposal_store import ProposalStore + + +def example(pool: AsyncConnectionPool[Any], outbox: PgTaskWebhookOutbox) -> None: + async def task_factory() -> tuple[PgTaskRegistry, PgTaskWebhookOutbox]: + return PgTaskRegistry(pool=pool, task_webhook_outbox=outbox), outbox + + factory: LazyTaskRegistryFactory = task_factory + registry = LazyTaskRegistry(factory) + minimal: TaskRegistry = registry + listable: ListableTaskRegistry = registry + + def metrics( + event: TaskTransition, + *, + task_id: str, + account_id: str, + task_type: str, + created_at: float, + updated_at: float, + ) -> None: + print(event, task_id, account_id, task_type, updated_at - created_at) + + registry.add_lifecycle_observer(metrics) + registry.remove_lifecycle_observer(metrics) + facade: LazyTaskWebhookOutbox = LazyTaskWebhookOutbox.from_registry(registry) + + def proposal_store() -> PgProposalStore: + return PgProposalStore(pool=pool) + + proposal_factory: LazyProposalStoreFactory = proposal_store + proposals: ProposalStore = LazyProposalStore(proposal_factory) + + def task_outbox() -> PgTaskWebhookOutbox: + return outbox + + outbox_factory: LazyTaskWebhookOutboxFactory = task_outbox + other_facade = LazyTaskWebhookOutbox(outbox_factory) + print(minimal, listable, facade, proposals, other_facade) + + +async def consume(registry: LazyTaskRegistry, facade: LazyTaskWebhookOutbox) -> None: + concrete: PgTaskRegistry = await registry.resolve() + box: PgTaskWebhookOutbox = await facade.resolve() + original: PgTaskRegistry | None = registry.resolved + tasks: dict[str, Any] = await registry.list(account_id="acct") + assert concrete.task_webhook_outbox is box + print(original, tasks)