langgraph-checkpoint-postgres 3.0.4__tar.gz → 3.1.0a1__tar.gz

{langgraph_checkpoint_postgres-3.0.4 → langgraph_checkpoint_postgres-3.1.0a1}/.gitignore

@@ -100,3 +100,4 @@ dmypy.json
 .turbo
 .editorconfig
 .scratch
+.worktrees/

{langgraph_checkpoint_postgres-3.0.4 → langgraph_checkpoint_postgres-3.1.0a1}/Makefile

@@ -1,4 +1,4 @@
-.PHONY: test test_watch lint format
+.PHONY: test test_watch lint type format
 
 ######################
 # TESTING AND COVERAGE
@@ -61,6 +61,9 @@ lint lint_diff lint_package lint_tests:
 	[ "$(PYTHON_FILES)" = "" ] || mkdir -p $(MYPY_CACHE)
 	[ "$(PYTHON_FILES)" = "" ] || uv run mypy $(PYTHON_FILES) --cache-dir $(MYPY_CACHE)
 
+type:
+	mkdir -p $(MYPY_CACHE) && uv run mypy $(PYTHON_FILES) --cache-dir $(MYPY_CACHE)
+
 format format_diff:
 	uv run ruff format $(PYTHON_FILES)
 	uv run ruff check --select I --fix $(PYTHON_FILES)

{langgraph_checkpoint_postgres-3.0.4 → langgraph_checkpoint_postgres-3.1.0a1}/PKG-INFO

@@ -1,16 +1,16 @@
 Metadata-Version: 2.4
 Name: langgraph-checkpoint-postgres
-Version: 3.0.4
+Version: 3.1.0a1
 Summary: Library with a Postgres implementation of LangGraph checkpoint saver.
 Project-URL: Source, https://github.com/langchain-ai/langgraph/tree/main/libs/checkpoint-postgres
-Project-URL: Twitter, https://x.com/LangChain
+Project-URL: Twitter, https://x.com/langchain_oss
 Project-URL: Slack, https://www.langchain.com/join-community
 Project-URL: Reddit, https://www.reddit.com/r/LangChain/
 License-Expression: MIT
 License-File: LICENSE
 Requires-Python: >=3.10
-Requires-Dist: langgraph-checkpoint<5.0.0,>=2.1.2
-Requires-Dist: orjson>=3.10.1
+Requires-Dist: langgraph-checkpoint<5.0.0,>=4.1.0a1
+Requires-Dist: orjson>=3.11.5
 Requires-Dist: psycopg-pool>=3.2.0
 Requires-Dist: psycopg>=3.2.0
 Description-Content-Type: text/markdown
@@ -23,6 +23,11 @@ Implementation of LangGraph CheckpointSaver that uses Postgres.
 
 By default `langgraph-checkpoint-postgres` installs `psycopg` (Psycopg 3) without any extras. However, you can choose a specific installation that best suits your needs [here](https://www.psycopg.org/psycopg3/docs/basic/install.html) (for example, `psycopg[binary]`).
 
+## Security
+
+> [!IMPORTANT]
+> Set `LANGGRAPH_STRICT_MSGPACK=true` or pass an explicit `allowed_msgpack_modules` list when creating your checkpointer. This restricts checkpoint deserialization to known-safe types, preventing code execution if the database is compromised. See the [langgraph-checkpoint README](https://github.com/langchain-ai/langgraph/tree/main/libs/checkpoint#serde) for details.
+
 ## Usage
 
 > [!IMPORTANT]

{langgraph_checkpoint_postgres-3.0.4 → langgraph_checkpoint_postgres-3.1.0a1}/README.md

@@ -6,6 +6,11 @@ Implementation of LangGraph CheckpointSaver that uses Postgres.
 
 By default `langgraph-checkpoint-postgres` installs `psycopg` (Psycopg 3) without any extras. However, you can choose a specific installation that best suits your needs [here](https://www.psycopg.org/psycopg3/docs/basic/install.html) (for example, `psycopg[binary]`).
 
+## Security
+
+> [!IMPORTANT]
+> Set `LANGGRAPH_STRICT_MSGPACK=true` or pass an explicit `allowed_msgpack_modules` list when creating your checkpointer. This restricts checkpoint deserialization to known-safe types, preventing code execution if the database is compromised. See the [langgraph-checkpoint README](https://github.com/langchain-ai/langgraph/tree/main/libs/checkpoint#serde) for details.
+
 ## Usage
 
 > [!IMPORTANT]

{langgraph_checkpoint_postgres-3.0.4 → langgraph_checkpoint_postgres-3.1.0a1}/langgraph/checkpoint/postgres/__init__.py

@@ -4,15 +4,17 @@ import threading
 from collections import defaultdict
 from collections.abc import Iterator, Sequence
 from contextlib import contextmanager
-from typing import Any
+from typing import Any, cast
 
 from langchain_core.runnables import RunnableConfig
 from langgraph.checkpoint.base import (
+    DELTA_SENTINEL,
     WRITES_IDX_MAP,
     ChannelVersions,
     Checkpoint,
     CheckpointMetadata,
     CheckpointTuple,
+    _ChannelWritesHistory,
     get_checkpoint_id,
     get_serializable_checkpoint_metadata,
 )
@@ -23,7 +25,11 @@ from psycopg.types.json import Jsonb
 from psycopg_pool import ConnectionPool
 
 from langgraph.checkpoint.postgres import _internal
-from langgraph.checkpoint.postgres.base import BasePostgresSaver
+from langgraph.checkpoint.postgres.base import (
+    SELECT_DELTA_COMBINED_SQL,
+    BasePostgresSaver,
+    _DeltaCombinedRow,
+)
 from langgraph.checkpoint.postgres.shallow import ShallowPostgresSaver
 
 Conn = _internal.Conn  # For backward compatibility
@@ -430,6 +436,48 @@ class PostgresSaver(BasePostgresSaver):
                 with conn.cursor(binary=True, row_factory=dict_row) as cur:
                     yield cur
 
+    def _get_channel_writes_history(
+        self, config: RunnableConfig, channel: str
+    ) -> _ChannelWritesHistory:
+        """Fast-path override of `BaseCheckpointSaver._get_channel_writes_history`.
+
+        One combined UNION ALL query (`SELECT_DELTA_COMBINED_SQL`) fetches rows
+        from `checkpoints`, `checkpoint_writes`, and `checkpoint_blobs` in a
+        single roundtrip; the ancestor walk runs in Python.
+        """
+        thread_id = config["configurable"]["thread_id"]
+        checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
+        checkpoint_id = get_checkpoint_id(config)
+        if checkpoint_id is None:
+            # Caller didn't specify a target — resolve to the latest
+            # checkpoint on the thread. `get_tuple` without `checkpoint_id`
+            # returns the newest; its config carries the resolved id.
+            target = self.get_tuple(config)
+            if target is None:
+                return _ChannelWritesHistory(seed=DELTA_SENTINEL, writes=[])
+            checkpoint_id = target.config["configurable"]["checkpoint_id"]
+        with self._cursor() as cur:
+            cur.execute(
+                SELECT_DELTA_COMBINED_SQL,
+                (
+                    channel,
+                    thread_id,
+                    checkpoint_ns,
+                    thread_id,
+                    checkpoint_ns,
+                    channel,
+                    thread_id,
+                    checkpoint_ns,
+                    channel,
+                ),
+            )
+            rows = cur.fetchall()
+        return self._build_delta_channel_writes_history(
+            channel=channel,
+            target_id=checkpoint_id,
+            rows=cast("list[_DeltaCombinedRow]", rows),
+        )
+
     def _load_checkpoint_tuple(self, value: DictRow) -> CheckpointTuple:
         """
         Convert a database row into a CheckpointTuple object.

{langgraph_checkpoint_postgres-3.0.4 → langgraph_checkpoint_postgres-3.1.0a1}/langgraph/checkpoint/postgres/aio.py

@@ -4,15 +4,17 @@ import asyncio
 from collections import defaultdict
 from collections.abc import AsyncIterator, Iterator, Sequence
 from contextlib import asynccontextmanager
-from typing import Any
+from typing import Any, cast
 
 from langchain_core.runnables import RunnableConfig
 from langgraph.checkpoint.base import (
+    DELTA_SENTINEL,
     WRITES_IDX_MAP,
     ChannelVersions,
     Checkpoint,
     CheckpointMetadata,
     CheckpointTuple,
+    _ChannelWritesHistory,
     get_checkpoint_id,
     get_serializable_checkpoint_metadata,
 )
@@ -23,7 +25,11 @@ from psycopg.types.json import Jsonb
 from psycopg_pool import AsyncConnectionPool
 
 from langgraph.checkpoint.postgres import _ainternal
-from langgraph.checkpoint.postgres.base import BasePostgresSaver
+from langgraph.checkpoint.postgres.base import (
+    SELECT_DELTA_COMBINED_SQL,
+    BasePostgresSaver,
+    _DeltaCombinedRow,
+)
 from langgraph.checkpoint.postgres.shallow import AsyncShallowPostgresSaver
 
 Conn = _ainternal.Conn  # For backward compatibility
@@ -391,6 +397,46 @@ class AsyncPostgresSaver(BasePostgresSaver):
                 async with conn.cursor(binary=True, row_factory=dict_row) as cur:
                     yield cur
 
+    async def _aget_channel_writes_history(
+        self, config: RunnableConfig, channel: str
+    ) -> _ChannelWritesHistory:
+        """Fast-path override of `BaseCheckpointSaver._aget_channel_writes_history`.
+
+        One combined UNION ALL query (`SELECT_DELTA_COMBINED_SQL`) fetches rows
+        from `checkpoints`, `checkpoint_writes`, and `checkpoint_blobs` in a
+        single roundtrip; rows are assembled by the shared pure helper on
+        `BasePostgresSaver`.
+        """
+        thread_id = config["configurable"]["thread_id"]
+        checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
+        checkpoint_id = get_checkpoint_id(config)
+        if checkpoint_id is None:
+            target = await self.aget_tuple(config)
+            if target is None:
+                return _ChannelWritesHistory(seed=DELTA_SENTINEL, writes=[])
+            checkpoint_id = target.config["configurable"]["checkpoint_id"]
+        async with self._cursor() as cur:
+            await cur.execute(
+                SELECT_DELTA_COMBINED_SQL,
+                (
+                    channel,
+                    thread_id,
+                    checkpoint_ns,
+                    thread_id,
+                    checkpoint_ns,
+                    channel,
+                    thread_id,
+                    checkpoint_ns,
+                    channel,
+                ),
+            )
+            rows = await cur.fetchall()
+        return self._build_delta_channel_writes_history(
+            channel=channel,
+            target_id=checkpoint_id,
+            rows=cast("list[_DeltaCombinedRow]", rows),
+        )
+
     async def _load_checkpoint_tuple(self, value: DictRow) -> CheckpointTuple:
         """
         Convert a database row into a CheckpointTuple object.

{langgraph_checkpoint_postgres-3.0.4 → langgraph_checkpoint_postgres-3.1.0a1}/langgraph/checkpoint/postgres/base.py

@@ -4,13 +4,16 @@ import random
 import warnings
 from collections.abc import Sequence
 from importlib.metadata import version as get_version
-from typing import Any, cast
+from typing import Any, TypedDict, cast
 
 from langchain_core.runnables import RunnableConfig
 from langgraph.checkpoint.base import (
+    DELTA_SENTINEL,
     WRITES_IDX_MAP,
     BaseCheckpointSaver,
     ChannelVersions,
+    PendingWrite,
+    _ChannelWritesHistory,
     get_checkpoint_id,
 )
 from langgraph.checkpoint.serde.types import TASKS
@@ -153,6 +156,62 @@ INSERT_CHECKPOINT_WRITES_SQL = """
 """
 
 
+class _DeltaCombinedRow(TypedDict, total=False):
+    """One row from `SELECT_DELTA_COMBINED_SQL` (a UNION ALL of three tables).
+
+    Every row carries `_kind` ("p" / "w" / "b") plus whichever columns are
+    relevant for that kind; irrelevant columns are NULL and typed as `None`.
+    """
+
+    _kind: str  # always present: "p", "w", or "b"
+    # checkpoint row ("p")
+    checkpoint_id: str | None
+    parent_checkpoint_id: str | None
+    ver: str | None
+    # write / blob rows ("w", "b")
+    type: str | None
+    blob: bytes | None
+    # write row only ("w")
+    task_id: str | None
+    idx: int | None
+    # blob row only ("b")
+    version: str | None
+
+
+# DeltaChannel reconstruction: one UNION ALL query fetches checkpoints,
+# writes, and blobs for `channel` in one roundtrip; the ancestor walk runs
+# in Python in `_build_delta_channel_writes_history`.
+#
+# Parameter order: (channel, thread_id, checkpoint_ns,
+#                   thread_id, checkpoint_ns, channel,
+#                   thread_id, checkpoint_ns, channel)
+SELECT_DELTA_COMBINED_SQL = """
+    SELECT 'p'::text       AS _kind,
+           checkpoint_id,
+           parent_checkpoint_id,
+           checkpoint -> 'channel_versions' ->> %s AS ver,
+           NULL::text       AS type,
+           NULL::bytea      AS blob,
+           NULL::text       AS task_id,
+           NULL::int        AS idx,
+           NULL::text       AS version
+    FROM checkpoints
+    WHERE thread_id = %s AND checkpoint_ns = %s
+    UNION ALL
+    SELECT 'w',
+           checkpoint_id, NULL, NULL,
+           type, blob, task_id, idx, NULL
+    FROM checkpoint_writes
+    WHERE thread_id = %s AND checkpoint_ns = %s AND channel = %s
+    UNION ALL
+    SELECT 'b',
+           NULL, NULL, NULL,
+           type, blob, NULL, NULL, version
+    FROM checkpoint_blobs
+    WHERE thread_id = %s AND checkpoint_ns = %s AND channel = %s
+"""
+
+
 class BasePostgresSaver(BaseCheckpointSaver[str]):
     SELECT_SQL = SELECT_SQL
     SELECT_PENDING_SENDS_SQL = SELECT_PENDING_SENDS_SQL
@@ -195,6 +254,83 @@ class BasePostgresSaver(BaseCheckpointSaver[str]):
             if t.decode() != "empty"
         }
 
+    def _build_delta_channel_writes_history(
+        self,
+        *,
+        channel: str,
+        target_id: str,
+        rows: Sequence[_DeltaCombinedRow],
+    ) -> _ChannelWritesHistory:
+        """Reconstruct one delta channel's history from the combined UNION ALL rows.
+
+        Pure data transform shared by sync (`PostgresSaver`) and async
+        (`AsyncPostgresSaver`); both paths run `SELECT_DELTA_COMBINED_SQL`
+        and feed the tagged rows here.
+
+        Walk is newest → oldest from the target's parent. A non-sentinel
+        blob in `checkpoint_blobs` (a pre-delta snapshot) terminates the
+        walk and is returned as the seed so replay starts from it.
+
+        Writes stored at `target_id` itself are pending writes for the next
+        step and are excluded — the walk begins at the target's parent.
+        """
+        parent_of: dict[str, str | None] = {}
+        ver_of: dict[str, str | None] = {}
+        writes_by_cid: dict[str, list[tuple[str, bytes, str, int]]] = {}
+        blob_by_ver: dict[str, tuple[str, bytes]] = {}
+
+        for r in rows:
+            kind = r["_kind"]
+            if kind == "p":
+                cid = cast(str, r["checkpoint_id"])
+                parent_of[cid] = r["parent_checkpoint_id"]
+                ver_of[cid] = r["ver"]
+            elif kind == "w":
+                cid = cast(str, r["checkpoint_id"])
+                writes_by_cid.setdefault(cid, []).append(
+                    cast(
+                        "tuple[str, bytes, str, int]",
+                        (r["type"], r["blob"], r["task_id"], r["idx"]),
+                    )
+                )
+            else:  # kind == "b"
+                blob_by_ver[cast(str, r["version"])] = cast(
+                    "tuple[str, bytes]", (r["type"], r["blob"])
+                )
+
+        # newest write first per ancestor (task_id DESC, idx DESC)
+        for ws in writes_by_cid.values():
+            ws.sort(key=lambda w: (w[2], w[3]), reverse=True)
+
+        ancestors: list[str] = []
+        cur_cid: str | None = parent_of.get(target_id)
+        while cur_cid is not None:
+            ancestors.append(cur_cid)
+            cur_cid = parent_of.get(cur_cid)
+        if not ancestors:
+            return _ChannelWritesHistory(seed=DELTA_SENTINEL, writes=[])
+
+        collected: list[PendingWrite] = []  # newest first; reversed at the end
+        for cid in ancestors:
+            # Collect writes first — they encode the transition FROM this
+            # ancestor's state to its child's and must be included even if
+            # this ancestor is also the seed checkpoint.
+            for type_tag, write_blob, task_id, _idx in writes_by_cid.get(cid, []):
+                val = self.serde.loads_typed((type_tag, write_blob))
+                collected.append((task_id, channel, val))
+            # Then check seed terminator.
+            ver = ver_of.get(cid)
+            if ver is not None:
+                seed_blob = blob_by_ver.get(ver)
+                if seed_blob is not None and seed_blob[0] != "empty":
+                    blob_value = self.serde.loads_typed(seed_blob)
+                    if blob_value is not DELTA_SENTINEL:
+                        collected.reverse()
+                        return _ChannelWritesHistory(seed=blob_value, writes=collected)
+
+        collected.reverse()  # oldest → newest
+        return _ChannelWritesHistory(seed=DELTA_SENTINEL, writes=collected)
+
     def _dump_blobs(
         self,
         thread_id: str,

{langgraph_checkpoint_postgres-3.0.4 → langgraph_checkpoint_postgres-3.1.0a1}/langgraph/store/postgres/aio.py

@@ -164,12 +164,11 @@ class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Con
         grouped_ops, num_ops = _group_ops(ops)
         results: list[Result] = [None] * num_ops
 
-        async with _ainternal.get_connection(self.conn) as conn:
-            if self.pipe:
-                async with self.pipe:
-                    await self._execute_batch(grouped_ops, results, conn)
-            else:
-                await self._execute_batch(grouped_ops, results, conn)
+        if self.pipe:
+            async with self.pipe:
+                await self._execute_batch(grouped_ops, results)
+        else:
+            await self._execute_batch(grouped_ops, results)
 
         return results
 
@@ -410,8 +409,10 @@ class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Con
         self,
         grouped_ops: dict,
         results: list[Result],
-        conn: AsyncConnection[DictRow],
+        conn: AsyncConnection[DictRow] | None = None,
     ) -> None:
+        # Keep `conn` for compatibility with subclasses overriding this private hook.
+        # All database I/O goes through `_cursor()`, which owns connection acquisition.
         async with self._cursor(pipeline=True) as cur:
             if GetOp in grouped_ops:
                 await self._batch_get_ops(

{langgraph_checkpoint_postgres-3.0.4 → langgraph_checkpoint_postgres-3.1.0a1}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "langgraph-checkpoint-postgres"
-version = "3.0.4"
+version = "3.1.0a1"
 description = "Library with a Postgres implementation of LangGraph checkpoint saver."
 authors = []
 requires-python = ">=3.10"
@@ -12,15 +12,15 @@ readme = "README.md"
 license = "MIT"
 license-files = ['LICENSE']
 dependencies = [
-  "langgraph-checkpoint>=2.1.2,<5.0.0",
-  "orjson>=3.10.1",
+  "langgraph-checkpoint>=4.1.0a1,<5.0.0",
+  "orjson>=3.11.5",
   "psycopg>=3.2.0",
   "psycopg-pool>=3.2.0",
 ]
 
 [project.urls]
 Source = "https://github.com/langchain-ai/langgraph/tree/main/libs/checkpoint-postgres"
-Twitter = "https://x.com/LangChain"
+Twitter = "https://x.com/langchain_oss"
 Slack = "https://www.langchain.com/join-community"
 Reddit = "https://www.reddit.com/r/LangChain/"
 

{langgraph_checkpoint_postgres-3.0.4 → langgraph_checkpoint_postgres-3.1.0a1}/tests/test_async.py

@@ -361,9 +361,9 @@ async def test_get_checkpoint_no_channel_values(
 
         load_checkpoint_tuple = saver._load_checkpoint_tuple
 
-        def patched_load_checkpoint_tuple(value):
+        async def patched_load_checkpoint_tuple(value):
             value["checkpoint"].pop("channel_values", None)
-            return load_checkpoint_tuple(value)
+            return await load_checkpoint_tuple(value)
 
         monkeypatch.setattr(
             saver, "_load_checkpoint_tuple", patched_load_checkpoint_tuple
@@ -371,3 +371,47 @@ async def test_get_checkpoint_no_channel_values(
 
         checkpoint = await saver.aget_tuple(config)
         assert checkpoint.checkpoint["channel_values"] == {}
+
+
+@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe"])
+async def test_delta_channel_chain_reconstruction(saver_name: str) -> None:
+    """AsyncPostgresSaver reconstructs DeltaChannel chain via point-lookup traversal."""
+    pytest.importorskip(
+        "langgraph.channels.delta", reason="langgraph core not installed"
+    )
+
+    from typing import Annotated
+
+    from langchain_core.messages import AIMessage, HumanMessage
+    from langgraph.channels.delta import DeltaChannel
+    from langgraph.graph import START, StateGraph
+    from langgraph.graph.message import _messages_delta_reducer
+    from typing_extensions import TypedDict
+
+    class State(TypedDict):
+        messages: Annotated[list, DeltaChannel(_messages_delta_reducer)]
+
+    def respond(state: State) -> dict:
+        n = len(state["messages"])
+        return {"messages": [AIMessage(content=f"reply-{n}", id=f"ai-{n}")]}
+
+    builder = StateGraph(State)
+    builder.add_node("respond", respond)
+    builder.add_edge(START, "respond")
+
+    async with _saver(saver_name) as saver:
+        graph = builder.compile(checkpointer=saver)
+        config = {"configurable": {"thread_id": "diff-channel-test-1"}}
+
+        await graph.ainvoke({"messages": [HumanMessage(content="hi", id="h1")]}, config)
+        await graph.ainvoke(
+            {"messages": [HumanMessage(content="there", id="h2")]}, config
+        )
+
+        state = await graph.aget_state(config)
+        msgs = state.values["messages"]
+        assert len(msgs) == 4, f"expected 4, got {len(msgs)}: {msgs}"
+        assert msgs[0].content == "hi"
+        assert msgs[1].content == "reply-1"
+        assert msgs[2].content == "there"
+        assert msgs[3].content == "reply-3"

{langgraph_checkpoint_postgres-3.0.4 → langgraph_checkpoint_postgres-3.1.0a1}/tests/test_async_store.py

@@ -20,6 +20,7 @@ from langgraph.store.base import (
 )
 from psycopg import AsyncConnection
 
+from langgraph.checkpoint.postgres import _ainternal
 from langgraph.store.postgres import AsyncPostgresStore
 from tests.conftest import (
     DEFAULT_URI,
@@ -346,6 +347,59 @@ async def test_batch_list_namespaces_ops(store: AsyncPostgresStore) -> None:
     assert ("test", "namespace2") in results[0]
 
 
+@asynccontextmanager
+async def _create_pool_store() -> AsyncIterator[AsyncPostgresStore]:
+    database = f"test_{uuid.uuid4().hex[:16]}"
+    uri_parts = DEFAULT_URI.split("/")
+    uri_base = "/".join(uri_parts[:-1])
+    query_params = ""
+    if "?" in uri_parts[-1]:
+        _, query_params = uri_parts[-1].split("?", 1)
+        query_params = "?" + query_params
+
+    conn_string = f"{uri_base}/{database}{query_params}"
+    admin_conn_string = DEFAULT_URI
+    async with await AsyncConnection.connect(
+        admin_conn_string, autocommit=True
+    ) as conn:
+        await conn.execute(f"CREATE DATABASE {database}")
+    try:
+        async with AsyncPostgresStore.from_conn_string(
+            conn_string, pool_config={"min_size": 1, "max_size": 1}
+        ) as store:
+            await store.setup()
+            yield store
+    finally:
+        async with await AsyncConnection.connect(
+            admin_conn_string, autocommit=True
+        ) as conn:
+            await conn.execute(f"DROP DATABASE {database}")
+
+
+async def test_abatch_uses_single_pool_checkout(monkeypatch) -> None:
+    async with _create_pool_store() as store:
+        await store.aput(("test",), "key1", {"data": "value1"})
+
+        original_get_connection = _ainternal.get_connection
+        checkout_count = 0
+
+        @asynccontextmanager
+        async def counting_get_connection(conn):
+            nonlocal checkout_count
+            checkout_count += 1
+            async with original_get_connection(conn) as checked_out_conn:
+                yield checked_out_conn
+
+        monkeypatch.setattr(_ainternal, "get_connection", counting_get_connection)
+
+        results = await store.abatch([GetOp(namespace=("test",), key="key1")])
+
+        assert len(results) == 1
+        assert results[0] is not None
+        assert results[0].value == {"data": "value1"}
+        assert checkout_count == 1
+
+
 @asynccontextmanager
 async def _create_vector_store(
     vector_type: str,