community-agent-framework-postgres 0.1.1__py3-none-any.whl

agent_framework_community_postgres/_migrations.py

@@ -0,0 +1,235 @@
+"""Numbered schema migrations. The list index plus one is the version, as in LangGraph's checkpointer."""
+
+from __future__ import annotations
+
+import json
+from collections.abc import Callable, Sequence
+from dataclasses import dataclass
+from datetime import timedelta
+
+from psycopg import sql
+
+from ._client import ClientHandle, PostgresStorageError, TableNames
+
+Migration = Callable[[TableNames], Sequence[sql.Composed]]
+
+LOCK_TIMEOUT = timedelta(seconds=30)
+
+
+def _lock_timeout_literal(lock_timeout: timedelta) -> sql.Literal:
+    milliseconds = lock_timeout // timedelta(milliseconds=1)
+    if milliseconds < 1:
+        # PostgreSQL reads a lock_timeout of 0 as "wait forever", the opposite of what a caller means.
+        raise ValueError("lock_timeout must be at least one millisecond.")
+    return sql.Literal(f"{milliseconds}ms")
+
+
+def _v1(names: TableNames) -> list[sql.Composed]:
+    t = names.table
+    i = names.index
+    return [
+        sql.SQL(
+            "CREATE TABLE {history} ("
+            " id bigint GENERATED ALWAYS AS IDENTITY PRIMARY KEY,"
+            " application_id text NOT NULL,"
+            " tenant_id text NOT NULL DEFAULT '',"
+            " agent_id text NOT NULL DEFAULT '',"
+            " source_id text NOT NULL,"
+            " session_id text NOT NULL,"
+            " message jsonb NOT NULL,"
+            " created_at timestamptz NOT NULL DEFAULT now(),"
+            " expires_at timestamptz)"
+        ).format(history=t("history_messages")),
+        sql.SQL(
+            "CREATE INDEX {idx} ON {history} (application_id, tenant_id, agent_id, source_id, session_id, id)"
+        ).format(idx=i("history_messages_session_idx"), history=t("history_messages")),
+        # Serves the purge's per-session "is any row still live" probe.
+        sql.SQL(
+            "CREATE INDEX {idx} ON {history} (application_id, tenant_id, agent_id, source_id, session_id, expires_at)"
+        ).format(idx=i("history_messages_session_expiry_idx"), history=t("history_messages")),
+        sql.SQL("CREATE INDEX {idx} ON {history} (expires_at) WHERE expires_at IS NOT NULL").format(
+            idx=i("history_messages_expires_idx"), history=t("history_messages")
+        ),
+        sql.SQL(
+            "CREATE TABLE {sessions} ("
+            " application_id text NOT NULL,"
+            " session_id text NOT NULL,"
+            " snapshot jsonb,"
+            " snapshot_version text NOT NULL,"
+            " created_at timestamptz NOT NULL DEFAULT now(),"
+            " updated_at timestamptz NOT NULL DEFAULT now(),"
+            " expires_at timestamptz,"
+            " purged_at timestamptz,"
+            " PRIMARY KEY (application_id, session_id))"
+        ).format(sessions=t("sessions")),
+        sql.SQL("CREATE INDEX {idx} ON {sessions} (expires_at) WHERE expires_at IS NOT NULL").format(
+            idx=i("sessions_expires_idx"), sessions=t("sessions")
+        ),
+        # "Latest" is write order within a scope (sequence), not the producer's clock; see docs/checkpoints.md.
+        sql.SQL(
+            "CREATE TABLE {checkpoints} ("
+            " application_id text NOT NULL,"
+            " tenant_id text NOT NULL DEFAULT '',"
+            " scope text NOT NULL DEFAULT '',"
+            " payload_format text NOT NULL DEFAULT 'agent-framework.python.checkpoint.v1',"
+            " workflow_name text NOT NULL,"
+            " checkpoint_id text NOT NULL,"
+            " previous_checkpoint_id text,"
+            " checkpoint_timestamp timestamptz NOT NULL,"
+            " sequence bigint GENERATED ALWAYS AS IDENTITY,"
+            " payload jsonb NOT NULL,"
+            " created_at timestamptz NOT NULL DEFAULT now(),"
+            " expires_at timestamptz,"
+            " PRIMARY KEY (application_id, tenant_id, scope, checkpoint_id))"
+        ).format(checkpoints=t("checkpoints")),
+        sql.SQL(
+            "CREATE INDEX {idx} ON {checkpoints} (application_id, tenant_id, scope, workflow_name, sequence)"
+        ).format(idx=i("checkpoints_scope_idx"), checkpoints=t("checkpoints")),
+        sql.SQL("CREATE INDEX {idx} ON {checkpoints} (expires_at) WHERE expires_at IS NOT NULL").format(
+            idx=i("checkpoints_expires_idx"), checkpoints=t("checkpoints")
+        ),
+        sql.SQL(
+            "CREATE TABLE {snapshots} ("
+            " application_id text NOT NULL,"
+            " scope text NOT NULL,"
+            " thread_id text NOT NULL,"
+            " messages jsonb,"
+            " state jsonb,"
+            " interrupt jsonb,"
+            " session_state jsonb,"
+            " created_at timestamptz NOT NULL DEFAULT now(),"
+            " updated_at timestamptz NOT NULL DEFAULT now(),"
+            " expires_at timestamptz,"
+            " purged_at timestamptz,"
+            " PRIMARY KEY (application_id, scope, thread_id))"
+        ).format(snapshots=t("thread_snapshots")),
+        sql.SQL("CREATE INDEX {idx} ON {snapshots} (expires_at) WHERE expires_at IS NOT NULL").format(
+            idx=i("thread_snapshots_expires_idx"), snapshots=t("thread_snapshots")
+        ),
+        sql.SQL(
+            "CREATE TABLE {leases} ("
+            " application_id text NOT NULL,"
+            " resource text NOT NULL,"
+            " owner text NOT NULL,"
+            " token text NOT NULL,"
+            " expires_at timestamptz NOT NULL,"
+            " PRIMARY KEY (application_id, resource))"
+        ).format(leases=t("leases")),
+        sql.SQL("CREATE INDEX {idx} ON {leases} (expires_at)").format(idx=i("leases_expires_idx"), leases=t("leases")),
+    ]
+
+
+MIGRATIONS: tuple[Migration, ...] = (_v1,)
+
+
+@dataclass(frozen=True)
+class MigrationReport:
+    """What ``migrate`` did: the versions it applied this run and the version now current."""
+
+    applied: tuple[int, ...]
+    current: int
+
+
+def _version_table(names: TableNames) -> sql.Composed:
+    return sql.SQL(
+        "CREATE TABLE IF NOT EXISTS {migrations} ("
+        " version integer PRIMARY KEY,"
+        " applied_at timestamptz NOT NULL DEFAULT now())"
+    ).format(migrations=names.table("migrations"))
+
+
+def render(names: TableNames, versions: Sequence[int] | None = None) -> str:
+    """The SQL ``migrate`` would run, as text, for review or manual application by a DBA.
+
+    Each version is one transaction (``BEGIN;`` ... ``COMMIT;``) under ``migrate``'s default
+    ``lock_timeout``, as in ``migrate``. No versions render as an empty string.
+    """
+    chosen = list(range(1, len(MIGRATIONS) + 1)) if versions is None else list(versions)
+    if any(not 1 <= version <= len(MIGRATIONS) for version in chosen):
+        raise ValueError(f"versions must be between 1 and {len(MIGRATIONS)}.")
+    if not chosen:
+        return ""
+    timeout = sql.SQL("SET LOCAL lock_timeout = {}").format(_lock_timeout_literal(LOCK_TIMEOUT)).as_string()
+    parts = [_version_table(names).as_string() + ";"]
+    for version in chosen:
+        parts.append(f"-- version {version}")
+        parts.append("BEGIN;")
+        parts.append(timeout + ";")
+        parts.extend(statement.as_string() + ";" for statement in MIGRATIONS[version - 1](names))
+        parts.append(
+            sql.SQL("INSERT INTO {migrations} (version) VALUES ({version})")  # noqa: S608
+            .format(migrations=names.table("migrations"), version=sql.Literal(version))
+            .as_string()
+            + ";"
+        )
+        parts.append("COMMIT;")
+    return "\n".join(parts) + "\n"
+
+
+async def current_version(client: ClientHandle, names: TableNames) -> int:
+    """The highest applied version, or 0 when nothing has been applied."""
+    async with client.connection(transactional=False) as connection:
+        cursor = await connection.execute("SELECT to_regclass(%s)", (names.qualified("migrations"),))
+        found = await cursor.fetchone()
+        if found is None or found[0] is None:
+            return 0
+        cursor = await connection.execute(
+            sql.SQL("SELECT coalesce(max(version), 0) FROM {migrations}").format(migrations=names.table("migrations"))
+        )
+        row = await cursor.fetchone()
+        return int(row[0]) if row else 0
+
+
+def ensure_supported(current: int) -> None:
+    """Raise ``PostgresStorageError`` when the database was migrated by a newer release of this package."""
+    if current > len(MIGRATIONS):
+        raise PostgresStorageError(
+            f"Database schema version {current} is newer than this package supports ({len(MIGRATIONS)});"
+            " upgrade the package."
+        )
+
+
+async def pending_versions(client: ClientHandle, names: TableNames) -> list[int]:
+    """Versions ``migrate`` would apply now; raise when the database is newer than this package."""
+    current = await current_version(client, names)
+    ensure_supported(current)
+    return list(range(current + 1, len(MIGRATIONS) + 1))
+
+
+async def migrate(
+    client: ClientHandle, names: TableNames, *, lock_timeout: timedelta = LOCK_TIMEOUT
+) -> MigrationReport:
+    """Apply every pending migration, one transaction each, serialized across processes by an advisory lock.
+
+    A version's DDL runs under ``lock_timeout`` (``SET LOCAL``, after the migration lock is held), so a
+    migration never queues indefinitely behind application traffic; a timeout raises a retryable
+    ``PostgresStorageError`` and a rerun resumes. Replicas waiting for another replica's migration
+    wait on the advisory lock without a timeout.
+    """
+    lock_key = json.dumps(["agent-framework-postgres:migrations", names.schema, names.prefix])
+    timeout = _lock_timeout_literal(lock_timeout)
+    applied: list[int] = []
+    # READ COMMITTED in both blocks: a statement after the lock must see what another replica committed while
+    # this one waited, which a REPEATABLE READ snapshot taken by the lock statement itself would hide.
+    async with client.connection(read_committed=True) as connection:
+        # The lock also covers the CREATE TABLE IF NOT EXISTS, which races on pg_type when two processes start fresh.
+        await connection.execute("SELECT pg_advisory_xact_lock(hashtextextended(%s, 0))", (lock_key,))
+        await connection.execute(_version_table(names))
+    for version in await pending_versions(client, names):
+        async with client.connection(read_committed=True) as connection:
+            await connection.execute("SELECT pg_advisory_xact_lock(hashtextextended(%s, 0))", (lock_key,))
+            cursor = await connection.execute(
+                sql.SQL("SELECT 1 FROM {migrations} WHERE version = %s").format(migrations=names.table("migrations")),
+                (version,),
+            )
+            if await cursor.fetchone() is not None:
+                continue  # another process applied it while we waited for the lock
+            await connection.execute(sql.SQL("SET LOCAL lock_timeout = {}").format(timeout))
+            for statement in MIGRATIONS[version - 1](names):
+                await connection.execute(statement)
+            await connection.execute(
+                sql.SQL("INSERT INTO {migrations} (version) VALUES (%s)").format(migrations=names.table("migrations")),
+                (version,),
+            )
+            applied.append(version)
+    return MigrationReport(applied=tuple(applied), current=await current_version(client, names))

