pyoq-sql 1.0.0__py3-none-any.whl

pyoq/query/sqlite/commands.py

@@ -0,0 +1,201 @@
+"""SQLite cursor execution and cooperative interruption."""
+
+from __future__ import annotations
+
+import sqlite3
+from collections.abc import Callable
+from time import monotonic
+from typing import NoReturn, TypeVar, cast
+
+from pyoq.errors import QueryCancelledError, QueryExecutionError, QueryTimeoutError
+from pyoq.query.execution import ExecutionControl
+from pyoq.query.sqlite.health import is_connection_failure
+from pyoq.query.sqlite.parameters import SQLiteParameter
+
+OperationResult = TypeVar("OperationResult")
+Invalidator = Callable[[], None]
+
+
+class SQLiteExecutionMonitor:
+    __slots__ = ("_control", "_installed", "_started")
+
+    def __init__(self, control: ExecutionControl | None) -> None:
+        self._control = control or ExecutionControl()
+        self._started = monotonic() if self._control.timeout is not None else 0.0
+        self._installed = False
+
+    def install(
+        self,
+        connection: sqlite3.Connection,
+        invalidate: Invalidator,
+    ) -> None:
+        self.checkpoint()
+        if not self._requires_progress_handler():
+            return
+        try:
+            connection.set_progress_handler(
+                self._progress_handler,
+                self._control.progress_steps,
+            )
+            self._installed = True
+        except sqlite3.Error as error:
+            _raise_driver_failure(error, invalidate)
+
+    def checkpoint(self) -> None:
+        token = self._control.cancellation_token
+        if token is not None and token.cancelled:
+            raise QueryCancelledError("SQLite query execution was cancelled")
+        timeout = self._control.timeout
+        if timeout is not None and monotonic() - self._started >= timeout:
+            raise QueryTimeoutError("SQLite query execution exceeded its timeout")
+
+    def translate_driver_error(
+        self,
+        error: sqlite3.Error,
+        invalidate: Invalidator,
+    ) -> NoReturn:
+        if is_connection_failure(error):
+            invalidate()
+        self.checkpoint()
+        _raise_query_execution_error(error)
+
+    def close(
+        self,
+        connection: sqlite3.Connection,
+        invalidate: Invalidator,
+    ) -> None:
+        if not self._installed:
+            return
+        self._installed = False
+        try:
+            connection.set_progress_handler(None, 0)
+        except sqlite3.Error:
+            invalidate()
+
+    def _progress_handler(self) -> int:
+        return int(self._stopped())
+
+    def _stopped(self) -> bool:
+        token = self._control.cancellation_token
+        if token is not None and token.cancelled:
+            return True
+        timeout = self._control.timeout
+        return timeout is not None and monotonic() - self._started >= timeout
+
+    def _requires_progress_handler(self) -> bool:
+        return (
+            self._control.timeout is not None
+            or self._control.cancellation_token is not None
+        )
+
+
+class SQLiteCursorSession:
+    __slots__ = ("_closed", "_connection", "_cursor", "_invalidate", "_monitor")
+
+    def __init__(
+        self,
+        connection: sqlite3.Connection,
+        cursor: sqlite3.Cursor,
+        monitor: SQLiteExecutionMonitor,
+        invalidate: Invalidator,
+    ) -> None:
+        self._connection = connection
+        self._cursor = cursor
+        self._monitor = monitor
+        self._invalidate = invalidate
+        self._closed = False
+
+    @property
+    def cursor(self) -> sqlite3.Cursor:
+        return self._cursor
+
+    def checkpoint(self) -> None:
+        self._monitor.checkpoint()
+
+    def translate_driver_error(self, error: sqlite3.Error) -> NoReturn:
+        self._monitor.translate_driver_error(error, self._invalidate)
+
+    def fetch_batch(self, size: int) -> list[tuple[object, ...]]:
+        try:
+            self._monitor.checkpoint()
+            rows = cast(list[tuple[object, ...]], self.cursor.fetchmany(size))
+            self._monitor.checkpoint()
+        except sqlite3.Error as error:
+            self._monitor.translate_driver_error(error, self._invalidate)
+        return rows
+
+    def close(self) -> None:
+        if self._closed:
+            return
+        self._closed = True
+        _close_cursor(self._cursor, self._invalidate)
+        self._monitor.close(self._connection, self._invalidate)
+
+
+def run_sql(
+    connection: sqlite3.Connection,
+    sql: str,
+    parameters: tuple[SQLiteParameter, ...],
+    operation: Callable[[sqlite3.Cursor], OperationResult],
+    control: ExecutionControl | None,
+    invalidate: Invalidator,
+) -> OperationResult:
+    session = open_cursor(connection, sql, parameters, control, invalidate)
+    try:
+        result = operation(session.cursor)
+        session.checkpoint()
+        return result
+    except sqlite3.Error as error:
+        session.translate_driver_error(error)
+    finally:
+        session.close()
+
+
+def open_cursor(
+    connection: sqlite3.Connection,
+    sql: str,
+    parameters: tuple[SQLiteParameter, ...],
+    control: ExecutionControl | None,
+    invalidate: Invalidator,
+) -> SQLiteCursorSession:
+    monitor = SQLiteExecutionMonitor(control)
+    cursor: sqlite3.Cursor | None = None
+    try:
+        monitor.install(connection, invalidate)
+        cursor = connection.execute(sql, parameters)
+        monitor.checkpoint()
+        return SQLiteCursorSession(connection, cursor, monitor, invalidate)
+    except sqlite3.Error as error:
+        if cursor is not None:
+            _close_cursor(cursor, invalidate)
+        monitor.close(connection, invalidate)
+        monitor.translate_driver_error(error, invalidate)
+    except BaseException:
+        if cursor is not None:
+            _close_cursor(cursor, invalidate)
+        monitor.close(connection, invalidate)
+        raise
+
+
+def _close_cursor(cursor: sqlite3.Cursor, invalidate: Invalidator) -> None:
+    try:
+        cursor.close()
+    except sqlite3.Error:
+        invalidate()
+
+
+def _raise_driver_failure(
+    error: sqlite3.Error,
+    invalidate: Invalidator,
+) -> NoReturn:
+    if is_connection_failure(error):
+        invalidate()
+    _raise_query_execution_error(error)
+
+
+def _raise_query_execution_error(error: sqlite3.Error) -> NoReturn:
+    message = f"SQLite query execution failed: {error}"
+    raise QueryExecutionError(message) from error
+
+
+__all__ = ("SQLiteCursorSession", "open_cursor", "run_sql")

