Refactor code to enhance selection and operations (#9)

* WIP: refactoring enhancements

* Update delete method

* Fix query filter parse

* Update tests

* Update multiple_conditions to and

* fix parse filters method

* fix tests

* Add arithmetic filters

Author: wu-clan

sqlalchemy_crud_plus/crud.py

@@ -1,184 +1,130 @@
 #!/usr/bin/env python3
 # -*- coding: utf-8 -*-
-from typing import Any, Generic, Iterable, Literal, Sequence, Type, TypeVar
+from typing import Any, Generic, Iterable, Sequence, Type
 
-from pydantic import BaseModel
-from sqlalchemy import Row, RowMapping, and_, asc, desc, or_, select
+from sqlalchemy import Row, RowMapping, select
 from sqlalchemy import delete as sa_delete
 from sqlalchemy import update as sa_update
 from sqlalchemy.ext.asyncio import AsyncSession
 
-from sqlalchemy_crud_plus.errors import ModelColumnError, SelectExpressionError
+from sqlalchemy_crud_plus.errors import MultipleResultsError
+from sqlalchemy_crud_plus.types import CreateSchema, Model, UpdateSchema
+from sqlalchemy_crud_plus.utils import apply_sorting, count, parse_filters
 
-_Model = TypeVar('_Model')
-_CreateSchema = TypeVar('_CreateSchema', bound=BaseModel)
-_UpdateSchema = TypeVar('_UpdateSchema', bound=BaseModel)
 
-
-class CRUDPlus(Generic[_Model]):
-    def __init__(self, model: Type[_Model]):
+class CRUDPlus(Generic[Model]):
+    def __init__(self, model: Type[Model]):
         self.model = model
 
-    async def create_model(self, session: AsyncSession, obj: _CreateSchema, commit: bool = False, **kwargs) -> _Model:
+    async def create_model(self, session: AsyncSession, obj: CreateSchema, commit: bool = False, **kwargs) -> Model:
         """
         Create a new instance of a model
 
-        :param session:
-        :param obj:
-        :param commit:
-        :param kwargs:
+        :param session: The SQLAlchemy async session.
+        :param obj: The Pydantic schema containing data to be saved.
+        :param commit: If `True`, commits the transaction immediately. Default is `False`.
+        :param kwargs: Additional model data not included in the pydantic schema.
         :return:
         """
-        if kwargs:
-            ins = self.model(**obj.model_dump(), **kwargs)
-        else:
+        if not kwargs:
             ins = self.model(**obj.model_dump())
+        else:
+            ins = self.model(**obj.model_dump(), **kwargs)
         session.add(ins)
         if commit:
             await session.commit()
         return ins
 
     async def create_models(
-        self, session: AsyncSession, obj: Iterable[_CreateSchema], commit: bool = False
-    ) -> list[_Model]:
+        self, session: AsyncSession, obj: Iterable[CreateSchema], commit: bool = False
+    ) -> list[Model]:
         """
         Create new instances of a model
 
-        :param session:
-        :param obj:
-        :param commit:
+        :param session: The SQLAlchemy async session.
+        :param obj: The Pydantic schema list containing data to be saved.
+        :param commit: If `True`, commits the transaction immediately. Default is `False`.
         :return:
         """
         ins_list = []
-        for i in obj:
-            ins_list.append(self.model(**i.model_dump()))
+        for ins in obj:
+            ins_list.append(self.model(**ins.model_dump()))
         session.add_all(ins_list)
         if commit:
             await session.commit()
         return ins_list
 
-    async def select_model_by_id(self, session: AsyncSession, pk: int) -> _Model | None:
+    async def select_model(self, session: AsyncSession, pk: int) -> Model | None:
         """
         Query by ID
 
-        :param session:
-        :param pk:
+        :param session: The SQLAlchemy async session.
+        :param pk: The database primary key value.
         :return:
         """
         stmt = select(self.model).where(self.model.id == pk)
         query = await session.execute(stmt)
         return query.scalars().first()
 
-    async def select_model_by_column(self, session: AsyncSession, column: str, column_value: Any) -> _Model | None:
+    async def select_model_by_column(self, session: AsyncSession, **kwargs) -> Model | None:
         """
         Query by column
 
-        :param session:
-        :param column:
-        :param column_value:
-        :return:
-        """
-        if hasattr(self.model, column):
-            model_column = getattr(self.model, column)
-            stmt = select(self.model).where(model_column == column_value)  # type: ignore
-            query = await session.execute(stmt)
-            return query.scalars().first()
-        else:
-            raise ModelColumnError(f'Column {column} is not found in {self.model}')
-
-    async def select_model_by_columns(
-        self, session: AsyncSession, expression: Literal['and', 'or'] = 'and', **conditions
-    ) -> _Model | None:
-        """
-        Query by columns
-
-        :param session:
-        :param expression:
-        :param conditions: Query conditions, format：column1=value1, column2=value2
+        :param session: The SQLAlchemy async session.
+        :param kwargs: Query expressions.
         :return:
         """
