harness-sdk-python 0.25.1__tar.gz → 0.25.2__tar.gz

{harness_sdk_python-0.25.1 → harness_sdk_python-0.25.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: harness-sdk-python
-Version: 0.25.1
+Version: 0.25.2
 Summary: RunManager: the harness-sdk runs subsystem for Python Statewire hosts
 Project-URL: Repository, https://github.com/assistant-ui/harness-sdk
 License-Expression: MIT

{harness_sdk_python-0.25.1 → harness_sdk_python-0.25.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "harness-sdk-python"
-version = "0.25.1"
+version = "0.25.2"
 description = "RunManager: the harness-sdk runs subsystem for Python Statewire hosts"
 readme = "README.md"
 license = "MIT"

{harness_sdk_python-0.25.1 → harness_sdk_python-0.25.2}/src/harness_sdk/checkpoints.py

@@ -11,6 +11,7 @@ Rows and entries are insert-only; ``last_checkpoint_id`` is the only update.
 """
 
 import asyncio
+from contextlib import asynccontextmanager
 from dataclasses import dataclass, field
 from typing import Any, Callable, Protocol
 
@@ -158,50 +159,69 @@ class MemoryCheckpointIndex:
 
 
 class PostgresCheckpointIndex:
-    """psycopg-backed CheckpointIndex; ``setup()`` creates the two tables."""
+    """psycopg-backed CheckpointIndex over an ``AsyncConnection`` or ``AsyncConnectionPool``; ``setup()`` creates the two tables."""
 
     def __init__(self, conn: Any) -> None:
         try:
-            import psycopg  # noqa: F401
+            import psycopg
+            import psycopg_pool
         except ImportError as exc:
             raise ImportError(
                 "harness_sdk.checkpoints.PostgresCheckpointIndex requires the "
                 "'postgres' extra: pip install 'harness-sdk-python[postgres]'"
             ) from exc
+        if not isinstance(
+            conn, (psycopg.AsyncConnection, psycopg_pool.AsyncConnectionPool)
+        ):
+            raise TypeError(
+                f"conn must be an AsyncConnection or AsyncConnectionPool, got {conn!r}"
+            )
         self._conn = conn
 
+    @asynccontextmanager
+    async def _connection(self):
+        import psycopg
+
+        if isinstance(self._conn, psycopg.AsyncConnection):
+            yield self._conn
+        else:
+            async with self._conn.connection() as conn:
+                yield conn
+
     async def setup(self) -> None:
-        await self._conn.execute(
-            "CREATE TABLE IF NOT EXISTS checkpoint_threads ("
-            " thread_id TEXT NOT NULL,"
-            " checkpoint_ns TEXT NOT NULL,"
-            " branches JSONB NOT NULL,"
-            " PRIMARY KEY (thread_id, checkpoint_ns))"
-        )
-        await self._conn.execute(
-            "CREATE TABLE IF NOT EXISTS checkpoint_messages ("
-            " thread_id TEXT NOT NULL,"
-            " checkpoint_ns TEXT NOT NULL,"
-            " message_id TEXT NOT NULL,"
-            " anchor_checkpoint_id TEXT,"
-            " branch_id INTEGER NOT NULL,"
-            " PRIMARY KEY (thread_id, checkpoint_ns, message_id))"
-        )
+        async with self._connection() as conn:
+            await conn.execute(
+                "CREATE TABLE IF NOT EXISTS checkpoint_threads ("
+                " thread_id TEXT NOT NULL,"
+                " checkpoint_ns TEXT NOT NULL,"
+                " branches JSONB NOT NULL,"
+                " PRIMARY KEY (thread_id, checkpoint_ns))"
+            )
+            await conn.execute(
+                "CREATE TABLE IF NOT EXISTS checkpoint_messages ("
+                " thread_id TEXT NOT NULL,"
+                " checkpoint_ns TEXT NOT NULL,"
+                " message_id TEXT NOT NULL,"
+                " anchor_checkpoint_id TEXT,"
+                " branch_id INTEGER NOT NULL,"
+                " PRIMARY KEY (thread_id, checkpoint_ns, message_id))"
+            )
 
     async def read_thread(
         self, thread_id: str, checkpoint_ns: str
     ) -> list[dict[str, Any]]:
-        cur = await self._conn.execute(
-            "SELECT branches FROM checkpoint_threads WHERE thread_id = %s AND checkpoint_ns = %s",
-            (thread_id, checkpoint_ns),
-        )
-        row = await cur.fetchone()
+        async with self._connection() as conn:
+            cur = await conn.execute(
+                "SELECT branches FROM checkpoint_threads WHERE thread_id = %s AND checkpoint_ns = %s",
+                (thread_id, checkpoint_ns),
+            )
+            row = await cur.fetchone()
         return row[0] if row is not None else []
 
-    async def _seed_root(self, thread_id: str, checkpoint_ns: str) -> None:
+    async def _seed_root(self, conn: Any, thread_id: str, checkpoint_ns: str) -> None:
         from psycopg.types.json import Jsonb
 
-        await self._conn.execute(
+        await conn.execute(
             "INSERT INTO checkpoint_threads (thread_id, checkpoint_ns, branches)"
             " VALUES (%s, %s, %s) ON CONFLICT DO NOTHING",
             (thread_id, checkpoint_ns, Jsonb([_ROOT_ENTRY])),
@@ -223,9 +243,9 @@ class PostgresCheckpointIndex:
             "fork_branch": fork_branch,
             "last_checkpoint_id": last_checkpoint_id,
         }
-        async with self._conn.transaction():
-            await self._seed_root(thread_id, checkpoint_ns)
-            cur = await self._conn.execute(
+        async with self._connection() as conn, conn.transaction():
+            await self._seed_root(conn, thread_id, checkpoint_ns)
+            cur = await conn.execute(
                 "UPDATE checkpoint_threads SET branches = branches || %s"
                 " WHERE thread_id = %s AND checkpoint_ns = %s"
                 " RETURNING jsonb_array_length(branches)",
@@ -240,26 +260,27 @@ class PostgresCheckpointIndex:
         branch_ids: list[int],
         checkpoint_id: str,
     ) -> None:
-        await self._seed_root(thread_id, checkpoint_ns)
-        cur = await self._conn.execute(
-            "SELECT jsonb_array_length(branches) FROM checkpoint_threads"
-            " WHERE thread_id = %s AND checkpoint_ns = %s",
-            (thread_id, checkpoint_ns),
-        )
-        row = await cur.fetchone()
-        _check_branch_ids(branch_ids, row[0] if row is not None else 0)
-        await self._conn.execute(
-            "UPDATE checkpoint_threads SET branches = (SELECT jsonb_agg("
-            " CASE WHEN idx::int - 1 = ANY(%s)"
-            " THEN elem || jsonb_build_object('last_checkpoint_id', %s::text)"
-            " ELSE elem END ORDER BY idx)"
-            " FROM jsonb_array_elements(branches) WITH ORDINALITY AS t(elem, idx))"
-            " WHERE thread_id = %s AND checkpoint_ns = %s",
-            (branch_ids, checkpoint_id, thread_id, checkpoint_ns),
-        )
+        async with self._connection() as conn:
+            await self._seed_root(conn, thread_id, checkpoint_ns)
+            cur = await conn.execute(
+                "SELECT jsonb_array_length(branches) FROM checkpoint_threads"
+                " WHERE thread_id = %s AND checkpoint_ns = %s",
+                (thread_id, checkpoint_ns),
+            )
+            row = await cur.fetchone()
+            _check_branch_ids(branch_ids, row[0] if row is not None else 0)
+            await conn.execute(
+                "UPDATE checkpoint_threads SET branches = (SELECT jsonb_agg("
+                " CASE WHEN idx::int - 1 = ANY(%s)"
+                " THEN elem || jsonb_build_object('last_checkpoint_id', %s::text)"
+                " ELSE elem END ORDER BY idx)"
+                " FROM jsonb_array_elements(branches) WITH ORDINALITY AS t(elem, idx))"
+                " WHERE thread_id = %s AND checkpoint_ns = %s",
+                (branch_ids, checkpoint_id, thread_id, checkpoint_ns),
+            )
 
     async def upsert_messages(self, rows: list[dict[str, Any]]) -> None:
-        async with self._conn.cursor() as cur:
+        async with self._connection() as conn, conn.cursor() as cur:
             await cur.executemany(
                 "INSERT INTO checkpoint_messages"
                 " (thread_id, checkpoint_ns, message_id, anchor_checkpoint_id, branch_id)"
@@ -272,12 +293,13 @@ class PostgresCheckpointIndex:
     async def get_message(
         self, thread_id: str, checkpoint_ns: str, message_id: str
     ) -> dict[str, Any] | None:
-        cur = await self._conn.execute(
-            "SELECT anchor_checkpoint_id, branch_id FROM checkpoint_messages"
-            " WHERE thread_id = %s AND checkpoint_ns = %s AND message_id = %s",
-            (thread_id, checkpoint_ns, message_id),
-        )
-        row = await cur.fetchone()
+        async with self._connection() as conn:
+            cur = await conn.execute(
+                "SELECT anchor_checkpoint_id, branch_id FROM checkpoint_messages"
+                " WHERE thread_id = %s AND checkpoint_ns = %s AND message_id = %s",
+                (thread_id, checkpoint_ns, message_id),
+            )
+            row = await cur.fetchone()
         if row is None:
             return None
         return {"anchor_checkpoint_id": row[0], "branch_id": row[1]}

{harness_sdk_python-0.25.1 → harness_sdk_python-0.25.2}/tests/test_checkpoints.py

@@ -663,11 +663,12 @@ def postgres_dsn():
         subprocess.run([docker, "rm", "-f", container], capture_output=True)
 
 
-@pytest.fixture
-async def pg_index(postgres_dsn):
+@pytest.fixture(params=["conn", "pool"])
+async def pg_index(request, postgres_dsn):
     import uuid
 
     import psycopg
+    from psycopg_pool import AsyncConnectionPool
 
     from harness_sdk import PostgresCheckpointIndex
 
@@ -675,15 +676,29 @@ async def pg_index(postgres_dsn):
     with psycopg.connect(postgres_dsn, autocommit=True) as admin:
         admin.execute(f'CREATE DATABASE "{dbname}"')
     dsn = postgres_dsn.rsplit("/", 1)[0] + f"/{dbname}"
-    conn = await psycopg.AsyncConnection.connect(dsn, autocommit=True)
-    index = PostgresCheckpointIndex(conn)
-    await index.setup()
-    yield index
-    await conn.close()
+    if request.param == "conn":
+        conn = await psycopg.AsyncConnection.connect(dsn, autocommit=True)
+        index = PostgresCheckpointIndex(conn)
+        await index.setup()
+        yield index
+        await conn.close()
+    else:
+        async with AsyncConnectionPool(dsn, open=False) as pool:
+            index = PostgresCheckpointIndex(pool)
+            await index.setup()
+            yield index
     with psycopg.connect(postgres_dsn, autocommit=True) as admin:
         admin.execute(f'DROP DATABASE "{dbname}" WITH (FORCE)')
 
 
+def test_postgres_index_rejects_other_conn_types():
+    pytest.importorskip("psycopg")
+    from harness_sdk import PostgresCheckpointIndex
+
+    with pytest.raises(TypeError, match="AsyncConnection or AsyncConnectionPool"):
+        PostgresCheckpointIndex(object())
+
+
 async def test_postgres_branches_round_trip(pg_index):
     assert await pg_index.read_thread(THREAD, "") == []
     first = await pg_index.append_branch(