langgraph-checkpoint-oceanbase-es 3.0.0__py3-none-any.whl

langgraph/checkpoint/oceanbase/__init__.py

@@ -0,0 +1,26 @@
+"""OceanBase checkpoint saver - MySQL compatible with OceanBase naming."""
+from __future__ import annotations
+
+from langgraph.checkpoint.oceanbase.pymysql import (
+    Conn,
+    OceanBaseSaver,
+    ShallowOceanBaseSaver,
+)
+from langgraph.checkpoint.oceanbase.aio import (
+    AsyncOceanBaseSaver,
+    ShallowAsyncOceanBaseSaver,
+)
+from langgraph.checkpoint.oceanbase.asyncmy import (
+    AsyncMyOceanBaseSaver,
+    ShallowAsyncMyOceanBaseSaver,
+)
+
+__all__ = [
+    "OceanBaseSaver",
+    "AsyncOceanBaseSaver",
+    "AsyncMyOceanBaseSaver",
+    "ShallowOceanBaseSaver",
+    "ShallowAsyncOceanBaseSaver",
+    "ShallowAsyncMyOceanBaseSaver",
+    "Conn",
+]

langgraph/checkpoint/oceanbase/_ainternal.py

@@ -0,0 +1,91 @@
+"""Shared utility functions for asynchronous OceanBase connections."""
+from __future__ import annotations
+
+from collections.abc import AsyncIterator, Mapping, Sequence
+from contextlib import AbstractAsyncContextManager, asynccontextmanager
+from typing import (
+    Any,
+    Generic,
+    Protocol,
+    TypeVar,
+    Union,
+    cast,
+)
+
+
+class AsyncDictCursor(AbstractAsyncContextManager, Protocol):
+    """Protocol that an asynchronous OceanBase/MySQL cursor should implement."""
+
+    async def execute(
+        self,
+        operation: str,
+        parameters: Sequence[Any] | Mapping[str, Any] = ...,
+        /,
+    ) -> object: ...
+    async def executemany(
+        self, operation: str, seq_of_parameters: Sequence[Sequence[Any]], /
+    ) -> object: ...
+    async def fetchone(self) -> dict[str, Any] | None: ...
+    async def fetchall(self) -> Sequence[dict[str, Any]]: ...
+
+    def __aiter__(self) -> AsyncIterator[dict[str, Any]]: ...
+
+
+R = TypeVar("R", bound=AsyncDictCursor)  # cursor type
+
+
+class AsyncConnection(AbstractAsyncContextManager, Protocol):
+    """Protocol that an asynchronous OceanBase/MySQL connection should implement."""
+
+    async def begin(self) -> None:
+        """Begin transaction."""
+        ...
+
+    async def commit(self) -> None:
+        """Commit changes to stable storage."""
+        ...
+
+    async def rollback(self) -> None:
+        """Roll back the current transaction."""
+        ...
+
+    async def set_charset(self, charset: str) -> None:
+        """Sets the character set for the current connection"""
+        ...
+
+
+C = TypeVar("C", bound=AsyncConnection)  # connection type
+COut = TypeVar("COut", bound=AsyncConnection, covariant=True)  # connection type
+
+
+class AsyncPool(Protocol, Generic[COut]):
+    def acquire(self) -> COut:
+        """Gets a connection from the connection pool."""
+        ...
+
+
+Conn = Union[C, AsyncPool[C]]
+
+
+@asynccontextmanager
+async def get_connection(
+    conn: Conn[C],
+) -> AsyncIterator[C]:
+    """Get a connection from a connection object or pool.
+    
+    Args:
+        conn: Either a connection object or an async pool
+    
+    Yields:
+        The connection object
+    """
+    if hasattr(conn, "cursor"):
+        yield cast(C, conn)
+    elif hasattr(conn, "acquire"):
+        async with cast(AsyncPool[C], conn).acquire() as _conn:
+            yield _conn
+    else:
+        raise TypeError(f"Invalid connection type: {type(conn)}")
+
+
+__all__ = ["get_connection", "Conn", "AsyncConnection", "AsyncDictCursor", "AsyncPool"]

langgraph/checkpoint/oceanbase/_internal.py