agent_framework_community_postgres/_persistence.py

@@ -0,0 +1,184 @@
+"""One pool, one schema, every store: the entry point most applications want."""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from datetime import timedelta
+from typing import TYPE_CHECKING, Any, Self
+
+from agent_framework import SecretString
+from psycopg import sql
+
+from ._checkpoint_storage import PostgresCheckpointStorage
+from ._client import PostgresClient, TableNames, create_client, require_text
+from ._entra import EntraCredential
+from ._framework import log_untested_once
+from ._history_provider import PostgresHistoryProvider, whole_session_expired
+from ._leases import PostgresLeases
+from ._migrations import LOCK_TIMEOUT, MigrationReport, current_version, migrate, pending_versions
+from ._retention import PurgeReport, RetentionMode, RetentionPolicy, purge_rows
+from ._session_store import PostgresSessionStore
+
+if TYPE_CHECKING:
+    from ._thread_snapshot_store import PostgresAGUIThreadSnapshotStore
+
+# Another application id, schema or prefix would escape the hub's migrations or purge; another client
+# or credential would escape its pool and its close.
+_HUB_OWNED = ("application_id", "client", "credential", "pool_kwargs", "schema", "table_prefix")
+
+
+class PostgresPersistence:
+    """Own one connection pool and hand it to every store this package provides.
+
+    Enter it (``async with``) or call ``open()`` to open the pool up front; otherwise the
+    pool opens on first use. Stores it creates borrow the pool and never close it, and
+    once the hub is closed they refuse to run. With ``credential``, the pool signs in to Azure
+    Database for PostgreSQL with Microsoft Entra ID tokens.
+    """
+
+    def __init__(
+        self,
+        *,
+        application_id: str,
+        connection_string: str | SecretString | None = None,
+        client: PostgresClient | None = None,
+        credential: EntraCredential | None = None,
+        env_file_path: str | None = None,
+        env_file_encoding: str | None = None,
+        schema: str = "public",
+        table_prefix: str = "af_",
+        retention: RetentionPolicy | None = None,
+        pool_kwargs: Mapping[str, Any] | None = None,
+    ) -> None:
+        log_untested_once()
+        self.application_id = require_text(application_id, "application_id")
+        self.names = TableNames(schema=schema, prefix=table_prefix)
+        self.retention = retention or RetentionPolicy()
+        self._client = create_client(
+            connection_string,
+            client=client,
+            env_file_path=env_file_path,
+            env_file_encoding=env_file_encoding,
+            credential=credential,
+            pool_kwargs=pool_kwargs,
+        )
+
+    @property
+    def pool(self) -> PostgresClient:
+        """The pool or connection every store created here shares."""
+        return self._client.client
+
+    async def open(self) -> None:
+        """Open the owned pool (a no-op for a borrowed client)."""
+        await self._client.open()
+
+    async def close(self) -> None:
+        """Close the owned pool."""
+        await self._client.close()
+
+    async def __aenter__(self) -> Self:
+        await self.open()
+        return self
+
+    async def __aexit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
+        await self.close()
+
+    async def migrate(self, *, lock_timeout: timedelta = LOCK_TIMEOUT) -> MigrationReport:
+        """Create or upgrade the schema; each version's DDL gives up after ``lock_timeout`` waiting on a lock."""
+        return await migrate(self._client, self.names, lock_timeout=lock_timeout)
+
+    async def pending_migrations(self) -> list[int]:
+        """Versions ``migrate`` would apply."""
+        return await pending_versions(self._client, self.names)
+
+    async def current_version(self) -> int:
+        """The schema version the database is at; 0 before the first migration."""
+        return await current_version(self._client, self.names)
+
+    def _common(self, overrides: dict[str, Any]) -> dict[str, Any]:
+        for key in _HUB_OWNED:
+            if key in overrides:
+                raise TypeError(f"{key} is set by the hub and cannot be overridden.")
+        base: dict[str, Any] = {
+            "application_id": self.application_id,
+            "client": self._client,
+            "schema": self.names.schema,
+            "table_prefix": self.names.prefix,
+            "retention": self.retention,
+        }
+        return base | overrides
+
+    def history_provider(
+        self, source_id: str = PostgresHistoryProvider.DEFAULT_SOURCE_ID, **kwargs: Any
+    ) -> PostgresHistoryProvider:
+        """A history provider on the shared pool; keyword arguments as on the class."""
+        return PostgresHistoryProvider(source_id, **self._common(kwargs))
+
+    def session_store(self, **kwargs: Any) -> PostgresSessionStore:
+        """A session store on the shared pool."""
+        return PostgresSessionStore(**self._common(kwargs))
+
+    def checkpoint_storage(
+        self,
+        *,
+        tenant_id: str | None = None,
+        scope: str | None = None,
+        allowed_checkpoint_types: list[str] | None = None,
+        **kwargs: Any,
+    ) -> PostgresCheckpointStorage:
+        """A checkpoint storage on the shared pool, optionally scoped to one tenant and one conversation."""
+        return PostgresCheckpointStorage(
+            tenant_id=tenant_id, scope=scope, allowed_checkpoint_types=allowed_checkpoint_types, **self._common(kwargs)
+        )
+
+    def thread_snapshot_store(self, **kwargs: Any) -> PostgresAGUIThreadSnapshotStore:
+        """An AG-UI thread snapshot store on the shared pool (needs the ``ag-ui`` extra)."""
+        from ._thread_snapshot_store import PostgresAGUIThreadSnapshotStore
+
+        return PostgresAGUIThreadSnapshotStore(**self._common(kwargs))
+
+    def leases(self) -> PostgresLeases:
+        """The lease table on the shared pool."""
+        return PostgresLeases(
+            application_id=self.application_id,
+            client=self._client,
+            schema=self.names.schema,
+            table_prefix=self.names.prefix,
+        )
+
+    async def purge(self, mode: RetentionMode | None = None) -> PurgeReport:
+        """Purge every expired row of this application in every table.
+
+        A row is expired when its ``expires_at``, stamped by whichever policy wrote it, has
+        passed; the hub needs no policy of its own. ``mode`` defaults to the hub policy's mode.
+        Each table is purged in batches, one transaction per batch. A history session is purged
+        only once every one of its messages has expired.
+        """
+        chosen = RetentionPolicy(mode=mode).mode if mode is not None else self.retention.mode
+        where = sql.SQL("application_id = %(app)s")
+        params = {"app": self.application_id}
+        plan: list[tuple[str, sql.Composable, sql.Composable | None, bool]] = [
+            ("history_messages", whole_session_expired(self.names.table("history_messages")), None, False),
+            ("sessions", where, sql.SQL("snapshot = NULL"), True),
+            ("checkpoints", where, None, False),
+            (
+                "thread_snapshots",
+                where,
+                sql.SQL("messages = NULL, state = NULL, interrupt = NULL, session_state = NULL"),
+                True,
+            ),
+            ("leases", where, None, False),
+        ]
+        counts: dict[str, int] = {}
+        # One transaction per batch: row locks are never held across tables, so a long purge cannot
+        # deadlock with, say, a history trim on the same session.
+        for table, table_where, tombstone, tombstone_allowed in plan:
+            counts[f"{self.names.prefix}{table}"] = await purge_rows(
+                self._client,
+                table=self.names.table(table),
+                where=table_where,
+                params=params,
+                mode=chosen if tombstone_allowed else "delete",
+                tombstone=tombstone,
+            )
+        return PurgeReport(counts)

