pyoq-sql 1.0.2__py3-none-any.whl

pyoq/diagnostics/observation.py

@@ -0,0 +1,227 @@
+"""Watching what one scope executed, so a repeated query becomes visible.
+
+A query issued once per row of a previous result is the shape of an N+1 access,
+and it is invisible from inside the loop that causes it. Recording shapes for
+the length of a scope makes it visible from outside.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from inspect import currentframe
+from threading import Lock
+from types import FrameType
+
+from pyoq.diagnostics.fingerprint import QueryShape, query_shape
+from pyoq.errors import QueryValidationError
+from pyoq.query.execution import CompiledQuery
+from pyoq.query.execution.results import StatementKind
+
+DEFAULT_SHAPE_LIMIT = 256
+DEFAULT_SITE_LIMIT = 5
+DEFAULT_REPEAT_THRESHOLD = 2
+
+_OWN_PACKAGE = "pyoq."
+
+
+@dataclass(frozen=True, slots=True)
+class CallSite:
+    """Where in the caller's own code a query was issued."""
+
+    file: str
+    line: int
+    function: str
+
+    def __str__(self) -> str:
+        return f"{self.file}:{self.line} in {self.function}"
+
+
+@dataclass(frozen=True, slots=True)
+class RepeatedQuery:
+    """One shape a scope executed more than once, and where from."""
+
+    shape: QueryShape
+    executions: int
+    kinds: frozenset[StatementKind]
+    sites: tuple[CallSite, ...]
+
+    def describe(self) -> str:
+        places = "; ".join(str(site) for site in self.sites)
+        suffix = f" from {places}" if places else ""
+        return f"{self.executions} executions of {self.shape.sql}{suffix}"
+
+
+def _no_kinds() -> set[StatementKind]:
+    return set()
+
+
+def _no_sites() -> dict[CallSite, None]:
+    """An ordered set of places, so a report reads in the order they occurred."""
+    return {}
+
+
+@dataclass(slots=True)
+class _Observed:
+    shape: QueryShape
+    executions: int = 0
+    kinds: set[StatementKind] = field(default_factory=_no_kinds)
+    sites: dict[CallSite, None] = field(default_factory=_no_sites)
+
+
+class QueryObserver:
+    """What one scope executed, kept within a fixed budget of memory.
+
+    Diagnostics must not become the thing that exhausts a process, so the number
+    of distinct shapes and the number of places recorded per shape are both
+    capped. Executions keep being counted after the caps are reached.
+
+    An observer is safe to share. A pool hands connections to whichever thread
+    asks, so an observer watching a whole application is written to from several
+    at once, and reading a report while that happens must not fail.
+    """
+
+    __slots__ = (
+        "_capture_sites",
+        "_executions",
+        "_lock",
+        "_observed",
+        "_shape_limit",
+        "_site_limit",
+        "_unrecorded_shapes",
+    )
+
+    def __init__(
+        self,
+        *,
+        shape_limit: int = DEFAULT_SHAPE_LIMIT,
+        site_limit: int = DEFAULT_SITE_LIMIT,
+        capture_sites: bool = True,
+    ) -> None:
+        _require_positive(shape_limit, "shape limit")
+        _require_positive(site_limit, "site limit")
+        self._shape_limit = shape_limit
+        self._site_limit = site_limit
+        self._capture_sites = capture_sites
+        self._observed: dict[str, _Observed] = {}
+        self._executions = 0
+        self._unrecorded_shapes = 0
+        self._lock = Lock()
+
+    @property
+    def executions(self) -> int:
+        """Every statement this scope executed, including ones not recorded."""
+        return self._executions
+
+    @property
+    def shapes(self) -> int:
+        with self._lock:
+            return len(self._observed)
+
+    @property
+    def unrecorded_shapes(self) -> int:
+        """Distinct shapes seen after the limit, counted but not kept."""
+        return self._unrecorded_shapes
+
+    def record(self, statement: CompiledQuery, /) -> int:
+        """Record a statement and report how often its shape has now run.
+
+        A shape seen after the cap reports nothing, because it is counted but
+        not kept. A query repeated is one shape repeated, so a scope with more
+        distinct shapes than the cap has a different problem.
+        """
+        shape = query_shape(statement.sql)
+        site = self._site()
+        with self._lock:
+            self._executions += 1
+            observed = self._observed.get(shape.digest)
+            if observed is None:
+                if len(self._observed) >= self._shape_limit:
+                    self._unrecorded_shapes += 1
+                    return 0
+                observed = _Observed(shape)
+                self._observed[shape.digest] = observed
+            observed.executions += 1
+            observed.kinds.add(statement.statement_kind)
+            if site is not None and len(observed.sites) < self._site_limit:
+                observed.sites[site] = None
+            return observed.executions
+
+    def report(self, digest: str, /) -> RepeatedQuery | None:
+        """One shape's report, for a caller that already knows which."""
+        with self._lock:
+            observed = self._observed.get(digest)
+            if observed is None:
+                return None
+            return _report(observed)
+
+    def repeated(
+        self,
+        *,
+        threshold: int = DEFAULT_REPEAT_THRESHOLD,
+    ) -> tuple[RepeatedQuery, ...]:
+        """Shapes executed at least ``threshold`` times, busiest first."""
+        _require_positive(threshold, "repeat threshold")
+        with self._lock:
+            matching = [
+                _report(observed)
+                for observed in self._observed.values()
+                if observed.executions >= threshold
+            ]
+        matching.sort(key=lambda entry: (-entry.executions, entry.shape.digest))
+        return tuple(matching)
+
+    def _site(self) -> CallSite | None:
+        """Read the caller's frame outside the lock, where it still is theirs."""
+        return calling_site() if self._capture_sites else None
+
+
+def _report(observed: _Observed) -> RepeatedQuery:
+    return RepeatedQuery(
+        observed.shape,
+        observed.executions,
+        frozenset(observed.kinds),
+        tuple(observed.sites),
+    )
+
+
+def calling_site() -> CallSite | None:
+    """The nearest frame that is not PyOQ's own.
+
+    A query is issued from inside this library, so the frame that matters is
+    the first one above it, which is the caller's own code. An interpreter that
+    does not offer frames reports no place rather than refusing to run.
+    """
+    frame = currentframe()
+    while frame is not None:
+        if not _is_own_frame(frame):
+            return CallSite(
+                frame.f_code.co_filename,
+                frame.f_lineno,
+                frame.f_code.co_qualname,
+            )
+        frame = frame.f_back
+    return None
+
+
+def _is_own_frame(frame: FrameType) -> bool:
+    module = frame.f_globals.get("__name__", "")
+    return isinstance(module, str) and (
+        module == "pyoq" or module.startswith(_OWN_PACKAGE)
+    )
+
+
+def _require_positive(value: int, label: str) -> None:
+    if value < 1:
+        message = f"{label} must be at least one"
+        raise QueryValidationError(message)
+
+
+__all__ = (
+    "DEFAULT_REPEAT_THRESHOLD",
+    "DEFAULT_SHAPE_LIMIT",
+    "DEFAULT_SITE_LIMIT",
+    "CallSite",
+    "QueryObserver",
+    "RepeatedQuery",
+    "calling_site",
+)

