feat: Add identity Mixin for Primary Keys (#473)

The sequences based BigInt key offers the most compatibility, but many would prefer to use the Identity column when the database supports it.

This changes implements a basic Identity primary key mixin

Author: cofin

.pre-commit-config.yaml

@@ -2,7 +2,7 @@ default_language_version:
   python: "3"
 repos:
   - repo: https://github.com/compilerla/conventional-pre-commit
-    rev: v4.1.0
+    rev: v4.2.0
     hooks:
       - id: conventional-pre-commit
         stages: [commit-msg]

advanced_alchemy/base.py

@@ -27,6 +27,7 @@
 from advanced_alchemy.mixins import (
     AuditColumns,
     BigIntPrimaryKey,
+    IdentityPrimaryKey,
     NanoIDPrimaryKey,
     UUIDPrimaryKey,
     UUIDv6PrimaryKey,
@@ -49,6 +50,9 @@
     "BigIntBase",
     "BigIntBaseT",
     "CommonTableAttributes",
+    "IdentityAuditBase",
+    "IdentityBase",
+    "IdentityBaseT",
     "ModelProtocol",
     "NanoIDAuditBase",
     "NanoIDBase",
@@ -77,6 +81,8 @@
 """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")
@@ -475,6 +481,39 @@ class BigIntAuditBase(CommonTableAttributes, BigIntPrimaryKey, AuditColumns, Adv
     __abstract__ = True
 
 
+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
+
+
+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
+
+
 class DefaultBase(CommonTableAttributes, AdvancedDeclarativeBase, AsyncAttrs):
     """Base for all SQLAlchemy declarative models.  No primary key is added.
 

advanced_alchemy/mixins/__init__.py

@@ -1,5 +1,5 @@
 from advanced_alchemy.mixins.audit import AuditColumns
-from advanced_alchemy.mixins.bigint import BigIntPrimaryKey
+from advanced_alchemy.mixins.bigint import BigIntPrimaryKey, IdentityPrimaryKey
 from advanced_alchemy.mixins.nanoid import NanoIDPrimaryKey
 from advanced_alchemy.mixins.sentinel import SentinelMixin
 from advanced_alchemy.mixins.slug import SlugKey
@@ -9,6 +9,7 @@
 __all__ = (
     "AuditColumns",
     "BigIntPrimaryKey",
+    "IdentityPrimaryKey",
     "NanoIDPrimaryKey",
     "SentinelMixin",
     "SlugKey",

advanced_alchemy/mixins/bigint.py

@@ -36,3 +36,21 @@ def id(cls) -> Mapped[int]:
             Sequence(f"{cls.__tablename__}_id_seq", **seq_kwargs),  # type: ignore[attr-defined]
             primary_key=True,
         )
+
+
+@declarative_mixin
+class IdentityPrimaryKey:
+    """Primary Key Field Mixin using database IDENTITY feature.
+
+    This mixin uses the database's native IDENTITY feature rather than a sequence.
+    This can be more efficient for databases that support IDENTITY natively.
+    """
+
+    @declared_attr
+    def id(cls) -> Mapped[int]:
+        """Primary key column using IDENTITY."""
+        return mapped_column(
+            BigIntIdentity,
+            primary_key=True,
+            autoincrement=True,
+        )

docs/usage/modeling.rst

@@ -20,6 +20,10 @@ Advanced Alchemy provides several base classes optimized for different use cases
      - BIGINT primary keys for tables
    * - ``BigIntAuditBase``
      - BIGINT primary keys for tables, Automatic created_at/updated_at timestamps
+   * - ``IdentityBase``
+     - Primary keys using database IDENTITY feature instead of sequences
+   * - ``IdentityAuditBase``
+     - Primary keys using database IDENTITY feature, Automatic created_at/updated_at timestamps
    * - ``UUIDBase``
      - UUID primary keys
    * - ``UUIDv6Base``
@@ -53,6 +57,10 @@ Additionally, Advanced Alchemy provides mixins to enhance model functionality:
    * - ``AuditColumns``
      - | Automatic created_at/updated_at timestamps
        | Tracks record modifications
+   * - ``BigIntPrimaryKey``
+     - | Adds BigInt primary key with sequence
+   * - ``IdentityPrimaryKey``
+     - | Adds primary key using database IDENTITY feature
    * - ``UniqueMixin``
      - | Automatic Select or Create for many-to-many relationships
 

tests/integration/test_bigint_identity.py

@@ -0,0 +1,153 @@
+from __future__ import annotations
+
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+import pytest
+from sqlalchemy import Column, ForeignKey, String, Table, create_engine, select
+from sqlalchemy.ext.associationproxy import AssociationProxy, association_proxy
+from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
+from sqlalchemy.orm import Mapped, Session, mapped_column, relationship, sessionmaker
+
+if TYPE_CHECKING:
+    from pytest import MonkeyPatch
+
+
+@pytest.mark.xdist_group("loader")
+def test_ap_sync(monkeypatch: MonkeyPatch, tmp_path: Path) -> None:
+    from sqlalchemy.orm import DeclarativeBase
+
+    from advanced_alchemy import base, mixins
+
+    orm_registry = base.create_registry()
+
+    class NewIdentityBase(mixins.IdentityPrimaryKey, base.CommonTableAttributes, DeclarativeBase):
+        registry = orm_registry
+
+    monkeypatch.setattr(base, "IdentityBase", NewIdentityBase)
+
+    product_tag_table = Table(
+        "product_tag",
+        orm_registry.metadata,
+        Column("product_id", ForeignKey("product.id", ondelete="CASCADE"), primary_key=True),  # pyright: ignore[reportUnknownArgumentType]
+        Column("tag_id", ForeignKey("tag.id", ondelete="CASCADE"), primary_key=True),  # pyright: ignore[reportUnknownArgumentType]
+    )
+
+    class Tag(NewIdentityBase):
+        name: Mapped[str] = mapped_column(index=True)
+        products: Mapped[list[Product]] = relationship(
+            secondary=lambda: product_tag_table,
+            back_populates="product_tags",
+            cascade="all, delete",
+            passive_deletes=True,
+            lazy="noload",
+        )
+
+    class Product(NewIdentityBase):
+        name: Mapped[str] = mapped_column(String(length=50))  # pyright: ignore
+        product_tags: Mapped[list[Tag]] = relationship(
+            secondary=lambda: product_tag_table,
+            back_populates="products",
+            cascade="all, delete",
+            passive_deletes=True,
+            lazy="joined",
+        )
+        tags: AssociationProxy[list[str]] = association_proxy(
+            "product_tags",
+            "name",
+            creator=lambda name: Tag(name=name),  # pyright: ignore[reportUnknownArgumentType,reportUnknownLambdaType]
+        )
+
+    engine = create_engine(f"sqlite:///{tmp_path}/test.sqlite1.db", echo=True)
+    session_factory: sessionmaker[Session] = sessionmaker(engine, expire_on_commit=False)
+
+    with engine.begin() as conn:
+        Product.metadata.create_all(conn)
+
+    with session_factory() as db_session:
+        product_1 = Product(name="Product 1", tags=["a new tag", "second tag"])
+        db_session.add(product_1)
+
+        tags = db_session.execute(select(Tag)).unique().fetchall()
+        assert len(tags) == 2
+
+        product_2 = Product(name="Product 2", tags=["third tag"])
+        db_session.add(product_2)
+        tags = db_session.execute(select(Tag)).unique().fetchall()
+        assert len(tags) == 3
+
+        product_2.tags = []
+        db_session.add(product_2)
+
+        product_2_validate = db_session.execute(select(Product).where(Product.name == "Product 2")).unique().fetchone()
+        assert product_2_validate
+        tags_2 = db_session.execute(select(Tag)).unique().fetchall()
+        assert len(product_2_validate[0].product_tags) == 0
+        assert len(tags_2) == 3
+        # add more assertions
+
+
+@pytest.mark.xdist_group("loader")
+async def test_ap_async(monkeypatch: MonkeyPatch, tmp_path: Path) -> None:
+    from sqlalchemy.orm import DeclarativeBase
+
+    from advanced_alchemy import base, mixins
+
+    orm_registry = base.create_registry()
+
+    class NewIdentityBase(mixins.IdentityPrimaryKey, base.CommonTableAttributes, DeclarativeBase):
+        registry = orm_registry
+
+    monkeypatch.setattr(base, "IdentityBase", NewIdentityBase)
+
+    product_tag_table = Table(
+        "product_tag",
+        orm_registry.metadata,
+        Column("product_id", ForeignKey("product.id", ondelete="CASCADE"), primary_key=True),  # pyright: ignore[reportUnknownArgumentType]
+        Column("tag_id", ForeignKey("tag.id", ondelete="CASCADE"), primary_key=True),  # pyright: ignore[reportUnknownArgumentType]
+    )
+
+    class Tag(NewIdentityBase):
+        name: Mapped[str] = mapped_column(index=True)
+        products: Mapped[list[Product]] = relationship(
+            secondary=lambda: product_tag_table,
+            back_populates="product_tags",
+            cascade="all, delete",
+            passive_deletes=True,
+            lazy="noload",
+        )
+
+    class Product(NewIdentityBase):
+        name: Mapped[str] = mapped_column(String(length=50))  # pyright: ignore
+        product_tags: Mapped[list[Tag]] = relationship(
+            secondary=lambda: product_tag_table,
+            back_populates="products",
+            cascade="all, delete",
+            passive_deletes=True,
+            lazy="joined",
+        )
+        tags: AssociationProxy[list[str]] = association_proxy(
+            "product_tags",
+            "name",
+            creator=lambda name: Tag(name=name),  # pyright: ignore[reportUnknownArgumentType,reportUnknownLambdaType]
+        )
+
+    engine = create_async_engine(f"sqlite+aiosqlite:///{tmp_path}/test.sqlite2.db", echo=True)
+    session_factory: async_sessionmaker[AsyncSession] = async_sessionmaker(engine, expire_on_commit=False)
+
+    async with engine.begin() as conn:
+        await conn.run_sync(Tag.metadata.create_all)
+
+    async with session_factory() as db_session:
+        product_1 = Product(name="Product 1 async", tags=["a new tag", "second tag"])
+        db_session.add(product_1)
+
+        tags = await db_session.execute(select(Tag))
+        assert len(tags.unique().fetchall()) == 2
+
+        product_2 = Product(name="Product 2 async", tags=["third tag"])
+        db_session.add(product_2)
+        tags = await db_session.execute(select(Tag))
+        assert len(tags.unique().fetchall()) == 3
+
+        # add more assertions