chore: import upstream snapshot with attribution

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

@@ -0,0 +1,116 @@
+from __future__ import annotations
+
+import json
+import re
+from collections.abc import Sequence
+from typing import Any
+
+from langchain_core.runnables import RunnableConfig
+from langgraph.checkpoint.base import get_checkpoint_id
+
+_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)