Schema class¶

def get_default_schema(self) -> str | None:
    """
    Retrieves the default schema name from the underlying database dialect.

    This method caches the default schema name after its first retrieval
    to optimize subsequent calls.

    Returns:
        The name of the default schema, or None if not applicable or found.
    """
    # Check if the _default_schema attribute has already been set.
    if not hasattr(self, "_default_schema"):
        # If not set, retrieve the default schema name from the database URL's
        # SQLAlchemy dialect and store it.
        self._default_schema = self.database.url.sqla_url.get_dialect(True).default_schema_name
    # Return the cached default schema name.
    return self._default_schema

async def create_schema(
    self,
    schema: str | None,
    if_not_exists: bool = False,
    init_models: bool = False,
    init_tenant_models: bool = False,
    update_cache: bool = True,
    databases: Sequence[str | None] = (None,),
) -> None:
    """
    Creates a new database schema and optionally initializes models within it.

    This method handles the creation of a new schema and can populate it
    with tables for both regular models and tenant-specific models,
    respecting global field constraints.

    Args:
        schema: The name of the schema to be created or None for main.
        if_not_exists: If True, the schema will only be created if it does
                       not already exist, preventing an error. Defaults to False.
        init_models: If True, all models registered with the registry will have
                     their tables created within the new schema. Defaults to False.
        init_tenant_models: If True, tenant-specific models will have their
                           tables created within the new schema. This operation
                           temporarily bypasses global field constraints. Defaults to False.
        update_cache: If True, the model's schema cache will be updated.
                      Defaults to True.
        databases: A sequence of database names (keys from `registry.extra`)
                   or None for the default database, on which the schema
                   should be created. Defaults to `(None,)`, meaning only
                   the default database.
    Raises:
        SchemaError: If there is an issue during schema creation or table
                     initialization within the schema.
    """
    tenant_tables: list[sqlalchemy.Table] = []
    # If init_models is True, iterate through all registered models and
    # update their table schema and cache.
    if init_models:
        for model_class in self.registry.models.values():
            # only init if no schema is set. This prevents problems with e.g. reflection.
            if model_class.__using_schema__ is Undefined:
                model_class.table_schema(schema=schema, update_cache=update_cache)

    # If init_tenant_models is True, handle the creation of tenant-specific model tables.
    if init_tenant_models:
        # Temporarily disable global field constraints for tenant model building.
        token = NO_GLOBAL_FIELD_CONSTRAINTS.set(True)
        try:
            # Iterate through tenant models and build their tables with the specified schema.
            for model_class in self.registry.tenant_models.values():
                tenant_tables.append(model_class.build(schema=schema))
        finally:
            # Ensure global field constraints are re-enabled after processing tenant models.
            NO_GLOBAL_FIELD_CONSTRAINTS.reset(token)

        # Perform a second pass to add global field constraints to tenant models.
        for model_class in self.registry.tenant_models.values():
            model_class.add_global_field_constraints(schema=schema)

    def execute_create(connection: sqlalchemy.Connection, name: str | None) -> None:
        """
        Internal helper function to execute the schema and table creation
        within a given database connection.
        """
        if schema is not None:
            try:
                # Attempt to create the schema.
                connection.execute(
                    sqlalchemy.schema.CreateSchema(name=schema, if_not_exists=if_not_exists)
                )
            except ProgrammingError as e:
                # Raise a SchemaError if there's a programming error during schema creation.
                raise SchemaError(detail=e.orig.args[0]) from e

        # If tenant_tables exist, create them within the schema.
        if tenant_tables:
            self.registry.metadata_by_name[name].create_all(
                connection, checkfirst=if_not_exists, tables=tenant_tables
            )
        # If init_models is True, create all registered model tables within the schema.
        if init_models:
            self.registry.metadata_by_name[name].create_all(
                connection, checkfirst=if_not_exists
            )

    # Iterate through the specified databases to perform schema creation.
    for database_name in databases:
        # Determine which database instance to use based on database_name.
        db = (
            self.registry.database
            if database_name is None
            else self.registry.extra[database_name]
        )
        # Enter an asynchronous context for the database connection, disabling rollback.
        # prevents warning of inperformance
        async with db as db:
            with db.force_rollback(False):
                # run the operation sequencially because we need the right context and
                # a connected database
                await db.run_sync(execute_create, database_name)

