Repository: igorbenav/fastcrud
Pull request #111 (2024-06-24 17:13:36)

Add returning clause to update
## Description

This change allows CRUD users to retrieve the updated record after performing an `update` operation by passing `return_columns` and optionally the model to load the response into.

## Changes

closes: #106

- add `return_columns` arg to the `update` function
- add `schema_to_select` arg to the `update` function
- add `return_as_model` arg to the `update` function
- add `one_or_none` arg to the `update` function
- add `Optional[Union[dict, BaseModel]]` return type as it now can return data

We use the same args that we also use in other CRUD API functions to make it easier for users to understand and use them.

## Tests

I have added tests for update operations returning a single row, multiple rows, with model and without (i.e. `dict`).

## Checklist
- [x] I have read the [CONTRIBUTING](CONTRIBUTING.md) document.
- [x] My code follows the code style of this project.
- [x] I have added necessary documentation (if appropriate).
- [x] I have added tests that cover my changes (if applicable).
- [x] All new and existing tests passed.



Patch:
diff --git a/docs/advanced/crud.md b/docs/advanced/crud.md
index 8d94eab..c1d2c0b 100644
--- a/docs/advanced/crud.md
+++ b/docs/advanced/crud.md
@@ -178,6 +178,46 @@ await crud_items.delete(
 # this will not actually delete until you run a db.commit()
 ```
 
+## Returning clause in `update`
+
+In `update` method, you can pass `return_columns` parameter containing a list of columns you want to return after the update.
+
+```python
+from fastcrud import FastCRUD
+
+from .models.item import Item
+from .database import session as db
+
+crud_items = FastCRUD(Item)
+item = await crud_items.update(
+    db=db,
+    object={"price": 9.99},
+    price__lt=10
+    return_columns=["price"]
+)
+# this will return the updated price
+```
+
+You can also pass `schema_to_select` parameter and `return_as_model` to return the updated data in the form of a Pydantic schema.
+
+```python
+from fastcrud import FastCRUD
+
+from .models.item import Item
+from .schemas.item import ItemSchema
+from .database import session as db
+
+crud_items = FastCRUD(Item)
+item = await crud_items.update(
+    db=db,
+    object={"price": 9.99},
+    price__lt=10
+    schema_to_select=ItemSchema,
+    return_as_model=True
+)
+# this will return the updated data in the form of ItemSchema
+```
+
 ## Unpaginated `get_multi` and `get_multi_joined`
 
 If you pass `None` to `limit` in `get_multi` and `get_multi_joined`, you get the whole unpaginated set of data that matches the filters. Use this with caution.
diff --git a/fastcrud/crud/fast_crud.py b/fastcrud/crud/fast_crud.py
index 82c3e3d..1e44a2e 100644
--- a/fastcrud/crud/fast_crud.py
+++ b/fastcrud/crud/fast_crud.py
@@ -2,7 +2,7 @@ from typing import Any, Dict, Generic, TypeVar, Union, Optional, Callable
 from datetime import datetime, timezone
 
 from pydantic import BaseModel, ValidationError
-from sqlalchemy import select, update, delete, func, inspect, asc, desc, or_
+from sqlalchemy import Result, select, update, delete, func, inspect, asc, desc, or_, column
 from sqlalchemy.exc import ArgumentError, MultipleResultsFound, NoResultFound
 from sqlalchemy.sql import Join
 from sqlalchemy.ext.asyncio import AsyncSession
@@ -1595,8 +1595,12 @@ class FastCRUD(
         object: Union[UpdateSchemaType, dict[str, Any]],
         allow_multiple: bool = False,
         commit: bool = True,
+        return_columns: Optional[list[str]] = None,
+        schema_to_select: Optional[type[BaseModel]] = None,
+        return_as_model: bool = False,
+        one_or_none: bool = False,
         **kwargs: Any,
-    ) -> None:
+    ) -> Optional[Union[dict, BaseModel]]:
         """
         Updates an existing record or multiple records in the database based on specified filters. This method allows for precise targeting of records to update.
         For filtering details see:
@@ -1607,14 +1611,19 @@ class FastCRUD(
             object: A Pydantic schema or dictionary containing the update data.
             allow_multiple: If True, allows updating multiple records that match the filters. If False, raises an error if more than one record matches the filters.
             commit: If True, commits the transaction immediately. Default is True.
+            return_columns: A list of column names to return after the update. If return_as_model is True, all columns are returned.
+            schema_to_select: Pydantic schema for selecting specific columns from the updated record(s). Required if `return_as_model` is True.
+            return_as_model: If True, returns the updated record(s) as Pydantic model instances based on `schema_to_select`. Default is False.
+            one_or_none: If True, returns a single record if only one record matches the filters. Default is False.
             **kwargs: Filters to identify the record(s) to update, supporting advanced comparison operators for refined querying.
 
         Returns:
-            None
+            The updated record(s) as a dictionary or Pydantic model instance or None, depending on the value of `return_as_model` and `return_columns`.
 
         Raises:
             MultipleResultsFound: If `allow_multiple` is False and more than one record matches the filters.
             ValueError: If extra fields not present in the model are provided in the update data.
+            ValueError: If `return_as_model` is True but `schema_to_select` is not provided.
 
         Examples:
             Update a user's email based on their ID:
@@ -1631,6 +1640,16 @@ class FastCRUD(
             ```python
             update(db, {'username': 'new_username'}, id__ne=1, allow_multiple=False)
             ```
+
+            Update a user's email and return the updated record as a Pydantic model instance:
+            ```python
+            update(db, {'email': 'new_email@example.com'}, id=1, schema_to_select=UserSchema, return_as_model=True)
+            ```
+
+            Update a user's email and return the updated record as a dictionary:
+            ```python
+            update(db, {'email': 'new_email@example.com'}, id=1, return_columns=['id', 'email'])
+            ```
         """
         if not allow_multiple and (total_count := await self.count(db, **kwargs)) > 1:
             raise MultipleResultsFound(
@@ -1655,9 +1674,74 @@ class FastCRUD(
         filters = self._parse_filters(**kwargs)
         stmt = update(self.model).filter(*filters).values(update_data)
 
+        if return_as_model:
+            # All columns are returned to ensure the model can be constructed
+            return_columns = self.model_col_names
+
+        if return_columns:
+            stmt = stmt.returning(*[column(name) for name in return_columns])
+            db_row = await db.execute(stmt)
+            if allow_multiple:
+                return self._as_multi_response(
+                    db_row,
+                    schema_to_select=schema_to_select,
+                    return_as_model=return_as_model,
+                )
+            return self._as_single_response(
+                db_row,
+                schema_to_select=schema_to_select,
+                return_as_model=return_as_model,
+                one_or_none=one_or_none,
+            )
+
         await db.execute(stmt)
         if commit:
             await db.commit()
+        return None
+
+    def _as_single_response(
+        self,
+        db_row: Result,
+        schema_to_select: Optional[type[BaseModel]] = None,
+        return_as_model: bool = False,
+        one_or_none: bool = False,
+    ) -> Optional[Union[dict, BaseModel]]:
+        result: Optional[Row] = db_row.one_or_none() if one_or_none else db_row.first()
+        if result is None:
+            return None
+        out: dict = dict(result._mapping)
+        if not return_as_model:
+            return out
+        if not schema_to_select:
+            raise ValueError(
+                "schema_to_select must be provided when return_as_model is True."
+            )
+        return schema_to_select(**out)
+
+    def _as_multi_response(
+        self,
+        db_row: Result,
+        schema_to_select: Optional[type[BaseModel]] = None,
+        return_as_model: bool = False,
+    ) -> dict:
+        data = [dict(row) for row in db_row.mappings()]
+
+        response: dict[str, Any] = {"data": data}
+
+        if return_as_model:
+            if not schema_to_select:
+                raise ValueError(
+                    "schema_to_select must be provided when return_as_model is True."
+                )
+            try:
+                model_data = [schema_to_select(**row) for row in data]
+                response["data"] = model_data
+            except ValidationError as e:
+                raise ValueError(
+                    f"Data validation error for schema {schema_to_select.__name__}: {e}"
+                )
+
+        return response
 
     async def db_delete(
         self,


Test patch:
diff --git a/tests/sqlalchemy/crud/test_update.py b/tests/sqlalchemy/crud/test_update.py
index 65cf190..c265370 100644
--- a/tests/sqlalchemy/crud/test_update.py
+++ b/tests/sqlalchemy/crud/test_update.py
@@ -215,3 +215,64 @@ async def test_update_auto_updates_updated_at(async_session, test_data):
     assert (
         updated.updated_at > initial_time
     ), "updated_at should be later than the initial timestamp."
+
+
+@pytest.mark.parametrize(
+    ["update_kwargs", "expected_result"],
+    [
+        pytest.param(
+            {"return_columns": ["id", "name"]},
+            {
+                "id": 1,
+                "name": "Updated Name",
+            },
+            id="dict",
+        ),
+        pytest.param(
+            {"schema_to_select": UpdateSchemaTest, "return_as_model": True},
+            UpdateSchemaTest(id=1, name="Updated Name"),
+            id="model",
+        ),
+        pytest.param(
+            {"allow_multiple": True, "return_columns": ["id", "name"]},
+            {
+                "data": [
+                    {
+                        "id": 1,
+                        "name": "Updated Name",
+                    }
+                ]
+            },
+            id="multiple_dict",
+        ),
+        pytest.param(
+            {
+                "allow_multiple": True,
+                "schema_to_select": UpdateSchemaTest,
+                "return_as_model": True,
+            },
+            {"data": [UpdateSchemaTest(id=1, name="Updated Name")]},
+            id="multiple_model",
+        ),
+    ],
+)
+@pytest.mark.asyncio
+async def test_update_with_returning(
+    async_session, test_data, update_kwargs, expected_result
+):
+    for item in test_data:
+        async_session.add(ModelTest(**item))
+    await async_session.commit()
+
+    crud = FastCRUD(ModelTest)
+    target_id = test_data[0]["id"]
+    updated_data = {"name": "Updated Name"}
+
+    updated_record = await crud.update(
+        db=async_session,
+        object=updated_data,
+        id=target_id,
+        **update_kwargs,
+    )
+
+    assert updated_record == expected_result