xtr-orm 1.5.0__py3-none-any.whl

xtr_orm/messenger/close_connection_middleware.py

@@ -0,0 +1,75 @@
+"""``orm_close_connection``: a worker gives its database connections back after each message."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import final
+
+from typing_extensions import override
+from xtr_messenger import (
+    Envelope,
+    MiddlewareInterface,
+    ReceivedStamp,
+    StackInterface,
+    as_middleware,
+)
+
+from xtr_orm.connection_registry import ConnectionRegistry
+from xtr_orm.exception import UnknownConnectionError
+
+__all__ = ["CloseConnectionMiddleware"]
+
+
+@final
+@as_middleware("orm_close_connection")
+class CloseConnectionMiddleware(MiddlewareInterface):
+    """Closes connections once a message a worker consumed is handled, saving connections.
+
+    A worker waiting for its next message then holds none open; the next
+    one connects again. A message dispatched in the process that sends it
+    closes nothing.
+
+    A worker handling several messages at once closes them once the last of
+    those still being handled is done: closing drops the pool, and dropping
+    it after each message would make every other one in flight connect anew.
+    """
+
+    __slots__ = ("_connection_names", "_connections", "_consumed")
+
+    def __init__(
+        self,
+        connections: ConnectionRegistry,
+        connection_names: str | Sequence[str] = (),
+    ) -> None:
+        """Close ``connection_names`` — every connection when none is given."""
+        self._connections = connections
+        self._consumed = 0
+        self._connection_names = (
+            (connection_names,) if isinstance(connection_names, str) else tuple(connection_names)
+        )
+
+    @override
+    async def handle(self, envelope: Envelope, stack: StackInterface, /) -> Envelope:
+        """Run the rest of the chain, then close the connections if a worker consumed it.
+
+        Closing waits for every consumed message this middleware is handling
+        to be done with.
+
+        Raises:
+            UnknownConnectionError: When a connection named is not registered.
+        """
+        names = self._connection_names or self._connections.names()
+        for name in names:
+            if not self._connections.has(name):
+                raise UnknownConnectionError(name, self._connections.names())
+        if envelope.last(ReceivedStamp) is None:
+            return await stack.next().handle(envelope, stack)
+        # Counted per instance, on one event loop: no lock needed.
+        self._consumed += 1
+        try:
+            return await stack.next().handle(envelope, stack)
+        finally:
+            self._consumed -= 1
+            if self._consumed == 0:
+                for name in names:
+                    await self._connections.close(name)

xtr_orm/messenger/open_transaction_logger_middleware.py

@@ -0,0 +1,90 @@
+"""``orm_open_transaction_logger``: say so when a handler leaves a transaction open."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from contextvars import ContextVar
+from typing import TYPE_CHECKING, Final, final
+
+from sqlalchemy.orm import SessionTransactionOrigin
+from typing_extensions import override
+from xtr_logging_contracts import LoggerInterface, NullLogger
+from xtr_messenger import Envelope, MiddlewareInterface, StackInterface, as_middleware
+
+from xtr_orm.connection_registry import ConnectionRegistry
+
+if TYPE_CHECKING:
+    from sqlalchemy.ext.asyncio import AsyncSession
+
+__all__ = ["OpenTransactionLoggerMiddleware"]
+
+_OPENED: Final = frozenset({SessionTransactionOrigin.BEGIN, SessionTransactionOrigin.BEGIN_NESTED})
+"""How a transaction a handler opened itself began; the session's own autobegin is not one."""
+
+_DISCARD: Final = NullLogger()
+"""Where findings go without a logger."""
+
+# Set while a message is checked, so a message dispatched during it is not checked twice.
+_handling: ContextVar[bool] = ContextVar("xtr_orm_open_transaction_logging", default=False)
+
+
+@final
+@as_middleware("orm_open_transaction_logger")
+class OpenTransactionLoggerMiddleware(MiddlewareInterface):
+    """Logs an error when a handler opened a transaction and did not close it.
+
+    A transaction counts when a handler began it — ``begin()`` or
+    ``begin_nested()`` on the unit of work's session — not when the session
+    began one by itself for a query.
+    """
+
+    __slots__ = ("_connection_names", "_connections", "_logger")
+
+    def __init__(
+        self,
+        connections: ConnectionRegistry,
+        logger: LoggerInterface = _DISCARD,
+        connection_names: str | Sequence[str] = (),
+    ) -> None:
+        """Check ``connection_names`` — every connection when none is given — logging to ``logger``.
+
+        Under a container, ``logger`` is the application's logger when the
+        logging bundle is active; otherwise what is found goes nowhere.
+        """
+        self._connections = connections
+        self._logger = logger
+        self._connection_names = (
+            (connection_names,) if isinstance(connection_names, str) else tuple(connection_names)
+        )
+
+    @override
+    async def handle(self, envelope: Envelope, stack: StackInterface, /) -> Envelope:
+        """Run the rest of the chain, then log each connection left with a transaction open."""
+        if _handling.get():
+            return await stack.next().handle(envelope, stack)
+        names = self._connection_names or self._connections.names()
+        sessions = {name: await self._connections.session(name) for name in names}
+        initial = {name: _opened(session) for name, session in sessions.items()}
+        token = _handling.set(True)
+        try:
+            return await stack.next().handle(envelope, stack)
+        finally:
+            _handling.reset(token)
+            left = [name for name, session in sessions.items() if _opened(session) > initial[name]]
+            if left:
+                self._logger.error(
+                    "A handler opened a transaction but did not close it.",
+                    {"connections": left, "message": envelope.message},
+                )
+
+
+def _opened(session: AsyncSession) -> int:
+    """Count the transactions open on ``session`` that someone began explicitly."""
+    sync = session.sync_session
+    transaction = sync.get_nested_transaction() or sync.get_transaction()
+    count = 0
+    while transaction is not None:
+        if transaction.is_active and transaction.origin in _OPENED:
+            count += 1
+        transaction = transaction.parent
+    return count