-        where_list = []
-        for column, value in conditions.items():
-            if hasattr(self.model, column):
-                model_column = getattr(self.model, column)
-                where_list.append(model_column == value)
-            else:
-                raise ModelColumnError(f'Column {column} is not found in {self.model}')
-        match expression:
-            case 'and':
-                stmt = select(self.model).where(and_(*where_list))
-                query = await session.execute(stmt)
-            case 'or':
-                stmt = select(self.model).where(or_(*where_list))
-                query = await session.execute(stmt)
-            case _:
-                raise SelectExpressionError(
-                    f'Select expression {expression} is not supported, only supports `and`, `or`'
-                )
+        filters = await parse_filters(self.model, **kwargs)
+        stmt = select(self.model).where(*filters)
+        query = await session.execute(stmt)
         return query.scalars().first()
 
-    async def select_models(self, session: AsyncSession) -> Sequence[Row[Any] | RowMapping | Any]:
+    async def select_models(self, session: AsyncSession, **kwargs) -> Sequence[Row[Any] | RowMapping | Any]:
         """
         Query all rows
 
-        :param session:
+        :param session: The SQLAlchemy async session.
+        :param kwargs: Query expressions.
         :return:
         """
-        stmt = select(self.model)
+        filters = await parse_filters(self.model, **kwargs)
+        stmt = select(self.model).where(*filters)
         query = await session.execute(stmt)
         return query.scalars().all()
 
     async def select_models_order(
-        self,
-        session: AsyncSession,
-        *columns,
-        model_sort: Literal['asc', 'desc'] = 'desc',
+        self, session: AsyncSession, sort_columns: str | list[str], sort_orders: str | list[str] | None = None, **kwargs
     ) -> Sequence[Row | RowMapping | Any] | None:
         """
-        Query all rows asc or desc
+        Query all rows and sort by columns
 
-        :param session:
-        :param columns:
-        :param model_sort:
+        :param session: The SQLAlchemy async session.
+        :param sort_columns: more details see apply_sorting
+        :param sort_orders: more details see apply_sorting
         :return:
         """
-        sort_list = []
-        for column in columns:
-            if hasattr(self.model, column):
-                model_column = getattr(self.model, column)
-                sort_list.append(model_column)
-            else:
-                raise ModelColumnError(f'Column {column} is not found in {self.model}')
-        match model_sort:
-            case 'asc':
-                query = await session.execute(select(self.model).order_by(asc(*sort_list)))
-            case 'desc':
-                query = await session.execute(select(self.model).order_by(desc(*sort_list)))
-            case _:
-                raise SelectExpressionError(
-                    f'Select sort expression {model_sort} is not supported, only supports `asc`, `desc`'
-                )
+        filters = await parse_filters(self.model, **kwargs)
+        stmt = select(self.model).where(*filters)
+        stmt_sort = await apply_sorting(self.model, stmt, sort_columns, sort_orders)
+        query = await session.execute(stmt_sort)
         return query.scalars().all()
 
     async def update_model(
-        self, session: AsyncSession, pk: int, obj: _UpdateSchema | dict[str, Any], commit: bool = False, **kwargs
+        self, session: AsyncSession, pk: int, obj: UpdateSchema | dict[str, Any], commit: bool = False
     ) -> int:
         """
-        Update an instance of model's primary key
+        Update an instance by model's primary key
 
-        :param session:
-        :param pk:
-        :param obj:
-        :param commit:
-        :param kwargs:
+        :param session: The SQLAlchemy async session.
+        :param pk: The database primary key value.
+        :param obj: A pydantic schema or dictionary containing the update data
+        :param commit: If `True`, commits the transaction immediately. Default is `False`.
         :return:
         """
         if isinstance(obj, dict):
             instance_data = obj
         else:
             instance_data = obj.model_dump(exclude_unset=True)
-        if kwargs:
-            instance_data.update(kwargs)
         stmt = sa_update(self.model).where(self.model.id == pk).values(**instance_data)
         result = await session.execute(stmt)
         if commit:
@@ -188,55 +134,80 @@ async def update_model(
     async def update_model_by_column(
         self,
         session: AsyncSession,
-        column: str,
-        column_value: Any,
-        obj: _UpdateSchema | dict[str, Any],
+        obj: UpdateSchema | dict[str, Any],
+        allow_multiple: bool = False,
         commit: bool = False,
         **kwargs,
     ) -> int:
         """
-        Update an instance of model column
+        Update an instance by model column
 
-        :param session:
-        :param column:
-        :param column_value:
-        :param obj:
-        :param commit:
-        :param kwargs:
+        :param session: The SQLAlchemy async session.
+        :param obj: A pydantic schema or dictionary containing the update data
+        :param allow_multiple: If `True`, allows updating multiple records that match the filters.
+        :param commit: If `True`, commits the transaction immediately. Default is `False`.
+        :param kwargs: Query expressions.
         :return:
         """
+        filters = await parse_filters(self.model, **kwargs)
+        total_count = await count(session, self.model, 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)
-        if kwargs:
-            instance_data.update(kwargs)
-        if hasattr(self.model, column):
-            model_column = getattr(self.model, column)
-        else:
-            raise ModelColumnError(f'Column {column} is not found in {self.model}')
-        stmt = sa_update(self.model).where(model_column == column_value).values(**instance_data)  # type: ignore
+        stmt = sa_update(self.model).where(*filters).values(**instance_data)  # type: ignore
         result = await session.execute(stmt)
         if commit:
             await session.commit()
         return result.rowcount  # type: ignore
 
-    async def delete_model(self, session: AsyncSession, pk: int, commit: bool = False, **kwargs) -> int:
+    async def delete_model(self, session: AsyncSession, pk: int, commit: bool = False) -> int:
         """
-        Delete an instance of a model
+        Delete an instance by model's primary key
 
-        :param session:
-        :param pk:
-        :param commit:
-        :param kwargs: for soft deletion only
+        :param session: The SQLAlchemy async session.
+        :param pk: The database primary key value.
+        :param commit: If `True`, commits the transaction immediately. Default is `False`.
         :return:
         """
-        if not kwargs:
-            stmt = sa_delete(self.model).where(self.model.id == pk)
-            result = await session.execute(stmt)
-        else:
-            stmt = sa_update(self.model).where(self.model.id == pk).values(**kwargs)
-            result = await session.execute(stmt)
+        stmt = sa_delete(self.model).where(self.model.id == pk)
+        result = await session.execute(stmt)
         if commit:
             await session.commit()
         return result.rowcount  # type: ignore
+
+    async def delete_model_by_column(
+        self,
+        session: AsyncSession,
+        allow_multiple: bool = False,
+        logical_deletion: bool = False,
+        deleted_flag_column: str = 'del_flag',
+        commit: bool = False,
+        **kwargs,
+    ) -> int:
+        """
+        Delete
+
+        :param session: The SQLAlchemy async session.
+        :param commit: If `True`, commits the transaction immediately. Default is `False`.
+        :param kwargs: Query expressions.
+        :param allow_multiple: If `True`, allows deleting multiple records that match the filters.
+        :param logical_deletion: If `True`, enable logical deletion instead of physical deletion
+        :param deleted_flag_column: Specify the flag column for logical deletion
+        :return:
+        """
+        filters = await parse_filters(self.model, **kwargs)
+        total_count = await count(session, self.model, 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 = sa_update(self.model).where(*filters).values(**deleted_flag)
+        else:
+            stmt = sa_delete(self.model).where(*filters)
+        await session.execute(stmt)
+        if commit:
+            await session.commit()
+        return total_count

sqlalchemy_crud_plus/errors.py

@@ -17,8 +17,22 @@ def __init__(self, msg: str) -> None:
         super().__init__(msg)
 
 
-class SelectExpressionError(SQLAlchemyCRUDPlusException):
+class SelectOperatorError(SQLAlchemyCRUDPlusException):
     """Error raised when a select expression is invalid."""
 
     def __init__(self, msg: str) -> None:
         super().__init__(msg)
+
+
+class ColumnSortError(SQLAlchemyCRUDPlusException):
+    """Error raised when a column sorting is invalid."""
+
+    def __init__(self, msg: str) -> None:
+        super().__init__(msg)
+
+
+class MultipleResultsError(SQLAlchemyCRUDPlusException):
+    """Error raised when multiple results are invalid."""
+
+    def __init__(self, msg: str) -> None:
+        super().__init__(msg)

sqlalchemy_crud_plus/types.py

@@ -0,0 +1,10 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+from typing import TypeVar
+
+from pydantic import BaseModel
+
+Model = TypeVar('Model')
+
+CreateSchema = TypeVar('CreateSchema', bound=BaseModel)
+UpdateSchema = TypeVar('UpdateSchema', bound=BaseModel)