@@ -0,0 +1,87 @@
+"""Shared utility functions for synchronous OceanBase connections."""
+from __future__ import annotations
+
+from collections.abc import Callable, Iterator, Mapping, Sequence
+from contextlib import AbstractContextManager, closing, contextmanager
+from typing import (
+    Any,
+    Protocol,
+    TypeVar,
+    Union,
+    cast,
+)
+
+
+class DictCursor(AbstractContextManager, Protocol):
+    """Protocol that a synchronous OceanBase/MySQL cursor should implement."""
+
+    def execute(
+        self,
+        operation: str,
+        parameters: Sequence[Any] | Mapping[str, Any] = ...,
+        /,
+    ) -> object: ...
+    def executemany(
+        self, operation: str, seq_of_parameters: Sequence[Sequence[Any]], /
+    ) -> object: ...
+    def fetchone(self) -> dict[str, Any] | None: ...
+    def fetchall(self) -> Sequence[dict[str, Any]]: ...
+
+
+R = TypeVar("R", bound=DictCursor)  # cursor type
+
+
+class Connection(AbstractContextManager, Protocol):
+    """Protocol that a synchronous OceanBase/MySQL connection should implement."""
+
+    def begin(self) -> None:
+        """Begin transaction."""
+        ...
+
+    def commit(self) -> None:
+        """Commit changes to stable storage."""
+        ...
+
+    def rollback(self) -> None:
+        """Roll back the current transaction."""
+        ...
+
+
+C = TypeVar("C", bound=Connection)  # connection type
+
+
+ConnectionFactory = Callable[[], Any]
+Conn = Union[C, ConnectionFactory]
+
+
+@contextmanager
+def get_connection(conn: Conn[C]) -> Iterator[C]:
+    """Get a connection from a connection object or factory.
+    
+    Args:
+        conn: Either a connection object, a callable that returns a connection,
+              or an object with a connect() method (like SQLAlchemy pool)
+    
+    Yields:
+        The connection object
+    """
+    if hasattr(conn, "cursor"):
+        yield cast(C, conn)
+    elif callable(conn):
+        _conn = conn()
+        if isinstance(_conn, AbstractContextManager):
+            yield cast(C, _conn)
+        else:
+            with closing(_conn) as __conn:
+                yield __conn
+    # This is kept for backwards compatibility
+    elif hasattr(conn, "connect"):
+        # sqlalchemy pool
+        factory: ConnectionFactory = getattr(conn, "connect")  # noqa: B009
+        with get_connection(factory) as _conn:
+            yield _conn
+    else:
+        raise TypeError(f"Invalid connection or pool type: {type(conn)}")
+
+
+__all__ = ["get_connection", "Conn", "Connection", "DictCursor"]

langgraph/checkpoint/oceanbase/aio.py

@@ -0,0 +1,105 @@
+"""OceanBase checkpoint saver using aiomysql (asynchronous)."""
+from __future__ import annotations
+
+import warnings
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
+from typing import Any, cast
+
+import aiomysql  # type: ignore
+from typing_extensions import Self, override
+
+from langgraph.checkpoint.oceanbase.aio_base import BaseAsyncOceanBaseSaver
+from langgraph.checkpoint.oceanbase.shallow import BaseShallowSyncOceanBaseSaver
+from langgraph.checkpoint.oceanbase.utils import parse_conn_string_async
+from langgraph.checkpoint.serde.base import SerializerProtocol
+
+# Type alias for backward compatibility
+Conn = aiomysql.Connection
+
+
+class AsyncOceanBaseSaver(BaseAsyncOceanBaseSaver[aiomysql.Connection, aiomysql.DictCursor]):
+    """OceanBase checkpoint saver using aiomysql (async).
+    
+    OceanBase is MySQL-compatible, so this inherits from BaseAsyncOceanBaseSaver.
+    Default port is 2881 (OceanBase) instead of 3306 (MySQL).
+    """
+
+    @staticmethod
+    def parse_conn_string(conn_string: str) -> dict[str, Any]:
+        """Parse connection string with OceanBase default port (2881)."""
+        return parse_conn_string_async(conn_string, default_port=2881)
+
+    @classmethod
+    @asynccontextmanager
+    async def from_conn_string(
+        cls,
+        conn_string: str,
+        *,
+        serde: SerializerProtocol | None = None,
+    ) -> AsyncIterator[Self]:
+        """Create a new AsyncOceanBaseSaver instance from a connection string.
+
+        Args:
+            conn_string: The OceanBase/MySQL connection info string.
+            serde: Optional serializer protocol.
+
+        Returns:
+            AsyncOceanBaseSaver: A new AsyncOceanBaseSaver instance.
+
+        Example:
+            conn_string=mysql+aiomysql://user:***@localhost:2881/db?unix_socket=/path/to/socket
+        """
+        async with aiomysql.connect(
+            **cls.parse_conn_string(conn_string),
+            autocommit=True,
+        ) as conn:
+            yield cls(conn=conn, serde=serde)
+
+    @override
+    @staticmethod
+    def _get_cursor_from_connection(conn: aiomysql.Connection) -> aiomysql.DictCursor:
+        return cast(aiomysql.DictCursor, conn.cursor(aiomysql.DictCursor))
+
+
+class ShallowAsyncOceanBaseSaver(BaseShallowSyncOceanBaseSaver):
+    """Deprecated shallow async OceanBase saver.
+    
+    Use AsyncOceanBaseSaver instead with `await graph.ainvoke(..., durability='exit')`.
+    """
+    
+    def __init__(
+        self,
+        conn: aiomysql.Connection,
+        serde: SerializerProtocol | None = None,
+    ) -> None:
+        warnings.warn(
+            "ShallowAsyncOceanBaseSaver is deprecated as of version 2.0.15 and will be removed in 3.0.0. "
+            "Use AsyncOceanBaseSaver instead, and invoke the graph with `await graph.ainvoke(..., durability='exit')`.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        super().__init__(conn, serde=serde)
+
+    @classmethod
+    @asynccontextmanager
+    async def from_conn_string(
+        cls,
+        conn_string: str,
+        *,
+        serde: SerializerProtocol | None = None,
+    ) -> AsyncIterator[Self]:
+        """Create a new ShallowAsyncOceanBaseSaver instance from a connection string."""
+        async with aiomysql.connect(
+            **AsyncOceanBaseSaver.parse_conn_string(conn_string),
+            autocommit=True,
+        ) as conn:
+            yield cls(conn)
+
+    @override
+    @staticmethod
+    def _get_cursor_from_connection(conn: aiomysql.Connection) -> aiomysql.DictCursor:
+        return cast(aiomysql.DictCursor, conn.cursor(aiomysql.DictCursor))
+
+
+__all__ = ["AsyncOceanBaseSaver", "ShallowAsyncOceanBaseSaver", "Conn"]

langgraph/checkpoint/oceanbase/aio_base.py

@@ -0,0 +1,382 @@
+"""OceanBase async checkpoint saver - async base class."""
+from __future__ import annotations
+
+import asyncio
+import json
+from collections import defaultdict
+from collections.abc import AsyncIterator, Iterator, Sequence
+from contextlib import asynccontextmanager
+from typing import Any, Generic
+
+from langchain_core.runnables import RunnableConfig
+
+from langgraph.checkpoint.base import (
+    WRITES_IDX_MAP,
+    ChannelVersions,
+    Checkpoint,
+    CheckpointMetadata,
+    CheckpointTuple,
+    get_checkpoint_id,
+    get_serializable_checkpoint_metadata,
+)
+from langgraph.checkpoint.oceanbase import _ainternal
+from langgraph.checkpoint.oceanbase.base import BaseOceanBaseSaver
+from langgraph.checkpoint.oceanbase.utils import (
+    deserialize_channel_values,
+    deserialize_pending_sends,
+    deserialize_pending_writes,
+)
+from langgraph.checkpoint.serde.base import SerializerProtocol
+
+
+class BaseAsyncOceanBaseSaver(BaseOceanBaseSaver, Generic[_ainternal.C, _ainternal.R]):
+    """Base async OceanBase checkpoint saver."""
+    
+    lock: asyncio.Lock
+
+    def __init__(
+        self,
+        conn: _ainternal.Conn[_ainternal.C],
+        serde: SerializerProtocol | None = None,
+    ) -> None:
+        super().__init__(serde=serde)
+        self.conn = conn
+        self.lock = asyncio.Lock()
+        self.loop = asyncio.get_running_loop()
+
+    @staticmethod
+    def _get_cursor_from_connection(conn: _ainternal.C) -> _ainternal.R:
+        raise NotImplementedError
+
+    async def setup(self) -> None:
+        """Set up the checkpoint database asynchronously."""
+        async with self._cursor() as cur:
+            await cur.execute(self.MIGRATIONS[0])
+            await cur.execute(
+                "SELECT v FROM checkpoint_migrations ORDER BY v DESC LIMIT 1"
+            )
+            row = await cur.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 :],
+            ):
+                await cur.execute(migration)
+                await cur.execute(
+                    "INSERT INTO checkpoint_migrations (v) VALUES (%s)", (v,)
+                )
+
+    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."""
+        where, args = self._search_where(config, filter, before)
+        query = self._select_sql(where) + " ORDER BY checkpoint_id DESC"
+        if limit is not None:
+            query += " LIMIT %(limit)s"
+            args = {**args, "limit": int(limit)}
+        
+        async with self._cursor() as cur:
+            await cur.execute(query, args)
+            values = await cur.fetchall()
+            if not values:
+                return
+            for value in values:
+                value["checkpoint"] = json.loads(value["checkpoint"])
+                value["channel_values"] = deserialize_channel_values(
+                    value["channel_values"]
+                )
+            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(len(to_migrate)),
+                    (
+                        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(
+                            deserialize_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."""
+        thread_id = config["configurable"]["thread_id"]
+        checkpoint_id = get_checkpoint_id(config)
+        checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
+        if checkpoint_id:
+            args = {
+                "thread_id": thread_id,
+                "checkpoint_ns": checkpoint_ns,
+                "checkpoint_id": checkpoint_id,
+            }
+            where = "WHERE thread_id = %(thread_id)s AND checkpoint_ns_hash = UNHEX(MD5(%(checkpoint_ns)s)) AND checkpoint_id = %(checkpoint_id)s"
+        else:
+            args = {
+                "thread_id": thread_id,
+                "checkpoint_ns": checkpoint_ns,
+            }
+            where = "WHERE thread_id = %(thread_id)s AND checkpoint_ns_hash = UNHEX(MD5(%(checkpoint_ns)s))"
+
+        query = self._select_sql(where)
+        if not checkpoint_id:
+            query += " ORDER BY checkpoint_id DESC LIMIT 1"
+        async with self._cursor() as cur:
+            await cur.execute(query, args)
+            value = await cur.fetchone()
+            if value is None:
+                return None
+
+            value["checkpoint"] = json.loads(value["checkpoint"])
+            value["channel_values"] = deserialize_channel_values(
+                value["channel_values"]
+            )
+
+            if value["checkpoint"]["v"] < 4 and value["parent_checkpoint_id"]:
+                await cur.execute(
+                    self._select_pending_sends_sql(1),
+                    (thread_id, value["parent_checkpoint_id"]),
+                )
+                if sends := await cur.fetchone():
+                    if value["channel_values"] is None:
+                        value["channel_values"] = []
+                    self._migrate_pending_sends(
+                        deserialize_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."""
+        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"],
+            }
+        }
+
+        blob_values = {}
+        for k, v in checkpoint["channel_values"].items():
+            if 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_ns,
+                    checkpoint["id"],
+                    checkpoint_id,
+                    json.dumps(copy),
+                    json.dumps(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."""
+        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."""
+        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[_ainternal.R]:
+        """Create a database cursor as a context manager."""
+        async with _ainternal.get_connection(self.conn) as conn:
+            if pipeline:
+                async with self.lock:
+                    await conn.begin()
+                    try:
+                        async with self._get_cursor_from_connection(conn) as cur:
+                            yield cur
+                        await conn.commit()
+                    except:
+                        await conn.rollback()
+                        raise
+            else:
+                async with (
+                    self.lock,
+                    self._get_cursor_from_connection(conn) as cur,
+                ):
+                    yield cur
+
+    async def _load_checkpoint_tuple(self, value: dict[str, Any]) -> CheckpointTuple:
+        """Convert a database row into a CheckpointTuple object."""
+        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"),
+                    **self._load_blobs(value["channel_values"]),
+                },
+            },
+            json.loads(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,
+                deserialize_pending_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 (sync wrapper)."""
+        try:
+            if asyncio.get_running_loop() is self.loop:
+                raise asyncio.InvalidStateError(
+                    "Synchronous calls to AsyncOceanBaseSaver are only allowed from a "
+                    "different thread. From the main thread, use the async interface."
+                )
+        except RuntimeError:
+            pass
+        aiter_ = self.alist(config, filter=filter, before=before, limit=limit)
+        while True:
+            try:
+                yield asyncio.run_coroutine_threadsafe(
+                    anext(aiter_),
+                    self.loop,
+                ).result()
+            except StopAsyncIteration:
+                break
+
+    def get_tuple(self, config: RunnableConfig) -> CheckpointTuple | None:
+        """Get a checkpoint tuple from the database (sync wrapper)."""
+        try:
+            if asyncio.get_running_loop() is self.loop:
+                raise asyncio.InvalidStateError(
+                    "Synchronous calls to AsyncOceanBaseSaver are only allowed from a "
+                    "different thread. From the main thread, use the async interface."
+                )
+        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 (sync wrapper)."""
+        return asyncio.run_coroutine_threadsafe(
+            self.aput(config, checkpoint, metadata, new_versions), self.loop
+        ).result()
+
+
+__all__ = ["BaseAsyncOceanBaseSaver"]