langgraph-checkpoint-dmdb 0.1.0__py3-none-any.whl

langgraph/checkpoint/dmdb/__init__.py

@@ -0,0 +1,12 @@
+"""DM Database checkpoint savers for LangGraph."""
+
+from .aio import AsyncDMDBSaver, AsyncShallowDMDBSaver
+from .dmpython import DMDBSaver
+from .shallow import ShallowDMDBSaver
+
+__all__ = [
+    "AsyncDMDBSaver",
+    "AsyncShallowDMDBSaver",
+    "DMDBSaver",
+    "ShallowDMDBSaver",
+]

langgraph/checkpoint/dmdb/_internal.py

@@ -0,0 +1,173 @@
+"""Internal DB-API helpers for the synchronous DMDB saver."""
+
+from __future__ import annotations
+
+import threading
+from collections.abc import Callable, Iterator, Sequence
+from contextlib import contextmanager
+from typing import Any, Protocol, TypeVar, Union, cast
+
+
+class Cursor(Protocol):
+    description: Sequence[Sequence[Any]] | None
+    rowcount: int
+
+    def execute(
+        self,
+        operation: str,
+        parameters: Sequence[Any] = ...,
+        /,
+    ) -> object: ...
+
+    def executemany(
+        self,
+        operation: str,
+        seq_of_parameters: Sequence[Sequence[Any]],
+        /,
+    ) -> object: ...
+
+    def fetchone(self) -> Sequence[Any] | None: ...
+
+    def fetchall(self) -> Sequence[Sequence[Any]]: ...
+
+    def fetchmany(self, size: int = ...) -> Sequence[Sequence[Any]]: ...
+
+    def setinputsizes(self, *sizes: object) -> None: ...
+
+    def close(self) -> None: ...
+
+
+class Connection(Protocol):
+    def cursor(self) -> Cursor: ...
+
+    def commit(self) -> None: ...
+
+    def rollback(self) -> None: ...
+
+    def close(self) -> None: ...
+
+
+C = TypeVar("C", bound=Connection)
+ConnectionFactory = Callable[[], C]
+Conn = Union[C, ConnectionFactory[C]]
+
+
+def is_connection(value: object) -> bool:
+    try:
+        return callable(value.cursor)  # type: ignore[attr-defined]
+    except AttributeError:
+        return False
+
+
+@contextmanager
+def get_connection(
+    source: Conn[C],
+    *,
+    direct_lock: threading.RLock,
+    direct_thread_id: int | None,
+) -> Iterator[C]:
+    """Yield a borrowed direct connection or an owned factory connection."""
+    if is_connection(source):
+        if direct_thread_id != threading.get_ident():
+            raise RuntimeError(
+                "A direct dmPython connection cannot be used from another thread; "
+                "construct DMDBSaver with a connection factory instead."
+            )
+        with direct_lock:
+            yield cast(C, source)
+        return
+
+    if not callable(source):
+        raise TypeError(f"Invalid DMDB connection source: {type(source)!r}")
+
+    conn = source()
+    if not is_connection(conn):
+        raise TypeError(f"Connection factory returned an invalid value: {type(conn)!r}")
+    try:
+        yield conn
+    finally:
+        conn.close()
+
+
+def column_names(cursor: Cursor) -> list[str]:
+    return [str(column[0]).lower() for column in cursor.description or ()]
+
+
+def row_to_dict(columns: Sequence[str], row: Sequence[Any]) -> dict[str, Any]:
+    return dict(zip(columns, row))
+
+
+def fetchone_dict(cursor: Cursor) -> dict[str, Any] | None:
+    columns = column_names(cursor)
+    row = cursor.fetchone()
+    return row_to_dict(columns, row) if row is not None else None
+
+
+def fetchall_dict(cursor: Cursor) -> list[dict[str, Any]]:
+    columns = column_names(cursor)
+    return [row_to_dict(columns, row) for row in cursor.fetchall()]
+
+
+def fetchmany_dict(cursor: Cursor, size: int) -> list[dict[str, Any]]:
+    columns = column_names(cursor)
+    return [row_to_dict(columns, row) for row in cursor.fetchmany(size)]
+
+
+def execute_batch(
+    cursor: Cursor,
+    operation: str,
+    parameters: Sequence[Sequence[Any]],
+) -> None:
+    """Execute DM MERGE rows individually within the caller's transaction.
+
+    dmPython 2.5.32 may size executemany buffers from an earlier row and raise
+    ``-70005 String truncated`` when later serialized BLOB values are larger.
+    """
+    for row in parameters:
+        cursor.execute(operation, row)
+
+
+def configure_input_sizes(cursor: Cursor, *sizes: int | str) -> None:
+    """Configure dmPython bind sizes without importing the driver in test fakes."""
+    try:
+        import dmPython  # type: ignore[import-not-found]
+    except ImportError:
+        return
+    try:
+        resolved = [
+            getattr(dmPython, size.upper()) if isinstance(size, str) else size
+            for size in sizes
+        ]
+    except AttributeError:
+        return
+    cursor.setinputsizes(*resolved)
+
+
+def _read_lob(value: object) -> object:
+    read = getattr(value, "read", None)
+    return _read_lob(read()) if callable(read) else value
+
+
+def read_clob(value: object) -> str:
+    value = _read_lob(value)
+    if isinstance(value, str):
+        return value
+    if isinstance(value, (bytes, bytearray, memoryview)):
+        return bytes(value).decode("utf-8")
+    raise TypeError(f"Unsupported CLOB value: {type(value)!r}")
+
+
+def read_blob(value: object) -> bytes:
+    value = _read_lob(value)
+    if isinstance(value, str):
+        return value.encode("utf-8")
+    if isinstance(value, (bytes, bytearray, memoryview)):
+        return bytes(value)
+    raise TypeError(f"Unsupported BLOB value: {type(value)!r}")
+
+
+def rollback_quietly(conn: Connection) -> None:
+    try:
+        conn.rollback()
+    except Exception:
+        pass

