langmigrate 1.0.0__py3-none-any.whl

langmigrate/__init__.py

@@ -0,0 +1,62 @@
+"""LangMigrate — declarative schema migrations for LangGraph state persistence.
+
+See ``CLAUDE.md`` for architecture and conventions.
+"""
+
+from .core.engine import HEAD, MigrationEngine
+from .core.exceptions import (
+    ChannelRemovalUnsupportedError,
+    CyclicHistoryError,
+    DuplicateRevisionError,
+    IrreversibleMigrationError,
+    LangMigrateError,
+    MissingRequiredFieldError,
+    MultipleHeadsError,
+    RevisionNotAncestorError,
+    RevisionNotFoundError,
+    TopologyMismatchError,
+    UnsafeMigrationError,
+)
+from .core.migration import BaseMigration, FunctionMigration, migration
+from .core.registry import MigrationRegistry, new_revision_id
+from .core.topology import NodeRemap
+from .core.types import REVISION_METADATA_KEY, RevisionMeta, StateEnvelope
+from .integrations.state import migrate_state_update
+from .runtime.batch import BatchResult, run_batch_downgrade, run_batch_upgrade
+from .runtime.factory import setup_langmigrate
+from .runtime.interceptor import MigrationInterceptor
+
+__version__ = "1.0.0"
+
+__all__ = [
+    "__version__",
+    "HEAD",
+    "BaseMigration",
+    "FunctionMigration",
+    "migration",
+    "MigrationEngine",
+    "MigrationRegistry",
+    "MigrationInterceptor",
+    "setup_langmigrate",
+    "new_revision_id",
+    "NodeRemap",
+    "StateEnvelope",
+    "RevisionMeta",
+    "REVISION_METADATA_KEY",
+    "BatchResult",
+    "run_batch_upgrade",
+    "run_batch_downgrade",
+    "migrate_state_update",
+    # exceptions
+    "LangMigrateError",
+    "UnsafeMigrationError",
+    "MissingRequiredFieldError",
+    "RevisionNotFoundError",
+    "RevisionNotAncestorError",
+    "DuplicateRevisionError",
+    "CyclicHistoryError",
+    "MultipleHeadsError",
+    "IrreversibleMigrationError",
+    "TopologyMismatchError",
+    "ChannelRemovalUnsupportedError",
+]

langmigrate/adapters/__init__.py

@@ -0,0 +1,5 @@
+"""Database-specific bulk adapters for the proactive batch CLI.
+
+Database client imports are confined to this package and loaded lazily so that
+importing ``langmigrate`` never requires an optional backend extra to be installed.
+"""

langmigrate/adapters/base.py

@@ -0,0 +1,48 @@
+"""The adapter contract for proactive (batch) migration.
+
+An adapter exposes a database's checkpoints to the batch CLI: it enumerates the
+checkpoints whose stored revision is behind the target (ideally via an indexed
+metadata query) and provides the underlying saver used to read/write them.
+
+This module is pure — it declares a :class:`Protocol` only. Concrete adapters
+(``postgres``, ``redis``) live alongside and import their DB client lazily.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Iterator
+from typing import Protocol, runtime_checkable
+
+from langchain_core.runnables import RunnableConfig
+from langgraph.checkpoint.base import BaseCheckpointSaver
+
+
+@runtime_checkable
+class CheckpointAdapter(Protocol):
+    """Backend-specific access used by the batch migration runner."""
+
+    @property
+    def saver(self) -> BaseCheckpointSaver:
+        """The underlying LangGraph checkpointer for reads/writes."""
+        ...
+
+    def count_stale(self, head: str) -> int:
+        """Number of checkpoints whose stored revision differs from ``head``."""
+        ...
+
+    def iter_stale_configs(self, head: str) -> Iterator[RunnableConfig]:
+        """Yield a full ``RunnableConfig`` (incl. ``checkpoint_id``) per stale checkpoint."""
+        ...
+
+
+@runtime_checkable
+class BatchCheckpointAdapter(CheckpointAdapter, Protocol):
+    """A :class:`CheckpointAdapter` that can also enumerate *all* checkpoints.
+
+    Needed for downgrades, whose target sits below the current head (so the
+    stale-only enumeration is insufficient).
+    """
+
+    def iter_all_configs(self) -> Iterator[RunnableConfig]:
+        """Yield a full ``RunnableConfig`` for every checkpoint in the store."""
+        ...

langmigrate/adapters/postgres.py

@@ -0,0 +1,126 @@
+"""PostgreSQL adapter for proactive batch migration.
+
+Finds stale checkpoints with a single indexed query against the ``metadata``
+JSONB column — no need to deserialize every row to discover its revision, which
+is what makes the batch path scale to large databases.
+
+The ``psycopg`` / ``langgraph-checkpoint-postgres`` imports are done lazily so the
+rest of LangMigrate stays importable without the ``[postgres]`` extra installed.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Iterator
+from typing import TYPE_CHECKING, Any
+
+from langchain_core.runnables import RunnableConfig
+
+from ..core.types import REVISION_METADATA_KEY
+
+if TYPE_CHECKING:  # pragma: no cover
+    from langgraph.checkpoint.postgres import PostgresSaver
+
+# Stale = the stored revision tag differs from (or is missing relative to) the head.
+_STALE_WHERE = f"metadata->>'{REVISION_METADATA_KEY}' IS DISTINCT FROM %s"
+
+
+class PostgresAdapter:
+    """Adapter over a ``PostgresSaver`` connection for batch enumeration."""
+
+    def __init__(self, conn: Any, saver: PostgresSaver) -> None:
+        self._conn = conn
+        self._saver = saver
+
+    @classmethod
+    def from_conn_string(cls, conn_string: str) -> PostgresAdapter:
+        """Open a connection and build the adapter (and its ``PostgresSaver``)."""
+        import psycopg
+        from langgraph.checkpoint.postgres import PostgresSaver
+        from psycopg.rows import dict_row
+
+        conn = psycopg.connect(conn_string, autocommit=True, row_factory=dict_row)
+        return cls(conn, PostgresSaver(conn))
+
+    @property
+    def saver(self) -> PostgresSaver:
+        return self._saver
+
+    def setup(self) -> None:
+        """Create the checkpoint tables if they do not yet exist."""
+        self._saver.setup()
+
+    def count_stale(self, head: str) -> int:
+        with self._conn.cursor() as cur:
+            cur.execute(f"SELECT count(*) AS c FROM checkpoints WHERE {_STALE_WHERE}", (head,))
+            return int(cur.fetchone()["c"])
+
+    def iter_stale_configs(self, head: str) -> Iterator[RunnableConfig]:
+        # Materialize first so the cursor is closed before the saver reuses the
+        # connection during migration of each checkpoint.
+        with self._conn.cursor() as cur:
+            cur.execute(
+                "SELECT thread_id, checkpoint_ns, checkpoint_id "
+                f"FROM checkpoints WHERE {_STALE_WHERE} "
+                "ORDER BY thread_id, checkpoint_ns, checkpoint_id",
+                (head,),
+            )
+            rows = cur.fetchall()
+        for row in rows:
+            yield {
+                "configurable": {
+                    "thread_id": row["thread_id"],
+                    "checkpoint_ns": row["checkpoint_ns"],
+                    "checkpoint_id": row["checkpoint_id"],
+                }
+            }
+
+    def iter_all_configs(self) -> Iterator[RunnableConfig]:
+        with self._conn.cursor() as cur:
+            cur.execute(
+                "SELECT thread_id, checkpoint_ns, checkpoint_id FROM checkpoints "
+                "ORDER BY thread_id, checkpoint_ns, checkpoint_id"
+            )
+            rows = cur.fetchall()
+        for row in rows:
+            yield {
+                "configurable": {
+                    "thread_id": row["thread_id"],
+                    "checkpoint_ns": row["checkpoint_ns"],
+                    "checkpoint_id": row["checkpoint_id"],
+                }
+            }
+
+    def stamp_all(self, revision: str) -> int:
+        """Set the revision tag on every checkpoint without running migrations.
+
+        Returns the number of rows updated. Use when adopting LangMigrate on a
+        database whose state already matches a known revision. ``COALESCE`` guards
+        against a row whose ``metadata`` is SQL/JSON ``null`` (``jsonb_set`` of a
+        null base returns null and would silently drop the tag).
+        """
+        with self._conn.cursor() as cur:
+            cur.execute(
+                "UPDATE checkpoints SET metadata = "
+                "jsonb_set(COALESCE(NULLIF(metadata, 'null'::jsonb), '{}'::jsonb), "
+                f"'{{{REVISION_METADATA_KEY}}}', to_jsonb(%s::text))",
+                (revision,),
+            )
+            return int(cur.rowcount or 0)
+
+    def revision_counts(self) -> dict[str, int]:
+        """Distribution of stored revision tags across all checkpoints (for ``current --db``)."""
+        with self._conn.cursor() as cur:
+            cur.execute(
+                f"SELECT metadata->>'{REVISION_METADATA_KEY}' AS rev, count(*) AS c "
+                "FROM checkpoints GROUP BY rev"
+            )
+            return {(row["rev"] or "<untagged>"): int(row["c"]) for row in cur.fetchall()}
+
+    def close(self) -> None:
+        self._conn.close()
+
+    def __enter__(self) -> PostgresAdapter:
+        return self
+
+    def __exit__(self, *exc: object) -> None:
+        self.close()

langmigrate/adapters/redis.py