pyoq/diagnostics/scoped.py

@@ -0,0 +1,103 @@
+"""Counting what one scope's database work actually does.
+
+A budget is only worth holding if something records against it. Every read and
+write funnels through one place before it reaches a driver, so that is where a
+statement is counted, and nothing about a dialect has to know a budget exists.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, TypeVar
+
+from pyoq.query.execution import QueryOperations
+from pyoq.query.execution.aio import AsyncQueryOperations
+
+if TYPE_CHECKING:
+    from collections.abc import Awaitable, Callable
+
+    from pyoq.diagnostics.budget import QueryScope
+    from pyoq.query.execution import (
+        BulkPlan,
+        CompiledQuery,
+        DatabaseCursor,
+        ExecutionControl,
+        WriteProvider,
+    )
+    from pyoq.query.execution.aio import AsyncDatabaseCursor
+    from pyoq.query.statements import StatementCompiler
+    from pyoq.query.write_nodes import WriteNode
+
+OperationResult = TypeVar("OperationResult")
+
+
+class ScopedOperations(QueryOperations):
+    """A database that reports every statement to a scope before running it.
+
+    Recording happens after compilation and before the driver sees anything, so
+    a statement that a budget refuses never reaches the database at all.
+    """
+
+    __slots__ = ("_inner", "_scope")
+
+    def __init__(self, inner: QueryOperations, scope: QueryScope) -> None:
+        self._inner = inner
+        self._scope = scope
+
+    @property
+    def scope(self) -> QueryScope:
+        return self._scope
+
+    @property
+    def compiler(self) -> StatementCompiler:
+        return self._inner.compiler
+
+    def plan_bulk(self, statement: WriteProvider | WriteNode, /) -> BulkPlan:
+        return self._inner.plan_bulk(statement)
+
+    def last_inserted_id(self, cursor: DatabaseCursor) -> int | None:
+        return self._inner.last_inserted_id(cursor)
+
+    def _run(
+        self,
+        statement: CompiledQuery,
+        operation: Callable[[DatabaseCursor], OperationResult],
+        control: ExecutionControl | None,
+    ) -> OperationResult:
+        self._scope.record(statement)
+        return self._inner._run(statement, operation, control)
+
+
+class AsyncScopedOperations(AsyncQueryOperations):
+    """The asynchronous counterpart, counting at the same point."""
+
+    __slots__ = ("_inner", "_scope")
+
+    def __init__(self, inner: AsyncQueryOperations, scope: QueryScope) -> None:
+        self._inner = inner
+        self._scope = scope
+
+    @property
+    def scope(self) -> QueryScope:
+        return self._scope
+
+    @property
+    def compiler(self) -> StatementCompiler:
+        return self._inner.compiler
+
+    def plan_bulk(self, statement: WriteProvider | WriteNode, /) -> BulkPlan:
+        return self._inner.plan_bulk(statement)
+
+    def last_inserted_id(self, cursor: AsyncDatabaseCursor) -> int | None:
+        return self._inner.last_inserted_id(cursor)
+
+    async def _run(
+        self,
+        statement: CompiledQuery,
+        operation: Callable[[AsyncDatabaseCursor], Awaitable[OperationResult]],
+        control: ExecutionControl | None,
+    ) -> OperationResult:
+        self._scope.record(statement)
+        return await self._inner._run(statement, operation, control)
+
+
+__all__ = ("AsyncScopedOperations", "ScopedOperations")

pyoq/django/__init__.py

@@ -0,0 +1,15 @@
+"""Running PyOQ inside a Django project."""
+
+from pyoq.django.execution import Dialect, DjangoOperations, alias_for, dialect_for
+from pyoq.django.generation import GenerationRequest, generate
+from pyoq.django.schema import MigrationStateSchemaSource
+
+__all__ = (
+    "Dialect",
+    "DjangoOperations",
+    "GenerationRequest",
+    "MigrationStateSchemaSource",
+    "alias_for",
+    "dialect_for",
+    "generate",
+)

pyoq/django/apps.py

@@ -0,0 +1,17 @@
+"""The application entry a project adds to install PyOQ's commands."""
+
+from __future__ import annotations
+
+from django.apps import AppConfig
+
+
+class PyoqConfig(AppConfig):
+    """Named for the toolkit rather than for the package it lives under.
+
+    The label Django would infer from `pyoq.django` is `django`, which reads as
+    the framework itself everywhere a label is shown.
+    """
+
+    name = "pyoq.django"
+    label = "pyoq"
+    verbose_name = "PyOQ"

