path: root/libs/flask_sqlalchemy/model.py

from __future__ import annotations

import re
import typing as t

import sqlalchemy as sa
import sqlalchemy.orm

from .query import Query

if t.TYPE_CHECKING:
    from .extension import SQLAlchemy


class _QueryProperty:
    """A class property that creates a query object for a model.

    :meta private:
    """

    @t.overload
    def __get__(self, obj: None, cls: type[Model]) -> Query:
        ...

    @t.overload
    def __get__(self, obj: Model, cls: type[Model]) -> Query:
        ...

    def __get__(self, obj: Model | None, cls: type[Model]) -> Query:
        return cls.query_class(
            cls, session=cls.__fsa__.session()  # type: ignore[arg-type]
        )


class Model:
    """The base class of the :attr:`.SQLAlchemy.Model` declarative model class.

    To define models, subclass :attr:`db.Model <.SQLAlchemy.Model>`, not this. To
    customize ``db.Model``, subclass this and pass it as ``model_class`` to
    :class:`.SQLAlchemy`. To customize ``db.Model`` at the metaclass level, pass an
    already created declarative model class as ``model_class``.
    """

    __fsa__: t.ClassVar[SQLAlchemy]
    """Internal reference to the extension object.

    :meta private:
    """

    query_class: t.ClassVar[type[Query]] = Query
    """Query class used by :attr:`query`. Defaults to :attr:`.SQLAlchemy.Query`, which
    defaults to :class:`.Query`.
    """

    query: t.ClassVar[Query] = _QueryProperty()  # type: ignore[assignment]
    """A SQLAlchemy query for a model. Equivalent to ``db.session.query(Model)``. Can be
    customized per-model by overriding :attr:`query_class`.

    .. warning::
        The query interface is considered legacy in SQLAlchemy. Prefer using
        ``session.execute(select())`` instead.
    """

    def __repr__(self) -> str:
        state = sa.inspect(self)
        assert state is not None

        if state.transient:
            pk = f"(transient {id(self)})"
        elif state.pending:
            pk = f"(pending {id(self)})"
        else:
            pk = ", ".join(map(str, state.identity))

        return f"<{type(self).__name__} {pk}>"


class BindMetaMixin(type):
    """Metaclass mixin that sets a model's ``metadata`` based on its ``__bind_key__``.

    If the model sets ``metadata`` or ``__table__`` directly, ``__bind_key__`` is
    ignored. If the ``metadata`` is the same as the parent model, it will not be set
    directly on the child model.
    """

    __fsa__: SQLAlchemy
    metadata: sa.MetaData

    def __init__(
        cls, name: str, bases: tuple[type, ...], d: dict[str, t.Any], **kwargs: t.Any
    ) -> None:
        if not ("metadata" in cls.__dict__ or "__table__" in cls.__dict__):
            bind_key = getattr(cls, "__bind_key__", None)
            parent_metadata = getattr(cls, "metadata", None)
            metadata = cls.__fsa__._make_metadata(bind_key)

            if metadata is not parent_metadata:
                cls.metadata = metadata

        super().__init__(name, bases, d, **kwargs)


class NameMetaMixin(type):
    """Metaclass mixin that sets a model's ``__tablename__`` by converting the
    ``CamelCase`` class name to ``snake_case``. A name is set for non-abstract models
    that do not otherwise define ``__tablename__``. If a model does not define a primary
    key, it will not generate a name or ``__table__``, for single-table inheritance.
    """

    metadata: sa.MetaData
    __tablename__: str
    __table__: sa.Table

    def __init__(
        cls, name: str, bases: tuple[type, ...], d: dict[str, t.Any], **kwargs: t.Any
    ) -> None:
        if should_set_tablename(cls):
            cls.__tablename__ = camel_to_snake_case(cls.__name__)

        super().__init__(name, bases, d, **kwargs)

        # __table_cls__ has run. If no table was created, use the parent table.
        if (
            "__tablename__" not in cls.__dict__
            and "__table__" in cls.__dict__
            and cls.__dict__["__table__"] is None
        ):
            del cls.__table__

    def __table_cls__(cls, *args: t.Any, **kwargs: t.Any) -> sa.Table | None:
        """This is called by SQLAlchemy during mapper setup. It determines the final
        table object that the model will use.

        If no primary key is found, that indicates single-table inheritance, so no table
        will be created and ``__tablename__`` will be unset.
        """
        schema = kwargs.get("schema")

        if schema is None:
            key = args[0]
        else:
            key = f"{schema}.{args[0]}"

        # Check if a table with this name already exists. Allows reflected tables to be
        # applied to models by name.
        if key in cls.metadata.tables:
            return sa.Table(*args, **kwargs)

        # If a primary key is found, create a table for joined-table inheritance.
        for arg in args:
            if (isinstance(arg, sa.Column) and arg.primary_key) or isinstance(
                arg, sa.PrimaryKeyConstraint
            ):
                return sa.Table(*args, **kwargs)

        # If no base classes define a table, return one that's missing a primary key
        # so SQLAlchemy shows the correct error.
        for base in cls.__mro__[1:-1]:
            if "__table__" in base.__dict__:
                break
        else:
            return sa.Table(*args, **kwargs)

        # Single-table inheritance, use the parent table name. __init__ will unset
        # __table__ based on this.
        if "__tablename__" in cls.__dict__:
            del cls.__tablename__

        return None


def should_set_tablename(cls: type) -> bool:
    """Determine whether ``__tablename__`` should be generated for a model.

    -   If no class in the MRO sets a name, one should be generated.
    -   If a declared attr is found, it should be used instead.
    -   If a name is found, it should be used if the class is a mixin, otherwise one
        should be generated.
    -   Abstract models should not have one generated.

    Later, ``__table_cls__`` will determine if the model looks like single or
    joined-table inheritance. If no primary key is found, the name will be unset.
    """
    if cls.__dict__.get("__abstract__", False) or not any(
        isinstance(b, sa.orm.DeclarativeMeta) for b in cls.__mro__[1:]
    ):
        return False

    for base in cls.__mro__:
        if "__tablename__" not in base.__dict__:
            continue

        if isinstance(base.__dict__["__tablename__"], sa.orm.declared_attr):
            return False

        return not (
            base is cls
            or base.__dict__.get("__abstract__", False)
            or not isinstance(base, sa.orm.DeclarativeMeta)
        )

    return True


def camel_to_snake_case(name: str) -> str:
    """Convert a ``CamelCase`` name to ``snake_case``."""
    name = re.sub(r"((?<=[a-z0-9])[A-Z]|(?!^)[A-Z](?=[a-z]))", r"_\1", name)
    return name.lower().lstrip("_")


class DefaultMeta(BindMetaMixin, NameMetaMixin, sa.orm.DeclarativeMeta):
    """SQLAlchemy declarative metaclass that provides ``__bind_key__`` and
    ``__tablename__`` support.
    """