sqlakit 0.1.0__py3-none-any.whl

sqlakit/sql.py

@@ -0,0 +1,211 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, TypeVar, cast
+
+import sqlalchemy as sa
+
+from ._sql import (
+    BaseSQLQuery,
+    Templates,
+    require_pydantic,
+    templates_of,
+)
+
+if TYPE_CHECKING:
+    from collections.abc import Iterator, Sequence
+
+    from sqlalchemy.engine import Result, ScalarResult
+    from sqlalchemy.sql import Executable
+
+    from ._db import Database
+
+__all__ = ["SQL", "SQLQuery", "SQLRows", "Templates"]
+
+RowT = TypeVar("RowT")
+OtherT = TypeVar("OtherT")
+
+
+class SQL:
+    """The SQL templates of one database, ready to run.
+
+    Reached as `db.sql`, and where the templates are is the database's own
+    `templates=`:
+
+    ```python
+    db = Database(DB_URL, templates="app/sql")
+
+    db.sql("reports/by_team.sql", since=since).typed(TeamReport).all()
+    db.sql.from_string("SELECT count(*) FROM users").scalars().one()
+    ```
+    """
+
+    def __init__(self, db: Database) -> None:
+        self.db = db
+
+    def __repr__(self) -> str:
+        return f"{type(self).__name__}({self.db!r})"
+
+    def __call__(self, template: str, **context: Any) -> SQLQuery:  # noqa: ANN401
+        """Read the rows of a template. Short for `from_file`.
+
+        ```python
+        db.sql("users/active.sql", team="red").all()
+        ```
+        """
+        return self.from_file(template, **context)
+
+    def from_file(self, template: str, **context: Any) -> SQLQuery:  # noqa: ANN401
+        """Read the rows of a template kept under the database's ``templates=``.
+
+        ```python
+        db.sql.from_file("users/active.sql", team="red").all()
+        ```
+
+        The keyword arguments are the template's context.
+        """
+        return SQLQuery(self.db, template, context)
+
+    def from_string(self, source: str, **context: Any) -> SQLQuery:  # noqa: ANN401
+        """Read the rows of SQL written out here rather than kept in a file.
+
+        ```python
+        db.sql.from_string("SELECT id FROM users WHERE team = {{ team }}", team="red")
+        ```
+
+        Values are named in `{{ }}` and passed by keyword, as in a template. A
+        `:name` or a `?` binds nothing here, and rendering says so rather than
+        reaching the driver. It needs no ``templates=``.
+        """
+        return SQLQuery(self.db, source, context, inline=True)
+
+    def from_statement(self, statement: Executable) -> SQLQuery:
+        """Read the rows of a statement built with SQLAlchemy.
+
+        ```python
+        db.sql.from_statement(sa.text("SELECT ...").bindparams(id=1)).all()
+        ```
+
+        Nothing is rendered: the statement is the one that runs, parameters
+        and all. What this adds is the reading, `typed` and `chunks` included.
+        """
+        return SQLQuery(self.db, statement, {})
+
+    @property
+    def templates(self) -> Templates:
+        """Where this database looks for its templates."""
+        return templates_of(self.db)
+
+    def check(self) -> None:
+        """Compile every `.sql` template, so a broken one fails where deploys do.
+
+        Call it at startup, next to the rest of the wiring: a template is read
+        when something asks for it, and that is a poor time to find a typo.
+        """
+        self.templates.check()
+
+
+class SQLRows(BaseSQLQuery[RowT, "Database"]):
+    """The rows of a SQL template, on the connection of the block it runs in.
+
+    What the rows are is settled: reading them is all that is left.
+    """
+
+    def all(self) -> Sequence[RowT]:
+        """Return every row."""
+        return cast("Sequence[RowT]", self._shaped(self._rows().all()))
+
+    def first(self) -> RowT | None:
+        """Return the first row, or None."""
+        return cast("RowT | None", self._shaped_one(self._rows().first()))
+
+    def one(self) -> RowT:
+        """Return the single row.
+
+        Raises:
+            NoResultFound: if there is none.
+            MultipleResultsFound: if there is more than one.
+
+        """
+        return cast("RowT", self._shaped_one(self._rows().one()))
+
+    def one_or_none(self) -> RowT | None:
+        """Return the single row, or None.
+
+        Raises:
+            MultipleResultsFound: if there is more than one.
+
+        """
+        return cast("RowT | None", self._shaped_one(self._rows().one_or_none()))
+
+    def chunks(self, size: int) -> Iterator[Sequence[RowT]]:
+        """Read every row, ``size`` of them at a time.
+
+        One statement, fetched in batches, for a job that walks a table too large to
+        hold. The rows come off a cursor the database holds open, so the whole walk
+        is one transaction.
+        """
+        for batch in self._rows(size=size).partitions(size):
+            yield cast("Sequence[RowT]", self._shaped(batch))
+
+    def execute(self) -> int:
+        """Run it for what it writes, and return how many rows it touched.
+
+        For a template that inserts, updates or deletes. Inside a transaction the
+        write is part of it and the block decides. In a block with no transaction
+        the call commits for itself, as ORM writes do.
+        """
+        connection = self._connection()
+        result = connection.execute(self.statement)
+        if not self.db.in_transaction():
+            connection.commit()
+        return result.rowcount
+
+    def _rows(self, *, size: int | None = None) -> Result[Any] | ScalarResult[Any]:
+        result = self._connection().execute(self._executable(size=size))
+        return result.scalars() if self.scalar else result
+
+    def _connection(self) -> sa.Connection:
+        """Return this block's connection, with any pending ORM writes on it."""
+        if self.db.in_session():
+            self.db.session.flush()
+        return self.db.connection
+
+
+class SQLQuery(SQLRows[sa.Row[Any]]):
+    """The rows a SQL template returns, as `db.sql(...)` hands them over.
+
+    `typed` and `scalars` say what one row is; both return rows that read the
+    same way and carry no further say, so each is asked once.
+    """
+
+    def typed(self, type_: type[OtherT], /) -> SQLRows[OtherT]:
+        """Read the rows as this type, one row at a time.
+
+        ```python
+        db.sql("reports/by_team.sql", since=since).typed(TeamReport).all()
+        ```
+
+        The type is what one row becomes, and the terminal decides the container.
+        Anything pydantic can validate works: a model, a dataclass, a
+        `TypedDict`. The type also says how much of the row it takes: one built
+        from columns is given the whole row, and anything else is given the
+        first column, so `SELECT count(*)` with `typed(int)` reads as an `int`.
+
+        Raises:
+            MissingDependencyError: if pydantic is not installed.
+
+        """
+        require_pydantic()
+        return cast("SQLRows[OtherT]", self._as(SQLRows, type_=type_))
+
+    def scalars(self) -> SQLRows[Any]:
+        """Read the first column of each row instead of whole rows.
+
+        ```python
+        db.sql.from_string("SELECT count(*) FROM users").scalars().one()
+        ```
+
+        `typed()` does the same when the type is worth naming; this is for when it is
+        not, and needs no pydantic.
+        """
+        return self._as(SQLRows, scalar=True)

