Database¶

def __copy__(self) -> Database:
    """Return a new runtime object with the same configuration.

    The copy shares URL and option values but not live SQLAlchemy engines,
    connections, transactions, or context state. This is used by registry
    preparation paths that need an isolated runtime handle.
    """
    return self.__class__(self)

@staticmethod
def _extract_options(
    database_url: DatabaseURL,
    **options: Any,
) -> tuple[DatabaseURL, dict[str, Any]]:
    """Separate SQLAlchemy engine options from URL query parameters.

    Saffier accepts selected engine options in the URL for compatibility
    with existing configuration. Known values are removed from the URL and
    passed directly to ``create_async_engine()`` so SQLAlchemy owns pool,
    isolation, echo, and JSON behavior.
    """
    options.setdefault("pool_reset_on_return", None)
    options.setdefault("json_serializer", _json_serializer)
    options.setdefault("json_deserializer", _json_deserializer)
    query_options = dict(database_url.options)
    for param in ("ssl", "echo", "echo_pool"):
        if param in query_options:
            assert param not in options
            value = cast("str", query_options.pop(param))
            options[param] = value.lower() in {"true", ""}
    if "isolation_level" in query_options:
        assert "isolation_level" not in options
        options["isolation_level"] = cast("str", query_options.pop("isolation_level"))
    for param in ("pool_size", "max_overflow"):
        if param in query_options:
            assert param not in options
            options[param] = int(cast("str", query_options.pop(param)))
    if "pool_recycle" in query_options:
        assert "pool_recycle" not in options
        options["pool_recycle"] = float(cast("str", query_options.pop("pool_recycle")))
    options.setdefault("isolation_level", "AUTOCOMMIT")
    return database_url.replace(options=query_options), options

async def session(self) -> AsyncSession:
    """Create a native SQLAlchemy ``AsyncSession``.

    The database must be connected so the session factory is bound to the
    current ``AsyncEngine``. Sessions use ``expire_on_commit=False`` to keep
    ORM-facing values usable after transaction boundaries.
    """
    await self._ensure_connected()
    assert self._sessionmaker is not None
    return self._sessionmaker()

async def connect_hook(self) -> None:
    """Run subclass setup before the SQLAlchemy engine is created.

    ``DatabaseTestClient`` uses this hook to create or verify the test
    database before normal connection proceeds. The base implementation is
    intentionally empty so regular runtimes connect directly.
    """
    pass

async def disconnect_hook(self) -> None:
    """Run subclass cleanup after the SQLAlchemy engine is disposed.

    Hooks run after Saffier has closed its force-rollback transaction and
    disposed the engine, which lets test clients safely drop databases or
    release resources outside the live connection pool.
    """
    pass

async def connect(self) -> bool:
    """Create the SQLAlchemy async engine for this database lifecycle.

    Connection is reference-counted because registries and nested async
    contexts may ask for the same database to connect more than once. The
    first caller creates ``AsyncEngine`` and ``async_sessionmaker``; later
    callers reuse that live SQLAlchemy runtime.
    """
    async with self.ref_lock:
        self.ref_counter += 1
        if self.ref_counter > 1:
            if bool(self.force_rollback):
                await self._ensure_global_force_rollback()
            return False
        try:
            await self.connect_hook()
            self._engine = create_async_engine(_async_url(self.url), **self.options)
            self._sessionmaker = async_sessionmaker(
                self._engine,
                expire_on_commit=False,
            )
            self.is_connected = True
            if bool(self.force_rollback):
                await self._ensure_global_force_rollback()
        except BaseException:
            self.ref_counter = 0
            self.is_connected = False
            self._engine = None
            self._sessionmaker = None
            raise
        return True

async def disconnect(self, force: bool = False) -> bool:
    """Dispose the SQLAlchemy async engine when lifecycle ownership ends.

    Normal disconnect decrements the reference counter and only disposes the
    engine for the last owner. ``force=True`` is reserved for teardown paths
    that must close the pool regardless of outstanding references.
    """
    async with self.ref_lock:
        if force:
            self.ref_counter = 0
        elif self.ref_counter > 0:
            self.ref_counter -= 1
        if self.ref_counter > 0:
            return False
        if not self.is_connected:
            return False
        try:
            await self._close_global_force_rollback()
        finally:
            engine, self._engine = self._engine, None
            self._sessionmaker = None
            self.is_connected = False
            if engine is not None:
                await engine.dispose(close=True)
            await self.disconnect_hook()
        return True

async def _ensure_connected(self) -> None:
    """Validate that the database has an active SQLAlchemy engine.

    Execution helpers call this before reaching for ``AsyncEngine``. Raising
    here gives callers a clear lifecycle error instead of a later attribute
    failure on a missing connection pool.
    """
    if not self.is_connected or self._engine is None:
        raise RuntimeError("Database is not connected")

def _current_connection(self) -> AsyncConnection | None:
    """Return the connection currently bound to this execution context.

    Transaction contexts take precedence because nested operations must use
    the same SQLAlchemy connection. If no transaction is active and forced
    rollback is enabled, the reusable rollback connection becomes the
    current connection for test isolation.
    """
    stack = _active_connections().get(id(self), ())
    if stack:
        return stack[-1]
    if bool(self.force_rollback):
        return self._global_connection
    return None

def _push_connection(
    self,
    connection: AsyncConnection,
) -> Token[dict[int, tuple[AsyncConnection, ...]] | None]:
    """Bind a SQLAlchemy connection to this database in the current context.

    Bindings are stored as a per-database stack so nested transactions can
    push savepoint connections and then restore the previous state exactly
    when they finish.
    """
    connections = _active_connections()
    copied = dict(connections)
    copied[id(self)] = (*copied.get(id(self), ()), connection)
    return _ACTIVE_CONNECTIONS.set(copied)

def _push_loop_bound_database(self) -> None:
    """Mark this database as using loop-bound SQLAlchemy resources.

    Force-rollback mode keeps one SQLAlchemy connection open across many ORM
    operations. The sync bridge reads this marker to avoid executing lazy
    loads on a different event loop from that connection.
    """
    databases = _LOOP_BOUND_DATABASES.get()
    token = _LOOP_BOUND_DATABASES.set((*databases, id(self)))
    token_stacks = dict(_loop_bound_token_stacks())
    token_stacks[id(self)] = (*token_stacks.get(id(self), ()), token)
    _LOOP_BOUND_TOKEN_STACKS.set(token_stacks)

def _pop_loop_bound_database(self) -> None:
    """Remove this database's loop-bound marker for the current context.

    The marker token must be reset in the same context that created it.
    Keeping the stack in a context variable prevents concurrent tasks using
    the same ``Database`` instance from popping and resetting each other's
    SQLAlchemy loop markers.
    """
    token_stacks = _loop_bound_token_stacks()
    stack = token_stacks.get(id(self), ())
    if not stack:
        return
    token = stack[-1]
    copied = dict(token_stacks)
    if len(stack) == 1:
        copied.pop(id(self), None)
    else:
        copied[id(self)] = stack[:-1]
    _LOOP_BOUND_TOKEN_STACKS.set(copied)
    _LOOP_BOUND_DATABASES.reset(token)

async def _ensure_global_force_rollback(self) -> AsyncConnection:
    """Open the reusable rollback transaction used for test isolation.

    The connection and transaction are created directly through SQLAlchemy
    and held for the database lifecycle. ORM statements run against this
    connection, and disconnect rolls the outer transaction back.
    """
    await self._ensure_connected()
    if self._global_connection is not None:
        return self._global_connection
    assert self._engine is not None
    connection = await self._engine.connect()
    await connection.execution_options(isolation_level="SERIALIZABLE")
    transaction = await connection.begin()
    self._global_connection = connection
    self._global_transaction = transaction
    return connection

