atomic operations: single session, rollback on error

Author: mahenzon

fastapi_jsonapi/atomic/atomic.py

@@ -20,12 +20,12 @@
     OperationRelationshipSchema,
 )
 from fastapi_jsonapi.utils.dependency_helper import DependencyHelper
+from fastapi_jsonapi.views.list_view import ListViewBase
 from fastapi_jsonapi.views.utils import HTTPMethodConfig
 from fastapi_jsonapi.views.view_base import ViewBase
 
 if TYPE_CHECKING:
     from fastapi_jsonapi.data_layers.base import BaseDataLayer
-    from fastapi_jsonapi.views.list_view import ListViewBase
 
 
 @dataclass
@@ -105,15 +105,18 @@ async def view_atomic(
 
         results = []
 
+        previous_dl: Optional["BaseDataLayer"] = None
         for operation in prepared_operations:
             dl = operation.data_layer
+            await dl.atomic_start(previous_dl=previous_dl)
+            previous_dl = dl
             if operation.action == "add":
                 data = operation.jsonapi.schema_in_post(data=operation.data)
                 created_object = await dl.create_object(
                     data_create=data.data,
                     view_kwargs={},
                 )
-                # assert isinstance(operation.view, ListViewBase)
+                assert isinstance(operation.view, ListViewBase)
                 view: "ListViewBase" = operation.view
                 response = await view.response_for_created_object(
                     dl=operation.data_layer,
@@ -129,6 +132,9 @@ async def view_atomic(
                 msg = f"unknown action {operation.action!r}"
                 raise ValueError(msg)
 
+        if previous_dl:
+            await previous_dl.atomic_end(success=True)
+
         return {"atomic:results": results}
 
     def _register_view(self):

fastapi_jsonapi/data_layers/base.py

@@ -63,6 +63,13 @@ def __init__(
         self.id_name_field = id_name_field
         self.disable_collection_count: bool = disable_collection_count
         self.default_collection_count: int = default_collection_count
+        self.is_atomic = False
+
+    async def atomic_start(self, previous_dl: Optional["BaseDataLayer"] = None):
+        self.is_atomic = True
+
+    async def atomic_end(self, success: bool = True):
+        raise NotImplementedError
 
     async def create_object(self, data_create: BaseJSONAPIItemInSchema, view_kwargs: dict) -> TypeModel:
         """

fastapi_jsonapi/data_layers/sqla_orm.py

@@ -3,13 +3,14 @@
 from typing import TYPE_CHECKING, Any, Iterable, List, Optional, Tuple, Type
 
 from sqlalchemy import delete, func, select
-from sqlalchemy.exc import DBAPIError, NoResultFound
-from sqlalchemy.ext.asyncio import AsyncSession
+from sqlalchemy.exc import DBAPIError, IntegrityError, NoResultFound
+from sqlalchemy.ext.asyncio import AsyncSession, AsyncSessionTransaction
 from sqlalchemy.inspection import inspect
 from sqlalchemy.orm import joinedload, selectinload
 from sqlalchemy.orm.attributes import InstrumentedAttribute
 from sqlalchemy.orm.collections import InstrumentedList
 
+from fastapi_jsonapi import BadRequest
 from fastapi_jsonapi.data_layers.base import BaseDataLayer
 from fastapi_jsonapi.data_layers.filtering.sqlalchemy import create_filters
 from fastapi_jsonapi.data_layers.sorting.sqlalchemy import create_sorts
@@ -86,6 +87,30 @@ def __init__(
         self.eagerload_includes_ = eagerload_includes
         self._query = query
         self.auto_convert_id_to_column_type = auto_convert_id_to_column_type
+        self.transaction: Optional[AsyncSessionTransaction] = None
+
+    async def atomic_start(self, previous_dl: Optional["SqlalchemyDataLayer"] = None):
+        self.is_atomic = True
+        if previous_dl:
+            self.session = previous_dl.session
+            if previous_dl.transaction:
+                self.transaction = previous_dl.transaction
+                return
+
+        self.transaction = self.session.begin()
+        await self.transaction.start()
+
+    async def atomic_end(self, success: bool = True):
+        if success:
+            await self.transaction.commit()
+        else:
+            await self.transaction.rollback()
+
+    async def save(self):
+        if self.is_atomic:
+            await self.session.flush()
+        else:
+            await self.session.commit()
 
     def prepare_id_value(self, col: InstrumentedAttribute, value: Any) -> Any:
         """
@@ -199,7 +224,11 @@ async def create_object(self, data_create: BaseJSONAPIItemInSchema, view_kwargs:
 
         self.session.add(obj)
         try:
-            await self.session.commit()
+            await self.save()
+        except IntegrityError:
+            log.exception("Could not create object with data create %s", data_create)
+            msg = "Object creation error"
+            raise BadRequest(msg, pointer="/data")
         except DBAPIError:
             log.exception("Could not create object with data create %s", data_create)
             msg = "Object creation error"
@@ -331,7 +360,7 @@ async def update_object(
                 setattr(obj, field_name, new_value)
                 has_updated = True
         try:
-            await self.session.commit()
+            await self.save()
         except DBAPIError as e:
             await self.session.rollback()
 
@@ -357,7 +386,7 @@ async def delete_object(self, obj: TypeModel, view_kwargs: dict):
         await self.before_delete_object(obj, view_kwargs)
         try:
             await self.session.delete(obj)
-            await self.session.commit()
+            await self.save()
         except DBAPIError as e:
             await self.session.rollback()
 
@@ -377,7 +406,7 @@ async def delete_objects(self, objects: List[TypeModel], view_kwargs: dict):
 
         try:
             await self.session.execute(query)
-            await self.session.commit()
+            await self.save()
         except DBAPIError as e:
             await self.session.rollback()
             raise InternalServerError(

tests/conftest.py

@@ -54,6 +54,9 @@
 
 def configure_logging():
     logging.getLogger("faker.factory").setLevel(logging.INFO)
+    logging.getLogger("aiosqlite").setLevel(logging.INFO)
+    # logging.getLogger("sqlalchemy.engine").setLevel(logging.INFO)
+    logging.basicConfig(level=logging.DEBUG)
 
 
 configure_logging()

tests/fixtures/app.py

@@ -132,6 +132,7 @@ def add_routers(app_plain: FastAPI):
         schema_in_post=ChildInSchema,
         model=Child,
     )
+
     RoutersJSONAPI(
         router=router,
         path="/computers",

tests/test_atomic/test_create_objects.py

@@ -5,6 +5,7 @@
 from sqlalchemy import and_, or_, select
 from sqlalchemy.engine import Result
 from sqlalchemy.ext.asyncio import AsyncSession
+from sqlalchemy.sql.functions import count
 from starlette import status
 
 from tests.misc.utils import fake
@@ -135,3 +136,65 @@ async def test_create_two_objects(
                     "type": "user",
                 },
             }
+
+    async def test_atomic_rollback_on_create_error(
+        self,
+        client: AsyncClient,
+        async_session: AsyncSession,
+    ):
+        """
+        User name is unique
+
+        - create first user
+        - create second user with the same name
+        - catch exc
+        - rollback all changes
+
+        :param client:
+        :param async_session:
+        :return:
+        """
+        user_data_1 = UserAttributesBaseSchema(
+            name=fake.name(),
+            age=fake.pyint(min_value=13, max_value=99),
+            email=fake.email(),
+        )
+        user_data_2 = UserAttributesBaseSchema(
+            name=user_data_1.name,
+            age=fake.pyint(min_value=13, max_value=99),
+            email=fake.email(),
+        )
+        users_data = [user_data_1, user_data_2]
+        data_atomic_request = {
+            "atomic:operations": [
+                {
+                    "op": "add",
+                    "data": {
+                        "type": "user",
+                        "attributes": user_data.dict(),
+                    },
+                }
+                for user_data in users_data
+            ],
+        }
+        response = await client.post("/operations", json=data_atomic_request)
+        assert response.status_code == status.HTTP_400_BAD_REQUEST, response.text
+        response_data = response.json()
+        assert "errors" in response_data, response_data
+        errors = response_data["errors"]
+        assert errors, response_data
+        error = errors[0]
+        assert error == {
+            "detail": "Object creation error",
+            "source": {"pointer": "/data"},
+            "status_code": status.HTTP_400_BAD_REQUEST,
+            "title": "Bad Request",
+        }
+        stmt = select(count(User.id)).where(
+            or_(
+                User.name == user_data_1.name,
+                User.name == user_data_2.name,
+            ),
+        )
+        result: Result = await async_session.execute(stmt)
+        assert result.scalar_one() == 0