pyoq/query/sqlite/compiler.py

@@ -0,0 +1,139 @@
+"""SQLite query and statement compiler."""
+
+from __future__ import annotations
+
+from typing import Protocol
+
+from pyoq.errors import UnsupportedQueryError
+from pyoq.query.execution import CompiledQuery, StatementKind
+from pyoq.query.rendering.identifiers import IdentifierQuoting
+from pyoq.query.rendering.queries import QueryRenderer
+from pyoq.query.select_nodes import (
+    JoinKind,
+    LockNode,
+    NullOrder,
+    QueryNode,
+    SelectNode,
+    SetNode,
+)
+from pyoq.query.sqlite.context import SQLiteCompilationContext
+from pyoq.query.sqlite.expressions import SQLiteExpressionCompiler
+from pyoq.query.sqlite.identifiers import QUOTING
+from pyoq.query.sqlite.model import SQLiteCapabilities
+from pyoq.query.sqlite.writes import SQLiteWriteCompilation
+from pyoq.query.statements import StatementNode, StatementProvider
+from pyoq.query.write_nodes import CallNode, DeleteNode, InsertNode, UpdateNode
+
+_STATEMENT_NODES = (
+    SelectNode,
+    SetNode,
+    InsertNode,
+    UpdateNode,
+    DeleteNode,
+    CallNode,
+)
+
+
+class QueryProvider(Protocol):
+    @property
+    def node(self) -> QueryNode: ...
+
+
+class SQLiteCompiler:
+    __slots__ = ("_capabilities",)
+
+    def __init__(
+        self,
+        capabilities: SQLiteCapabilities | None = None,
+    ) -> None:
+        self._capabilities = capabilities or SQLiteCapabilities()
+
+    @property
+    def capabilities(self) -> SQLiteCapabilities:
+        return self._capabilities
+
+    def compile(
+        self,
+        statement: StatementProvider | StatementNode,
+        /,
+    ) -> CompiledQuery:
+        node = statement if isinstance(statement, _STATEMENT_NODES) else statement.node
+        context = SQLiteCompilationContext(self.capabilities)
+        if isinstance(node, CallNode):
+            message = "SQLite keeps no stored procedures; run the statements themselves"
+            raise UnsupportedQueryError(message)
+        if isinstance(node, (InsertNode, UpdateNode, DeleteNode)):
+            sql, statement_kind = SQLiteWriteCompilation(context).compile(node)
+            return context.result(sql, statement_kind, node)
+        sql = _SQLiteQueryCompilation(context).compile(node)
+        return context.result(sql, StatementKind.SELECT, node)
+
+
+class _SQLiteQueryCompilation(QueryRenderer):
+    __slots__ = ("_context",)
+
+    def __init__(self, context: SQLiteCompilationContext) -> None:
+        self._context = context
+        super().__init__(SQLiteExpressionCompiler(context))
+
+    @property
+    def dialect(self) -> str:
+        return "SQLite"
+
+    @property
+    def quoting(self) -> IdentifierQuoting:
+        return QUOTING
+
+    def require_join_capability(self, kind: JoinKind) -> None:
+        if kind is JoinKind.RIGHT and not self._context.capabilities.right_join:
+            message = "SQLite RIGHT JOIN is disabled by compiler capabilities"
+            raise UnsupportedQueryError(message)
+        if kind is JoinKind.FULL and not self._context.capabilities.full_join:
+            message = "SQLite FULL JOIN is disabled by compiler capabilities"
+            raise UnsupportedQueryError(message)
+
+    def locking(self, node: LockNode) -> str:
+        """SQLite locks a database rather than rows, so it cannot do this."""
+        del node
+        message = (
+            "SQLite locks the whole database rather than rows, so it cannot "
+            "hold the rows a query read; use a transaction instead"
+        )
+        raise UnsupportedQueryError(message)
+
+    def null_ordered(
+        self,
+        expression: str,
+        direction: str,
+        nulls: NullOrder,
+    ) -> str:
+        if not self._context.capabilities.null_ordering:
+            message = "SQLite null ordering is disabled by compiler capabilities"
+            raise UnsupportedQueryError(message)
+        return super().null_ordered(expression, direction, nulls)
+
+    def require_recursive_common_tables(self) -> None:
+        if not self._context.capabilities.recursive_common_tables:
+            message = "SQLite recursive common tables are disabled"
+            raise UnsupportedQueryError(message)
+
+    def require_lateral_sources(self) -> None:
+        if not self._context.capabilities.lateral_sources:
+            message = "SQLite has no lateral sources"
+            raise UnsupportedQueryError(message)
+
+    def append_pagination(
+        self,
+        fragments: list[str],
+        limit: int | None,
+        offset: int | None,
+    ) -> None:
+        if limit is not None:
+            fragments.append(f"LIMIT {self._context.bind_page_value(limit)}")
+        elif offset is not None:
+            fragments.append("LIMIT -1")
+        if offset is not None:
+            fragments.append(f"OFFSET {self._context.bind_page_value(offset)}")
+
+
+__all__ = ("QueryProvider", "SQLiteCompiler")