xtr_orm/messenger/transaction_middleware.py

@@ -0,0 +1,104 @@
+"""``orm_transaction``: every handler of a message in one database transaction."""
+
+from __future__ import annotations
+
+from contextvars import ContextVar
+from typing import TYPE_CHECKING, final
+
+from typing_extensions import override
+from xtr_messenger import (
+    Envelope,
+    HandledStamp,
+    HandlersFailedError,
+    MiddlewareInterface,
+    StackInterface,
+    as_middleware,
+)
+
+from xtr_orm.connection_registry import ConnectionRegistry
+
+if TYPE_CHECKING:
+    from sqlalchemy.ext.asyncio import AsyncSession
+
+__all__ = ["TransactionMiddleware"]
+
+# The sessions a transaction middleware has begun in this context: a message
+# dispatched while another is handled runs inside the outer one's transaction.
+_begun: ContextVar[frozenset[int]] = ContextVar("xtr_orm_transaction_begun", default=frozenset())
+
+
+@final
+@as_middleware("orm_transaction")
+class TransactionMiddleware(MiddlewareInterface):
+    """Wraps every handler of a message in a single transaction on one connection.
+
+    The transaction is the session of the message's unit of work, the one
+    every handler is given: it is committed — flushing what the handlers
+    added — once they all succeeded, and rolled back otherwise. Handlers
+    leave committing to it.
+
+    When a handler fails, the handled stamps of the others are dropped from
+    the error's envelope: their work was rolled back too, so a retry runs
+    every handler again.
+
+    A message dispatched while another is handled runs inside that one's
+    transaction, in a savepoint of its own: if it fails, what its handlers
+    wrote is rolled back to the savepoint — so a handler that catches the
+    failure and carries on commits only its own work — and the failure goes
+    on to the handler that dispatched it.
+    """
+
+    __slots__ = ("_connection_name", "_connections")
+
+    def __init__(self, connections: ConnectionRegistry, connection_name: str | None = None) -> None:
+        """Run transactions on ``connection_name``, the default connection when ``None``."""
+        self._connections = connections
+        self._connection_name = connection_name
+
+    @override
+    async def handle(self, envelope: Envelope, stack: StackInterface, /) -> Envelope:
+        """Run the rest of the chain inside a transaction, committed once it all succeeded."""
+        session = await self._connections.session(self._connection_name)
+        begun = _begun.get()
+        if id(session) in begun:
+            return await _in_savepoint(session, envelope, stack)
+        if not session.in_transaction():
+            _ = await session.begin()
+        token = _begun.set(begun | {id(session)})
+        try:
+            handled = await stack.next().handle(envelope, stack)
+            await session.commit()
+        except HandlersFailedError as error:
+            await session.rollback()
+            raise _unhandled(error) from error
+        except BaseException:
+            await session.rollback()
+            raise
+        finally:
+            _begun.reset(token)
+        return handled
+
+
+async def _in_savepoint(
+    session: AsyncSession, envelope: Envelope, stack: StackInterface
+) -> Envelope:
+    """Run the rest of the chain in a savepoint of the transaction under way."""
+    savepoint = await session.begin_nested()
+    try:
+        handled = await stack.next().handle(envelope, stack)
+    except HandlersFailedError as error:
+        if savepoint.is_active:
+            await savepoint.rollback()
+        raise _unhandled(error) from error
+    except BaseException:
+        if savepoint.is_active:
+            await savepoint.rollback()
+        raise
+    if savepoint.is_active:
+        await savepoint.commit()
+    return handled
+
+
+def _unhandled(error: HandlersFailedError) -> HandlersFailedError:
+    """Return ``error`` without the handled stamps of handlers whose work was rolled back."""
+    return HandlersFailedError(error.envelope.without_stamps(HandledStamp), error.errors)

