161 lines
6.8 KiB
Python
161 lines
6.8 KiB
Python
"""Bound JSON predicates for SQLite, PostgreSQL and MySQL.
|
|
|
|
Adapted from upstream #2865. MySQL business tables store JSON in TEXT, so
|
|
these predicates use JSON functions rather than assuming a native JSON column.
|
|
Unsupported keys/values must stay in the repository's residual filter; never
|
|
silently discard a predicate (in particular a shared task's taskId filter).
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import math
|
|
import re
|
|
from typing import Any
|
|
|
|
from sqlalchemy import BigInteger, Float, String, bindparam
|
|
from sqlalchemy.ext.compiler import compiles
|
|
from sqlalchemy.sql.elements import ColumnElement
|
|
from sqlalchemy.sql.visitors import InternalTraversal
|
|
from sqlalchemy.types import Boolean
|
|
|
|
_SAFE_KEY = re.compile(r"[A-Za-z0-9_-]+\Z")
|
|
_STRIP_CHARS = "".join(chr(i) for i in range(0x3001) if chr(i).isspace())
|
|
|
|
|
|
def supports_json_filter(key: Any, value: Any) -> bool:
|
|
if not isinstance(key, str) or not _SAFE_KEY.fullmatch(key):
|
|
return False
|
|
if value is None or isinstance(value, (str, bool)):
|
|
return True
|
|
if isinstance(value, int):
|
|
return -(2**63) <= value < 2**63
|
|
return isinstance(value, float) and math.isfinite(value)
|
|
|
|
|
|
class JsonPredicate(ColumnElement):
|
|
"""A scalar equality, Python truthiness, or nonblank-string predicate."""
|
|
|
|
inherit_cache = True
|
|
type = Boolean()
|
|
_is_implicitly_boolean = True
|
|
_traverse_internals = [
|
|
("column", InternalTraversal.dp_clauseelement),
|
|
("key", InternalTraversal.dp_string),
|
|
("operation", InternalTraversal.dp_string),
|
|
("value", InternalTraversal.dp_plain_obj),
|
|
("value_kind", InternalTraversal.dp_string),
|
|
]
|
|
|
|
def __init__(self, column, key: str, operation: str = "equal", value: Any = None):
|
|
if not supports_json_filter(key, value):
|
|
raise ValueError("Unsupported JSON filter key or value")
|
|
if operation not in {"equal", "truthy", "nonblank"}:
|
|
raise ValueError("Unsupported JSON predicate operation")
|
|
self.column = column
|
|
self.key = key
|
|
self.operation = operation
|
|
self.value = value
|
|
# Keep Python value types separate in SQLAlchemy's compiled cache.
|
|
self.value_kind = type(value).__name__
|
|
super().__init__()
|
|
|
|
|
|
def _compile(element, compiler, family, **kw):
|
|
if not _SAFE_KEY.fullmatch(element.key):
|
|
raise ValueError("Unsafe JSON key")
|
|
column = compiler.process(element.column, **kw)
|
|
|
|
def bound(value, type_=String()):
|
|
return compiler.process(bindparam(None, value, type_=type_), **kw)
|
|
|
|
if family == "postgres":
|
|
raw = f"({column} -> '{element.key}')"
|
|
kind = f"json_typeof({raw})"
|
|
text = f"({column} ->> '{element.key}')"
|
|
string_kind, number_kinds = "string", "'number'"
|
|
else:
|
|
path = f"'$.\"{element.key}\"'"
|
|
raw = f"json_extract({column}, {path})"
|
|
if family == "mysql":
|
|
kind, text = f"lower(json_type({raw}))", f"json_unquote({raw})"
|
|
string_kind, number_kinds = "string", "'integer', 'double', 'decimal'"
|
|
else:
|
|
kind, text = f"json_type({column}, {path})", raw
|
|
string_kind, number_kinds = "text", "'integer', 'real'"
|
|
|
|
if element.operation == "nonblank":
|
|
if family == "mysql":
|
|
# MySQL 5.7 regexes are byte-oriented; a Unicode whitespace regex
|
|
# can misclassify CJK text. Removing each complete whitespace
|
|
# character preserves str.strip()'s nonempty test on 5.7 and 8.x.
|
|
stripped = text
|
|
for char in _STRIP_CHARS:
|
|
stripped = f"replace({stripped}, {bound(char)}, '')"
|
|
has_text = f"length({stripped}) > 0"
|
|
else:
|
|
trim = "btrim" if family == "postgres" else "trim"
|
|
# SQLite length(TEXT) stops at NUL, unlike Python str.strip().
|
|
has_text = f"{trim}({text}, {bound(_STRIP_CHARS)}) <> ''"
|
|
clause = f"({kind} = '{string_kind}' AND {has_text})"
|
|
elif element.operation == "truthy":
|
|
if family == "sqlite":
|
|
containers = f"json({raw}) NOT IN ('[]', '{{}}')"
|
|
boolean = f"{kind} = 'true'"
|
|
elif family == "mysql":
|
|
containers = f"json_length({raw}) > 0"
|
|
boolean = f"{text} = 'true'"
|
|
else:
|
|
containers = f"CAST({raw} AS jsonb) NOT IN ('[]'::jsonb, '{{}}'::jsonb)"
|
|
boolean = f"{text} = 'true'"
|
|
cast = "NUMERIC" if family == "postgres" else "DOUBLE PRECISION"
|
|
number = f"({text} + 0e0)" if family == "mysql" else f"CAST({text} AS {cast})"
|
|
nonempty = f"{text} <> ''" if family == "sqlite" else f"length({text}) > 0"
|
|
clause = (
|
|
f"CASE WHEN {kind} IN ({number_kinds}) THEN {number} <> 0 "
|
|
f"WHEN {kind} = '{string_kind}' THEN {nonempty} "
|
|
f"WHEN {kind} IN ('array', 'object') THEN {containers} "
|
|
f"WHEN {kind} IN ('boolean', 'true', 'false') THEN {boolean} ELSE FALSE END"
|
|
)
|
|
elif element.value is None:
|
|
# Preserve this project's md.get(key) == None semantics for old rows.
|
|
clause = f"({kind} IS NULL OR {kind} = 'null')"
|
|
elif isinstance(element.value, (bool, int, float)):
|
|
if family != "sqlite":
|
|
raise NotImplementedError("Server numeric metadata requires the exact residual filter")
|
|
value = int(element.value) if isinstance(element.value, bool) else element.value
|
|
# Preserve md.get(key) == value: Python considers True == 1 == 1.0.
|
|
# Text ids remain distinct from numbers. SQLite compares its extracted
|
|
# integers/reals natively; server numeric filters use the residual path.
|
|
param = bound(value, Float() if isinstance(value, float) else BigInteger())
|
|
clause = f"({kind} IN ({number_kinds}, 'true', 'false') AND {text} = {param})"
|
|
else:
|
|
# Binary comparison is essential for task/agent ids: MySQL TEXT can
|
|
# inherit a case-insensitive collation and ignores trailing spaces.
|
|
extract = f"CAST({text} AS BINARY)" if family == "mysql" else text
|
|
param = bound(element.value)
|
|
if family == "mysql":
|
|
param = f"CAST({param} AS BINARY)"
|
|
clause = f"({kind} = '{string_kind}' AND {extract} = {param})"
|
|
# Missing/null values must remain false under NOT (system exclusions).
|
|
return f"COALESCE(({clause}), FALSE)"
|
|
|
|
|
|
@compiles(JsonPredicate, "sqlite")
|
|
def _sqlite(element, compiler, **kw):
|
|
return _compile(element, compiler, "sqlite", **kw)
|
|
|
|
|
|
@compiles(JsonPredicate, "postgresql")
|
|
def _postgres(element, compiler, **kw):
|
|
return _compile(element, compiler, "postgres", **kw)
|
|
|
|
|
|
@compiles(JsonPredicate, "mysql")
|
|
def _mysql(element, compiler, **kw):
|
|
return _compile(element, compiler, "mysql", **kw)
|
|
|
|
|
|
@compiles(JsonPredicate)
|
|
def _unsupported(element, compiler, **kw):
|
|
raise NotImplementedError(f"JSON predicates unsupported for {compiler.dialect.name}")
|