langgraph/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/utils.py

from __future__ import annotations

import json
import re
from collections.abc import Iterable, Sequence
from typing import Any

from langchain_core.runnables import RunnableConfig
from langgraph.checkpoint.base import PendingWrite, get_checkpoint_id, writes_sort_key
from langgraph.checkpoint.serde.base import SerializerProtocol

_FILTER_PATTERN = re.compile(r"^[a-zA-Z0-9_.-]+$")


def _validate_filter_key(key: str) -> None:
    """Validate that a filter key is safe for use in SQL queries.

    Args:
        key: The filter key to validate

    Raises:
        ValueError: If the key contains invalid characters that could enable SQL injection
    """
    # Allow alphanumeric characters, underscores, dots, and hyphens
    # This covers typical JSON property names while preventing SQL injection
    if not _FILTER_PATTERN.match(key):
        raise ValueError(
            f"Invalid filter key: '{key}'. Filter keys must contain only alphanumeric characters, underscores, dots, and hyphens."
        )


def _metadata_predicate(
    metadata_filter: dict[str, Any],
) -> tuple[Sequence[str], Sequence[Any]]:
    """Return WHERE clause predicates for (a)search() given metadata filter.

    This method returns a tuple of a string and a tuple of values. The string
    is the parametered WHERE clause predicate (excluding the WHERE keyword):
    "column1 = ? AND column2 IS ?". The tuple of values contains the values
    for each of the corresponding parameters.
    """

    def _where_value(query_value: Any) -> tuple[str, Any]:
        """Return tuple of operator and value for WHERE clause predicate."""
        if query_value is None:
            return ("IS ?", None)
        elif (
            isinstance(query_value, str)
            or isinstance(query_value, int)
            or isinstance(query_value, float)
        ):
            return ("= ?", query_value)
        elif isinstance(query_value, bool):
            return ("= ?", 1 if query_value else 0)
        elif isinstance(query_value, dict) or isinstance(query_value, list):
            # query value for JSON object cannot have trailing space after separators (, :)
            # SQLite json_extract() returns JSON string without whitespace
            return ("= ?", json.dumps(query_value, separators=(",", ":")))
        else:
            return ("= ?", str(query_value))

    predicates = []
    param_values = []

    # process metadata query
    for query_key, query_value in metadata_filter.items():
        _validate_filter_key(query_key)
        operator, param_value = _where_value(query_value)
        predicates.append(
            f"json_extract(CAST(metadata AS TEXT), '$.{query_key}') {operator}"
        )
        param_values.append(param_value)

    return (predicates, param_values)


def search_where(
    config: RunnableConfig | None,
    filter: dict[str, Any] | None,
    before: RunnableConfig | None = None,
) -> tuple[str, Sequence[Any]]:
    """Return WHERE clause predicates for (a)search() given metadata filter
    and `before` config.

    This method returns a tuple of a string and a tuple of values. The string
    is the parametered WHERE clause predicate (including the WHERE keyword):
    "WHERE column1 = ? AND column2 IS ?". The tuple of values contains the
    values for each of the corresponding parameters.
    """
    wheres = []
    param_values = []

    # construct predicate for config filter
    if config is not None:
        wheres.append("thread_id = ?")
        param_values.append(config["configurable"]["thread_id"])
        checkpoint_ns = config["configurable"].get("checkpoint_ns")
        if checkpoint_ns is not None:
            wheres.append("checkpoint_ns = ?")
            param_values.append(checkpoint_ns)

        if checkpoint_id := get_checkpoint_id(config):
            wheres.append("checkpoint_id = ?")
            param_values.append(checkpoint_id)

    # construct predicate for metadata filter
    if filter:
        metadata_predicates, metadata_values = _metadata_predicate(filter)
        wheres.extend(metadata_predicates)
        param_values.extend(metadata_values)

    # construct predicate for `before`
    if before is not None:
        wheres.append("checkpoint_id < ?")
        param_values.append(get_checkpoint_id(before))

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


def pending_writes_sql(has_task_path: bool) -> str:
    task_path = "task_path" if has_task_path else "''"
    return (
        f"SELECT task_id, channel, type, value, {task_path}, idx FROM writes "
        "WHERE thread_id = ? AND checkpoint_ns = ? AND checkpoint_id = ?"
    )


def load_pending_writes(
    rows: Iterable[Any], serde: SerializerProtocol
) -> list[PendingWrite]:
    """Deserialize `pending_writes_sql` rows in `writes_sort_key` order."""
    return [
        (task_id, channel, serde.loads_typed((type_, value)))
        for task_id, channel, type_, value, _, _ in sorted(
            rows, key=lambda r: writes_sort_key(r[4], r[0], r[5])
        )
    ]