async def drop_schema(
    self,
    schema: str | None,
    cascade: bool = False,
    if_exists: bool = False,
    databases: Sequence[str | None] = (None,),
) -> None:
    """
    Drops an existing database schema, optionally cascading the drop
    to all contained objects.

    Args:
        schema: The name of the schema to be dropped. If None fallback to the classic drop_all logic.
        cascade: If True, all objects (tables, views, etc.) within the
                 schema will also be dropped. Defaults to False.
        if_exists: If True, the schema will only be dropped if it exists,
                   preventing an error if it does not. Defaults to False.
        databases: A sequence of database names (keys from `registry.extra`)
                   or None for the default database, from which the schema
                   should be dropped. Defaults to `(None,)`, meaning only
                   the default database.
    Raises:
        SchemaError: If there is an issue during schema drop operation.
    """

    def execute_drop(connection: sqlalchemy.Connection) -> None:
        """
        Internal helper function to execute the schema drop
        within a given database connection.
        """
        try:
            # just remove the schema
            if schema is not None:
                # Attempt to drop the schema.
                connection.execute(
                    sqlalchemy.schema.DropSchema(
                        name=schema, cascade=cascade, if_exists=if_exists
                    )
                )
            else:
                # if no schema is found remove all registered models in meta
                self.registry.metadata_by_name[database_name].drop_all(
                    connection, checkfirst=if_exists
                )
        except DBAPIError as e:
            # Raise a SchemaError if there's a database API error during schema drop.
            raise SchemaError(detail=e.orig.args[0]) from e

    # Iterate through the specified databases to perform schema drop.
    for database_name in databases:
        # Determine which database instance to use based on database_name.
        db = (
            self.registry.database
            if database_name is None
            else self.registry.extra[database_name]
        )
        # Enter an asynchronous context for the database connection, disabling rollback.
        # prevents warning of inperformance
        async with db as db:
            with db.force_rollback(False):
                # run the operation sequencially because we need the right context and
                # a connected database
                await db.run_sync(execute_drop)

async def get_metadata_of_all_schemes(
    self, database: Database, *, no_reflect: bool = False
) -> tuple[sqlalchemy.MetaData, list[str]]:
    """
    Retrieves metadata and a list of all schema names for a given database.

    This method reflects the table structures for registered models
    within each discovered schema if `no_reflect` is False.

    Args:
        database: The database instance from which to retrieve schema metadata.
        no_reflect: If True, tables will not be reflected into the metadata.
                    Defaults to False.

    Returns:
        A tuple containing:
            - sqlalchemy.MetaData: An SQLAlchemy MetaData object populated
                                   with reflected tables (if `no_reflect` is False).
            - list[str]: A list of schema names found in the database.
    """
    # Get a set of all table names registered in the registry.
    tablenames = self.registry.get_tablenames()

    async with database as database:
        # Initialize an empty list to store schema names.
        list_schemes: list[str] = []
        # Create a new SQLAlchemy MetaData object.
        metadata = sqlalchemy.MetaData()
        # Force no rollback for the database connection.
        with database.force_rollback(False):

            def wrapper(connection: sqlalchemy.Connection) -> None:
                """
                Internal wrapper function to perform synchronous database
                inspection and reflection.
                """
                nonlocal list_schemes
                # Create an inspector from the database connection.
                inspector = sqlalchemy.inspect(connection)
                # Get the default schema name from the inspector.
                default_schema_name = inspector.default_schema_name
                # Get all schema names and replace the default schema name with an empty string.
                list_schemes = [
                    "" if default_schema_name == schema else schema
                    for schema in inspector.get_schema_names()
                ]
                # If no_reflect is False, reflect table metadata for each schema.
                if not no_reflect:
                    for schema in list_schemes:
                        metadata.reflect(
                            connection, schema=schema, only=lambda name, _: name in tablenames
                        )

            # Run the synchronous wrapper function on the database.
            await database.run_sync(wrapper)
            # Return the populated metadata and the list of schema names.
            return metadata, list_schemes

async def get_schemes_tree(
    self, *, no_reflect: bool = False
) -> dict[str | None, tuple[str, sqlalchemy.MetaData, list[str]]]:
    """
    Builds a comprehensive tree-like structure of schemas across all
    registered databases.

    Each entry in the resulting dictionary represents a database (identified
    by its name or None for the default), containing its URL, its SQLAlchemy
    MetaData, and a list of schema names.

    Args:
        no_reflect: If True, tables will not be reflected into the metadata
                    for any schema. Defaults to False.

    Returns:
        A dictionary where keys are database names (or None for the default
        database) and values are tuples containing:
            - str: The URL of the database.
            - sqlalchemy.MetaData: The MetaData object for the database.
            - list[str]: A list of schema names found in that database.
    """
    # Initialize the schemes_tree dictionary.
    schemes_tree: dict[str | None, tuple[str, sqlalchemy.MetaData, list[str]]] = {
        None: (
            str(self.database.url),
            # Get metadata and schemes for the default database.
            *(
                await self.get_metadata_of_all_schemes(
                    self.registry.database, no_reflect=no_reflect
                )
            ),
        )
    }
    # Iterate through extra databases registered in the registry.
    for key, val in self.registry.extra.items():
        # Populate schemes_tree for each extra database.
        schemes_tree[key] = (
            str(val.url),
            # Get metadata and schemes for the current extra database.
            *(await self.get_metadata_of_all_schemes(val, no_reflect=no_reflect)),
        )
    # Return the complete schemes_tree.
    return schemes_tree