Database class¶

def __copy__(self) -> Database:
    return self.__class__(self)

async def inc_refcount(self) -> bool:
    """
    Internal method to bump the ref_count.

    Return True if ref_count is 0, False otherwise.

    Should not be used outside of tests. Use connect and hooks instead.
    Not multithreading safe!
    """
    async with self.ref_lock:
        self.ref_counter += 1
        # on the first call is count is 1 because of the former +1
        if self.ref_counter == 1:
            return True
    return False

async def decr_refcount(self) -> bool:
    """
    Internal method to decrease the ref_count.

    Return True if ref_count drops to 0, False otherwise.

    Should not be used outside of tests. Use disconnect and hooks instead.
    Not multithreading safe!
    """
    async with self.ref_lock:
        self.ref_counter -= 1
        # on the last call, the count is 0
        if self.ref_counter == 0:
            return True
    return False

async def connect_hook(self) -> None:
    """Refcount protected connect hook. Executed before engine and global connection setup."""

async def connect(self) -> bool:
    """
    Establish the connection pool.
    """
    loop = asyncio.get_running_loop()
    if self._loop is not None and loop != self._loop:
        if self.poll_interval < 0:
            raise RuntimeError("Subdatabases and polling are disabled")
        # copy when not in map
        if loop not in self._databases_map:
            assert self._global_connection is not None, (
                "global connection should have been set"
            )
            # correctly initialize force_rollback with parent value
            database = self.__class__(
                self, force_rollback=bool(self.force_rollback), full_isolation=False
            )
            # prevent side effects of connect_hook
            database._call_hooks = False
            database._global_connection = self._global_connection
            self._databases_map[loop] = database
        # forward call
        return await self._databases_map[loop].connect()

    if not await self.inc_refcount():
        assert self.is_connected, "ref_count < 0"
        return False
    if self._call_hooks:
        try:
            await self.connect_hook()
        except BaseException as exc:
            await self.decr_refcount()
            raise exc
    self._loop = asyncio.get_event_loop()

    await self.backend.connect(self.url, **self.options)
    self.is_connected = True

    if self._global_connection is None:
        connection = Connection(self, force_rollback=True, full_isolation=self._full_isolation)
        self._global_connection = connection
    else:
        self._remove_global_connection = False
    return True

async def disconnect_hook(self) -> None:
    """Refcount protected disconnect hook. Executed after connection, engine cleanup."""

@multiloop_protector(True, inject_parent=True)
async def disconnect(
    self, force: bool = False, *, parent_database: Database | None = None
) -> bool:
    """
    Close all connections in the connection pool.
    """
    # parent_database is injected and should not be specified manually
    if not await self.decr_refcount() or force:
        if not self.is_connected:
            logger.debug("Already disconnected, skip disconnecting")
            return False
        if force:
            logger.warning("Force disconnect, despite refcount not 0")
        else:
            return False
    if parent_database is not None:
        loop = asyncio.get_running_loop()
        del parent_database._databases_map[loop]
    if force and self._databases_map:
        assert not self._databases_map, "sub databases still active, force terminate them"
        for sub_database in self._databases_map.values():
            await arun_coroutine_threadsafe(
                sub_database.disconnect(True),
                sub_database._loop,
                self.poll_interval,
            )
        self._databases_map = {}
    assert not self._databases_map, "sub databases still active"

    try:
        assert self._global_connection is not None
        if self._remove_global_connection:
            await self._global_connection.__aexit__()
            self._global_connection = None
        self._connection = None
    finally:
        self.is_connected = False
        await self.backend.disconnect()
        self._loop = None
        if self._call_hooks:
            await self.disconnect_hook()
    return True

async def fetch_all(
    self,
    query: ClauseElement | str,
    values: dict | None = None,
    timeout: float | None = None,
) -> list[interfaces.Record]:
    async with self.connection() as connection:
        return await connection.fetch_all(query, values, timeout=timeout)

async def fetch_one(
    self,
    query: ClauseElement | str,
    values: dict | None = None,
    pos: int = 0,
    timeout: float | None = None,
) -> interfaces.Record | None:
    async with self.connection() as connection:
        return await connection.fetch_one(query, values, pos=pos, timeout=timeout)

async def fetch_val(
    self,
    query: ClauseElement | str,
    values: dict | None = None,
    column: Any = 0,
    pos: int = 0,
    timeout: float | None = None,
) -> Any:
    async with self.connection() as connection:
        return await connection.fetch_val(
            query,
            values,
            column=column,
            pos=pos,
            timeout=timeout,
        )

async def execute(
    self,
    query: ClauseElement | str,
    values: Any = None,
    timeout: float | None = None,
) -> interfaces.Record | int:
    async with self.connection() as connection:
        return await connection.execute(query, values, timeout=timeout)

async def execute_many(
    self,
    query: ClauseElement | str,
    values: Any = None,
    timeout: float | None = None,
) -> Sequence[interfaces.Record] | int:
    async with self.connection() as connection:
        return await connection.execute_many(query, values, timeout=timeout)

async def iterate(
    self,
    query: ClauseElement | str,
    values: dict | None = None,
    chunk_size: int | None = None,
    timeout: float | None = None,
) -> AsyncGenerator[interfaces.Record, None]:
    async with self.connection() as connection:
        async for record in connection.iterate(query, values, chunk_size, timeout=timeout):
            yield record