async def _close_global_force_rollback(self) -> None:
    """Roll back and close the force-rollback connection.

    Teardown suppresses rollback errors so disposal can still close the
    underlying SQLAlchemy connection. This mirrors test cleanup expectations
    where isolation should be best-effort during exceptional shutdown.
    """
    transaction, self._global_transaction = self._global_transaction, None
    connection, self._global_connection = self._global_connection, None
    try:
        if transaction is not None:
            with contextlib.suppress(Exception):
                await transaction.rollback()
    finally:
        if connection is not None:
            await connection.close()

async def _transaction_connection(self) -> AsyncConnection:
    """Choose the SQLAlchemy connection for a new transaction scope.

    Existing context connections are reused for nesting, forced rollback
    scopes reuse the lifecycle-wide rollback connection, and ordinary
    transactions open a fresh connection from the async engine's pool.
    """
    current = self._current_connection()
    if current is not None:
        return current
    if bool(self.force_rollback):
        return await self._ensure_global_force_rollback()
    await self._ensure_connected()
    assert self._engine is not None
    return await self._engine.connect()

@contextlib.asynccontextmanager
async def connection(self) -> AsyncGenerator[AsyncConnection, None]:
    """Yield a raw SQLAlchemy ``AsyncConnection`` for direct operations.

    The yielded connection is the current transaction-bound connection when
    one exists; otherwise a short-lived connection is opened from the
    ``AsyncEngine`` pool. This keeps direct SQLAlchemy usage aligned with
    Saffier's transaction context instead of bypassing it.
    """
    current = self._current_connection()
    if current is not None:
        yield current
        return
    await self._ensure_connected()
    assert self._engine is not None
    async with self._engine.connect() as connection:
        yield connection

@contextlib.asynccontextmanager
async def _execution_connection(self) -> AsyncGenerator[AsyncConnection, None]:
    """Yield the SQLAlchemy connection for one ORM execution helper.

    ORM helpers should not decide whether they are inside a transaction,
    forced rollback block, or normal engine checkout. This helper centralizes
    that choice so fetch, execute, and iteration paths all share the same
    connection ownership rules.
    """
    current = self._current_connection()
    if current is not None:
        yield current
        return
    async with self.connection() as connection:
        yield connection

async def fetch_all(
    self,
    query: sqlalchemy.ClauseElement | str,
    values: dict[str, Any] | None = None,
    timeout: float | None = None,
) -> list[sqlalchemy.Row[Any]]:
    """Execute a statement and return all rows from SQLAlchemy.

    The method preserves Saffier's public ``fetch_all`` API while routing
    execution directly through ``AsyncConnection.execute()``. Results are
    materialized before the cursor is closed so callers can inspect rows
    after the connection context exits.
    """
    del timeout
    statement = _coerce_statement(query)
    async with self._execution_connection() as connection:
        result = await connection.execute(statement, values or {})
        try:
            return list(result.fetchall())
        finally:
            result.close()

async def fetch_one(
    self,
    query: sqlalchemy.ClauseElement | str,
    values: dict[str, Any] | None = None,
    pos: int = 0,
    timeout: float | None = None,
) -> sqlalchemy.Row[Any] | None:
    """Execute a statement and return one row by logical position.

    Positive positions are translated into SQLAlchemy ``offset`` and
    ``limit`` calls when the statement supports them. ``-1`` keeps the
    historic "last row" behavior by materializing the result and returning
    the final entry.
    """
    del timeout
    statement = _coerce_statement(query)
    if pos > 0 and hasattr(statement, "offset"):
        statement = statement.offset(pos)  # type: ignore[assignment,union-attr]
    if pos >= 0 and hasattr(statement, "limit"):
        statement = statement.limit(1)  # type: ignore[assignment,union-attr]
    rows = await self.fetch_all(statement, values)
    if not rows:
        return None
    if pos == -1:
        return rows[-1]
    if pos < -1:
        raise NotImplementedError(
            f"Only positive numbers and -1 for the last result are currently supported: {pos}"
        )
    return rows[0]

