pyoq-sql 1.0.0__py3-none-any.whl

pyoq/diagnostics/budget.py

@@ -0,0 +1,136 @@
+"""What one scope is allowed to ask a database to do.
+
+A repeated query that is only reported is a repeated query that still ships. A
+budget turns the same observation into a refusal, at the point where the scope
+exceeds what it was allowed rather than after the fact.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import cast
+
+from pyoq.diagnostics.fingerprint import QueryShape, query_shape
+from pyoq.diagnostics.observation import QueryObserver, RepeatedQuery
+from pyoq.errors import QueryBudgetExceededError, QueryValidationError
+from pyoq.query.execution import CompiledQuery
+
+
+@dataclass(frozen=True, slots=True)
+class QueryBudget:
+    """A ceiling on one scope's database work.
+
+    ``maximum_repeats`` is the one that catches an N+1 access, because the shape
+    executed once per row of an earlier result is the shape that repeats.
+    """
+
+    maximum_queries: int | None = None
+    maximum_repeats: int | None = None
+
+    def __post_init__(self) -> None:
+        _require_optional_positive(self.maximum_queries, "maximum queries")
+        _require_optional_positive(self.maximum_repeats, "maximum repeats")
+
+    @property
+    def unlimited(self) -> bool:
+        return self.maximum_queries is None and self.maximum_repeats is None
+
+
+class QueryScope:
+    """One request's worth of database work, watched and bounded.
+
+    The scope holds an observer so that a refusal can say which shape ran too
+    often and where from, which is what makes the refusal actionable rather than
+    merely correct.
+    """
+
+    __slots__ = ("_budget", "_observer")
+
+    def __init__(
+        self,
+        budget: QueryBudget | None = None,
+        *,
+        observer: QueryObserver | None = None,
+    ) -> None:
+        self._budget = budget or QueryBudget()
+        self._observer = observer or QueryObserver()
+
+    @property
+    def budget(self) -> QueryBudget:
+        return self._budget
+
+    @property
+    def observer(self) -> QueryObserver:
+        return self._observer
+
+    @property
+    def executions(self) -> int:
+        return self._observer.executions
+
+    def record(self, statement: CompiledQuery, /) -> None:
+        """Record a statement, refusing the one that exceeds the budget.
+
+        The statement is recorded before it is judged, so a report taken after
+        a refusal includes the execution that caused it.
+        """
+        executions = self._observer.record(statement)
+        if self._budget.unlimited:
+            return
+        self._require_within_total()
+        self._require_within_repeats(statement, executions)
+
+    def repeated(self, *, threshold: int = 2) -> tuple[RepeatedQuery, ...]:
+        return self._observer.repeated(threshold=threshold)
+
+    def _require_within_total(self) -> None:
+        limit = self._budget.maximum_queries
+        if limit is None or self._observer.executions <= limit:
+            return
+        message = (
+            f"this scope executed {self._observer.executions} statements, beyond "
+            f"the {limit} it was allowed"
+        )
+        raise QueryBudgetExceededError(message)
+
+    def _require_within_repeats(
+        self,
+        statement: CompiledQuery,
+        executions: int,
+    ) -> None:
+        """Judge only the shape just recorded, which is the only one that moved.
+
+        Scanning every shape on every statement would make the cost of holding a
+        budget grow with the variety of a scope's queries.
+        """
+        limit = self._budget.maximum_repeats
+        if limit is None or executions <= limit:
+            return
+        digest = query_shape(statement.sql).digest
+        # A count above the limit means the observer kept this shape, and it
+        # never lets one go, so the report is there to be read.
+        repeated = cast("RepeatedQuery", self._observer.report(digest))
+        raise QueryBudgetExceededError(_repeat_message(repeated, limit))
+
+
+def _repeat_message(repeated: RepeatedQuery, limit: int) -> str:
+    places = "; ".join(str(site) for site in repeated.sites)
+    origin = f" from {places}" if places else ""
+    return (
+        f"one statement ran {repeated.executions} times in this scope, beyond "
+        f"the {limit} it was allowed{origin}: {_shape_text(repeated.shape)}"
+    )
+
+
+def _shape_text(shape: QueryShape) -> str:
+    return shape.sql
+
+
+def _require_optional_positive(value: int | None, label: str) -> None:
+    if value is None:
+        return
+    if isinstance(value, bool) or value < 1:
+        message = f"{label} must be a positive integer or None"
+        raise QueryValidationError(message)
+
+
+__all__ = ("QueryBudget", "QueryScope")

pyoq/diagnostics/events.py

@@ -0,0 +1,137 @@
+"""What a statement did, described without what it was asked about.
+
+A shape carries the SQL with its values taken out, which is safe to log. A bound
+value is not, and neither is the message a driver raises: PostgreSQL states the
+offending value inside it, so a failure that says only its type is the one that
+can be reported anywhere.
+
+A project that has decided otherwise says so through a policy, one field at a
+time. A value marked sensitive when the statement was compiled is never included
+whatever the policy says.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Protocol, TypeAlias
+
+if TYPE_CHECKING:
+    from pyoq.diagnostics.fingerprint import QueryShape
+    from pyoq.query.execution import CompiledQuery, StatementKind
+
+
+@dataclass(frozen=True, slots=True)
+class EventPolicy:
+    """What an event may carry beyond the shape of a statement.
+
+    Every field is off, so instrumentation added without a decision reports
+    nothing a query was asked about.
+    """
+
+    include_values: bool = False
+    include_failure_detail: bool = False
+    slow_after: float | None = None
+
+    def __post_init__(self) -> None:
+        if self.slow_after is not None and self.slow_after <= 0:
+            message = "a slow statement threshold must be positive"
+            raise ValueError(message)
+
+
+@dataclass(frozen=True, slots=True)
+class StatementStarted:
+    """A statement about to reach a driver."""
+
+    shape: QueryShape
+    kind: StatementKind
+    parameters: int
+    sensitive_parameters: int
+    values: tuple[object, ...] = ()
+
+
+@dataclass(frozen=True, slots=True)
+class StatementFinished:
+    """A statement the database answered."""
+
+    shape: QueryShape
+    kind: StatementKind
+    parameters: int
+    sensitive_parameters: int
+    duration: float
+    rows: int | None = None
+    slow: bool = False
+    values: tuple[object, ...] = ()
+
+
+@dataclass(frozen=True, slots=True)
+class StatementFailed:
+    """A statement the database refused.
+
+    `failure` names the error's type. Its message is only carried when a policy
+    says so, because a driver states the value that caused the failure in it.
+    """
+
+    shape: QueryShape
+    kind: StatementKind
+    parameters: int
+    sensitive_parameters: int
+    duration: float
+    failure: str
+    detail: str | None = None
+    values: tuple[object, ...] = ()
+
+
+StatementEvent: TypeAlias = StatementStarted | StatementFinished | StatementFailed
+
+
+class EventSink(Protocol):
+    """Somewhere for events to go."""
+
+    def record(self, event: StatementEvent, /) -> None: ...
+
+
+class CollectingSink:
+    """Keeps what it is given, for a caller that wants to look afterwards."""
+
+    __slots__ = ("events",)
+    events: list[StatementEvent]
+
+    def __init__(self) -> None:
+        self.events = []
+
+    def record(self, event: StatementEvent, /) -> None:
+        self.events.append(event)
+
+    def clear(self) -> None:
+        self.events.clear()
+
+
+def reportable_values(
+    statement: CompiledQuery,
+    policy: EventPolicy,
+    /,
+) -> tuple[object, ...]:
+    """The values an event may carry, which is none of them by default.
+
+    A value the compiler marked sensitive is left out even when a policy asks
+    for values, because marking it was the decision that it must not be shown.
+    """
+    if not policy.include_values:
+        return ()
+    return tuple(
+        value
+        for index, value in enumerate(statement.parameters)
+        if index not in statement.sensitive_parameter_indexes
+    )
+
+
+__all__ = (
+    "CollectingSink",
+    "EventPolicy",
+    "EventSink",
+    "StatementEvent",
+    "StatementFailed",
+    "StatementFinished",
+    "StatementStarted",
+    "reportable_values",
+)