async def batched_iterate(
    self,
    query: ClauseElement | str,
    values: dict | None = None,
    batch_size: int | None = None,
    batch_wrapper: BatchCallable = tuple,
    timeout: float | None = None,
) -> AsyncGenerator[BatchCallableResult, None]:
    async with self.connection() as connection:
        async for batch in cast(
            AsyncGenerator["BatchCallableResult", None],
            connection.batched_iterate(
                query,
                values,
                batch_wrapper=batch_wrapper,
                batch_size=batch_size,
                timeout=timeout,
            ),
        ):
            yield batch

def transaction(self, *, force_rollback: bool = False, **kwargs: Any) -> Transaction:
    return Transaction(self.connection, force_rollback=force_rollback, **kwargs)

async def run_sync(
    self,
    fn: Callable[..., Any],
    *args: Any,
    timeout: float | None = None,
    **kwargs: Any,
) -> Any:
    async with self.connection() as connection:
        return await connection.run_sync(fn, *args, **kwargs, timeout=timeout)

async def create_all(
    self, meta: MetaData, timeout: float | None = None, **kwargs: Any
) -> None:
    async with self.connection() as connection:
        await connection.create_all(meta, **kwargs, timeout=timeout)

async def drop_all(self, meta: MetaData, timeout: float | None = None, **kwargs: Any) -> None:
    async with self.connection() as connection:
        await connection.drop_all(meta, **kwargs, timeout=timeout)

@multiloop_protector(False)
def _non_global_connection(
    self,
    timeout: (
        float | None
    ) = None,  # stub for type checker, multiloop_protector handles timeout
) -> Connection:
    if self._connection is None:
        _connection = self._connection = Connection(self)
        return _connection
    return self._connection

def connection(self, timeout: float | None = None) -> Connection:
    if not self.is_connected:
        raise RuntimeError("Database is not connected")
    if self.force_rollback:
        return cast(Connection, self._global_connection)
    return self._non_global_connection(timeout=timeout)

def asgi(
    self,
    app: ASGIApp | None = None,
    handle_lifespan: bool = False,
) -> ASGIApp | Callable[[ASGIApp], ASGIApp]:
    """Return wrapper for asgi integration."""

    async def setup() -> contextlib.AsyncExitStack:
        cleanupstack = contextlib.AsyncExitStack()
        await self.connect()
        cleanupstack.push_async_callback(self.disconnect)
        return cleanupstack

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

@classmethod
def get_backends(
    cls,
    # let scheme empty for direct imports
    scheme: str = "",
    *,
    overwrite_paths: Sequence[str] = ["databasez.overwrites"],
    database_name: str = "Database",
    connection_name: str = "Connection",
    transaction_name: str = "Transaction",
    database_class: type[interfaces.DatabaseBackend] | None = None,
    connection_class: type[interfaces.ConnectionBackend] | None = None,
    transaction_class: type[interfaces.TransactionBackend] | None = None,
) -> tuple[
    type[interfaces.DatabaseBackend],
    type[interfaces.ConnectionBackend],
    type[interfaces.TransactionBackend],
]:
    module: Any = None
    # when not using utils...., the constant cannot be changed at runtime for debug purposes
    more_debug = utils.DATABASEZ_OVERWRITE_LOGGING
    for overwrite_path in overwrite_paths:
        imp_path = f"{overwrite_path}.{scheme.replace('+', '_')}" if scheme else overwrite_path
        try:
            module = importlib.import_module(imp_path)
        except ImportError as exc:
            logging.debug(
                f'Could not import "{imp_path}". Continue search.',
                exc_info=exc if more_debug else None,
            )
            if "+" in scheme:
                imp_path = f"{overwrite_path}.{scheme.split('+', 1)[0]}"
                try:
                    module = importlib.import_module(imp_path)
                except ImportError as exc:
                    logging.debug(
                        f'Could not import "{imp_path}". Continue search.',
                        exc_info=exc if more_debug else None,
                    )
        if module is not None:
            break
    if module is None:
        logging.debug(
            "No overwrites found. Use default.",
        )
    database_class = getattr(module, database_name, database_class)
    assert database_class is not None and issubclass(
        database_class, interfaces.DatabaseBackend
    )
    connection_class = getattr(module, connection_name, connection_class)
    assert connection_class is not None and issubclass(
        connection_class, interfaces.ConnectionBackend
    )
    transaction_class = getattr(module, transaction_name, transaction_class)
    assert transaction_class is not None and issubclass(
        transaction_class, interfaces.TransactionBackend
    )
    return database_class, connection_class, transaction_class

@classmethod
def apply_database_url_and_options(
    cls,
    url: DatabaseURL | str,
    *,
    overwrite_paths: Sequence[str] = ["databasez.overwrites"],
    **options: Any,
) -> tuple[interfaces.DatabaseBackend, DatabaseURL, dict[str, Any]]:
    url = DatabaseURL(url)
    database_class, connection_class, transaction_class = cls.get_backends(
        url.scheme,
        database_class=default_database,
        connection_class=default_connection,
        transaction_class=default_transaction,
        overwrite_paths=overwrite_paths,
    )

    backend = database_class(
        connection_class=connection_class, transaction_class=transaction_class
    )
    url, options = backend.extract_options(url, **options)
    # check against transformed url
    assert url.sqla_url.get_dialect(True).is_async, f'Dialect: "{url.scheme}" is not async.'

    return backend, url, options