pyoq/query/sqlite/context.py

@@ -0,0 +1,20 @@
+"""SQLite compilation context."""
+
+from __future__ import annotations
+
+from pyoq.query.execution import CompilationContext
+from pyoq.query.sqlite.model import SQLiteCapabilities
+
+
+class SQLiteCompilationContext(CompilationContext):
+    __slots__ = ("capabilities",)
+
+    def __init__(self, capabilities: SQLiteCapabilities) -> None:
+        super().__init__(
+            capabilities.maximum_parameters,
+            capabilities.parameter_style,
+        )
+        self.capabilities = capabilities
+
+
+__all__ = ("SQLiteCompilationContext",)

pyoq/query/sqlite/executor.py

@@ -0,0 +1,119 @@
+"""Synchronous typed SQLite query execution."""
+
+from __future__ import annotations
+
+import sqlite3
+from collections.abc import Callable
+from typing import TypeVar
+
+from pyoq.query.execution import ExecutionControl
+from pyoq.query.select import Query
+from pyoq.query.sqlite.commands import open_cursor, run_sql
+from pyoq.query.sqlite.compiler import SQLiteCompiler
+from pyoq.query.sqlite.model import CompiledQuery
+from pyoq.query.sqlite.operations import SQLiteQueryOperations
+from pyoq.query.sqlite.parameters import adapt_parameters
+from pyoq.query.sqlite.pool import SQLitePool
+from pyoq.query.sqlite.streaming import (
+    SQLiteRowStream,
+    SQLiteStreamPolicy,
+    SQLiteStreamResource,
+)
+from pyoq.query.sqlite.transactions import SQLiteTransaction, TransactionMode
+
+OperationResult = TypeVar("OperationResult")
+ResultRow = TypeVar("ResultRow")
+
+
+class SQLiteExecutor(SQLiteQueryOperations):
+    __slots__ = ("_pool",)
+
+    def __init__(
+        self,
+        pool: SQLitePool,
+        compiler: SQLiteCompiler | None = None,
+    ) -> None:
+        self._pool = pool
+        super().__init__(compiler or SQLiteCompiler())
+
+    @property
+    def pool(self) -> SQLitePool:
+        return self._pool
+
+    def stream(
+        self,
+        query: Query[ResultRow],
+        /,
+        *,
+        policy: SQLiteStreamPolicy | None = None,
+    ) -> SQLiteRowStream[ResultRow]:
+        stream_policy = policy or SQLiteStreamPolicy()
+        statement = self.compiler.compile(query)
+        return SQLiteRowStream(
+            lambda: self._open_stream(
+                statement.sql,
+                statement.parameters,
+                stream_policy,
+            ),
+            stream_policy,
+            statement.row_types,
+        )
+
+    def transaction(
+        self,
+        mode: TransactionMode = TransactionMode.DEFERRED,
+        *,
+        control: ExecutionControl | None = None,
+    ) -> SQLiteTransaction:
+        return SQLiteTransaction(
+            self.pool,
+            self.compiler,
+            mode,
+            control,
+        )
+
+    def _run(
+        self,
+        statement: CompiledQuery,
+        operation: Callable[[sqlite3.Cursor], OperationResult],
+        control: ExecutionControl | None,
+    ) -> OperationResult:
+        parameters = adapt_parameters(statement.parameters)
+        lease = self.pool.checkout()
+        with lease as connection:
+            return run_sql(
+                connection,
+                statement.sql,
+                parameters,
+                operation,
+                control,
+                lease.invalidate,
+            )
+
+    def _open_stream(
+        self,
+        sql: str,
+        parameters: tuple[object, ...],
+        policy: SQLiteStreamPolicy,
+    ) -> SQLiteStreamResource:
+        adapted_parameters = adapt_parameters(parameters)
+        lease = self.pool.checkout()
+        connection = lease.__enter__()
+        try:
+            session = open_cursor(
+                connection,
+                sql,
+                adapted_parameters,
+                policy.control,
+                lease.invalidate,
+            )
+        except BaseException:
+            lease.__exit__(None, None, None)
+            raise
+        return SQLiteStreamResource(
+            session,
+            lambda: lease.__exit__(None, None, None),
+        )
+
+
+__all__ = ("SQLiteExecutor",)