langgraph/checkpoint/dmdb/aio.py

@@ -0,0 +1,182 @@
+"""Async facades for dmPython-backed checkpoint savers."""
+
+from __future__ import annotations
+
+import asyncio
+from collections.abc import AsyncIterator, Callable, Iterator, Sequence
+from typing import Any
+
+from langchain_core.runnables import RunnableConfig
+from langgraph.checkpoint.base import (
+    ChannelVersions,
+    Checkpoint,
+    CheckpointMetadata,
+    CheckpointTuple,
+)
+from langgraph.checkpoint.serde.base import SerializerProtocol
+from typing_extensions import Self
+
+from . import _internal
+from .base import BaseDMDBSaver
+from .dmpython import (
+    DMDBSaver,
+    _connection_factory,
+    _validate_conn_params,
+)
+from .shallow import ShallowDMDBSaver
+
+
+class AsyncDMDBSaver(BaseDMDBSaver):
+    """Async CheckpointSaver using one factory-created connection per call."""
+
+    sync_saver_class: type[DMDBSaver] = DMDBSaver
+
+    def __init__(
+        self,
+        conn: Callable[[], _internal.Connection],
+        *,
+        serde: SerializerProtocol | None = None,
+        schema: str | None = None,
+    ) -> None:
+        if _internal.is_connection(conn):
+            raise TypeError(
+                "AsyncDMDBSaver requires a connection factory; a dmPython connection "
+                "cannot be shared with worker threads."
+            )
+        if not callable(conn):
+            raise TypeError("AsyncDMDBSaver requires a connection factory")
+        super().__init__(serde=serde)
+        self._sync = self.sync_saver_class(conn, serde=self.serde, schema=schema)
+
+    @classmethod
+    def from_conn_string(
+        cls,
+        conn_string: str,
+        *,
+        serde: SerializerProtocol | None = None,
+    ) -> Self:
+        params = DMDBSaver.parse_conn_string(conn_string)
+        return cls(_connection_factory(params), serde=serde, schema=params.schema)
+
+    @classmethod
+    def from_conn_params(
+        cls,
+        *,
+        host: str,
+        user: str,
+        password: str,
+        port: int = 5236,
+        schema: str | None = None,
+        serde: SerializerProtocol | None = None,
+    ) -> Self:
+        params = _validate_conn_params(
+            host=host,
+            user=user,
+            password=password,
+            port=port,
+            schema=schema,
+        )
+        return cls(_connection_factory(params), serde=serde, schema=params.schema)
+
+    async def __aenter__(self) -> Self:
+        self._sync._ensure_open()
+        return self
+
+    async def __aexit__(self, *exc_info: object) -> None:
+        await self.aclose()
+
+    def close(self) -> None:
+        self._sync.close()
+
+    async def aclose(self) -> None:
+        self.close()
+
+    async def setup(self) -> None:
+        await asyncio.to_thread(self._sync.setup)
+
+    def list(
+        self,
+        config: RunnableConfig | None,
+        *,
+        filter: dict[str, Any] | None = None,
+        before: RunnableConfig | None = None,
+        limit: int | None = None,
+    ) -> Iterator[CheckpointTuple]:
+        return self._sync.list(config, filter=filter, before=before, limit=limit)
+
+    async def alist(
+        self,
+        config: RunnableConfig | None,
+        *,
+        filter: dict[str, Any] | None = None,
+        before: RunnableConfig | None = None,
+        limit: int | None = None,
+    ) -> AsyncIterator[CheckpointTuple]:
+        values = await asyncio.to_thread(
+            lambda: list(
+                self._sync.list(config, filter=filter, before=before, limit=limit)
+            )
+        )
+        for value in values:
+            yield value
+
+    def get_tuple(self, config: RunnableConfig) -> CheckpointTuple | None:
+        return self._sync.get_tuple(config)
+
+    async def aget_tuple(self, config: RunnableConfig) -> CheckpointTuple | None:
+        return await asyncio.to_thread(self._sync.get_tuple, config)
+
+    def put(
+        self,
+        config: RunnableConfig,
+        checkpoint: Checkpoint,
+        metadata: CheckpointMetadata,
+        new_versions: ChannelVersions,
+    ) -> RunnableConfig:
+        return self._sync.put(config, checkpoint, metadata, new_versions)
+
+    async def aput(
+        self,
+        config: RunnableConfig,
+        checkpoint: Checkpoint,
+        metadata: CheckpointMetadata,
+        new_versions: ChannelVersions,
+    ) -> RunnableConfig:
+        return await asyncio.to_thread(
+            self._sync.put, config, checkpoint, metadata, new_versions
+        )
+
+    def put_writes(
+        self,
+        config: RunnableConfig,
+        writes: Sequence[tuple[str, Any]],
+        task_id: str,
+        task_path: str = "",
+    ) -> None:
+        self._sync.put_writes(config, writes, task_id, task_path)
+
+    async def aput_writes(
+        self,
+        config: RunnableConfig,
+        writes: Sequence[tuple[str, Any]],
+        task_id: str,
+        task_path: str = "",
+    ) -> None:
+        await asyncio.to_thread(
+            self._sync.put_writes, config, writes, task_id, task_path
+        )
+
+    def delete_thread(self, thread_id: str) -> None:
+        self._sync.delete_thread(thread_id)
+
+    async def adelete_thread(self, thread_id: str) -> None:
+        await asyncio.to_thread(self._sync.delete_thread, thread_id)
+
+
+class AsyncShallowDMDBSaver(AsyncDMDBSaver):
+    """Async facade for :class:`ShallowDMDBSaver`."""
+
+    sync_saver_class = ShallowDMDBSaver
+
+
+__all__ = ["AsyncDMDBSaver", "AsyncShallowDMDBSaver"]