@@ -0,0 +1,152 @@
+"""Redis adapter for proactive batch migration.
+
+``langgraph-checkpoint-redis`` stores each checkpoint as a RedisJSON document at
+``checkpoint:<thread>:<ns>:<id>`` with the LangGraph metadata kept as a serialized
+JSON string under ``$.metadata``. Our revision tag is *not* part of the RediSearch
+index, so batch enumeration scans the checkpoint keys and reads each document's
+metadata (an O(n) sweep — inherent to Redis without a custom index).
+
+Lazy *online* migration needs none of this: wrap a ``RedisSaver`` with
+:class:`~langmigrate.runtime.interceptor.MigrationInterceptor` and it works today.
+The ``redis`` client imports are done lazily so the rest of LangMigrate stays
+importable without the ``[redis]`` extra.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import json
+from collections.abc import Iterator
+from typing import TYPE_CHECKING, Any, cast
+
+from langchain_core.runnables import RunnableConfig
+from langgraph.checkpoint.base import CheckpointMetadata
+
+from ..core.version import read_revision, stamp_metadata
+
+if TYPE_CHECKING:  # pragma: no cover
+    from langgraph.checkpoint.redis import RedisSaver
+
+# Matches checkpoint docs only ("checkpoint:..."), not "checkpoint_write:...".
+_CHECKPOINT_MATCH = "checkpoint:*"
+
+
+def _first(fields: dict[str, Any], path: str) -> Any:
+    """First value at a RedisJSON ``$.path`` result (paths return a list)."""
+    value = fields.get(path)
+    return value[0] if value else None
+
+
+class RedisAdapter:
+    """Adapter over a ``RedisSaver`` for batch enumeration and stamping."""
+
+    def __init__(self, saver: RedisSaver) -> None:
+        self._saver = saver
+
+    @classmethod
+    def from_conn_string(cls, conn_string: str) -> RedisAdapter:
+        """Open a connection and build the adapter (and its ``RedisSaver``)."""
+        from langgraph.checkpoint.redis import RedisSaver
+
+        saver = RedisSaver(conn_string)
+        saver.setup()
+        return cls(saver)
+
+    @property
+    def saver(self) -> RedisSaver:
+        return self._saver
+
+    def setup(self) -> None:
+        self._saver.setup()
+
+    # -- enumeration --------------------------------------------------------
+
+    def _client(self) -> Any:
+        return self._saver._redis
+
+    def _iter_docs(self) -> Iterator[tuple[RunnableConfig, str | None]]:
+        """Yield ``(config, stored_revision)`` for every checkpoint document."""
+        from langgraph.checkpoint.redis.util import (
+            from_storage_safe_id,
+            from_storage_safe_str,
+            safely_decode,
+        )
+
+        client = self._client()
+        for raw_key in client.scan_iter(match=_CHECKPOINT_MATCH, count=200):
+            key = safely_decode(raw_key)
+            fields = client.json().get(
+                key, "$.thread_id", "$.checkpoint_ns", "$.checkpoint_id", "$.metadata"
+            )
+            if not fields:
+                continue
+
+            thread_id = from_storage_safe_id(_first(fields, "$.thread_id") or "")
+            checkpoint_ns = from_storage_safe_str(_first(fields, "$.checkpoint_ns") or "")
+            checkpoint_id = from_storage_safe_id(_first(fields, "$.checkpoint_id") or "")
+            revision = self._revision_from_metadata(_first(fields, "$.metadata"))
+            config: RunnableConfig = {
+                "configurable": {
+                    "thread_id": thread_id,
+                    "checkpoint_ns": checkpoint_ns,
+                    "checkpoint_id": checkpoint_id,
+                }
+            }
+            yield config, revision
+
+    @staticmethod
+    def _revision_from_metadata(metadata: Any) -> str | None:
+        # Stored as a serialized JSON string (occasionally already a dict).
+        if isinstance(metadata, str):
+            try:
+                metadata = json.loads(metadata)
+            except json.JSONDecodeError:
+                return None
+        return read_revision(metadata if isinstance(metadata, dict) else None)
+
+    def count_stale(self, head: str) -> int:
+        return sum(1 for _, rev in self._iter_docs() if rev != head)
+
+    def iter_stale_configs(self, head: str) -> Iterator[RunnableConfig]:
+        for config, rev in self._iter_docs():
+            if rev != head:
+                yield config
+
+    def iter_all_configs(self) -> Iterator[RunnableConfig]:
+        for config, _ in self._iter_docs():
+            yield config
+
+    def revision_counts(self) -> dict[str, int]:
+        counts: dict[str, int] = {}
+        for _, rev in self._iter_docs():
+            key = rev or "<untagged>"
+            counts[key] = counts.get(key, 0) + 1
+        return counts
+
+    def stamp_all(self, revision: str) -> int:
+        """Set the revision tag on every checkpoint without running migrations."""
+        updated = 0
+        for config, _ in self._iter_docs():
+            tup = self._saver.get_tuple(config)
+            if tup is None:
+                continue
+            metadata = stamp_metadata(dict(tup.metadata or {}), revision)
+            put_config = tup.parent_config or {
+                "configurable": {
+                    "thread_id": config["configurable"]["thread_id"],
+                    "checkpoint_ns": config["configurable"]["checkpoint_ns"],
+                }
+            }
+            self._saver.put(put_config, tup.checkpoint, cast(CheckpointMetadata, metadata), {})
+            updated += 1
+        return updated
+
+    def close(self) -> None:
+        with contextlib.suppress(Exception):  # best effort
+            self._client().close()
+
+    def __enter__(self) -> RedisAdapter:
+        return self
+
+    def __exit__(self, *exc: object) -> None:
+        self.close()

langmigrate/cli/__init__.py

@@ -0,0 +1 @@
+"""Typer-based command line interface."""