"""Common base classes for SQLAlchemy declarative models."""
import contextlib
import datetime
import re
from collections.abc import Iterator
from typing import TYPE_CHECKING, Any, Optional, Protocol, Union, cast, runtime_checkable
from uuid import UUID
from sqlalchemy import Date, MetaData, String
from sqlalchemy.ext.asyncio import AsyncAttrs
from sqlalchemy.orm import (
DeclarativeBase,
Mapper,
declared_attr,
)
from sqlalchemy.orm import (
registry as SQLAlchemyRegistry, # noqa: N812
)
from sqlalchemy.orm.decl_base import (
_TableArgsType as TableArgsType, # pyright: ignore[reportPrivateUsage]
)
from sqlalchemy.types import TypeEngine
from typing_extensions import Self, TypeVar
from advanced_alchemy.mixins import (
AuditColumns,
BigIntPrimaryKey,
IdentityPrimaryKey,
NanoIDPrimaryKey,
UUIDPrimaryKey,
UUIDv6PrimaryKey,
UUIDv7PrimaryKey,
)
from advanced_alchemy.types import GUID, DateTimeUTC, FileObject, FileObjectList, JsonB, StoredObject
from advanced_alchemy.utils.dataclass import DataclassProtocol
if TYPE_CHECKING:
from sqlalchemy.sql import FromClause
from sqlalchemy.sql.schema import (
_NamingSchemaParameter as NamingSchemaParameter, # pyright: ignore[reportPrivateUsage]
)
__all__ = (
"AdvancedDeclarativeBase",
"BasicAttributes",
"BigIntAuditBase",
"BigIntBase",
"BigIntBaseT",
"CommonTableAttributes",
"DefaultBase",
"IdentityAuditBase",
"IdentityBase",
"IdentityBaseT",
"ModelProtocol",
"NanoIDAuditBase",
"NanoIDBase",
"NanoIDBaseT",
"SQLQuery",
"TableArgsType",
"UUIDAuditBase",
"UUIDBase",
"UUIDBaseT",
"UUIDv6AuditBase",
"UUIDv6Base",
"UUIDv6BaseT",
"UUIDv7AuditBase",
"UUIDv7Base",
"UUIDv7BaseT",
"convention",
"create_registry",
"merge_table_arguments",
"metadata_registry",
"orm_registry",
"table_name_regexp",
)
UUIDBaseT = TypeVar("UUIDBaseT", bound="UUIDBase")
"""Type variable for :class:`UUIDBase`."""
BigIntBaseT = TypeVar("BigIntBaseT", bound="BigIntBase")
"""Type variable for :class:`BigIntBase`."""
IdentityBaseT = TypeVar("IdentityBaseT", bound="IdentityBase")
"""Type variable for :class:`IdentityBase`."""
UUIDv6BaseT = TypeVar("UUIDv6BaseT", bound="UUIDv6Base")
"""Type variable for :class:`UUIDv6Base`."""
UUIDv7BaseT = TypeVar("UUIDv7BaseT", bound="UUIDv7Base")
"""Type variable for :class:`UUIDv7Base`."""
NanoIDBaseT = TypeVar("NanoIDBaseT", bound="NanoIDBase")
"""Type variable for :class:`NanoIDBase`."""
convention: "NamingSchemaParameter" = {
"ix": "ix_%(column_0_label)s",
"uq": "uq_%(table_name)s_%(column_0_name)s",
"ck": "ck_%(table_name)s_%(constraint_name)s",
"fk": "fk_%(table_name)s_%(column_0_name)s_%(referred_table_name)s",
"pk": "pk_%(table_name)s",
}
"""Templates for automated constraint name generation."""
table_name_regexp = re.compile(r"((?<=[a-z0-9])[A-Z]|(?!^)[A-Z](?=[a-z]))")
"""Regular expression for table name"""
[docs]
def merge_table_arguments(cls: type[DeclarativeBase], table_args: Optional[TableArgsType] = None) -> TableArgsType:
"""Merge Table Arguments.
This function helps merge table arguments when using mixins that include their own table args,
making it easier to append additional information such as comments or constraints to the model.
Args:
cls (type[:class:`sqlalchemy.orm.DeclarativeBase`]): The model that will get the table args.
table_args (:class:`TableArgsType`, optional): Additional information to add to table_args.
Returns:
:class:`TableArgsType`: Merged table arguments.
"""
args: list[Any] = []
kwargs: dict[str, Any] = {}
mixin_table_args = (getattr(super(base_cls, cls), "__table_args__", None) for base_cls in cls.__bases__) # pyright: ignore[reportUnknownParameter,reportUnknownArgumentType,reportArgumentType]
for arg_to_merge in (*mixin_table_args, table_args):
if arg_to_merge:
if isinstance(arg_to_merge, tuple):
last_positional_arg = arg_to_merge[-1] # pyright: ignore[reportUnknownVariableType]
args.extend(arg_to_merge[:-1]) # pyright: ignore[reportUnknownArgumentType]
if isinstance(last_positional_arg, dict):
kwargs.update(last_positional_arg) # pyright: ignore[reportUnknownArgumentType]
else:
args.append(last_positional_arg)
else:
kwargs.update(arg_to_merge) # pyright: ignore
if args:
if kwargs:
return (*args, kwargs)
return tuple(args)
return kwargs
[docs]
@runtime_checkable
class ModelProtocol(Protocol):
"""The base SQLAlchemy model protocol.
Attributes:
__table__ (:class:`sqlalchemy.sql.FromClause`): The table associated with the model.
__mapper__ (:class:`sqlalchemy.orm.Mapper`): The mapper for the model.
__name__ (str): The name of the model.
"""
if TYPE_CHECKING:
__table__: FromClause
__mapper__: Mapper[Any]
__name__: str
[docs]
def to_dict(self, exclude: Optional[set[str]] = None) -> dict[str, Any]:
"""Convert model to dictionary.
Returns:
Dict[str, Any]: A dict representation of the model
"""
...
[docs]
class BasicAttributes:
"""Basic attributes for SQLAlchemy tables and queries.
Provides a method to convert the model to a dictionary representation.
Methods:
to_dict: Converts the model to a dictionary, excluding specified fields. :no-index:
"""
if TYPE_CHECKING:
__name__: str
__table__: FromClause
__mapper__: Mapper[Any]
[docs]
def to_dict(self, exclude: Optional[set[str]] = None) -> dict[str, Any]:
"""Convert model to dictionary.
Returns:
Dict[str, Any]: A dict representation of the model
"""
exclude = {"sa_orm_sentinel", "_sentinel"}.union(self._sa_instance_state.unloaded).union(exclude or []) # type: ignore[attr-defined]
return {
field: getattr(self, field)
for field in self.__mapper__.columns.keys() # noqa: SIM118
if field not in exclude
}
[docs]
class CommonTableAttributes(BasicAttributes):
"""Common attributes for SQLAlchemy tables.
Inherits from :class:`BasicAttributes` and provides a mechanism to infer table names from class names
while respecting SQLAlchemy's inheritance patterns.
This mixin supports all three SQLAlchemy inheritance patterns:
- **Single Table Inheritance (STI)**: Child classes automatically use parent's table
- **Joined Table Inheritance (JTI)**: Child classes have their own tables with foreign keys
- **Concrete Table Inheritance (CTI)**: Child classes have independent tables
Attributes:
__tablename__ (str | None): The inferred table name, or None for Single Table Inheritance children.
"""
[docs]
def __init_subclass__(cls, **kwargs: Any) -> None:
"""Hook called when a subclass is created.
This method intercepts class creation to correctly handle ``__tablename__`` for
Single Table Inheritance (STI) hierarchies. When a parent class explicitly
defines ``__tablename__``, subclasses would normally inherit that string value.
For STI, child classes must have ``__tablename__`` resolve to ``None`` to indicate
they share the parent's table. This hook enforces that rule.
The detection logic identifies STI children by checking:
1. Class has ``polymorphic_identity`` in ``__mapper_args__`` (explicit STI child marker)
2. AND doesn't have ``concrete=True`` (which would make it CTI)
3. AND doesn't have ``polymorphic_on`` itself (which would make it a base)
4. AND doesn't explicitly define ``__tablename__`` in its own ``__dict__``
For children without ``polymorphic_identity`` but with a parent that has
``polymorphic_on``, SQLAlchemy treats them as abstract intermediate classes
and will issue a warning. We don't modify ``__tablename__`` for these cases.
This allows both usage patterns:
1. Auto-generated names (don't set ``__tablename__`` on parent)
2. Explicit names (set ``__tablename__`` on parent, STI still works)
"""
# IMPORTANT: Modify the class BEFORE calling super().__init_subclass__()
# because super() triggers SQLAlchemy's declarative processing
mapper_args = getattr(cls, "__mapper_args__", {})
# Skip if this class explicitly defines its own __tablename__
if "__tablename__" in cls.__dict__:
super().__init_subclass__(**kwargs)
return
# Skip if this is CTI (concrete table inheritance)
if mapper_args.get("concrete", False):
super().__init_subclass__(**kwargs)
return
# Check if this class might be an STI child
# An STI child either has polymorphic_identity in its own __mapper_args__,
# or inherits from a parent with polymorphic_on
is_potential_sti_child = False
# Check if THIS class (not inherited) defines polymorphic_on
# If it does, it's a base class, not a child
if "__mapper_args__" in cls.__dict__:
own_mapper_args = cls.__dict__["__mapper_args__"]
if "polymorphic_on" in own_mapper_args:
# This is a base class, not a child - skip
super().__init_subclass__(**kwargs)
return
# Check if any parent has polymorphic_on (indicates we're in an STI hierarchy)
for parent in cls.__mro__[1:]:
if not hasattr(parent, "__mapper_args__"):
continue
parent_mapper_args = getattr(parent, "__mapper_args__", {})
if "polymorphic_on" in parent_mapper_args:
# We're inheriting from a polymorphic base, so we're an STI child
is_potential_sti_child = True
break
if is_potential_sti_child and "__tablename__" not in cls.__dict__:
# For STI children that inherited an explicit __tablename__ from a parent,
# we need to explicitly set it to None so SQLAlchemy knows to use the parent's table.
# This overrides the inherited string value.
cls.__tablename__ = None # type: ignore[misc]
# Now call super() which triggers SQLAlchemy's declarative system
super().__init_subclass__(**kwargs)
if TYPE_CHECKING:
__tablename__: Optional[str]
else:
@declared_attr.directive
@classmethod
def __tablename__(cls) -> Optional[str]:
"""Generate table name automatically for base models.
This is called for models that do not have an explicit ``__tablename__``.
For STI child models, ``__init_subclass__`` will have already set
``__tablename__ = None``, so this function returns ``None`` to indicate
the child should use the parent's table.
The generation logic:
1. If class explicitly defines ``__tablename__`` in its ``__dict__``, use that
2. Otherwise, generate from class name using snake_case conversion
Returns:
str | None: Table name generated from class name in snake_case, or None for STI children.
Example:
Single Table Inheritance (both patterns work)::
# Pattern 1: Auto-generated table name (recommended)
class Employee(UUIDBase):
# __tablename__ auto-generated as "employee"
type: Mapped[str]
__mapper_args__ = {
"polymorphic_on": "type",
"polymorphic_identity": "employee",
}
class Manager(Employee):
# __tablename__ = None (set by __init_subclass__)
department: Mapped[str | None]
__mapper_args__ = {"polymorphic_identity": "manager"}
# Pattern 2: Explicit table name on parent
class Employee(UUIDBase):
__tablename__ = "custom_employee" # Explicit!
type: Mapped[str]
__mapper_args__ = {
"polymorphic_on": "type",
"polymorphic_identity": "employee",
}
class Manager(Employee):
# __tablename__ = None (set by __init_subclass__)
# Still uses parent's "custom_employee" table
department: Mapped[str | None]
__mapper_args__ = {"polymorphic_identity": "manager"}
Joined Table Inheritance::
class Employee(UUIDBase):
__tablename__ = "employee"
type: Mapped[str]
__mapper_args__ = {"polymorphic_on": "type"}
class Manager(Employee):
__tablename__ = "manager" # Explicit - has own table
id: Mapped[int] = mapped_column(
ForeignKey("employee.id"), primary_key=True
)
department: Mapped[str]
__mapper_args__ = {"polymorphic_identity": "manager"}
Concrete Table Inheritance::
class Employee(UUIDBase):
__tablename__ = "employee"
id: Mapped[int] = mapped_column(primary_key=True)
class Manager(Employee):
__tablename__ = "manager" # Independent table
__mapper_args__ = {"concrete": True}
"""
# Check if class explicitly defines __tablename__ in its own __dict__
if "__tablename__" in cls.__dict__:
value = cls.__dict__["__tablename__"]
# If explicitly set to None (e.g., by __init_subclass__ for STI), return None
if value is None:
return None
return value
# Check if this is an STI child class that needs auto-detection
# This handles cases where the parent didn't explicitly set __tablename__
mapper_args = getattr(cls, "__mapper_args__", {})
# Skip STI detection if this class defines polymorphic_on (it's a base, not a child)
if "polymorphic_on" not in mapper_args:
is_sti_child = False
# Check explicit STI marker
if "polymorphic_identity" in mapper_args:
is_sti_child = True
else:
# Check if any parent has polymorphic_on (indicates STI hierarchy)
for parent in cls.__mro__[1:]:
if not hasattr(parent, "__mapper_args__"):
continue
parent_mapper_args = getattr(parent, "__mapper_args__", {})
if "polymorphic_on" in parent_mapper_args:
is_sti_child = True
break
if is_sti_child:
# This is an STI child - return None to use parent's table
return None
# Generate table name from class name using snake_case conversion
return table_name_regexp.sub(r"_\1", cls.__name__).lower()
[docs]
def create_registry(
custom_annotation_map: Optional[dict[Any, Union[type[TypeEngine[Any]], TypeEngine[Any]]]] = None,
) -> SQLAlchemyRegistry:
"""Create a new SQLAlchemy registry.
Args:
custom_annotation_map (dict, optional): Custom type annotations to use for the registry.
Returns:
:class:`sqlalchemy.orm.registry`: A new SQLAlchemy registry with the specified type annotations.
"""
import uuid as core_uuid
meta = MetaData(naming_convention=convention)
type_annotation_map: dict[Any, Union[type[TypeEngine[Any]], TypeEngine[Any]]] = {
UUID: GUID,
core_uuid.UUID: GUID,
datetime.datetime: DateTimeUTC,
datetime.date: Date,
dict: JsonB,
dict[str, Any]: JsonB,
dict[str, str]: JsonB,
DataclassProtocol: JsonB,
FileObject: StoredObject,
FileObjectList: StoredObject,
}
with contextlib.suppress(ImportError):
from pydantic import AnyHttpUrl, AnyUrl, EmailStr, IPvAnyAddress, IPvAnyInterface, IPvAnyNetwork, Json
type_annotation_map.update(
{
EmailStr: String,
AnyUrl: String,
AnyHttpUrl: String,
Json: JsonB,
IPvAnyAddress: String,
IPvAnyInterface: String,
IPvAnyNetwork: String,
}
)
with contextlib.suppress(ImportError):
from msgspec import Struct
type_annotation_map[Struct] = JsonB
if custom_annotation_map is not None:
type_annotation_map.update(custom_annotation_map)
return SQLAlchemyRegistry(metadata=meta, type_annotation_map=type_annotation_map)
orm_registry = create_registry()
class MetadataRegistry:
"""A registry for metadata.
Provides methods to get and set metadata for different bind keys.
Methods:
get: Retrieves the metadata for a given bind key.
set: Sets the metadata for a given bind key.
"""
_instance: Optional["MetadataRegistry"] = None
_registry: dict[Union[str, None], MetaData] = {None: orm_registry.metadata}
def __new__(cls) -> Self:
if cls._instance is None:
cls._instance = super().__new__(cls)
return cast("Self", cls._instance)
def get(self, bind_key: Optional[str] = None) -> MetaData:
"""Get the metadata for the given bind key.
Args:
bind_key (Optional[str]): The bind key for the metadata.
Returns:
:class:`sqlalchemy.MetaData`: The metadata for the given bind key.
"""
return self._registry.setdefault(bind_key, MetaData(naming_convention=convention))
def set(self, bind_key: Optional[str], metadata: MetaData) -> None:
"""Set the metadata for the given bind key.
Args:
bind_key (Optional[str]): The bind key for the metadata.
metadata (:class:`sqlalchemy.MetaData`): The metadata to set.
"""
self._registry[bind_key] = metadata
def __iter__(self) -> Iterator[Union[str, None]]:
return iter(self._registry)
def __getitem__(self, bind_key: Union[str, None]) -> MetaData:
return self._registry[bind_key]
def __setitem__(self, bind_key: Union[str, None], metadata: MetaData) -> None:
self._registry[bind_key] = metadata
def __contains__(self, bind_key: Union[str, None]) -> bool:
return bind_key in self._registry
metadata_registry = MetadataRegistry()
[docs]
class AdvancedDeclarativeBase(DeclarativeBase):
"""A subclass of declarative base that allows for overriding of the registry.
Inherits from :class:`sqlalchemy.orm.DeclarativeBase`.
Attributes:
registry (:class:`sqlalchemy.orm.registry`): The registry for the declarative base.
__metadata_registry__ (:class:`~advanced_alchemy.base.MetadataRegistry`): The metadata registry.
__bind_key__ (Optional[:class:`str`]): The bind key for the metadata.
"""
registry = orm_registry
__abstract__ = True
__metadata_registry__: MetadataRegistry = MetadataRegistry()
__bind_key__: Optional[str] = None
def __init_subclass__(cls, **kwargs: Any) -> None:
bind_key = getattr(cls, "__bind_key__", None)
if bind_key is not None:
cls.metadata = cls.__metadata_registry__.get(bind_key)
elif None not in cls.__metadata_registry__ and getattr(cls, "metadata", None) is not None:
cls.__metadata_registry__[None] = cls.metadata
super().__init_subclass__(**kwargs)
[docs]
class UUIDBase(UUIDPrimaryKey, CommonTableAttributes, AdvancedDeclarativeBase, AsyncAttrs):
"""Base for all SQLAlchemy declarative models with UUID v4 primary keys.
.. seealso::
:class:`CommonTableAttributes`
:class:`advanced_alchemy.mixins.UUIDPrimaryKey`
:class:`AdvancedDeclarativeBase`
:class:`AsyncAttrs`
"""
__abstract__ = True
[docs]
class UUIDAuditBase(CommonTableAttributes, UUIDPrimaryKey, AuditColumns, AdvancedDeclarativeBase, AsyncAttrs):
"""Base for declarative models with UUID v4 primary keys and audit columns.
.. seealso::
:class:`CommonTableAttributes`
:class:`advanced_alchemy.mixins.UUIDPrimaryKey`
:class:`advanced_alchemy.mixins.AuditColumns`
:class:`AdvancedDeclarativeBase`
:class:`AsyncAttrs`
"""
__abstract__ = True
[docs]
class UUIDv6Base(UUIDv6PrimaryKey, CommonTableAttributes, AdvancedDeclarativeBase, AsyncAttrs):
"""Base for all SQLAlchemy declarative models with UUID v6 primary keys.
.. seealso::
:class:`advanced_alchemy.mixins.UUIDv6PrimaryKey`
:class:`CommonTableAttributes`
:class:`AdvancedDeclarativeBase`
:class:`AsyncAttrs`
"""
__abstract__ = True
[docs]
class UUIDv6AuditBase(CommonTableAttributes, UUIDv6PrimaryKey, AuditColumns, AdvancedDeclarativeBase, AsyncAttrs):
"""Base for declarative models with UUID v6 primary keys and audit columns.
.. seealso::
:class:`CommonTableAttributes`
:class:`advanced_alchemy.mixins.UUIDv6PrimaryKey`
:class:`advanced_alchemy.mixins.AuditColumns`
:class:`AdvancedDeclarativeBase`
:class:`AsyncAttrs`
"""
__abstract__ = True
[docs]
class UUIDv7Base(UUIDv7PrimaryKey, CommonTableAttributes, AdvancedDeclarativeBase, AsyncAttrs):
"""Base for all SQLAlchemy declarative models with UUID v7 primary keys.
.. seealso::
:class:`advanced_alchemy.mixins.UUIDv7PrimaryKey`
:class:`CommonTableAttributes`
:class:`AdvancedDeclarativeBase`
:class:`AsyncAttrs`
"""
__abstract__ = True
[docs]
class UUIDv7AuditBase(CommonTableAttributes, UUIDv7PrimaryKey, AuditColumns, AdvancedDeclarativeBase, AsyncAttrs):
"""Base for declarative models with UUID v7 primary keys and audit columns.
.. seealso::
:class:`CommonTableAttributes`
:class:`advanced_alchemy.mixins.UUIDv7PrimaryKey`
:class:`advanced_alchemy.mixins.AuditColumns`
:class:`AdvancedDeclarativeBase`
:class:`AsyncAttrs`
"""
__abstract__ = True
[docs]
class NanoIDBase(NanoIDPrimaryKey, CommonTableAttributes, AdvancedDeclarativeBase, AsyncAttrs):
"""Base for all SQLAlchemy declarative models with Nano ID primary keys.
.. seealso::
:class:`advanced_alchemy.mixins.NanoIDPrimaryKey`
:class:`CommonTableAttributes`
:class:`AdvancedDeclarativeBase`
:class:`AsyncAttrs`
"""
__abstract__ = True
[docs]
class NanoIDAuditBase(CommonTableAttributes, NanoIDPrimaryKey, AuditColumns, AdvancedDeclarativeBase, AsyncAttrs):
"""Base for declarative models with Nano ID primary keys and audit columns.
.. seealso::
:class:`CommonTableAttributes`
:class:`advanced_alchemy.mixins.NanoIDPrimaryKey`
:class:`advanced_alchemy.mixins.AuditColumns`
:class:`AdvancedDeclarativeBase`
:class:`AsyncAttrs`
"""
__abstract__ = True
[docs]
class BigIntBase(BigIntPrimaryKey, CommonTableAttributes, AdvancedDeclarativeBase, AsyncAttrs):
"""Base for all SQLAlchemy declarative models with BigInt primary keys.
.. seealso::
:class:`advanced_alchemy.mixins.BigIntPrimaryKey`
:class:`CommonTableAttributes`
:class:`AdvancedDeclarativeBase`
:class:`AsyncAttrs`
"""
__abstract__ = True
[docs]
class BigIntAuditBase(CommonTableAttributes, BigIntPrimaryKey, AuditColumns, AdvancedDeclarativeBase, AsyncAttrs):
"""Base for declarative models with BigInt primary keys and audit columns.
.. seealso::
:class:`CommonTableAttributes`
:class:`advanced_alchemy.mixins.BigIntPrimaryKey`
:class:`advanced_alchemy.mixins.AuditColumns`
:class:`AdvancedDeclarativeBase`
:class:`AsyncAttrs`
"""
__abstract__ = True
[docs]
class IdentityBase(IdentityPrimaryKey, CommonTableAttributes, AdvancedDeclarativeBase, AsyncAttrs):
"""Base for all SQLAlchemy declarative models with database IDENTITY primary keys.
This model uses the database native IDENTITY feature for generating primary keys
instead of using database sequences.
.. seealso::
:class:`advanced_alchemy.mixins.IdentityPrimaryKey`
:class:`CommonTableAttributes`
:class:`AdvancedDeclarativeBase`
:class:`AsyncAttrs`
"""
__abstract__ = True
[docs]
class IdentityAuditBase(CommonTableAttributes, IdentityPrimaryKey, AuditColumns, AdvancedDeclarativeBase, AsyncAttrs):
"""Base for declarative models with database IDENTITY primary keys and audit columns.
This model uses the database native IDENTITY feature for generating primary keys
instead of using database sequences.
.. seealso::
:class:`CommonTableAttributes`
:class:`advanced_alchemy.mixins.IdentityPrimaryKey`
:class:`advanced_alchemy.mixins.AuditColumns`
:class:`AdvancedDeclarativeBase`
:class:`AsyncAttrs`
"""
__abstract__ = True
[docs]
class DefaultBase(CommonTableAttributes, AdvancedDeclarativeBase, AsyncAttrs):
"""Base for all SQLAlchemy declarative models. No primary key is added.
.. seealso::
:class:`CommonTableAttributes`
:class:`AdvancedDeclarativeBase`
:class:`AsyncAttrs`
"""
__abstract__ = True
[docs]
class SQLQuery(BasicAttributes, AdvancedDeclarativeBase, AsyncAttrs):
"""Base for all SQLAlchemy custom mapped objects.
.. seealso::
:class:`BasicAttributes`
:class:`AdvancedDeclarativeBase`
:class:`AsyncAttrs`
"""
__abstract__ = True
__allow_unmapped__ = True
