fix: IdentityPrimaryKey correctly generates an IDENTITY DDL (#523)

Ensure that the IdentityPrimaryKey correctly generates IDENTITY DDL across multiple database dialects, including PostgreSQL, Oracle, and SQL Server. Update dependencies and add tests to verify the functionality.

Author: cofin

.pre-commit-config.yaml

@@ -22,7 +22,7 @@ repos:
       - id: unasyncd
         additional_dependencies: ["ruff"]
   - repo: https://github.com/charliermarsh/ruff-pre-commit
-    rev: "v0.12.9"
+    rev: "v0.12.10"
     hooks:
       # Run the linter.
       - id: ruff

advanced_alchemy/mixins/bigint.py

@@ -1,6 +1,6 @@
 from typing import Any, Optional
 
-from sqlalchemy import Sequence
+from sqlalchemy import Identity, Sequence
 from sqlalchemy.orm import Mapped, declarative_mixin, declared_attr, mapped_column
 
 from advanced_alchemy.types import BigIntIdentity
@@ -51,6 +51,6 @@ def id(cls) -> Mapped[int]:
         """Primary key column using IDENTITY."""
         return mapped_column(
             BigIntIdentity,
+            Identity(start=1, increment=1),
             primary_key=True,
-            autoincrement=True,
         )

tests/unit/test_base.py

@@ -2,6 +2,12 @@
 from __future__ import annotations
 
 import warnings
+from typing import cast
+
+from sqlalchemy import Table, create_engine
+from sqlalchemy.dialects import mssql, oracle, postgresql
+from sqlalchemy.orm import declarative_mixin
+from sqlalchemy.schema import CreateTable
 
 from tests.helpers import purge_module
 
@@ -41,3 +47,115 @@ def test_deprecated_classes_functionality() -> None:
     assert hasattr(nanoid_pk, "_sentinel")
     assert hasattr(audit, "created_at")
     assert hasattr(audit, "updated_at")
+
+
+def test_identity_primary_key_generates_identity_ddl() -> None:
+    """Test that IdentityPrimaryKey generates proper IDENTITY DDL for PostgreSQL."""
+    from advanced_alchemy.base import BigIntBase
+    from advanced_alchemy.mixins.bigint import IdentityPrimaryKey
+
+    @declarative_mixin
+    class TestMixin(IdentityPrimaryKey):
+        __tablename__ = "test_identity"
+
+    class TestModel(TestMixin, BigIntBase):
+        pass
+
+    # Get the CREATE TABLE statement
+    create_stmt = CreateTable(cast(Table, TestModel.__table__))
+
+    # Test with PostgreSQL dialect
+    pg_ddl = str(create_stmt.compile(dialect=postgresql.dialect()))  # type: ignore[no-untyped-call,unused-ignore]
+
+    # Should contain GENERATED BY DEFAULT AS IDENTITY
+    assert "GENERATED BY DEFAULT AS IDENTITY" in pg_ddl
+    assert "BIGSERIAL" not in pg_ddl.upper()
+    assert "START WITH 1" in pg_ddl
+    assert "INCREMENT BY 1" in pg_ddl
+
+
+def test_identity_audit_base_generates_identity_ddl() -> None:
+    """Test that IdentityAuditBase generates proper IDENTITY DDL for PostgreSQL."""
+    from advanced_alchemy.base import IdentityAuditBase
+
+    class TestModel(IdentityAuditBase):
+        __tablename__ = "test_identity_audit"
+
+    # Get the CREATE TABLE statement
+    create_stmt = CreateTable(cast(Table, TestModel.__table__))
+
+    # Test with PostgreSQL dialect
+    pg_ddl = str(create_stmt.compile(dialect=postgresql.dialect()))  # type: ignore[no-untyped-call,unused-ignore]
+
+    # Should contain GENERATED BY DEFAULT AS IDENTITY
+    assert "GENERATED BY DEFAULT AS IDENTITY" in pg_ddl
+    assert "BIGSERIAL" not in pg_ddl.upper()
+
+
+def test_bigint_primary_key_still_uses_sequence() -> None:
+    """Test that BigIntPrimaryKey still uses sequences as before."""
+    from advanced_alchemy.base import BigIntBase
+    from advanced_alchemy.mixins.bigint import BigIntPrimaryKey
+
+    @declarative_mixin
+    class TestMixin(BigIntPrimaryKey):
+        __tablename__ = "test_bigint"
+
+    class TestModel(TestMixin, BigIntBase):
+        pass
+
+    # Get the CREATE TABLE statement
+    create_stmt = CreateTable(cast(Table, TestModel.__table__))
+
+    # Test with PostgreSQL dialect
+    pg_ddl = str(create_stmt.compile(dialect=postgresql.dialect()))  # type: ignore[no-untyped-call,unused-ignore]
+
+    # BigIntPrimaryKey should use a Sequence (not IDENTITY)
+    assert "GENERATED" not in pg_ddl
+    assert "IDENTITY" not in pg_ddl.upper()
+    # The sequence is defined on the column but rendered separately
+    assert TestModel.__table__.c.id.default is not None
+    assert TestModel.__table__.c.id.default.name == "test_bigint_id_seq"
+
+
+def test_identity_ddl_for_oracle() -> None:
+    """Test Identity DDL generation for Oracle."""
+    from advanced_alchemy.base import IdentityAuditBase
+
+    class TestModel(IdentityAuditBase):
+        __tablename__ = "test_oracle"
+
+    create_stmt = CreateTable(cast(Table, TestModel.__table__))
+    oracle_ddl = str(create_stmt.compile(dialect=oracle.dialect()))  # type: ignore[no-untyped-call,unused-ignore]
+
+    # Oracle should generate IDENTITY
+    assert "GENERATED BY DEFAULT AS IDENTITY" in oracle_ddl
+
+
+def test_identity_ddl_for_mssql() -> None:
+    """Test Identity DDL generation for SQL Server."""
+    from advanced_alchemy.base import IdentityAuditBase
+
+    class TestModel(IdentityAuditBase):
+        __tablename__ = "test_mssql"
+
+    create_stmt = CreateTable(cast(Table, TestModel.__table__))
+    mssql_ddl = str(create_stmt.compile(dialect=mssql.dialect()))  # type: ignore[no-untyped-call,unused-ignore]
+
+    # SQL Server should generate IDENTITY
+    assert "IDENTITY(1,1)" in mssql_ddl
+
+
+def test_identity_works_with_sqlite() -> None:
+    """Test that Identity columns work with SQLite (fallback to autoincrement)."""
+    from advanced_alchemy.base import IdentityAuditBase
+
+    class TestModel(IdentityAuditBase):
+        __tablename__ = "test_sqlite"
+
+    # Create an in-memory SQLite engine
+    engine = create_engine("sqlite:///:memory:")
+    cast(Table, TestModel.__table__).create(engine)
+
+    # Should not raise any errors
+    assert True  # If we get here, it worked