async def fetch_val(
    self,
    query: sqlalchemy.ClauseElement | str,
    values: dict[str, Any] | None = None,
    column: int | str = 0,
    pos: int = 0,
    timeout: float | None = None,
) -> Any:
    """Execute a statement and return a single value from one row.

    The value can be selected by integer position or by row mapping key.
    Returning ``None`` for empty results matches the existing Saffier query
    helper contract.
    """
    row = await self.fetch_one(query, values, pos=pos, timeout=timeout)
    if row is None:
        return None
    if isinstance(column, str):
        return row._mapping[column]
    return row[column]

@staticmethod
def _parse_execute_result(result: sqlalchemy.CursorResult[Any]) -> Any:
    """Convert SQLAlchemy cursor metadata into Saffier's execute result.

    Inserts historically return primary-key metadata when the dialect can
    provide it, while updates and deletes return ``rowcount``. Keeping this
    translation in one helper lets ``execute()`` stay SQLAlchemy-native
    without leaking dialect-specific cursor details throughout the ORM.
    """
    if result.is_insert:
        with contextlib.suppress(AttributeError):
            inserted_primary_key = result.inserted_primary_key
            if inserted_primary_key:
                return inserted_primary_key
        with contextlib.suppress(AttributeError):
            return result.lastrowid
    return result.rowcount

async def execute(
    self,
    query: sqlalchemy.ClauseElement | str,
    values: Any = None,
    timeout: float | None = None,
) -> Any:
    """Execute one SQLAlchemy statement and return the public result value.

    Statements are run through the current execution connection, which means
    they participate in active Saffier transactions and force-rollback test
    contexts. The returned value follows the legacy Saffier contract:
    insert metadata when available, otherwise affected row count.
    """
    del timeout
    statement = _coerce_statement(query)
    async with self._execution_connection() as connection:
        result = (
            await connection.execute(statement, values)
            if values is not None
            else await connection.execute(statement)
        )
        try:
            return self._parse_execute_result(result)
        finally:
            result.close()

async def execute_many(
    self,
    query: sqlalchemy.ClauseElement | str,
    values: Any = None,
    timeout: float | None = None,
) -> Any:
    """Execute one statement against many parameter dictionaries.

    SQLAlchemy handles executemany behavior when a sequence of parameter
    dictionaries is passed to ``execute()``. Saffier only adapts the cursor
    metadata back into the value expected by bulk insert and update code.
    """
    del timeout
    statement = _coerce_statement(query)
    async with self._execution_connection() as connection:
        result = (
            await connection.execute(statement, values)
            if values is not None
            else await connection.execute(statement)
        )
        try:
            if result.is_insert:
                with contextlib.suppress(AttributeError):
                    if result.inserted_primary_key_rows is not None:
                        return result.inserted_primary_key_rows
            return result.rowcount
        finally:
            result.close()

async def iterate(
    self,
    query: sqlalchemy.ClauseElement | str,
    values: dict[str, Any] | None = None,
    chunk_size: int | None = None,
    timeout: float | None = None,
) -> AsyncGenerator[sqlalchemy.Row[Any], None]:
    """Yield rows from a SQLAlchemy statement without hiding result objects.

    Dialects with server-side cursor support use ``AsyncConnection.stream``
    and ``yield_per`` when a transaction is already active. Dialects without
    streaming support, or autocommit connections where asyncpg cannot open a
    cursor, fall back to a materialized result while preserving the async
    generator API.
    """
    del timeout
    statement = _coerce_statement(query)
    chunk_size = chunk_size or self.default_batch_size
    async with self._execution_connection() as connection:
        if (
            not connection.dialect.supports_server_side_cursors
            or not connection.in_transaction()
        ):
            result = await connection.execute(statement, values or {})
            try:
                for row in result.fetchall():
                    yield row
            finally:
                result.close()
            return

        await connection.execution_options(yield_per=chunk_size)
        try:
            async with connection.stream(statement, values or {}) as result:
                async for row in result:
                    yield row
        finally:
            await connection.execution_options(yield_per=0)