agent_framework_community_postgres/_retention.py

@@ -0,0 +1,159 @@
+"""Retention: an optional time-to-live per store and the purge that acts on it."""
+
+from __future__ import annotations
+
+from collections.abc import Iterator, Mapping
+from dataclasses import dataclass, field
+from datetime import timedelta
+from typing import TYPE_CHECKING, Any, Literal, NoReturn, Self, cast
+
+from psycopg import sql
+
+if TYPE_CHECKING:
+    from ._client import ClientHandle
+
+RetentionMode = Literal["tombstone", "delete"]
+
+EXPIRES_AT = sql.SQL("CASE WHEN %(ttl)s::interval IS NULL THEN NULL ELSE now() + %(ttl)s::interval END")
+"""Value for an ``expires_at`` column on write; pass ``ttl`` (``timedelta | None``) in the parameters."""
+
+
+@dataclass(frozen=True)
+class RetentionPolicy:
+    """How long rows live after their last write and what ``purge`` does to expired ones.
+
+    ``ttl`` of ``None`` (the default) means nothing ever expires. ``tombstone`` keeps
+    the row with its ids and timestamps and drops the payload;
+    ``delete`` removes the row. History messages and checkpoints are always deleted, and so are
+    expired leases, which carry no policy.
+    """
+
+    ttl: timedelta | None = None
+    mode: RetentionMode = "tombstone"
+
+    def __post_init__(self) -> None:
+        if self.ttl is not None and self.ttl <= timedelta(0):
+            raise ValueError("ttl must be a positive timedelta or None.")
+        if self.mode not in ("tombstone", "delete"):
+            raise ValueError("mode must be 'tombstone' or 'delete'.")
+
+    @property
+    def enabled(self) -> bool:
+        """Whether writes set ``expires_at``."""
+        return self.ttl is not None
+
+
+class _ReadOnlyCounts(Mapping[str, int]):
+    """A read-only mapping over a private ``dict``; unlike a ``dict`` subclass, no ``dict`` method can change it.
+
+    It is built in ``__new__``, so calling ``__init__`` again changes nothing. ``deepcopy`` returns a plain
+    ``dict``, which is what ``dataclasses.asdict`` puts under ``counts``; ``PurgeReport`` copies itself through
+    its own ``__reduce__``, so a copied report stays read-only.
+    """
+
+    __slots__ = ("_data",)
+    _data: dict[str, int]
+
+    def __new__(cls, data: Mapping[str, int]) -> Self:
+        self = super().__new__(cls)
+        object.__setattr__(self, "_data", dict(data))
+        return self
+
+    def __setattr__(self, name: str, value: object) -> NoReturn:
+        raise TypeError("PurgeReport.counts is read-only.")
+
+    def __getitem__(self, key: str) -> int:
+        return self._data[key]
+
+    def __iter__(self) -> Iterator[str]:
+        return iter(self._data)
+
+    def __len__(self) -> int:
+        return len(self._data)
+
+    def __repr__(self) -> str:
+        return repr(self._data)
+
+    def __eq__(self, other: object) -> bool:
+        if isinstance(other, Mapping):
+            return self._data == dict(cast("Mapping[object, object]", other))
+        return NotImplemented
+
+    __hash__ = None  # type: ignore[assignment]  # equal to a dict, so unhashable like one
+
+    def __reduce__(self) -> tuple[type[_ReadOnlyCounts], tuple[dict[str, int]]]:
+        return (_ReadOnlyCounts, (dict(self._data),))
+
+    def __deepcopy__(self, memo: dict[int, object]) -> dict[str, int]:
+        return dict(self._data)
+
+
+@dataclass(frozen=True)
+class PurgeReport:
+    """Rows affected per table by one purge; ``counts`` is a read-only copy. ``sum(reports)`` works."""
+
+    counts: Mapping[str, int] = field(default_factory=dict[str, int])
+
+    def __post_init__(self) -> None:
+        object.__setattr__(self, "counts", _ReadOnlyCounts(self.counts))
+
+    def __reduce__(self) -> tuple[type[PurgeReport], tuple[dict[str, int]]]:
+        # Pickle, copy and deepcopy rebuild the report from a plain dict, so the copy is read-only too.
+        return (PurgeReport, (dict(self.counts),))
+
+    @property
+    def total(self) -> int:
+        return sum(self.counts.values())
+
+    def __add__(self, other: PurgeReport) -> PurgeReport:
+        merged = dict(self.counts)
+        for table, count in other.counts.items():
+            merged[table] = merged.get(table, 0) + count
+        return PurgeReport(merged)
+
+    def __radd__(self, other: int) -> PurgeReport:
+        if other == 0:  # the start value of sum()
+            return self
+        return NotImplemented
+
+
+async def purge_rows(
+    client: ClientHandle,
+    *,
+    table: sql.Identifier,
+    where: sql.Composable,
+    params: Mapping[str, Any],
+    mode: RetentionMode,
+    tombstone: sql.Composable | None,
+    batch_size: int = 1000,
+) -> int:
+    """Tombstone or delete the expired rows matching ``where``, ``batch_size`` rows per transaction.
+
+    ``where`` may qualify columns with ``c``, the candidate row's alias. Rows another purge holds
+    are skipped and taken by a later batch, so two purges never block each other.
+    """
+    if batch_size < 1:
+        raise ValueError("batch_size must be at least 1.")
+    if mode == "tombstone" and tombstone is not None:
+        candidates = sql.SQL(
+            "SELECT c.ctid FROM {table} c WHERE ({where}) AND c.expires_at < now() AND c.purged_at IS NULL"
+            " LIMIT %(batch)s FOR UPDATE SKIP LOCKED"
+        )
+        statement = sql.SQL("UPDATE {table} SET {tombstone}, purged_at = now() WHERE ctid IN ({candidates})")
+    else:
+        candidates = sql.SQL(
+            "SELECT c.ctid FROM {table} c WHERE ({where}) AND c.expires_at < now()"
+            " LIMIT %(batch)s FOR UPDATE SKIP LOCKED"
+        )
+        statement = sql.SQL("DELETE FROM {table} WHERE ctid IN ({candidates})")
+    composed = statement.format(
+        table=table, tombstone=tombstone or sql.SQL(""), candidates=candidates.format(table=table, where=where)
+    )
+    total = 0
+    while True:
+        async with client.connection() as connection:
+            cursor = await connection.execute(composed, dict(params) | {"batch": batch_size})
+            count = cursor.rowcount
+        total += count
+        if count < batch_size:
+            return total