xtr_orm/migrations/__init__.py

@@ -0,0 +1,29 @@
+"""Versioned schema changes: revision files, run and recorded against one database.
+
+:class:`Migrator` is the whole surface: it reads the revision files, plans a
+migration, runs it, writes new revisions — blank, diffed against the table
+definitions, or dumped from the database — and keeps the version and
+history tables in step.
+"""
+
+from __future__ import annotations
+
+from .available_migration import AvailableMigration
+from .direction import Direction
+from .executed_migration import ExecutedMigration
+from .execution_result import ExecutionResult
+from .migration_plan import MigrationPlan
+from .migration_status import MigrationStatus
+from .migrations_config import MigrationsConfig
+from .migrator import Migrator
+
+__all__ = [
+    "AvailableMigration",
+    "Direction",
+    "ExecutedMigration",
+    "ExecutionResult",
+    "MigrationPlan",
+    "MigrationStatus",
+    "MigrationsConfig",
+    "Migrator",
+]

xtr_orm/migrations/_environment.py

@@ -0,0 +1,188 @@
+"""The migration runner, set up in memory for one database.
+
+No ``env.py`` and no ``alembic.ini``: the revision directory is built from a
+:class:`MigrationsConfig`, and the runner is handed the connection, the
+table definitions and every option directly, the way a start-up script would
+hand them over.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import TYPE_CHECKING, Final, TypeAlias, cast, final
+
+from alembic.config import Config
+from alembic.runtime.environment import EnvironmentContext
+from alembic.runtime.migration import MigrationContext
+from alembic.script import ScriptDirectory
+
+if TYPE_CHECKING:
+    from collections.abc import Callable, Iterable, Sequence
+    from typing import IO
+
+    from alembic.autogenerate import RevisionContext
+    from alembic.autogenerate.api import AutogenContext
+    from alembic.config import PostWriteHookConfig
+    from alembic.runtime.migration import MigrationStep
+    from sqlalchemy import MetaData
+    from sqlalchemy.engine import Connection
+
+    from .migrations_config import MigrationsConfig
+
+    MigrationsFn: TypeAlias = Callable[[tuple[str, ...], MigrationContext], Iterable[MigrationStep]]
+    """What the runner asks for the steps to apply, given the current revisions."""
+
+__all__ = ["Environment"]
+
+_TEMPLATE_DIRECTORY: Final = Path(__file__).parent / "template"
+"""Where the package's own ``script.py.mako`` is."""
+
+
+@final
+class Environment:
+    """One database's revision directory and runner options."""
+
+    __slots__ = ("_config", "_metadata")
+
+    _config: MigrationsConfig
+    _metadata: MetaData | Sequence[MetaData] | None
+
+    def __init__(
+        self, config: MigrationsConfig, metadata: MetaData | Sequence[MetaData] | None
+    ) -> None:
+        """Set up for the revisions ``config`` describes, diffing against ``metadata``."""
+        self._config = config
+        self._metadata = metadata
+
+    def directory(self) -> ScriptDirectory:
+        """Read the revision directory afresh, so files written since are seen."""
+        config = self._config
+        hooks = cast(
+            "list[PostWriteHookConfig]",
+            [
+                {"_hook_name": f"hook_{index}", **hook}
+                for index, hook in enumerate(config.post_write_hooks)
+            ],
+        )
+        return ScriptDirectory(
+            config.template_directory or _TEMPLATE_DIRECTORY,
+            file_template=config.alembic_file_template(),
+            truncate_slug_length=config.truncate_slug_length,
+            version_locations=[config.directory],
+            timezone=config.timezone,
+            hooks=hooks,
+            recursive_version_locations=config.recursive_version_locations,
+            messaging_opts={"quiet": True},
+        )
+
+    def context(self, connection: Connection) -> MigrationContext:
+        """Return a context reading ``connection``'s version table, for reports."""
+        return MigrationContext.configure(
+            connection,
+            opts={
+                "version_table": self._config.version_table,
+                "version_table_schema": self._config.version_table_schema,
+            },
+        )
+
+    def run(  # noqa: PLR0913 — every runner option, each passed by keyword
+        self,
+        connection: Connection,
+        directory: ScriptDirectory,
+        steps: MigrationsFn,
+        *,
+        as_sql: bool = False,
+        output: IO[str] | None = None,
+        starting_rev: tuple[str, ...] | None = None,
+        purge: bool = False,
+        on_version_apply: Callable[..., None] | None = None,
+        revision_context: RevisionContext | None = None,
+        read_only: bool = False,
+    ) -> None:
+        """Run ``steps`` against ``connection``, or write their SQL to ``output``.
+
+        ``read_only`` leaves even a missing version table uncreated, for a run
+        that only reads — one writing a revision file.
+        """
+        config = self._config
+        extra: dict[str, object] = {}
+        if revision_context is not None:
+            extra["revision_context"] = revision_context
+            extra["template_args"] = revision_context.template_args
+        with EnvironmentContext(
+            Config(),
+            directory,
+            fn=steps,
+            as_sql=as_sql,
+            starting_rev=list(starting_rev) if as_sql and starting_rev else None,
+            purge=purge,
+            **extra,
+        ) as environment:
+            options: dict[str, object] = {
+                "connection": connection,
+                "target_metadata": self._metadata,
+                "version_table": config.version_table,
+                "version_table_schema": config.version_table_schema,
+                "transaction_per_migration": config.transaction_per_migration,
+                "compare_type": config.compare_type,
+                "compare_server_default": config.compare_server_default,
+                "render_as_batch": config.render_as_batch,
+                "include_schemas": config.include_schemas,
+                "include_name": _include_name(config.include_name, config.history_table),
+                "include_object": config.include_object,
+                "render_item": _render_item(config.render_item, config.user_module_prefix),
+                "process_revision_directives": config.process_revision_directives,
+                "user_module_prefix": config.user_module_prefix,
+                "literal_binds": as_sql,
+                "output_buffer": output,
+                "on_version_apply": on_version_apply,
+                "dont_mutate": read_only,
+            }
+            if config.transactional_ddl is not None:
+                options["transactional_ddl"] = config.transactional_ddl
+            options.update(config.context_options)
+            # Forwarded by name, the application's own options included; the
+            # runner checks each as it reads it.
+            configure = cast("Callable[..., None]", environment.configure)
+            configure(**options)
+            with environment.begin_transaction():
+                environment.run_migrations()
+
+
+def _include_name(
+    custom: Callable[..., bool] | None, history_table: str
+) -> Callable[[str | None, str, dict[str, str | None]], bool]:
+    """Keep the history table out of every diff, as the runner keeps its version table out.
+
+    Any configured filter decides for every other name.
+    """
+
+    def include_name(name: str | None, kind: str, parents: dict[str, str | None]) -> bool:
+        if kind == "table" and name == history_table:
+            return False
+        return custom is None or custom(name, kind, parents)
+
+    return include_name
+
+
+def _render_item(
+    custom: Callable[..., str | bool] | None, user_module_prefix: str | None
+) -> Callable[[str, object, AutogenContext], str | bool]:
+    """Chain the configured renderer with one importing each custom type's module.
+
+    A type from outside the database library is written under its own
+    module's name; the revision imports that module so the name resolves.
+    """
+
+    def render_item(kind: str, value: object, context: AutogenContext) -> str | bool:
+        if custom is not None:
+            rendered = custom(kind, value, context)
+            if rendered is not False:
+                return rendered
+        if kind == "type" and user_module_prefix is None:
+            module = type(value).__module__
+            if not module.startswith("sqlalchemy.") and module != "builtins":
+                context.imports.add(f"import {module}")
+        return False
+
+    return render_item