async def batched_iterate(
    self,
    query: sqlalchemy.ClauseElement | str,
    values: dict[str, Any] | None = None,
    batch_size: int | None = None,
    batch_wrapper: Callable[[Sequence[Any]], Any] = tuple,
    timeout: float | None = None,
) -> AsyncGenerator[Any, None]:
    """Yield result rows grouped into caller-configured batches.

    This compatibility helper materializes rows with ``fetch_all`` and then
    wraps each batch with ``batch_wrapper``. The execution itself still uses
    SQLAlchemy async connections and participates in active transactions.
    """
    del timeout
    batch_size = batch_size or self.default_batch_size
    rows = await self.fetch_all(query, values)
    for batch in _batch_rows(rows, batch_size):
        yield batch_wrapper(batch)

def transaction(self, *, force_rollback: bool = False, **kwargs: Any) -> Transaction:
    """Create a SQLAlchemy-backed Saffier transaction context.

    Keyword arguments are passed through as SQLAlchemy execution options
    before the transaction begins. The returned object supports ``async
    with``, manual ``await`` startup, and decorator usage.
    """
    return Transaction(self, force_rollback=force_rollback, **kwargs)

async def run_sync(
    self,
    fn: Callable[..., Any],
    *args: Any,
    timeout: float | None = None,
    **kwargs: Any,
) -> Any:
    """Run a synchronous SQLAlchemy callable against an async connection.

    SQLAlchemy exposes ``AsyncConnection.run_sync()`` for operations such as
    metadata creation and reflection that still use synchronous Core APIs.
    Saffier forwards directly to that method so schema helpers remain
    idiomatic SQLAlchemy 2.x async code.
    """
    del timeout
    async with self._execution_connection() as connection:
        return await connection.run_sync(fn, *args, **kwargs)

async def create_all(
    self,
    meta: sqlalchemy.MetaData,
    timeout: float | None = None,
    **kwargs: Any,
) -> None:
    """Create every table in a SQLAlchemy ``MetaData`` object.

    Table creation is executed through ``AsyncConnection.run_sync()`` so the
    synchronous ``MetaData.create_all`` API runs on the connection paired
    with the active async engine or transaction.
    """
    del timeout
    await self.run_sync(meta.create_all, **kwargs)

async def drop_all(
    self,
    meta: sqlalchemy.MetaData,
    timeout: float | None = None,
    **kwargs: Any,
) -> None:
    """Drop every table in a SQLAlchemy ``MetaData`` object.

    The method mirrors ``create_all`` and keeps Saffier schema teardown on
    SQLAlchemy's public metadata API. Any active transaction or
    force-rollback connection is respected by the execution connection.
    """
    del timeout
    await self.run_sync(meta.drop_all, **kwargs)

def asgi(
    self,
    app: ASGIApp | None = None,
    handle_lifespan: bool = False,
) -> ASGIApp | Callable[[ASGIApp], ASGIApp]:
    """Wrap an ASGI application with SQLAlchemy engine lifecycle hooks.

    The wrapper connects the database during ASGI startup and registers
    ``disconnect`` as cleanup. This keeps framework integration small while
    still ensuring SQLAlchemy pools are opened and disposed at application
    lifecycle boundaries.
    """
    if LifespanHook is None:
        raise RuntimeError("monkay.asgi is required for ASGI integration.")

    async def setup() -> contextlib.AsyncExitStack:
        """Connect the database and register ASGI lifespan cleanup.

        ``LifespanHook`` expects a stack-like object that owns async cleanup
        callbacks. Saffier connects the SQLAlchemy engine first and then
        pushes ``disconnect`` so pool disposal always runs on shutdown.
        """
        cleanupstack = contextlib.AsyncExitStack()
        await self.connect()
        cleanupstack.push_async_callback(self.disconnect)
        return cleanupstack

    return LifespanHook(app, setup=setup, do_forward=not handle_lifespan)