base.py

"""Base MySQL checkpoint saver with VARCHAR(32) hash for cross-db-sync compatibility."""

from __future__ import annotations

import asyncio
import json
import random
from collections.abc import Sequence
from typing import Any, cast

from langchain_core.runnables import RunnableConfig

from langgraph.checkpoint.base import (
    WRITES_IDX_MAP,
    BaseCheckpointSaver,
    ChannelVersions,
    Checkpoint,
    CheckpointMetadata,
    CheckpointTuple,
    get_checkpoint_id,
    get_serializable_checkpoint_metadata,
)
from langgraph.checkpoint.serde.base import SerializerProtocol
from langgraph.checkpoint.serde.types import TASKS

from src.checkpoint_mysql.serde import (
    deserialize_channel_values,
    deserialize_pending_sends,
    deserialize_pending_writes,
)


"""
MIGRATIONS: Sequential schema migrations.
The position of the migration in the list is the version number.
"""
MIGRATIONS = [
    """CREATE TABLE IF NOT EXISTS checkpoint_migrations (
    v INTEGER PRIMARY KEY
);""",
    """CREATE TABLE IF NOT EXISTS checkpoints (
    thread_id VARCHAR(150) NOT NULL,
    checkpoint_ns VARCHAR(2000) NOT NULL DEFAULT '',
    checkpoint_ns_hash VARCHAR(32) NOT NULL DEFAULT '',
    checkpoint_id VARCHAR(150) NOT NULL,
    parent_checkpoint_id VARCHAR(150),
    type VARCHAR(150),
    checkpoint JSON NOT NULL,
    metadata JSON NOT NULL DEFAULT ('{}'),
    PRIMARY KEY (thread_id, checkpoint_ns_hash, checkpoint_id)
);""",
    """CREATE TABLE IF NOT EXISTS checkpoint_blobs (
    thread_id VARCHAR(150) NOT NULL,
    checkpoint_ns VARCHAR(2000) NOT NULL DEFAULT '',
    checkpoint_ns_hash VARCHAR(32) NOT NULL DEFAULT '',
    channel VARCHAR(150) NOT NULL,
    version VARCHAR(150) NOT NULL,
    type VARCHAR(150) NOT NULL,
    `blob` LONGBLOB,
    PRIMARY KEY (thread_id, checkpoint_ns_hash, channel, version)
);""",
    """CREATE TABLE IF NOT EXISTS checkpoint_writes (
    thread_id VARCHAR(150) NOT NULL,
    checkpoint_ns VARCHAR(2000) NOT NULL DEFAULT '',
    checkpoint_ns_hash VARCHAR(32) NOT NULL DEFAULT '',
    checkpoint_id VARCHAR(150) NOT NULL,
    task_id VARCHAR(150) NOT NULL,
    idx INTEGER NOT NULL,
    channel VARCHAR(150) NOT NULL,
    type VARCHAR(150),
    `blob` LONGBLOB NOT NULL,
    PRIMARY KEY (thread_id, checkpoint_ns_hash, checkpoint_id, task_id, idx)
);""",
    """
    CREATE INDEX checkpoints_thread_id_idx ON checkpoints (thread_id);
    """,
    """
    CREATE INDEX checkpoint_blobs_thread_id_idx ON checkpoint_blobs (thread_id);
    """,
    """
    CREATE INDEX checkpoint_writes_thread_id_idx ON checkpoint_writes (thread_id);
    """,
    """
    CREATE INDEX checkpoints_checkpoint_id_idx ON checkpoints (checkpoint_id);
    """,
    """
    ALTER TABLE checkpoint_writes ADD COLUMN task_path VARCHAR(2000) NOT NULL DEFAULT '';
    """,
]