xtr_orm/migrations/_history.py

@@ -0,0 +1,178 @@
+"""The history table: a row per revision applied, beside the version table's current ones.
+
+The version table holds only the revisions a database is *at*; the migration
+runner rewrites it as it goes. The history table keeps what that table
+forgets — every revision applied, when, and how long it took — so a report
+can name each applied revision, including one whose file is gone.
+
+A row is written by the step that applies its revision, in the same
+transaction, so the two tables agree wherever the database's transactions
+reach.
+"""
+
+from __future__ import annotations
+
+import time
+from datetime import datetime
+from typing import TYPE_CHECKING, Final, cast, final
+
+from sqlalchemy import (
+    Column,
+    DateTime,
+    Integer,
+    MetaData,
+    String,
+    Table,
+    delete,
+    func,
+    insert,
+    select,
+)
+from sqlalchemy import inspect as inspect_connection
+from sqlalchemy.schema import CreateTable
+
+from .direction import Direction
+from .execution_result import ExecutionResult
+
+if TYPE_CHECKING:
+    from collections.abc import Iterable, Mapping
+
+    from alembic.runtime.migration import MigrationContext, MigrationInfo
+    from sqlalchemy.engine import Connection
+
+__all__ = ["HistoryRecorder", "HistoryTable"]
+
+_VERSION_LENGTH: Final = 255
+_MILLISECONDS: Final = 1000
+
+
+@final
+class HistoryTable:
+    """The history table of one database: its definition, and how it is read and written."""
+
+    __slots__ = ("table",)
+
+    table: Table
+
+    def __init__(self, name: str, schema: str | None) -> None:
+        """Describe the table ``name`` in ``schema``."""
+        self.table = Table(
+            name,
+            MetaData(),
+            Column("version", String(_VERSION_LENGTH), primary_key=True),
+            Column("executed_at", DateTime(timezone=True), nullable=True),
+            Column("execution_time", Integer, nullable=True),
+            schema=schema,
+        )
+
+    def exists(self, connection: Connection) -> bool:
+        """Whether the database has the table."""
+        return inspect_connection(connection).has_table(self.table.name, self.table.schema)
+
+    def read(self, connection: Connection) -> dict[str, tuple[datetime | None, float | None]]:
+        """Return when each recorded revision was applied and how long it took, in seconds."""
+        if not self.exists(connection):
+            return {}
+        columns = self.table.c
+        rows = connection.execute(
+            select(columns.version, columns.executed_at, columns.execution_time),
+        )
+        found: dict[str, tuple[datetime | None, float | None]] = {}
+        for row in rows:
+            version, executed_at, execution_time = cast(
+                "tuple[str, object, int | None]", tuple(row)
+            )
+            found[version] = (
+                _as_datetime(executed_at),
+                None if execution_time is None else execution_time / _MILLISECONDS,
+            )
+        return found
+
+    def ensure(self, context: MigrationContext, connection: Connection) -> None:
+        """Create the table through ``context`` unless ``connection`` shows it exists.
+
+        Through the context, so in SQL-only mode the statement is written
+        with the others instead of run.
+        """
+        if not self.exists(connection):
+            context.execute(CreateTable(self.table))
+
+    def add(self, context: MigrationContext, version: str, duration: float | None) -> None:
+        """Record ``version`` as applied now, having taken ``duration`` seconds."""
+        self.remove(context, (version,))
+        context.execute(
+            insert(self.table).values(
+                version=version,
+                executed_at=func.current_timestamp(),
+                execution_time=None if duration is None else round(duration * _MILLISECONDS),
+            ),
+        )
+
+    def remove(self, context: MigrationContext, versions: Iterable[str]) -> None:
+        """Forget ``versions``."""
+        for version in versions:
+            context.execute(delete(self.table).where(self.table.c.version == version))
+
+    def clear(self, context: MigrationContext) -> None:
+        """Forget every revision."""
+        context.execute(delete(self.table))
+
+
+@final
+class HistoryRecorder:
+    """Writes a history row for every step the runner applies, and times each one.
+
+    The runner calls it after each step, inside the step's transaction. The
+    steps it times are handed out through :meth:`started`, which the planned
+    steps' generator calls just before giving the runner each one.
+    """
+
+    __slots__ = ("_history", "_marking", "_results", "_started")
+
+    _history: HistoryTable
+    _marking: bool
+    _results: list[ExecutionResult]
+    _started: float
+
+    def __init__(self, history: HistoryTable, *, marking: bool = False) -> None:
+        """Record into ``history``; ``marking`` steps change the tables without running."""
+        self._history = history
+        self._marking = marking
+        self._results = []
+        self._started = time.perf_counter()
+
+    @property
+    def results(self) -> tuple[ExecutionResult, ...]:
+        """Every step recorded so far, in the order it ran."""
+        return tuple(self._results)
+
+    def started(self) -> None:
+        """Note that the next step starts now."""
+        self._started = time.perf_counter()
+
+    def __call__(
+        self,
+        *,
+        ctx: MigrationContext,
+        step: MigrationInfo,
+        heads: set[str],
+        run_args: Mapping[str, object],
+    ) -> None:
+        """Record ``step``, which the runner has just applied."""
+        del heads, run_args
+        duration = time.perf_counter() - self._started
+        versions = step.up_revision_ids
+        if step.is_upgrade:
+            for version in versions:
+                self._history.add(ctx, version, None if self._marking else duration)
+        else:
+            self._history.remove(ctx, versions)
+        direction = Direction.UP if step.is_upgrade else Direction.DOWN
+        self._results.extend(ExecutionResult(version, direction, duration) for version in versions)
+
+
+def _as_datetime(value: object) -> datetime | None:
+    """Read a timestamp column, which some drivers hand back as text."""
+    if value is None or isinstance(value, datetime):
+        return value
+    return datetime.fromisoformat(str(value))