pyoq/django/execution.py

@@ -0,0 +1,317 @@
+"""Running PyOQ statements on the connection Django already has open.
+
+Django owns its connections, their aliases, and whatever transaction is in
+progress. PyOQ borrows one rather than opening a pool beside it, because two
+pools against one database is two views of what has been committed.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable, Generator, Sequence
+from contextlib import AbstractContextManager, contextmanager
+from dataclasses import dataclass
+from typing import Protocol, TypeVar, cast
+
+from django.db import DEFAULT_DB_ALIAS, Error, connections, router, transaction
+from django.db.backends.base.base import BaseDatabaseWrapper
+from django.db.models import Model
+
+from pyoq.django.parameters import DjangoValue, adapt_parameters
+from pyoq.django.timeouts import (
+    MYSQL_TIMEOUT,
+    POSTGRES_TIMEOUT,
+    TimeoutSetting,
+    applied,
+)
+from pyoq.errors import (
+    OperationUnavailableError,
+    QueryCancelledError,
+    QueryExecutionError,
+)
+from pyoq.query.execution import (
+    BulkPlan,
+    CompiledQuery,
+    DatabaseCursor,
+    ExecutionControl,
+    QueryOperations,
+    WriteProvider,
+)
+from pyoq.query.statements import StatementCompiler
+from pyoq.query.write_nodes import WriteNode
+
+OperationResult = TypeVar("OperationResult")
+
+
+class BulkPlanner(Protocol):
+    def plan(self, statement: WriteProvider | WriteNode, /) -> BulkPlan: ...
+
+
+class LastRowCursor(Protocol):
+    @property
+    def lastrowid(self) -> object: ...
+
+
+@dataclass(frozen=True, slots=True)
+class Dialect:
+    """What the backend Django is already speaking needs from PyOQ.
+
+    Values are adapted the same way they are on PyOQ's own connections, because
+    a backend does not care whose connection it is: a decimal or a date still
+    has to reach it in the shape its driver accepts.
+    """
+
+    compiler: StatementCompiler
+    planner: BulkPlanner
+    timeout: TimeoutSetting | None = None
+
+
+def _sqlite() -> Dialect:
+    from pyoq.query.execution import ParameterStyle
+    from pyoq.query.sqlite import (
+        SQLiteBulkPlanner,
+        SQLiteCapabilities,
+        SQLiteCompiler,
+    )
+
+    compiler = SQLiteCompiler(SQLiteCapabilities(parameter_style=ParameterStyle.FORMAT))
+    return Dialect(compiler, SQLiteBulkPlanner(compiler))
+
+
+def _postgres() -> Dialect:
+    from pyoq.query.postgres import PostgresBulkPlanner, PostgresCompiler
+
+    compiler = PostgresCompiler()
+    return Dialect(compiler, PostgresBulkPlanner(compiler), POSTGRES_TIMEOUT)
+
+
+def _mysql() -> Dialect:
+    from pyoq.query.mysql import MySQLBulkPlanner, MySQLCompiler
+
+    compiler = MySQLCompiler()
+    return Dialect(compiler, MySQLBulkPlanner(compiler), MYSQL_TIMEOUT)
+
+
+_DIALECTS: dict[str, Callable[[], Dialect]] = {
+    "sqlite": _sqlite,
+    "postgresql": _postgres,
+    "mysql": _mysql,
+}
+"""How to build each dialect, called once a connection turns out to be one.
+
+A dialect package imports the driver it speaks to, so building these eagerly
+would make the Django extra need every driver PyOQ supports.
+"""
+
+_DRIVER_EXTRAS: dict[str, str] = {"postgresql": "postgres", "mysql": "mysql"}
+"""What a project installs to reach each backend.
+
+SQLite is absent because it needs nothing, and so can never be the one that
+is missing.
+"""
+
+
+def dialect_for(connection: BaseDatabaseWrapper, /) -> Dialect:
+    """Choose the dialect Django is already speaking.
+
+    The connection decides, not the caller. Speaking a different dialect than
+    the one on the other end of the socket is not a choice worth offering.
+    """
+    factory = _DIALECTS.get(connection.vendor)
+    if factory is None:
+        message = (
+            f"PyOQ has no dialect for the Django backend {connection.vendor!r}; "
+            f"supported backends: {', '.join(sorted(_DIALECTS))}"
+        )
+        raise OperationUnavailableError(message)
+    return _built(factory, connection.vendor)
+
+
+def _built(factory: Callable[[], Dialect], vendor: str) -> Dialect:
+    """Say which extra is missing, rather than which module is.
+
+    A project installs the drivers it uses, so reaching a backend it did not
+    install is an ordinary mistake and worth answering with the fix.
+    """
+    try:
+        return factory()
+    except ImportError as error:
+        extra = _DRIVER_EXTRAS[vendor]
+        message = (
+            f"the Django backend {vendor!r} needs a driver PyOQ does not "
+            f"install by default; add it with pip install "
+            f"'pyoq-sql[django,{extra}]'"
+        )
+        raise OperationUnavailableError(message) from error
+
+
+def alias_for(
+    model: type[Model],
+    /,
+    *,
+    write: bool = False,
+    **hints: object,
+) -> str:
+    """Ask Django's routers which database a model lives in.
+
+    Routing is a question about a model, because that is the only thing a
+    router is given to decide on. A caller with a model gets the answer its own
+    routers would give; a caller without one names the alias instead.
+    """
+    if write:
+        return str(router.db_for_write(model, **hints) or DEFAULT_DB_ALIAS)
+    return str(router.db_for_read(model, **hints) or DEFAULT_DB_ALIAS)
+
+
+class DjangoOperations(QueryOperations):
+    """Typed operations against one Django database alias.
+
+    The connection is resolved for each statement rather than held, because
+    Django hands a different connection to each thread and closes them between
+    requests, so holding one would outlive what it belongs to.
+    """
+
+    __slots__ = ("_alias", "_dialect")
+
+    def __init__(self, alias: str = DEFAULT_DB_ALIAS) -> None:
+        self._alias = alias
+        self._dialect = dialect_for(self.connection)
+
+    @classmethod
+    def for_model(
+        cls,
+        model: type[Model],
+        /,
+        *,
+        write: bool = False,
+        **hints: object,
+    ) -> DjangoOperations:
+        """Run against whichever database this project routes a model to."""
+        return cls(alias_for(model, write=write, **hints))
+
+    @property
+    def alias(self) -> str:
+        return self._alias
+
+    @property
+    def connection(self) -> BaseDatabaseWrapper:
+        """Whatever connection Django holds for this alias right now."""
+        return connections[self._alias]
+
+    @property
+    def compiler(self) -> StatementCompiler:
+        return self._dialect.compiler
+
+    @property
+    def in_transaction(self) -> bool:
+        """Whether this database is inside a transaction right now."""
+        return self.connection.in_atomic_block
+
+    def atomic(
+        self,
+        *,
+        savepoint: bool = True,
+        durable: bool = False,
+    ) -> AbstractContextManager[None]:
+        """A transaction on this database, rather than on whichever is default.
+
+        Django's own block covers the default alias unless told otherwise, so a
+        caller running here against another alias would get no transaction at
+        all and no indication of it. Nesting one of these is a savepoint, which
+        is what Django already does.
+
+        PyOQ never commits or rolls back. The block that opened a transaction is
+        the one that ends it, and here that block is Django's.
+        """
+        return transaction.atomic(
+            using=self._alias, savepoint=savepoint, durable=durable
+        )
+
+    def on_commit(
+        self,
+        callback: Callable[[], object],
+        /,
+        *,
+        robust: bool = False,
+    ) -> None:
+        """Run something once this database's transaction has committed.
+
+        Outside a transaction it runs immediately, which is Django's own rule
+        and the right one: there is nothing left to wait for.
+        """
+        transaction.on_commit(callback, using=self._alias, robust=robust)
+
+    def plan_bulk(self, statement: WriteProvider | WriteNode, /) -> BulkPlan:
+        return self._dialect.planner.plan(statement)
+
+    def last_inserted_id(self, cursor: DatabaseCursor) -> int | None:
+        identifier = _last_row_identifier(cursor)
+        if not isinstance(identifier, int) or identifier <= 0:
+            return None
+        return identifier
+
+    def _run(
+        self,
+        statement: CompiledQuery,
+        operation: Callable[[DatabaseCursor], OperationResult],
+        control: ExecutionControl | None,
+    ) -> OperationResult:
+        _checkpoint(control)
+        connection = self.connection
+        adapted = adapt_parameters(connection, statement.parameters)
+        parameters = cast("Sequence[DjangoValue]", adapted)
+        with self._bounded(control), connection.cursor() as cursor:
+            try:
+                # Django's stub omits the duration values its own interval
+                # backends accept, so the call is wider than the stub says.
+                cursor.execute(statement.sql, parameters)  # type: ignore[arg-type]
+                result = operation(cast("DatabaseCursor", cursor))
+            except Error as error:
+                message = f"Django query execution failed: {error}"
+                raise QueryExecutionError(message) from error
+        _checkpoint(control)
+        return result
+
+    @contextmanager
+    def _bounded(self, control: ExecutionControl | None) -> Generator[None]:
+        """Bound the statement where the backend offers a way to.
+
+        SQLite has no session setting for this, and MySQL's covers reads only,
+        so a caller is given what the backend can honour rather than a promise
+        it cannot keep.
+        """
+        setting = self._dialect.timeout
+        seconds = None if control is None else control.timeout
+        if setting is None or seconds is None:
+            yield
+            return
+        with applied(setting, self.connection, seconds):
+            yield
+
+
+def _last_row_identifier(cursor: DatabaseCursor) -> object:
+    """Not every backend offers one, and none of them promise a type."""
+    return getattr(cursor, "lastrowid", None)
+
+
+def _checkpoint(control: ExecutionControl | None) -> None:
+    """Honour a caller's own limits without touching Django's connection.
+
+    A server-side timeout is a change to session state, and the session belongs
+    to Django. What is left is refusing before and after a statement, which is
+    what a caller asking for cancellation can be given honestly.
+    """
+    if control is None:
+        return
+    token = control.cancellation_token
+    if token is not None and token.cancelled:
+        message = "Django query execution was cancelled"
+        raise QueryCancelledError(message)
+
+
+__all__ = (
+    "BulkPlanner",
+    "Dialect",
+    "DjangoOperations",
+    "alias_for",
+    "dialect_for",
+)