SELECT_SQL = """
WITH channel_versions AS (
    SELECT thread_id, checkpoint_ns_hash, checkpoint_id, channel,
           json_unquote(
               json_extract(checkpoint, concat('$.channel_versions.', '"', channel, '"'))
           ) AS version
    FROM checkpoints, JSON_TABLE(
        JSON_KEYS(checkpoint, '$.channel_versions'),
        '$[*]' COLUMNS (channel VARCHAR(150) CHARACTER SET utf8mb4 PATH '$')
    ) AS channels
)
SELECT
    thread_id,
    checkpoint_ns,
    checkpoint_ns_hash,
    checkpoint_id,
    parent_checkpoint_id,
    checkpoint,
    metadata,
    (
        SELECT JSON_ARRAYAGG(JSON_ARRAY(
            bl.channel,
            bl.type,
            IFNULL(TO_BASE64(bl.`blob`), '')
        ))
        FROM channel_versions
        INNER JOIN checkpoint_blobs bl
            ON bl.channel = channel_versions.channel
            AND bl.version = channel_versions.version
        WHERE bl.thread_id = checkpoints.thread_id
            AND bl.checkpoint_ns_hash = checkpoints.checkpoint_ns_hash
            AND channel_versions.thread_id = checkpoints.thread_id
            AND channel_versions.checkpoint_ns_hash = checkpoints.checkpoint_ns_hash
            AND channel_versions.checkpoint_id = checkpoints.checkpoint_id
    ) AS channel_values,
    (
        SELECT JSON_ARRAYAGG(JSON_ARRAY(
            cw.task_id,
            cw.channel,
            cw.type,
            TO_BASE64(cw.`blob`),
            cw.idx
        ))
        FROM checkpoint_writes cw
        WHERE cw.thread_id = checkpoints.thread_id
            AND cw.checkpoint_ns_hash = checkpoints.checkpoint_ns_hash
            AND cw.checkpoint_id = checkpoints.checkpoint_id
    ) AS pending_writes,
    (
        SELECT JSON_ARRAYAGG(JSON_ARRAY(
            cw.task_path,
            cw.task_id,
            cw.type,
            TO_BASE64(cw.`blob`),
            cw.idx
        ))
        FROM checkpoint_writes cw
        WHERE cw.thread_id = checkpoints.thread_id
            AND cw.checkpoint_ns_hash = checkpoints.checkpoint_ns_hash
            AND cw.checkpoint_id = checkpoints.checkpoint_id
            AND cw.channel = '__pregel_tasks'
    ) AS pending_sends
FROM checkpoints {WHERE}
"""

SELECT_SQL_LIGHT = """
WITH channel_versions AS (
    SELECT thread_id, checkpoint_ns_hash, checkpoint_id, channel,
           json_unquote(
               json_extract(checkpoint, concat('$.channel_versions.', '"', channel, '"'))
           ) AS version
    FROM checkpoints, JSON_TABLE(
        JSON_KEYS(checkpoint, '$.channel_versions'),
        '$[*]' COLUMNS (channel VARCHAR(150) CHARACTER SET utf8mb4 PATH '$')
    ) AS channels
    WHERE checkpoints.thread_id = %s
      AND checkpoints.checkpoint_ns_hash = %s
)
SELECT
    checkpoints.thread_id,
    checkpoints.checkpoint_ns,
    checkpoints.checkpoint_ns_hash,
    checkpoints.checkpoint_id,
    checkpoints.parent_checkpoint_id,
    checkpoints.checkpoint,
    checkpoints.metadata,
    (
        SELECT JSON_ARRAYAGG(JSON_ARRAY(
            bl.channel,
            bl.type,
            IFNULL(TO_BASE64(bl.`blob`), '')
        ))
        FROM channel_versions
        INNER JOIN checkpoint_blobs bl
            ON bl.channel = channel_versions.channel
            AND bl.version = channel_versions.version
        WHERE bl.thread_id = checkpoints.thread_id
            AND bl.checkpoint_ns_hash = checkpoints.checkpoint_ns_hash
    ) AS channel_values,
    (
        SELECT JSON_ARRAYAGG(JSON_ARRAY(
            cw.task_id,
            cw.channel,
            cw.type,
            TO_BASE64(cw.`blob`),
            cw.idx
        ))
        FROM checkpoint_writes cw
        WHERE cw.thread_id = checkpoints.thread_id
            AND cw.checkpoint_ns_hash = checkpoints.checkpoint_ns_hash
            AND cw.checkpoint_id = checkpoints.checkpoint_id
    ) AS pending_writes
FROM checkpoints
WHERE checkpoints.thread_id = %s
  AND checkpoints.checkpoint_ns_hash = %s
"""