sqlakit/testing.py

@@ -0,0 +1,91 @@
+"""Helpers a test needs of a database, beyond a rollback and a schema."""
+
+from __future__ import annotations
+
+from contextlib import ExitStack, contextmanager
+from typing import TYPE_CHECKING, Any
+
+from . import db as sync_db
+from ._recording import Recording, check, require_expectation
+from .asyncio import db as async_db
+from .exceptions import UnknownDatabaseError
+
+if TYPE_CHECKING:
+    from collections.abc import Iterator
+
+__all__ = ["assert_queries"]
+
+
+@contextmanager
+def assert_queries(
+    count: int | None = None,
+    *,
+    at_most: int | None = None,
+    duplicates: bool = True,
+    using: str | Any = None,  # noqa: ANN401
+) -> Iterator[Recording]:
+    """Assert what the block asks of the database.
+
+    ```python
+    with assert_queries(2):
+        User.query.order_by("name").page(limit=10)
+
+    with assert_queries(at_most=5):
+        render(dashboard)
+
+    with assert_queries(duplicates=False):
+        render(users)  # the N+1 test, without a number
+    ```
+
+    The three checks stand alone or together: a count, a ceiling, and whether a
+    statement may run twice. What fails prints the statements, numbered and
+    timed, with the repeated ones pointing at each other.
+
+    ``using`` is the database to watch, as an alias or as the database itself.
+    Left out, it watches the importable registries, `sqlakit.db` and
+    `sqlakit.asyncio.db`, and every database each of them has. Awaited work is
+    watched the same way, so the block stays `with`.
+
+    Transaction control is not counted: `BEGIN` and `COMMIT` reach a cursor on
+    some drivers and not others.
+
+    Raises:
+        TypeError: if there is nothing to assert, or nothing to watch.
+        UnknownDatabaseError: if no configured registry has that alias.
+
+    """
+    require_expectation(count, at_most, duplicates)
+
+    with ExitStack() as stack:
+        recording = Recording()
+        for db in _watched(using):
+            stack.enter_context(db.recording(into=recording))
+        yield recording
+
+    check(recording, count=count, at_most=at_most, duplicates=duplicates)
+
+
+def _watched(using: str | Any) -> list[Any]:  # noqa: ANN401
+    """Return the databases to record: the one named, or every configured one."""
+    if using is None:
+        return _configured()
+    if not isinstance(using, str):
+        return [using]
+    registries = _configured()
+    watched = [registry[using] for registry in registries if using in registry]
+    if not watched:
+        known = {alias for registry in registries for alias in registry.aliases}
+        raise UnknownDatabaseError(using, tuple(sorted(known)))
+    return watched
+
+
+def _configured() -> list[Any]:
+    """Return the importable registries an application has configured."""
+    registries = [db for db in (sync_db, async_db) if db.is_configured]
+    if not registries:
+        message = (
+            "assert_queries has no database to watch. Configure `sqlakit.db`, "
+            "or name one with `assert_queries(..., using=db)`."
+        )
+        raise TypeError(message)
+    return registries

