add base db interface

Author: NatalyaGrigoreva

fastapi_jsonapi/data_layers/sqla/base_model.py

@@ -0,0 +1,237 @@
+import logging
+from typing import Any, Literal, Optional, Type, Union
+
+from sqlalchemy import and_, delete, func, select
+from sqlalchemy.exc import IntegrityError
+from sqlalchemy.ext.asyncio import AsyncSession
+from sqlalchemy.sql import Select, column, distinct
+from sqlalchemy.sql.elements import UnaryExpression
+from sqlalchemy.sql.expression import BinaryExpression
+
+from fastapi_jsonapi.data_layers.sqla.query_building import RelationshipInfo
+from fastapi_jsonapi.data_typing import TypeModel
+from fastapi_jsonapi.exceptions import BadRequest, InternalServerError, ObjectNotFound
+
+log = logging.getLogger(__name__)
+
+
+class BaseSQLA:
+    @classmethod
+    def _check_field_exists(
+        cls,
+        model: TypeModel,
+        key: str,
+    ) -> None:
+        try:
+            getattr(model, key)
+        except AttributeError as ex:
+            err_message = f"No fields `{key}` on `{type(model).__name__}`. Make sure schema conforms model."
+            log.exception(err_message, exc_info=ex)
+            raise InternalServerError(
+                detail=err_message,
+                pointer="/data",
+            )
+
+    @classmethod
+    def _fill(
+        cls,
+        model: TypeModel,
+        relationships: list[tuple[str, TypeModel]],
+        **kwargs,
+    ) -> None:
+        for key, value in kwargs.items():
+            cls._check_field_exists(model, key)
+            setattr(model, key, value)
+        for relation_name, related_data in relationships:
+            cls._check_field_exists(model, relation_name)
+            setattr(model, relation_name, related_data)
+
+    @classmethod
+    async def _save(
+        cls,
+        session: AsyncSession,
+        model: TypeModel,
+        action_trigger: Literal["update", "create", "delete"],
+        resource_type: str,
+        commit: bool = True,
+        id_: Optional[str] = None,
+        **kwargs,
+    ) -> TypeModel:
+        try:
+            if not commit:
+                await session.flush()
+                return model
+
+            await session.commit()
+            return model
+        except IntegrityError as ex:
+            err_message = f"Could not {action_trigger} object"
+            log.exception("%s with data %s", err_message, kwargs, exc_info=ex)
+            raise BadRequest(
+                detail=err_message,
+                pointer="/data",
+                meta={
+                    "type": resource_type,
+                    "id": id_,
+                },
+            )
+        except Exception as ex:
+            err_message = f"Got an error {ex.__class__.__name__} during updating obj {kwargs} data in DB"
+            log.exception(err_message, exc_info=ex)
+            await session.rollback()
+            raise InternalServerError(
+                detail=err_message,
+                pointer="/data",
+                meta={
+                    "type": resource_type,
+                    "id": id_,
+                },
+            )
+
+    @classmethod
+    async def all(
+        cls,
+        session: AsyncSession,
+        stmt: Select,
+    ) -> Union[Type[TypeModel], Any]:
+        return (await session.execute(stmt)).unique().scalars().all()
+
+    @classmethod
+    async def count(
+        cls,
+        session: AsyncSession,
+        stmt: Select,
+    ) -> int:
+        stmt = select(func.count(distinct(column("id")))).select_from(stmt.subquery())
+        return (await session.execute(stmt)).scalar_one()
+
+    @classmethod
+    async def create(
+        cls,
+        session: AsyncSession,
+        model: TypeModel,
+        resource_type: str,
+        relationships: list[tuple[str, TypeModel]],
+        commit: bool = True,
+        id_: Optional[str] = None,
+        **kwargs,
+    ) -> TypeModel:
+        cls._fill(model, relationships, **kwargs)
+        session.add(model)
+        return await cls._save(
+            session=session,
+            model=model,
+            action_trigger="create",
+            resource_type=resource_type,
+            commit=commit,
+            id_=id_,
+            **kwargs,
+        )
+
+    @classmethod
+    async def delete(
+        cls,
+        session: AsyncSession,
+        model: TypeModel,
+        filters: list[Union[BinaryExpression, bool]],
+        resource_type: str,
+        commit: bool = True,
+        id_: Optional[str] = None,
+        **kwargs,
+    ) -> None:
+        await session.execute(delete(model).where(*filters))
+        await cls._save(
+            session=session,
+            model=model,
+            action_trigger="delete",
+            resource_type=resource_type,
+            commit=commit,
+            id_=id_,
+            **kwargs,
+        )
+
+    @classmethod
+    async def one_or_raise(
+        cls,
+        session: AsyncSession,
+        model: TypeModel,
+        filters: list[Union[BinaryExpression, bool]],
+        stmt: Select,
+    ) -> Union[TypeModel, Any]:
+        result = (await session.execute(stmt)).scalar_one_or_none()
+        if result is None:
+            compiled_conditions = and_(*filters).compile(
+                dialect=session.bind.dialect,
+                compile_kwargs={"literal_binds": True},
+            )
+            raise ObjectNotFound(
+                detail=f"Resource {model.__name__} `{compiled_conditions}` not found",
+            )
+        return result
+
+    @classmethod
+    async def query(
+        cls,
+        model: TypeModel,
+        distinct_: bool = False,
+        fields: Optional[list] = None,
+        filters: Optional[list[Union[BinaryExpression, bool]]] = None,
+        for_update: Optional[dict] = None,
+        join: Optional[list[RelationshipInfo]] = None,
+        number: Optional[int] = None,
+        options: set = (),
+        order: Optional[Union[str, UnaryExpression]] = None,
+        size: Optional[int] = None,
+        stmt: Optional[Select] = None,
+    ) -> Select:
+        if stmt is None:
+            stmt = select(model) if fields is None else select(*fields)
+
+        if filters is not None:
+            stmt = stmt.where(*filters)
+
+        if options:
+            stmt = stmt.options(*options)
+
+        if for_update is not None:
+            stmt = stmt.with_for_update(**for_update)
+
+        if order is not None:
+            stmt = stmt.order_by(*order)
+
+        if join:
+            for relationship_info in join:
+                stmt = stmt.join(relationship_info.aliased_model, relationship_info.join_column)
+
+        if size not in [0, None]:
+            stmt = stmt.limit(size)
+            if number:
+                stmt = stmt.offset((number - 1) * size)
+
+        if distinct_:
+            stmt = stmt.distinct()
+
+        return stmt
+
+    @classmethod
+    async def update(
+        cls,
+        session: AsyncSession,
+        model: TypeModel,
+        resource_type: str,
+        relationships: list[tuple[str, TypeModel]],
+        commit: bool = True,
+        id_: Optional[str] = None,
+        **kwargs,
+    ) -> TypeModel:
+        cls._fill(model, relationships, **kwargs)
+        session.add(model)
+        return await cls._save(
+            session=session,
+            model=model,
+            action_trigger="update",
+            resource_type=resource_type,
+            commit=commit,
+            id_=id_,
+            **kwargs,
+        )