SELECT_SQL_NO_NS_HASH = """
SELECT
    checkpoints.thread_id,
    checkpoints.checkpoint_ns,
    checkpoints.checkpoint_ns_hash,
    checkpoints.checkpoint_id,
    checkpoints.parent_checkpoint_id,
    checkpoints.checkpoint,
    checkpoints.metadata,
    (
        SELECT JSON_ARRAYAGG(JSON_ARRAY(
            cw.task_id,
            cw.channel,
            cw.type,
            TO_BASE64(cw.`blob`),
            cw.idx
        ))
        FROM checkpoint_writes cw
        WHERE cw.thread_id = checkpoints.thread_id
            AND cw.checkpoint_ns_hash = checkpoints.checkpoint_ns_hash
            AND cw.checkpoint_id = checkpoints.checkpoint_id
    ) AS pending_writes
FROM checkpoints {WHERE}
"""

SELECT_SQL_NO_NS_HASH_LIGHT = """
SELECT
    checkpoints.thread_id,
    checkpoints.checkpoint_ns,
    checkpoints.checkpoint_ns_hash,
    checkpoints.checkpoint_id,
    checkpoints.parent_checkpoint_id,
    checkpoints.checkpoint,
    checkpoints.metadata,
    (
        SELECT JSON_ARRAYAGG(JSON_ARRAY(
            cw.task_id,
            cw.channel,
            cw.type,
            TO_BASE64(cw.`blob`),
            cw.idx
        ))
        FROM checkpoint_writes cw
        WHERE cw.thread_id = checkpoints.thread_id
            AND cw.checkpoint_ns_hash = checkpoints.checkpoint_ns_hash
            AND cw.checkpoint_id = checkpoints.checkpoint_id
    ) AS pending_writes
FROM checkpoints
WHERE checkpoints.thread_id = %s
  AND checkpoints.checkpoint_ns_hash = %s
"""

SELECT_SQL_MARIADB = """
WITH channel_versions AS (
    SELECT thread_id, checkpoint_ns_hash, checkpoint_id, channel,
           json_unquote(
               json_extract(checkpoint, concat('$.channel_versions.', '"', channel, '"'))
           ) AS version
    FROM checkpoints, JSON_TABLE(
        JSON_KEYS(checkpoint, '$.channel_versions'),
        '$[*]' COLUMNS (channel VARCHAR(150) CHARACTER SET utf8mb4 PATH '$')
    ) AS channels
    {WHERE}
)
SELECT
    thread_id,
    checkpoint_ns,
    checkpoint_ns_hash,
    checkpoint_id,
    parent_checkpoint_id,
    checkpoint,
    metadata,
    (
        SELECT JSON_ARRAYAGG(JSON_ARRAY(
            bl.channel,
            bl.type,
            IFNULL(bl.`blob`, '')
        ))
        FROM channel_versions
        INNER JOIN checkpoint_blobs bl
            ON bl.channel = channel_versions.channel
            AND bl.version = channel_versions.version
        WHERE bl.thread_id = checkpoints.thread_id
            AND bl.checkpoint_ns_hash = checkpoints.checkpoint_ns_hash
            AND channel_versions.thread_id = checkpoints.thread_id
            AND channel_versions.checkpoint_ns_hash = checkpoints.checkpoint_ns_hash
            AND channel_versions.checkpoint_id = checkpoints.checkpoint_id
    ) AS channel_values,
    (
        SELECT JSON_ARRAYAGG(JSON_ARRAY(
            cw.task_id,
            cw.channel,
            cw.type,
            cw.`blob`,
            cw.idx
        ))
        FROM checkpoint_writes cw
        WHERE cw.thread_id = checkpoints.thread_id
            AND cw.checkpoint_ns_hash = checkpoints.checkpoint_ns_hash
            AND cw.checkpoint_id = checkpoints.checkpoint_id
    ) AS pending_writes,
    (
        SELECT JSON_ARRAYAGG(JSON_ARRAY(
            cw.task_path,
            cw.task_id,
            cw.type,
            cw.`blob`,
            cw.idx
        ))
        FROM checkpoint_writes cw
        WHERE cw.thread_id = checkpoints.thread_id
            AND cw.checkpoint_ns_hash = checkpoints.checkpoint_ns_hash
            AND cw.checkpoint_id = checkpoints.checkpoint_id
            AND cw.channel = '__pregel_tasks'
    ) AS pending_sends
FROM checkpoints {WHERE}
"""

SELECT_SQL_LIGHT_MARIADB = """
WITH channel_versions AS (
    SELECT thread_id, checkpoint_ns_hash, checkpoint_id, channel,
           json_unquote(
               json_extract(checkpoint, concat('$.channel_versions.', '"', channel, '"'))
           ) AS version
    FROM checkpoints, JSON_TABLE(
        JSON_KEYS(checkpoint, '$.channel_versions'),
        '$[*]' COLUMNS (channel VARCHAR(150) CHARACTER SET utf8mb4 PATH '$')
    ) AS channels
    WHERE checkpoints.thread_id = %s
      AND checkpoints.checkpoint_ns_hash = %s
)
SELECT
    checkpoints.thread_id,
    checkpoints.checkpoint_ns,
    checkpoints.checkpoint_ns_hash,
    checkpoints.checkpoint_id,
    checkpoints.parent_checkpoint_id,
    checkpoints.checkpoint,
    checkpoints.metadata,
    (
        SELECT JSON_ARRAYAGG(JSON_ARRAY(
            bl.channel,
            bl.type,
            IFNULL(bl.`blob`, '')
        ))
        FROM channel_versions
        INNER JOIN checkpoint_blobs bl
            ON bl.channel = channel_versions.channel
            AND bl.version = channel_versions.version
        WHERE bl.thread_id = checkpoints.thread_id
            AND bl.checkpoint_ns_hash = checkpoints.checkpoint_ns_hash
    ) AS channel_values,
    (
        SELECT JSON_ARRAYAGG(JSON_ARRAY(
            cw.task_id,
            cw.channel,
            cw.type,
            cw.`blob`,
            cw.idx
        ))
        FROM checkpoint_writes cw
        WHERE cw.thread_id = checkpoints.thread_id
            AND cw.checkpoint_ns_hash = checkpoints.checkpoint_ns_hash
            AND cw.checkpoint_id = checkpoints.checkpoint_id
    ) AS pending_writes
FROM checkpoints
WHERE checkpoints.thread_id = %s
  AND checkpoints.checkpoint_ns_hash = %s
"""

SELECT_SQL_NO_NS_HASH_MARIADB = """
SELECT
    checkpoints.thread_id,
    checkpoints.checkpoint_ns,
    checkpoints.checkpoint_ns_hash,
    checkpoints.checkpoint_id,
    checkpoints.parent_checkpoint_id,
    checkpoints.checkpoint,
    checkpoints.metadata,
    (
        SELECT JSON_ARRAYAGG(JSON_ARRAY(
            cw.task_id,
            cw.channel,
            cw.type,
            cw.`blob`,
            cw.idx
        ))
        FROM checkpoint_writes cw
        WHERE cw.thread_id = checkpoints.thread_id
            AND cw.checkpoint_ns_hash = checkpoints.checkpoint_ns_hash
            AND cw.checkpoint_id = checkpoints.checkpoint_id
    ) AS pending_writes
FROM checkpoints {WHERE}
"""

SELECT_SQL_NO_NS_HASH_LIGHT_MARIADB = """
SELECT
    checkpoints.thread_id,
    checkpoints.checkpoint_ns,
    checkpoints.checkpoint_ns_hash,
    checkpoints.checkpoint_id,
    checkpoints.parent_checkpoint_id,
    checkpoints.checkpoint,
    checkpoints.metadata,
    (
        SELECT JSON_ARRAYAGG(JSON_ARRAY(
            cw.task_id,
            cw.channel,
            cw.type,
            cw.`blob`,
            cw.idx
        ))
        FROM checkpoint_writes cw
        WHERE cw.thread_id = checkpoints.thread_id
            AND cw.checkpoint_ns_hash = checkpoints.checkpoint_ns_hash
            AND cw.checkpoint_id = checkpoints.checkpoint_id
    ) AS pending_writes
FROM checkpoints
WHERE checkpoints.thread_id = %s
  AND checkpoints.checkpoint_ns_hash = %s
"""


UPSERT_CHECKPOINT_BLOBS_SQL = """
    INSERT INTO checkpoint_blobs (thread_id, checkpoint_ns, checkpoint_ns_hash, channel, version, type, `blob`)
    VALUES (%s, %s, MD5(%s), %s, %s, %s, %s)
    AS new
    ON DUPLICATE KEY UPDATE
        type = new.type,
        `blob` = new.`blob`;
"""

UPSERT_CHECKPOINTS_SQL = """
    INSERT INTO checkpoints (thread_id, checkpoint_ns, checkpoint_ns_hash, checkpoint_id, parent_checkpoint_id, checkpoint, metadata)
    VALUES (%s, %s, MD5(%s), %s, %s, %s, %s)
    AS new
    ON DUPLICATE KEY UPDATE
        checkpoint = new.checkpoint,
        metadata = new.metadata;
"""

UPSERT_CHECKPOINT_WRITES_SQL = """
    INSERT INTO checkpoint_writes (thread_id, checkpoint_ns, checkpoint_ns_hash, checkpoint_id, task_id, task_path, idx, channel, type, `blob`)
    VALUES (%s, %s, MD5(%s), %s, %s, %s, %s, %s, %s, %s)
    AS new
    ON DUPLICATE KEY UPDATE
        channel = new.channel,
        type = new.type,
        `blob` = new.`blob`;
"""

INSERT_CHECKPOINT_WRITES_SQL = """
    INSERT IGNORE INTO checkpoint_writes (thread_id, checkpoint_ns, checkpoint_ns_hash, checkpoint_id, task_id, task_path, idx, channel, type, `blob`)
    VALUES (%s, %s, MD5(%s), %s, %s, %s, %s, %s, %s, %s)
"""