pyoq/diagnostics/fingerprint.py

@@ -0,0 +1,267 @@
+"""Recognizing one query shape across the many times it is executed.
+
+A query executed in a loop differs only in its bound values, a query executed
+against a growing list of keys differs only in how many placeholders it carries,
+and a write chunked to fit a parameter limit differs only in how many groups of
+them it carries. All three are one shape, and seeing them as one is what makes a
+repeated query visible.
+
+A literal written into the statement itself is taken out for the same reason,
+and for one more: a shape is reported and logged, so anything left in it is
+reported and logged too. PyOQ binds every value it is given, but raw SQL can
+carry one, and a shape that kept it would put it in a log.
+"""
+
+from __future__ import annotations
+
+import re
+from dataclasses import dataclass
+from functools import lru_cache
+from hashlib import blake2b
+
+_ESCAPED_PERCENT = "\x00escaped-percent\x00"
+# PostgreSQL opens a string with a tag it chooses, and closes it with the same.
+_DOLLAR_TAG = re.compile(r"\$[A-Za-z_]\w*\$|\$\$")
+# What quotes a name rather than a value, in the dialects PyOQ speaks.
+# What quotes a name rather than a value. A single quote is absent because a
+# single quoted string is a value, and is read as one before this is asked.
+_NAME_QUOTES = '"`'
+# What begins a comment that runs to the end of its line.
+_LINE_COMMENTS = ("--", "#")
+# A number written in another radix, and one written in the usual way.
+_RADIX_NUMBER = re.compile(r"0[xX][0-9a-fA-F]+|0[bB][01]+")
+_DECIMAL_NUMBER = re.compile(r"\d+(?:\.\d+)?(?:[eE][+-]?\d+)?")
+_PLACEHOLDER = re.compile(r"\$\d+|%s|\?")
+_PLACEHOLDER_RUN = re.compile(r"\?(?:\s*,\s*\?)+")
+_GROUP_RUN = re.compile(r"\(\?\)(?:\s*,\s*\(\?\))+")
+_WHITESPACE = re.compile(r"\s+")
+_DIGEST_BYTES = 16
+# What a shape may carry into a log line or a span attribute.
+_MAXIMUM_SHAPE = 4096
+_TRUNCATED = " ..."
+
+
+@dataclass(frozen=True, slots=True)
+class QueryShape:
+    """One query with its values and their number taken out.
+
+    The shape carries no bound value, so it can be reported and logged without
+    revealing anything a query was asked about.
+    """
+
+    sql: str
+    digest: str
+
+    def __str__(self) -> str:
+        return self.sql
+
+
+# An application runs the same statements over and over, so the shape of one is
+# asked for far more often than it changes. The cache is bounded, because the
+# statements a process runs are not.
+_SHAPE_CACHE_SIZE = 2048
+
+
+@lru_cache(maxsize=_SHAPE_CACHE_SIZE)
+def query_shape(sql: str, /) -> QueryShape:
+    """The shape of a statement, remembered while it keeps being asked for.
+
+    A shape is written to a log and set on a span, so what it carries is
+    bounded. The digest is taken from the whole statement, so two that differ
+    only past the bound are still told apart.
+    """
+    normalized = normalize_sql(sql)
+    return QueryShape(_within_bounds(normalized), _digest(normalized))
+
+
+def _within_bounds(normalized: str) -> str:
+    """A generated statement can be enormous, and a log line cannot."""
+    if len(normalized) <= _MAXIMUM_SHAPE:
+        return normalized
+    return normalized[:_MAXIMUM_SHAPE] + _TRUNCATED
+
+
+def shape_cache_metrics() -> CacheMetrics:
+    """How well the shapes being asked for are already known."""
+    info = query_shape.cache_info()
+    return CacheMetrics(
+        hits=info.hits,
+        misses=info.misses,
+        held=info.currsize,
+        capacity=_SHAPE_CACHE_SIZE,
+    )
+
+
+def forget_shapes() -> None:
+    """Empty the cache, which only a test measuring it should need."""
+    query_shape.cache_clear()
+
+
+@dataclass(frozen=True, slots=True)
+class CacheMetrics:
+    """What a cache has been asked for and how much of it it kept."""
+
+    hits: int
+    misses: int
+    held: int
+    capacity: int
+
+    @property
+    def hit_rate(self) -> float:
+        asked = self.hits + self.misses
+        return 0.0 if asked == 0 else self.hits / asked
+
+
+def normalize_sql(sql: str, /) -> str:
+    """Reduce a statement to the shape it shares with its repetitions.
+
+    The statement is read once, left to right, because what a character means
+    depends on what it is inside. A quote inside a comment starts nothing, and
+    two dashes inside a string are not a comment, so deciding either one without
+    tracking the other gets both wrong.
+    """
+    pieces: list[str] = []
+    position = 0
+    length = len(sql)
+    while position < length:
+        step = (
+            _skip_comment(sql, position)
+            or _take_literal(sql, position)
+            or _take_quoted_name(sql, position)
+        )
+        if step is None:
+            pieces.append(sql[position])
+            position += 1
+            continue
+        text, position = step
+        pieces.append(text)
+    return _collapse("".join(pieces))
+
+
+def _skip_comment(sql: str, position: int) -> tuple[str, int] | None:
+    """A comment is dropped, because it says nothing about the shape."""
+    if sql.startswith(_LINE_COMMENTS, position):
+        ending = sql.find("\n", position)
+        return (" ", len(sql) if ending == -1 else ending)
+    if sql.startswith("/*", position):
+        return (" ", _end_of_block_comment(sql, position))
+    return None
+
+
+def _end_of_block_comment(sql: str, position: int) -> int:
+    """Past the closing marker, counting the nesting PostgreSQL allows.
+
+    A dialect that does not nest would have ended at the first marker, so this
+    can drop more than that dialect needed. Dropping too much loses a little of
+    a shape; dropping too little puts whatever was written there in a log.
+    """
+    depth = 0
+    index = position
+    while index < len(sql):
+        if sql.startswith("/*", index):
+            depth += 1
+            index += 2
+            continue
+        if sql.startswith("*/", index):
+            depth -= 1
+            index += 2
+            if depth == 0:
+                return index
+            continue
+        index += 1
+    return len(sql)
+
+
+def _take_literal(sql: str, position: int) -> tuple[str, int] | None:
+    """Every kind of literal becomes one placeholder."""
+    dollar = _take_dollar_quoted(sql, position)
+    if dollar is not None:
+        return dollar
+    if sql[position] == "'":
+        return ("?", _end_of_quoted(sql, position, "'"))
+    if _starts_number(sql, position):
+        return ("?", _end_of_number(sql, position))
+    return None
+
+
+def _take_dollar_quoted(sql: str, position: int) -> tuple[str, int] | None:
+    """PostgreSQL quotes with a tag of its own choosing, so it is read out."""
+    if sql[position] != "$":
+        return None
+    opening = _DOLLAR_TAG.match(sql, position)
+    if opening is None:
+        return None
+    tag = opening.group(0)
+    ending = sql.find(tag, opening.end())
+    return ("?", len(sql) if ending == -1 else ending + len(tag))
+
+
+def _end_of_quoted(sql: str, position: int, quote: str) -> int:
+    """Past the closing quote, counting a doubled or escaped one as inside."""
+    index = position + 1
+    while index < len(sql):
+        character = sql[index]
+        if character == "\\":
+            index += 2
+            continue
+        if character == quote:
+            if sql.startswith(quote * 2, index):
+                index += 2
+                continue
+            return index + 1
+        index += 1
+    return len(sql)
+
+
+def _take_quoted_name(sql: str, position: int) -> tuple[str, int] | None:
+    """A quoted name is a name, so it is kept as it was written."""
+    quote = sql[position]
+    if quote not in _NAME_QUOTES:
+        return None
+    ending = _end_of_quoted(sql, position, quote)
+    return (sql[position:ending], ending)
+
+
+def _starts_number(sql: str, position: int) -> bool:
+    """A decimal digit, which is narrower than what `isdigit` accepts.
+
+    A superscript is a digit to `isdigit` and not to the pattern that reads the
+    rest of the number, and taking one out a character at a time would mean a
+    shape that changes every time it is taken.
+    """
+    if not sql[position].isdecimal():
+        return False
+    before = sql[position - 1] if position else ""
+    return not (before.isalnum() or before in "_.$")
+
+
+def _end_of_number(sql: str, position: int) -> int:
+    """Past the whole number, in any of the ways one can be written."""
+    radix = _RADIX_NUMBER.match(sql, position)
+    if radix is not None:
+        return radix.end()
+    decimal = _DECIMAL_NUMBER.match(sql, position)
+    return position + 1 if decimal is None else decimal.end()
+
+
+def _collapse(text: str) -> str:
+    text = text.replace("%%", _ESCAPED_PERCENT)
+    text = _PLACEHOLDER.sub("?", text)
+    text = _PLACEHOLDER_RUN.sub("?", text)
+    text = _GROUP_RUN.sub("(?)", text)
+    text = _WHITESPACE.sub(" ", text).strip()
+    return text.replace(_ESCAPED_PERCENT, "%%")
+
+
+def _digest(normalized: str) -> str:
+    return blake2b(normalized.encode(), digest_size=_DIGEST_BYTES).hexdigest()
+
+
+__all__ = (
+    "CacheMetrics",
+    "QueryShape",
+    "forget_shapes",
+    "normalize_sql",
+    "query_shape",
+    "shape_cache_metrics",
+)