coco-framework 0.1.0__py3-none-any.whl

coco_graph/checkpoint/postgres/aio.py

@@ -0,0 +1,691 @@
+from __future__ import annotations
+
+import asyncio
+from collections import defaultdict
+from collections.abc import AsyncIterator, Iterator, Mapping, Sequence
+from contextlib import asynccontextmanager
+from typing import Any, cast
+
+from coco_core.runnables import RunnableConfig
+from coco_graph.checkpoint.base import (
+    WRITES_IDX_MAP,
+    ChannelVersions,
+    Checkpoint,
+    CheckpointMetadata,
+    CheckpointTuple,
+    DeltaChannelHistory,
+    get_checkpoint_id,
+    get_serializable_checkpoint_metadata,
+)
+from coco_graph.checkpoint.serde.base import SerializerProtocol
+from coco_graph.checkpoint.serde.types import _DeltaSnapshot
+from psycopg import AsyncConnection, AsyncCursor, AsyncPipeline, Capabilities
+from psycopg.rows import DictRow, dict_row
+from psycopg.types.json import Jsonb
+from psycopg_pool import AsyncConnectionPool
+
+from coco_graph.checkpoint.postgres import _ainternal
+from coco_graph.checkpoint.postgres.base import (
+    _DELTA_PAGE_SIZE,
+    BasePostgresSaver,
+    _build_delta_stage1_sql,
+    _build_delta_stage2_sql,
+    _DeltaStage2Row,
+)
+from coco_graph.checkpoint.postgres.shallow import AsyncShallowPostgresSaver
+
+Conn = _ainternal.Conn  # For backward compatibility
+
+
+class AsyncPostgresSaver(BasePostgresSaver):
+    """Asynchronous checkpointer that stores checkpoints in a Postgres database."""
+
+    lock: asyncio.Lock
+
+    def __init__(
+        self,
+        conn: _ainternal.Conn,
+        pipe: AsyncPipeline | None = None,
+        serde: SerializerProtocol | None = None,
+    ) -> None:
+        super().__init__(serde=serde)
+        if isinstance(conn, AsyncConnectionPool) and pipe is not None:
+            raise ValueError(
+                "Pipeline should be used only with a single AsyncConnection, not AsyncConnectionPool."
+            )
+
+        self.conn = conn
+        self.pipe = pipe
+        self.lock = asyncio.Lock()
+        self.loop = asyncio.get_running_loop()
+        self.supports_pipeline = Capabilities().has_pipeline()
+
+    @classmethod
+    @asynccontextmanager
+    async def from_conn_string(
+        cls,
+        conn_string: str,
+        *,
+        pipeline: bool = False,
+        serde: SerializerProtocol | None = None,
+    ) -> AsyncIterator[AsyncPostgresSaver]:
+        """Create a new AsyncPostgresSaver instance from a connection string.
+
+        Args:
+            conn_string: The Postgres connection info string.
+            pipeline: whether to use AsyncPipeline
+
+        Returns:
+            AsyncPostgresSaver: A new AsyncPostgresSaver instance.
+        """
+        async with await AsyncConnection.connect(
+            conn_string, autocommit=True, prepare_threshold=0, row_factory=dict_row
+        ) as conn:
+            if pipeline:
+                async with conn.pipeline() as pipe:
+                    yield cls(conn=conn, pipe=pipe, serde=serde)
+            else:
+                yield cls(conn=conn, serde=serde)
+
+    async def setup(self) -> None:
+        """Set up the checkpoint database asynchronously.
+
+        This method creates the necessary tables in the Postgres database if they don't
+        already exist and runs database migrations. It MUST be called directly by the user
+        the first time checkpointer is used.
+        """
+        async with self._cursor() as cur:
+            await cur.execute(self.MIGRATIONS[0])
+            results = await cur.execute(
+                "SELECT v FROM checkpoint_migrations ORDER BY v DESC LIMIT 1"
+            )
+            row = await results.fetchone()
+            if row is None:
+                version = -1
+            else:
+                version = row["v"]
+            for v, migration in zip(
+                range(version + 1, len(self.MIGRATIONS)),
+                self.MIGRATIONS[version + 1 :],
+                strict=False,
+            ):
+                await cur.execute(migration)
+                await cur.execute(
+                    "INSERT INTO checkpoint_migrations (v) VALUES (%s)", (v,)
+                )
+        if self.pipe:
+            await self.pipe.sync()
+
+    async def alist(
+        self,
+        config: RunnableConfig | None,
+        *,
+        filter: dict[str, Any] | None = None,
+        before: RunnableConfig | None = None,
+        limit: int | None = None,
+    ) -> AsyncIterator[CheckpointTuple]:
+        """List checkpoints from the database asynchronously.
+
+        This method retrieves a list of checkpoint tuples from the Postgres database based
+        on the provided config. The checkpoints are ordered by checkpoint ID in descending order (newest first).
+
+        Args:
+            config: Base configuration for filtering checkpoints.
+            filter: Additional filtering criteria for metadata.
+            before: If provided, only checkpoints before the specified checkpoint ID are returned.
+            limit: Maximum number of checkpoints to return.
+
+        Yields:
+            An asynchronous iterator of matching checkpoint tuples.
+        """
+        where, args = self._search_where(config, filter, before)
+        query = self.SELECT_SQL + where + " ORDER BY checkpoint_id DESC"
+        params = list(args)
+        if limit is not None:
+            query += " LIMIT %s"
+            params.append(int(limit))
+        # if we change this to use .stream() we need to make sure to close the cursor
+        async with self._cursor() as cur:
+            await cur.execute(query, params, binary=True)
+            values = await cur.fetchall()
+            if not values:
+                return
+            # migrate pending sends if necessary
+            if to_migrate := [
+                v
+                for v in values
+                if v["checkpoint"]["v"] < 4 and v["parent_checkpoint_id"]
+            ]:
+                await cur.execute(
+                    self.SELECT_PENDING_SENDS_SQL,
+                    (
+                        values[0]["thread_id"],
+                        [v["parent_checkpoint_id"] for v in to_migrate],
+                    ),
+                )
+                grouped_by_parent = defaultdict(list)
+                for value in to_migrate:
+                    grouped_by_parent[value["parent_checkpoint_id"]].append(value)
+                async for sends in cur:
+                    for value in grouped_by_parent[sends["checkpoint_id"]]:
+                        if value["channel_values"] is None:
+                            value["channel_values"] = []
+                        self._migrate_pending_sends(
+                            sends["sends"],
+                            value["checkpoint"],
+                            value["channel_values"],
+                        )
+            for value in values:
+                yield await self._load_checkpoint_tuple(value)
+
+    async def aget_tuple(self, config: RunnableConfig) -> CheckpointTuple | None:
+        """Get a checkpoint tuple from the database asynchronously.
+
+        This method retrieves a checkpoint tuple from the Postgres database based on the
+        provided config. If the config contains a `checkpoint_id` key, the checkpoint with
+        the matching thread ID and "checkpoint_id" is retrieved. Otherwise, the latest checkpoint
+        for the given thread ID is retrieved.
+
+        Args:
+            config: The config to use for retrieving the checkpoint.
+
+        Returns:
+            The retrieved checkpoint tuple, or None if no matching checkpoint was found.
+        """
+        thread_id = config["configurable"]["thread_id"]
+        checkpoint_id = get_checkpoint_id(config)
+        checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
+        if checkpoint_id:
+            args: tuple[Any, ...] = (thread_id, checkpoint_ns, checkpoint_id)
+            where = "WHERE thread_id = %s AND checkpoint_ns = %s AND checkpoint_id = %s"
+        else:
+            args = (thread_id, checkpoint_ns)
+            where = "WHERE thread_id = %s AND checkpoint_ns = %s ORDER BY checkpoint_id DESC LIMIT 1"
+
+        async with self._cursor() as cur:
+            await cur.execute(
+                self.SELECT_SQL + where,
+                args,
+                binary=True,
+            )
+            value = await cur.fetchone()
+            if value is None:
+                return None
+
+            # migrate pending sends if necessary
+            if value["checkpoint"]["v"] < 4 and value["parent_checkpoint_id"]:
+                await cur.execute(
+                    self.SELECT_PENDING_SENDS_SQL,
+                    (thread_id, [value["parent_checkpoint_id"]]),
+                )
+                if sends := await cur.fetchone():
+                    if value["channel_values"] is None:
+                        value["channel_values"] = []
+                    self._migrate_pending_sends(
+                        sends["sends"],
+                        value["checkpoint"],
+                        value["channel_values"],
+                    )
+
+            return await self._load_checkpoint_tuple(value)
+
+    async def aput(
+        self,
+        config: RunnableConfig,
+        checkpoint: Checkpoint,
+        metadata: CheckpointMetadata,
+        new_versions: ChannelVersions,
+    ) -> RunnableConfig:
+        """Save a checkpoint to the database asynchronously.
+
+        This method saves a checkpoint to the Postgres database. The checkpoint is associated
+        with the provided config and its parent config (if any).
+
+        Args:
+            config: The config to associate with the checkpoint.
+            checkpoint: The checkpoint to save.
+            metadata: Additional metadata to save with the checkpoint.
+            new_versions: New channel versions as of this write.
+
+        Returns:
+            RunnableConfig: Updated configuration after storing the checkpoint.
+        """
+        configurable = config["configurable"].copy()
+        thread_id = configurable.pop("thread_id")
+        checkpoint_ns = configurable.pop("checkpoint_ns")
+        checkpoint_id = configurable.pop("checkpoint_id", None)
+
+        copy = checkpoint.copy()
+        copy["channel_values"] = copy["channel_values"].copy()
+        next_config = {
+            "configurable": {
+                "thread_id": thread_id,
+                "checkpoint_ns": checkpoint_ns,
+                "checkpoint_id": checkpoint["id"],
+            }
+        }
+
+        # inline primitive values in checkpoint table
+        # others are stored in blobs table
+        blob_values = {}
+        for k, v in checkpoint["channel_values"].items():
+            if isinstance(v, _DeltaSnapshot):
+                blob_values[k] = copy["channel_values"].pop(k)
+                copy["channel_values"][k] = True
+            elif v is None or isinstance(v, (str, int, float, bool)):
+                pass
+            else:
+                blob_values[k] = copy["channel_values"].pop(k)
+
+        async with self._cursor(pipeline=True) as cur:
+            if blob_versions := {
+                k: v for k, v in new_versions.items() if k in blob_values
+            }:
+                await cur.executemany(
+                    self.UPSERT_CHECKPOINT_BLOBS_SQL,
+                    await asyncio.to_thread(
+                        self._dump_blobs,
+                        thread_id,
+                        checkpoint_ns,
+                        blob_values,
+                        blob_versions,
+                    ),
+                )
+            await cur.execute(
+                self.UPSERT_CHECKPOINTS_SQL,
+                (
+                    thread_id,
+                    checkpoint_ns,
+                    checkpoint["id"],
+                    checkpoint_id,
+                    Jsonb(copy),
+                    Jsonb(get_serializable_checkpoint_metadata(config, metadata)),
+                ),
+            )
+        return next_config
+
+    async def aput_writes(
+        self,
+        config: RunnableConfig,
+        writes: Sequence[tuple[str, Any]],
+        task_id: str,
+        task_path: str = "",
+    ) -> None:
+        """Store intermediate writes linked to a checkpoint asynchronously.
+
+        This method saves intermediate writes associated with a checkpoint to the database.
+
+        Args:
+            config: Configuration of the related checkpoint.
+            writes: List of writes to store, each as (channel, value) pair.
+            task_id: Identifier for the task creating the writes.
+        """
+        query = (
+            self.UPSERT_CHECKPOINT_WRITES_SQL
+            if all(w[0] in WRITES_IDX_MAP for w in writes)
+            else self.INSERT_CHECKPOINT_WRITES_SQL
+        )
+        params = await asyncio.to_thread(
+            self._dump_writes,
+            config["configurable"]["thread_id"],
+            config["configurable"]["checkpoint_ns"],
+            config["configurable"]["checkpoint_id"],
+            task_id,
+            task_path,
+            writes,
+        )
+        async with self._cursor(pipeline=True) as cur:
+            await cur.executemany(query, params)
+
+    async def adelete_thread(self, thread_id: str) -> None:
+        """Delete all checkpoints and writes associated with a thread ID.
+
+        Args:
+            thread_id: The thread ID to delete.
+
+        Returns:
+            None
+        """
+        async with self._cursor(pipeline=True) as cur:
+            await cur.execute(
+                "DELETE FROM checkpoints WHERE thread_id = %s",
+                (str(thread_id),),
+            )
+            await cur.execute(
+                "DELETE FROM checkpoint_blobs WHERE thread_id = %s",
+                (str(thread_id),),
+            )
+            await cur.execute(
+                "DELETE FROM checkpoint_writes WHERE thread_id = %s",
+                (str(thread_id),),
+            )
+
+    @asynccontextmanager
+    async def _cursor(
+        self, *, pipeline: bool = False
+    ) -> AsyncIterator[AsyncCursor[DictRow]]:
+        """Create a database cursor as a context manager.
+
+        Args:
+            pipeline: whether to use pipeline for the DB operations inside the context manager.
+                Will be applied regardless of whether the AsyncPostgresSaver instance was initialized with a pipeline.
+                If pipeline mode is not supported, will fall back to using transaction context manager.
+        """
+        async with self.lock, _ainternal.get_connection(self.conn) as conn:
+            if self.pipe:
+                # a connection in pipeline mode can be used concurrently
+                # in multiple threads/coroutines, but only one cursor can be
+                # used at a time
+                try:
+                    async with conn.cursor(binary=True, row_factory=dict_row) as cur:
+                        yield cur
+                finally:
+                    if pipeline:
+                        await self.pipe.sync()
+            elif pipeline:
+                # a connection not in pipeline mode can only be used by one
+                # thread/coroutine at a time, so we acquire a lock
+                if self.supports_pipeline:
+                    async with (
+                        conn.pipeline(),
+                        conn.cursor(binary=True, row_factory=dict_row) as cur,
+                    ):
+                        yield cur
+                else:
+                    # Use connection's transaction context manager when pipeline mode not supported
+                    async with (
+                        conn.transaction(),
+                        conn.cursor(binary=True, row_factory=dict_row) as cur,
+                    ):
+                        yield cur
+            else:
+                async with conn.cursor(binary=True, row_factory=dict_row) as cur:
+                    yield cur
+
+    async def aget_delta_channel_history(
+        self, *, config: RunnableConfig, channels: Sequence[str]
+    ) -> Mapping[str, DeltaChannelHistory]:
+        """Fast-path override of `BaseCheckpointSaver.aget_delta_channel_history`.
+
+        See `PostgresSaver.get_delta_channel_history` for design notes; this is
+        the async equivalent with internal stage-1 paging and per-channel
+        UNION ALL stage-2.
+        """
+        if not channels:
+            return {}
+        channels = list(channels)
+        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 {ch: {"writes": []} for ch in channels}
+            checkpoint_id = target.config["configurable"]["checkpoint_id"]
+
+        stage1_sql = _build_delta_stage1_sql(channels, paged=True)
+        parent_of: dict[str, str | None] = {}
+        ver_by_i_by_cid: list[dict[str, str | None]] = [{} for _ in channels]
+        hb_by_i_by_cid: list[dict[str, bool]] = [{} for _ in channels]
+        inline_by_i_by_cid: list[dict[str, Any]] = [{} for _ in channels]
+        chain_by_ch: dict[str, list[str]] = {ch: [] for ch in channels}
+        seed_ver_by_ch: dict[str, str | None] = {ch: None for ch in channels}
+        seed_inline_by_ch: dict[str, Any] = {}
+        walk_cursor_by_ch: dict[str, str | None] = {}
+        seeded: set[str] = set()
+        cursor: str | None = None
+
+        async with self._cursor() as cur:
+            while True:
+                stage1_params: list[Any] = []
+                for ch in channels:
+                    # ver_i, blob channel, blob version, inline_i
+                    stage1_params.extend([ch, ch, ch, ch])
+                stage1_params.extend(
+                    [thread_id, checkpoint_ns, cursor, cursor, _DELTA_PAGE_SIZE]
+                )
+                await cur.execute(stage1_sql, stage1_params)
+                page = await cur.fetchall()
+                if not page:
+                    break
+                oldest = self._ingest_stage1_page(
+                    cast("list[Mapping[str, Any]]", page),
+                    channels,
+                    parent_of,
+                    ver_by_i_by_cid,
+                    hb_by_i_by_cid,
+                    inline_by_i_by_cid,
+                )
+                self._try_advance_walks(
+                    checkpoint_id,
+                    channels,
+                    parent_of,
+                    ver_by_i_by_cid,
+                    hb_by_i_by_cid,
+                    inline_by_i_by_cid,
+                    chain_by_ch,
+                    seed_ver_by_ch,
+                    seed_inline_by_ch,
+                    walk_cursor_by_ch,
+                    seeded,
+                )
+                if len(seeded) == len(channels) or len(page) < _DELTA_PAGE_SIZE:
+                    break
+                cursor = oldest
+
+        channels_with_chain = [ch for ch in channels if chain_by_ch[ch]]
+        channels_with_seed = [ch for ch in channels if seed_ver_by_ch[ch] is not None]
+        stage2_sql = _build_delta_stage2_sql(
+            channels_with_chain=channels_with_chain,
+            channels_with_seed=channels_with_seed,
+        )
+
+        if stage2_sql:
+            stage2_params: list[Any] = []
+            for ch in channels_with_chain:
+                stage2_params.extend([thread_id, checkpoint_ns, ch, chain_by_ch[ch]])
+            for ch in channels_with_seed:
+                stage2_params.extend([thread_id, checkpoint_ns, ch, seed_ver_by_ch[ch]])
+            async with self._cursor() as cur:
+                await cur.execute(stage2_sql, stage2_params)
+                stage2_rows = await cur.fetchall()
+        else:
+            stage2_rows = []
+
+        return self._build_delta_channels_writes_history(
+            channels=channels,
+            chain_by_ch=chain_by_ch,
+            seed_ver_by_ch=seed_ver_by_ch,
+            seed_inline_by_ch=seed_inline_by_ch,
+            stage2_rows=cast("list[_DeltaStage2Row]", stage2_rows),
+        )
+
+    async def _load_checkpoint_tuple(self, value: DictRow) -> CheckpointTuple:
+        """
+        Convert a database row into a CheckpointTuple object.
+
+        Args:
+            value: A row from the database containing checkpoint data.
+
+        Returns:
+            CheckpointTuple: A structured representation of the checkpoint,
+            including its configuration, metadata, parent checkpoint (if any),
+            and pending writes.
+        """
+        return CheckpointTuple(
+            {
+                "configurable": {
+                    "thread_id": value["thread_id"],
+                    "checkpoint_ns": value["checkpoint_ns"],
+                    "checkpoint_id": value["checkpoint_id"],
+                }
+            },
+            {
+                **value["checkpoint"],
+                "channel_values": {
+                    **(value["checkpoint"].get("channel_values") or {}),
+                    **self._load_blobs(value["channel_values"]),
+                },
+            },
+            value["metadata"],
+            (
+                {
+                    "configurable": {
+                        "thread_id": value["thread_id"],
+                        "checkpoint_ns": value["checkpoint_ns"],
+                        "checkpoint_id": value["parent_checkpoint_id"],
+                    }
+                }
+                if value["parent_checkpoint_id"]
+                else None
+            ),
+            await asyncio.to_thread(self._load_writes, value["pending_writes"]),
+        )
+
+    def list(
+        self,
+        config: RunnableConfig | None,
+        *,
+        filter: dict[str, Any] | None = None,
+        before: RunnableConfig | None = None,
+        limit: int | None = None,
+    ) -> Iterator[CheckpointTuple]:
+        """List checkpoints from the database.
+
+        This method retrieves a list of checkpoint tuples from the Postgres database based
+        on the provided config. The checkpoints are ordered by checkpoint ID in descending order (newest first).
+
+        Args:
+            config: Base configuration for filtering checkpoints.
+            filter: Additional filtering criteria for metadata.
+            before: If provided, only checkpoints before the specified checkpoint ID are returned.
+            limit: Maximum number of checkpoints to return.
+
+        Yields:
+            An iterator of matching checkpoint tuples.
+        """
+        try:
+            # check if we are in the main thread, only bg threads can block
+            # we don't check in other methods to avoid the overhead
+            if asyncio.get_running_loop() is self.loop:
+                raise asyncio.InvalidStateError(
+                    "Synchronous calls to AsyncPostgresSaver are only allowed from a "
+                    "different thread. From the main thread, use the async interface. "
+                    "For example, use `checkpointer.alist(...)` or `await "
+                    "graph.ainvoke(...)`."
+                )
+        except RuntimeError:
+            pass
+        aiter_ = self.alist(config, filter=filter, before=before, limit=limit)
+        while True:
+            try:
+                yield asyncio.run_coroutine_threadsafe(
+                    anext(aiter_),  # type: ignore[arg-type]
+                    self.loop,
+                ).result()
+            except StopAsyncIteration:
+                break
+
+    def get_tuple(self, config: RunnableConfig) -> CheckpointTuple | None:
+        """Get a checkpoint tuple from the database.
+
+        This method retrieves a checkpoint tuple from the Postgres database based on the
+        provided config. If the config contains a `checkpoint_id` key, the checkpoint with
+        the matching thread ID and "checkpoint_id" is retrieved. Otherwise, the latest checkpoint
+        for the given thread ID is retrieved.
+
+        Args:
+            config: The config to use for retrieving the checkpoint.
+
+        Returns:
+            The retrieved checkpoint tuple, or None if no matching checkpoint was found.
+        """
+        try:
+            # check if we are in the main thread, only bg threads can block
+            # we don't check in other methods to avoid the overhead
+            if asyncio.get_running_loop() is self.loop:
+                raise asyncio.InvalidStateError(
+                    "Synchronous calls to AsyncPostgresSaver are only allowed from a "
+                    "different thread. From the main thread, use the async interface. "
+                    "For example, use `await checkpointer.aget_tuple(...)` or `await "
+                    "graph.ainvoke(...)`."
+                )
+        except RuntimeError:
+            pass
+        return asyncio.run_coroutine_threadsafe(
+            self.aget_tuple(config), self.loop
+        ).result()
+
+    def put(
+        self,
+        config: RunnableConfig,
+        checkpoint: Checkpoint,
+        metadata: CheckpointMetadata,
+        new_versions: ChannelVersions,
+    ) -> RunnableConfig:
+        """Save a checkpoint to the database.
+
+        This method saves a checkpoint to the Postgres database. The checkpoint is associated
+        with the provided config and its parent config (if any).
+
+        Args:
+            config: The config to associate with the checkpoint.
+            checkpoint: The checkpoint to save.
+            metadata: Additional metadata to save with the checkpoint.
+            new_versions: New channel versions as of this write.
+
+        Returns:
+            RunnableConfig: Updated configuration after storing the checkpoint.
+        """
+        return asyncio.run_coroutine_threadsafe(
+            self.aput(config, checkpoint, metadata, new_versions), self.loop
+        ).result()
+
+    def put_writes(
+        self,
+        config: RunnableConfig,
+        writes: Sequence[tuple[str, Any]],
+        task_id: str,
+        task_path: str = "",
+    ) -> None:
+        """Store intermediate writes linked to a checkpoint.
+
+        This method saves intermediate writes associated with a checkpoint to the database.
+
+        Args:
+            config: Configuration of the related checkpoint.
+            writes: List of writes to store, each as (channel, value) pair.
+            task_id: Identifier for the task creating the writes.
+            task_path: Path of the task creating the writes.
+        """
+        return asyncio.run_coroutine_threadsafe(
+            self.aput_writes(config, writes, task_id, task_path), self.loop
+        ).result()
+
+    def delete_thread(self, thread_id: str) -> None:
+        """Delete all checkpoints and writes associated with a thread ID.
+
+        Args:
+            thread_id: The thread ID to delete.
+
+        Returns:
+            None
+        """
+        try:
+            # check if we are in the main thread, only bg threads can block
+            # we don't check in other methods to avoid the overhead
+            if asyncio.get_running_loop() is self.loop:
+                raise asyncio.InvalidStateError(
+                    "Synchronous calls to AsyncPostgresSaver are only allowed from a "
+                    "different thread. From the main thread, use the async interface. "
+                    "For example, use `await checkpointer.aget_tuple(...)` or `await "
+                    "graph.ainvoke(...)`."
+                )
+        except RuntimeError:
+            pass
+        return asyncio.run_coroutine_threadsafe(
+            self.adelete_thread(thread_id), self.loop
+        ).result()
+
+
+__all__ = ["AsyncPostgresSaver", "AsyncShallowPostgresSaver", "Conn"]