class BaseMySQLSaver(BaseCheckpointSaver[str]):
    MIGRATIONS = MIGRATIONS
    UPSERT_CHECKPOINT_BLOBS_SQL = UPSERT_CHECKPOINT_BLOBS_SQL
    UPSERT_CHECKPOINTS_SQL = UPSERT_CHECKPOINTS_SQL
    UPSERT_CHECKPOINT_WRITES_SQL = UPSERT_CHECKPOINT_WRITES_SQL
    INSERT_CHECKPOINT_WRITES_SQL = INSERT_CHECKPOINT_WRITES_SQL

    SELECT_SQL: str = SELECT_SQL
    SELECT_SQL_LIGHT: str = SELECT_SQL_LIGHT
    SELECT_SQL_MARIADB: str = SELECT_SQL_MARIADB
    SELECT_SQL_LIGHT_MARIADB: str = SELECT_SQL_LIGHT_MARIADB
    SELECT_SQL_NO_NS_HASH: str = SELECT_SQL_NO_NS_HASH
    SELECT_SQL_NO_NS_HASH_LIGHT: str = SELECT_SQL_NO_NS_HASH_LIGHT
    SELECT_SQL_NO_NS_HASH_MARIADB: str = SELECT_SQL_NO_NS_HASH_MARIADB
    SELECT_SQL_NO_NS_HASH_LIGHT_MARIADB: str = SELECT_SQL_NO_NS_HASH_LIGHT_MARIADB

    def _migrate_pending_sends(
        self,
        pending_sends: list[tuple[str, bytes]],
        checkpoint: dict[str, Any],
        channel_values: list[tuple[str, str, bytes | None]],
    ) -> None:
        if not pending_sends:
            return
        enc, blob = self.serde.dumps_typed(
            [self.serde.loads_typed((c, b)) for c, b in pending_sends],
        )
        channel_values.append((TASKS, enc, blob))
        checkpoint["channel_versions"][TASKS] = (
            max(checkpoint["channel_versions"].values())
            if checkpoint["channel_versions"]
            else self.get_next_version(None, None)
        )

    def _load_blobs(
        self, blob_values: list[tuple[str, str, bytes | None]]
    ) -> dict[str, Any]:
        if not blob_values:
            return {}
        return {
            k: self.serde.loads_typed((t, v))
            for k, t, v in blob_values
            if t != "empty" and v is not None
        }

    def _dump_blobs(
        self,
        thread_id: str,
        checkpoint_ns: str,
        values: dict[str, Any],
        versions: ChannelVersions,
    ) -> list[tuple[str, str, str, str, str, str, bytes | None]]:
        if not versions:
            return []

        return [
            (
                thread_id,
                checkpoint_ns,
                checkpoint_ns,
                k,
                cast(str, ver),
                *(
                    self.serde.dumps_typed(values[k])
                    if k in values
                    else ("empty", None)
                ),
            )
            for k, ver in versions.items()
        ]

    def _load_writes(
        self, writes: list[tuple[str, str, str, bytes]]
    ) -> list[tuple[str, str, Any]]:
        return (
            [
                (
                    tid,
                    channel,
                    self.serde.loads_typed((t, v)),
                )
                for tid, channel, t, v in writes
            ]
            if writes
            else []
        )

    def _dump_writes(
        self,
        thread_id: str,
        checkpoint_ns: str,
        checkpoint_id: str,
        task_id: str,
        task_path: str,
        writes: Sequence[tuple[str, Any]],
    ) -> list[tuple[str, str, str, str, str, str, int, str, str, bytes]]:
        return [
            (
                thread_id,
                checkpoint_ns,
                checkpoint_ns,
                checkpoint_id,
                task_id,
                task_path,
                WRITES_IDX_MAP.get(channel, idx),
                channel,
                *self.serde.dumps_typed(value),
            )
            for idx, (channel, value) in enumerate(writes)
        ]

    def get_next_version(self, current: str | None, channel: None) -> str:
        if current is None:
            current_v = 0
        elif isinstance(current, int):
            current_v = current
        else:
            current_v = int(current.split(".")[0])
        next_v = current_v + 1
        next_h = random.random()
        return f"{next_v:032}.{next_h:016}"

    def _search_where(
        self,
        config: RunnableConfig | None,
        filter: dict[str, Any] | None = None,
        before: RunnableConfig | None = None,
    ) -> tuple[str, dict[str, Any]]:
        wheres = []
        param_values: dict[str, Any] = {}

        if config:
            wheres.append("thread_id = %(thread_id)s ")
            param_values["thread_id"] = config["configurable"]["thread_id"]
            checkpoint_ns = config["configurable"].get("checkpoint_ns")
            if checkpoint_ns is not None:
                wheres.append("checkpoint_ns_hash = MD5(%(checkpoint_ns)s)")
                param_values["checkpoint_ns"] = checkpoint_ns

            if checkpoint_id := get_checkpoint_id(config):
                wheres.append("checkpoint_id = %(checkpoint_id)s ")
                param_values["checkpoint_id"] = checkpoint_id

        if filter:
            wheres.append("JSON_CONTAINS(metadata, %(filter)s) ")
            param_values["filter"] = json.dumps(filter)

        if before is not None:
            wheres.append("checkpoint_id < %(before)s ")
            param_values["before"] = get_checkpoint_id(before)

        return (
            "WHERE " + " AND ".join(wheres) if wheres else "",
            param_values,
        )

    @staticmethod
    def _select_sql(where: str, *, light: bool = False, mariadb: bool = False, no_ns_hash: bool = False) -> str:
        if no_ns_hash:
            return (BaseMySQLSaver.SELECT_SQL_NO_NS_HASH_LIGHT_MARIADB if (light and mariadb) else
                    BaseMySQLSaver.SELECT_SQL_NO_NS_HASH_LIGHT if light else
                    BaseMySQLSaver.SELECT_SQL_NO_NS_HASH_MARIADB if mariadb else
                    BaseMySQLSaver.SELECT_SQL_NO_NS_HASH).replace("{WHERE}", where)
        elif light:
            return BaseMySQLSaver.SELECT_SQL_LIGHT if not mariadb else BaseMySQLSaver.SELECT_SQL_LIGHT_MARIADB
        else:
            return (BaseMySQLSaver.SELECT_SQL_MARIADB if mariadb else BaseMySQLSaver.SELECT_SQL).replace("{WHERE}", where)