sqlakit/types.py

@@ -0,0 +1,104 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Literal, TypedDict
+
+if TYPE_CHECKING:
+    from collections.abc import Callable, Mapping, Sequence
+
+    import sqlalchemy as sa
+    from sqlalchemy.engine import Connection, Engine
+    from sqlalchemy.orm import Query, Session
+    from sqlalchemy.pool import Pool
+
+__all__ = ["DatabaseConfig", "EngineArgs", "SessionArgs", "UrlParts"]
+
+
+class EngineArgs(TypedDict, total=False):
+    """Keyword arguments accepted by [`sqlalchemy.create_engine`](https://docs.sqlalchemy.org/en/20/core/engines.html#sqlalchemy.create_engine)."""
+
+    connect_args: dict[str, Any]
+    echo: bool | Literal["debug"]
+    echo_pool: bool | Literal["debug"]
+    enable_from_linting: bool
+    execution_options: dict[str, Any]
+    hide_parameters: bool
+    insertmanyvalues_page_size: int
+    isolation_level: str
+    json_deserializer: Callable[[str], Any]
+    json_serializer: Callable[[Any], str]
+    label_length: int | None
+    logging_name: str
+    max_identifier_length: int | None
+    max_overflow: int
+    module: Any
+    paramstyle: Literal["qmark", "numeric", "named", "format", "pyformat"]
+    plugins: list[str]
+    pool: Pool
+    pool_logging_name: str
+    pool_pre_ping: bool
+    pool_recycle: int
+    pool_reset_on_return: Literal["rollback", "commit"] | None
+    pool_size: int
+    pool_timeout: float
+    pool_use_lifo: bool
+    poolclass: type[Pool]
+    query_cache_size: int
+    skip_autocommit_rollback: bool
+    use_insertmanyvalues: bool
+
+
+class SessionArgs(TypedDict, total=False):
+    """Keyword arguments accepted by `sqlalchemy.orm.sessionmaker`.
+
+    No ``bind``: sessions bind to the connection of the surrounding block.
+    """
+
+    autobegin: bool
+    autoflush: bool
+    binds: dict[Any, Engine | Connection]
+    class_: type[Session]
+    enable_baked_queries: bool
+    expire_on_commit: bool
+    info: dict[Any, Any]
+    join_transaction_mode: Literal[
+        "conditional_savepoint",
+        "rollback_only",
+        "control_fully",
+        "create_savepoint",
+    ]
+    query_cls: type[Query[Any]]
+    twophase: bool
+
+
+class UrlParts(TypedDict, total=False):
+    """A database URL spelled out, as `sqlalchemy.URL.create` takes it."""
+
+    drivername: str
+    username: str | None
+    password: str | None
+    host: str | None
+    port: int | None
+    database: str | None
+    query: Mapping[str, Sequence[str] | str]
+
+
+class QueryStats(TypedDict):
+    """The numbers a recording adds up to, for a log read by machine."""
+
+    queries: int
+    milliseconds: float
+    slowest_milliseconds: float
+    duplicated: int
+    databases: tuple[str, ...]
+    label: str | None
+
+
+class DatabaseConfig(UrlParts, total=False):
+    """One database in a configuration keyed by alias.
+
+    Give it a ``url`` or the parts to build one from, never both.
+    """
+
+    url: str | sa.URL
+    engine_args: EngineArgs
+    session_args: SessionArgs