pyoq/query/sqlite/expressions.py

@@ -0,0 +1,224 @@
+"""SQLite scalar expression compilation."""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+
+from pyoq.config import DatabaseDialect
+from pyoq.errors import UnsupportedQueryError
+from pyoq.query.nodes import (
+    BinaryNode,
+    BinaryOperator,
+    DatePart,
+    ExtractNode,
+    FieldNode,
+    FunctionName,
+    FunctionNode,
+    JsonContainsNode,
+    JsonKeysNode,
+    JsonNode,
+    JsonOperation,
+    JsonWrite,
+    JsonWriteNode,
+    ScalarFamily,
+)
+from pyoq.query.rendering.expressions import (
+    CollectionFunctions,
+    ExpressionRenderer,
+)
+from pyoq.query.rendering.identifiers import IdentifierQuoting
+from pyoq.query.results import expression_family
+from pyoq.query.select_nodes import JsonElementsSourceNode
+from pyoq.query.sqlite.context import SQLiteCompilationContext
+from pyoq.query.sqlite.identifiers import QUOTING
+from pyoq.schema.models import TypeKind
+
+_SQLITE_BINARY_SQL = {
+    BinaryOperator.IS_DISTINCT_FROM: "IS NOT",
+    BinaryOperator.IS_NOT_DISTINCT_FROM: "IS",
+}
+
+_DATE_FORMAT = {
+    DatePart.YEAR: "%Y",
+    DatePart.MONTH: "%m",
+    DatePart.DAY: "%d",
+    DatePart.HOUR: "%H",
+    DatePart.MINUTE: "%M",
+    DatePart.SECOND: "%S",
+}
+
+
+_SQLITE_CAST_SQL: Mapping[TypeKind, str] = {
+    TypeKind.BOOLEAN: "INTEGER",
+    TypeKind.SMALL_INTEGER: "INTEGER",
+    TypeKind.INTEGER: "INTEGER",
+    TypeKind.BIG_INTEGER: "INTEGER",
+    TypeKind.DECIMAL: "NUMERIC",
+    TypeKind.REAL: "REAL",
+    TypeKind.DOUBLE: "REAL",
+    TypeKind.STRING: "TEXT",
+    TypeKind.BINARY: "BLOB",
+}
+"""The five storage classes SQLite has, and nothing it would only tolerate.
+
+SQLite accepts `CAST(x AS DATE)` and answers with an integer, because an
+unknown type name falls back to affinity rather than being rejected. Naming a
+type it does not have would promise a date and return something else, so the
+kinds it cannot do are absent and refused.
+"""
+
+_SQLITE_CHOOSING_SQL: Mapping[FunctionName, str] = {
+    FunctionName.GREATEST: "MAX",
+    FunctionName.LEAST: "MIN",
+}
+
+
+_SQLITE_JSON_WRITES: Mapping[JsonWrite, str] = {
+    JsonWrite.SET: "JSON_SET",
+    JsonWrite.INSERT: "JSON_INSERT",
+    JsonWrite.REPLACE: "JSON_REPLACE",
+}
+
+
+class SQLiteExpressionCompiler(ExpressionRenderer[SQLiteCompilationContext]):
+    __slots__ = ()
+
+    @property
+    def collection_functions(self) -> CollectionFunctions:
+        return CollectionFunctions(
+            "json_group_array",
+            "json_object",
+            "json_array()",
+            "json_array",
+            "json_group_object",
+            ordered=True,
+        )
+
+    @property
+    def dialect(self) -> str:
+        return "SQLite"
+
+    def qualified_field(self, node: FieldNode) -> str:
+        return QUOTING.qualified(
+            node.name,
+            table=node.table,
+            schema=node.schema,
+            catalog=node.catalog,
+        )
+
+    def binary_operator_sql(self, operator: BinaryOperator) -> str | None:
+        override = _SQLITE_BINARY_SQL.get(operator)
+        if override is not None:
+            return override
+        return super().binary_operator_sql(operator)
+
+    def validate_binary(self, node: BinaryNode) -> None:
+        if node.operator not in {BinaryOperator.ADD, BinaryOperator.SUBTRACT}:
+            return
+        families = {expression_family(node.left), expression_family(node.right)}
+        if ScalarFamily.TEMPORAL in families:
+            message = "SQLite temporal arithmetic requires an explicit expression"
+            raise UnsupportedQueryError(message)
+
+    def render_json_elements(self, node: JsonElementsSourceNode) -> str:
+        """SQLite walks a JSON value and already calls the column it yields value."""
+        subject = self.compile(node.expression)
+        return f"JSON_EACH({subject}) AS {QUOTING.quote(node.alias)}"
+
+    def render_json_keys(self, node: JsonKeysNode) -> str:
+        """SQLite walks the members of an object one row at a time."""
+        subject = self.compile(node.expression)
+        return f"(SELECT JSON_GROUP_ARRAY(key) FROM JSON_EACH({subject}))"  # noqa: S608
+
+    def render_json(self, node: JsonNode) -> str:
+        """SQLite walks a written path."""
+        subject = self.compile(node.expression)
+        path = self.json_path_text(node.path)
+        if node.operation is JsonOperation.REMOVE:
+            return f"JSON_REMOVE({subject}, {path})"
+        if node.operation is JsonOperation.EXISTS:
+            return f"(JSON_TYPE({subject}, {path}) IS NOT NULL)"
+        if node.operation is JsonOperation.LENGTH:
+            return f"JSON_ARRAY_LENGTH({subject}, {path})"
+        if node.operation is JsonOperation.TEXT:
+            return f"JSON_EXTRACT({subject}, {path})"
+        return f"({subject} -> {path})"
+
+    def render_json_write(self, node: JsonWriteNode) -> str:
+        """Compiled in the order it is written, because binding follows that."""
+        subject = self.compile(node.expression)
+        if node.operation is JsonWrite.MERGE:
+            return f"JSON_PATCH({subject}, {self.compile(node.value)})"
+        if node.operation is JsonWrite.CONCAT:
+            raise UnsupportedQueryError(_NO_CONCAT)
+        path = self.json_path_text(node.path)
+        value = self.compile(node.value)
+        return f"{_SQLITE_JSON_WRITES[node.operation]}({subject}, {path}, {value})"
+
+    def render_json_contains(self, node: JsonContainsNode) -> str:
+        """SQLite has no containment operator, and will not pretend to."""
+        del node
+        message = "SQLite has no JSON containment; compare the members you mean instead"
+        raise UnsupportedQueryError(message)
+
+    @property
+    def quoting(self) -> IdentifierQuoting:
+        return QUOTING
+
+    @property
+    def database_dialect(self) -> DatabaseDialect:
+        return DatabaseDialect.SQLITE
+
+    def render_extract(self, node: ExtractNode) -> str:
+        date_format = _DATE_FORMAT.get(node.part)
+        if date_format is None:
+            message = f"unsupported SQLite date part: {node.part}"
+            raise UnsupportedQueryError(message)
+        return (
+            f"CAST(STRFTIME('{date_format}', {self.compile(node.expression)}) "
+            "AS INTEGER)"
+        )
+
+    @property
+    def cast_type_names(self) -> Mapping[TypeKind, str]:
+        return _SQLITE_CAST_SQL
+
+    def require_window_functions(self) -> None:
+        """Window functions need SQLite 3.25 or later."""
+        if not self._context.capabilities.window_functions:
+            message = "this SQLite build cannot measure a row against its neighbours"
+            raise UnsupportedQueryError(message)
+
+    def choosing_function_sql(self, name: FunctionName) -> str:
+        """SQLite spells the widest and narrowest of several as MAX and MIN."""
+        return _SQLITE_CHOOSING_SQL.get(name, super().choosing_function_sql(name))
+
+    def render_string_predicate(self, node: FunctionNode) -> str:
+        left_node, right_node = node.arguments
+        left = self.compile(left_node)
+        right = self.compile(right_node)
+        if node.name is FunctionName.CONTAINS:
+            return f"(INSTR({left}, {right}) > 0)"
+        if node.name is FunctionName.STARTS_WITH:
+            comparison = self.compile(right_node)
+            return f"(SUBSTR({left}, 1, LENGTH({right})) = {comparison})"
+        if node.name is FunctionName.ENDS_WITH:
+            comparison = self.compile(right_node)
+            return f"(SUBSTR({left}, -LENGTH({right})) = {comparison})"
+        message = f"unsupported SQLite function: {node.name}"
+        raise UnsupportedQueryError(message)
+
+    def require_raw_expressions(self) -> None:
+        if not self._context.capabilities.raw_expressions:
+            message = "SQLite raw expressions are disabled by compiler capabilities"
+            raise UnsupportedQueryError(message)
+
+
+__all__ = ("SQLiteExpressionCompiler", "expression_family")
+
+
+_NO_CONCAT = (
+    "SQLite merges a document all the way down and has no one-level "
+    "concatenation; json_merge says what it does"
+)
+"""Refused rather than emulated, because all the way down is not one level."""