Add composite primary key support

Author: wu-clan

sqlalchemy_crud_plus/crud.py

@@ -26,7 +26,7 @@ def __init__(self, model: Type[Model]):
         self.model = model
         self.primary_key = self._get_primary_key()
 
-    def _get_primary_key(self) -> Column:
+    def _get_primary_key(self) -> Column | list[Column]:
         """
         Dynamically retrieve the primary key column(s) for the model.
         """
@@ -35,7 +35,21 @@ def _get_primary_key(self) -> Column:
         if len(primary_key) == 1:
             return primary_key[0]
         else:
-            raise CompositePrimaryKeysError('Composite primary keys are not supported')
+            return list(primary_key)
+
+    def _get_pk_filter(self, pk: Any | Sequence[Any]) -> list[bool]:
+        """
+        Get the primary key filter(s).
+
+        :param pk:
+        :return:
+        """
+        if isinstance(self.primary_key, list):
+            if len(pk) != len(self.primary_key):
+                raise CompositePrimaryKeysError(f'Expected {len(self.primary_key)} values for composite primary key')
+            return [column == value for column, value in zip(self.primary_key, pk)]
+        else:
+            return [self.primary_key == pk]
 
     async def create_model(
         self,
@@ -55,10 +69,7 @@ async def create_model(
         :param kwargs: Additional model data not included in the pydantic schema.
         :return:
         """
-        if not kwargs:
-            ins = self.model(**obj.model_dump())
-        else:
-            ins = self.model(**obj.model_dump(), **kwargs)
+        ins = self.model(**obj.model_dump()) if not kwargs else self.model(**obj.model_dump(), **kwargs)
 
         session.add(ins)
 
@@ -89,10 +100,7 @@ async def create_models(
         """
         ins_list = []
         for obj in objs:
-            if not kwargs:
-                ins = self.model(**obj.model_dump())
-            else:
-                ins = self.model(**obj.model_dump(), **kwargs)
+            ins = self.model(**obj.model_dump()) if not kwargs else self.model(**obj.model_dump(), **kwargs)
             ins_list.append(ins)
 
         session.add_all(ins_list)
@@ -118,12 +126,12 @@ async def count(
         :param kwargs: Query expressions.
         :return:
         """
-        filter_list = list(whereclause)
+        filters = list(whereclause)
 
         if kwargs:
-            filter_list.extend(parse_filters(self.model, **kwargs))
+            filters.extend(parse_filters(self.model, **kwargs))
 
-        stmt = select(func.count()).select_from(self.model).where(*filter_list)
+        stmt = select(func.count()).select_from(self.model).where(*filters)
         query = await session.execute(stmt)
         total_count = query.scalar()
         return total_count if total_count is not None else 0
@@ -154,21 +162,20 @@ async def exists(
     async def select_model(
         self,
         session: AsyncSession,
-        pk: int,
+        pk: Any | Sequence[Any],
         *whereclause: ColumnExpressionArgument[bool],
     ) -> Model | None:
         """
-        Query by ID
+        Query by primary key(s)
 
         :param session: The SQLAlchemy async session.
-        :param pk: The database primary key value.
+        :param pk: Single value for simple primary key, or tuple for composite primary key.
         :param whereclause: The WHERE clauses to apply to the query.
         :return:
         """
-        filter_list = list(whereclause)
-        _filters = [self.primary_key == pk]
-        _filters.extend(filter_list)
-        stmt = select(self.model).where(*_filters)
+        filters = self._get_pk_filter(pk)
+        filters + list(whereclause)
+        stmt = select(self.model).where(*filters)
         query = await session.execute(stmt)
         return query.scalars().first()
 
@@ -186,10 +193,8 @@ async def select_model_by_column(
         :param kwargs: Query expressions.
         :return:
         """
-        filter_list = list(whereclause)
-        _filters = parse_filters(self.model, **kwargs)
-        _filters.extend(filter_list)
-        stmt = select(self.model).where(*_filters)
+        filters = parse_filters(self.model, **kwargs) + list(whereclause)
+        stmt = select(self.model).where(*filters)
         query = await session.execute(stmt)
         return query.scalars().first()
 
@@ -201,10 +206,8 @@ async def select(self, *whereclause: ColumnExpressionArgument[bool], **kwargs) -
         :param kwargs: Query expressions.
         :return:
         """
-        filter_list = list(whereclause)
-        _filters = parse_filters(self.model, **kwargs)
-        _filters.extend(filter_list)
-        stmt = select(self.model).where(*_filters)
+        filters = parse_filters(self.model, **kwargs) + list(whereclause)
+        stmt = select(self.model).where(*filters)
         return stmt
 
     async def select_order(
@@ -270,7 +273,7 @@ async def select_models_order(
     async def update_model(
         self,
         session: AsyncSession,
-        pk: int,
+        pk: Any | Sequence[Any],
         obj: UpdateSchema | dict[str, Any],
         flush: bool = False,
         commit: bool = False,
@@ -280,21 +283,17 @@ async def update_model(
         Update an instance by model's primary key
 
         :param session: The SQLAlchemy async session.
-        :param pk: The database primary key value.
+        :param pk: Single value for simple primary key, or tuple for composite primary key.
         :param obj: A pydantic schema or dictionary containing the update data
         :param flush: If `True`, flush all object changes to the database. Default is `False`.
         :param commit: If `True`, commits the transaction immediately. Default is `False`.
         :param kwargs: Additional model data not included in the pydantic schema.
         :return:
         """
-        if isinstance(obj, dict):
-            instance_data = obj
-        else:
-            instance_data = obj.model_dump(exclude_unset=True)
-        if kwargs:
-            instance_data.update(kwargs)
-
-        stmt = update(self.model).where(self.primary_key == pk).values(**instance_data)
+        filters = self._get_pk_filter(pk)
+        instance_data = obj if isinstance(obj, dict) else obj.model_dump(exclude_unset=True)
+        instance_data.update(kwargs)
+        stmt = update(self.model).where(*filters).values(**instance_data)
         result = await session.execute(stmt)
 
         if flush:
@@ -325,15 +324,13 @@ async def update_model_by_column(
         :return:
         """
         filters = parse_filters(self.model, **kwargs)
+
         total_count = await self.count(session, *filters)
         if not allow_multiple and total_count > 1:
             raise MultipleResultsError(f'Only one record is expected to be update, found {total_count} records.')
-        if isinstance(obj, dict):
-            instance_data = obj
-        else:
-            instance_data = obj.model_dump(exclude_unset=True)
 
-        stmt = update(self.model).where(*filters).values(**instance_data)  # type: ignore
+        instance_data = obj if isinstance(obj, dict) else obj.model_dump(exclude_unset=True)
+        stmt = update(self.model).where(*filters).values(**instance_data)
         result = await session.execute(stmt)
 
         if flush:
@@ -346,20 +343,22 @@ async def update_model_by_column(
     async def delete_model(
         self,
         session: AsyncSession,
-        pk: int,
+        pk: Any | Sequence[Any],
         flush: bool = False,
         commit: bool = False,
     ) -> int:
         """
         Delete an instance by model's primary key
 
         :param session: The SQLAlchemy async session.
-        :param pk: The database primary key value.
+        :param pk: Single value for simple primary key, or tuple for composite primary key.
         :param flush: If `True`, flush all object changes to the database. Default is `False`.
         :param commit: If `True`, commits the transaction immediately. Default is `False`.
         :return:
         """
-        stmt = delete(self.model).where(self.primary_key == pk)
+        filters = self._get_pk_filter(pk)
+
+        stmt = delete(self.model).where(*filters)
         result = await session.execute(stmt)
 
         if flush:
@@ -392,14 +391,16 @@ async def delete_model_by_column(
         :return:
         """
         filters = parse_filters(self.model, **kwargs)
+
         total_count = await self.count(session, *filters)
         if not allow_multiple and total_count > 1:
             raise MultipleResultsError(f'Only one record is expected to be delete, found {total_count} records.')
-        if logical_deletion:
-            deleted_flag = {deleted_flag_column: True}
-            stmt = update(self.model).where(*filters).values(**deleted_flag)
-        else:
-            stmt = delete(self.model).where(*filters)
+
+        stmt = (
+            update(self.model).where(*filters).values(**{deleted_flag_column: True})
+            if logical_deletion
+            else delete(self.model).where(*filters)
+        )
 
         result = await session.execute(stmt)
 

tests/conftest.py

@@ -5,7 +5,7 @@
 
 from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
 
-from tests.model import Base, Ins
+from tests.model import Base, Ins, InsPks
 
 _async_engine = create_async_engine('sqlite+aiosqlite:///:memory:', future=True)
 _async_session = async_sessionmaker(_async_engine, autoflush=False, expire_on_commit=False)
@@ -29,3 +29,12 @@ async def create_test_model():
     async with _async_session.begin() as session:
         data = [Ins(name=f'name_{i}') for i in range(1, 10)]
         session.add_all(data)
+
+
+@pytest_asyncio.fixture
+async def create_test_model_pks():
+    async with _async_session.begin() as session:
+        data = [InsPks(id=i, name=f'name_{i}', sex='men') for i in range(1, 5)]
+        session.add_all(data)
+        data = [InsPks(id=i, name=f'name_{i}', sex='women') for i in range(6, 10)]
+        session.add_all(data)

tests/model.py

@@ -20,3 +20,14 @@ class Ins(Base):
     del_flag: Mapped[bool] = mapped_column(default=False)
     created_time: Mapped[datetime] = mapped_column(init=False, default_factory=datetime.now)
     updated_time: Mapped[datetime | None] = mapped_column(init=False, onupdate=datetime.now)
+
+
+class InsPks(Base):
+    __tablename__ = 'ins_pks'
+
+    id: Mapped[int] = mapped_column(primary_key=True, index=True)
+    name: Mapped[str] = mapped_column(String(64))
+    sex: Mapped[str] = mapped_column(String(16), primary_key=True, index=True)
+    del_flag: Mapped[bool] = mapped_column(default=False)
+    created_time: Mapped[datetime] = mapped_column(init=False, default_factory=datetime.now)
+    updated_time: Mapped[datetime | None] = mapped_column(init=False, onupdate=datetime.now)

tests/schema.py

@@ -5,3 +5,9 @@
 
 class ModelTest(BaseModel):
     name: str
+
+
+class ModelTestPks(BaseModel):
+    id: int
+    name: str
+    sex: str

tests/test_create.py

@@ -1,12 +1,14 @@
 #!/usr/bin/env python3
 # -*- coding: utf-8 -*-
+from random import choice
+
 import pytest
 
 from sqlalchemy import select
 
 from sqlalchemy_crud_plus import CRUDPlus
-from tests.model import Ins
-from tests.schema import ModelTest
+from tests.model import Ins, InsPks
+from tests.schema import ModelTest, ModelTestPks
 
 
 @pytest.mark.asyncio
@@ -34,3 +36,30 @@ async def test_create_models(async_db_session):
         for i in range(1, 10):
             query = await session.scalar(select(Ins).where(Ins.id == i))
             assert query.name == f'name_{i}'
+
+
+@pytest.mark.asyncio
+async def test_create_model_pks(async_db_session):
+    async with async_db_session.begin() as session:
+        crud = CRUDPlus(InsPks)
+        for i in range(1, 10):
+            data = ModelTestPks(id=i, name=f'name_{i}', sex=choice(['men', 'women']))
+            await crud.create_model(session, data)
+    async with async_db_session() as session:
+        for i in range(1, 10):
+            query = await session.scalar(select(InsPks).where(InsPks.id == i))
+            assert query.name == f'name_{i}'
+
+
+@pytest.mark.asyncio
+async def test_create_models_pks(async_db_session):
+    async with async_db_session.begin() as session:
+        crud = CRUDPlus(InsPks)
+        data = []
+        for i in range(1, 10):
+            data.append(ModelTestPks(id=i, name=f'name_{i}', sex=choice(['men', 'women'])))
+        await crud.create_models(session, data)
+    async with async_db_session() as session:
+        for i in range(1, 10):
+            query = await session.scalar(select(InsPks).where(InsPks.id == i))
+            assert query.name == f'name_{i}'

tests/test_delete.py

@@ -3,7 +3,7 @@
 import pytest
 
 from sqlalchemy_crud_plus import CRUDPlus
-from tests.model import Ins
+from tests.model import Ins, InsPks
 
 
 @pytest.mark.asyncio
@@ -14,6 +14,14 @@ async def test_delete_model(create_test_model, async_db_session):
         assert result == 1
 
 
+@pytest.mark.asyncio
+async def test_delete_model_pks(create_test_model_pks, async_db_session):
+    async with async_db_session.begin() as session:
+        crud = CRUDPlus(InsPks)
+        result = await crud.delete_model(session, (1, 'men'))
+        assert result == 1
+
+
 @pytest.mark.asyncio
 async def test_delete_model_by_column(create_test_model, async_db_session):
     async with async_db_session.begin() as session:

tests/test_select.py

@@ -5,7 +5,7 @@
 from sqlalchemy import Select
 
 from sqlalchemy_crud_plus import CRUDPlus
-from tests.model import Ins
+from tests.model import Ins, InsPks
 
 
 @pytest.mark.asyncio
@@ -62,6 +62,18 @@ async def test_select_model(create_test_model, async_db_session):
             assert result.name == f'name_{i}'
 
 
+@pytest.mark.asyncio
+async def test_select_model_pks(create_test_model_pks, async_db_session):
+    async with async_db_session() as session:
+        crud = CRUDPlus(InsPks)
+        for i in range(1, 5):
+            result = await crud.select_model(session, (i, 'men'))
+            assert result.name == f'name_{i}'
+        for i in range(6, 10):
+            result = await crud.select_model(session, (i, 'women'))
+            assert result.name == f'name_{i}'
+
+
 @pytest.mark.asyncio
 async def test_select_model_by_column(create_test_model, async_db_session):
     async with async_db_session() as session: