Async Types

sanctumlabs_dbkit/sql/callbacks.py

@@ -4,8 +4,8 @@
 
 from typing import List, cast
 
-from sanctumlabs_dbkit.sql.session import Session
-from sanctumlabs_dbkit.sql.types import CommitCallback
+from sanctumlabs_dbkit.sql.session import Session, AsyncSession
+from sanctumlabs_dbkit.sql.types import CommitCallback, CommitCallbackAsync
 
 
 def on_commit(current_session: Session, callback: CommitCallback) -> None:
@@ -14,3 +14,14 @@ def on_commit(current_session: Session, callback: CommitCallback) -> None:
         List[CommitCallback], current_session.info.setdefault("on_commit_hooks", [])
     )
     commit_hooks.append(callback)
+
+
+def on_commit_async(
+    current_session: AsyncSession, callback: CommitCallbackAsync
+) -> None:
+    """Sets an async commit callback to the current session"""
+    commit_hooks = cast(
+        List[CommitCallbackAsync],
+        current_session.info.setdefault("on_commit_hooks", []),
+    )
+    commit_hooks.append(callback)

sanctumlabs_dbkit/sql/repository/__init__.py

@@ -0,0 +1,7 @@
+from sanctumlabs_dbkit.sql.repository.repository import Repository
+from sanctumlabs_dbkit.sql.repository.async_repository import AsyncRepository
+
+__all__ = [
+    "Repository",
+    "AsyncRepository",
+]

sanctumlabs_dbkit/sql/repository/async_repository.py

@@ -0,0 +1,130 @@
+"""
+Contains a generic base repository or DAO for access patterns to data for a given database model
+"""
+
+from datetime import datetime, UTC
+from typing import (
+    Generic,
+    Any,
+    Optional,
+    Sequence,
+    Type,
+    cast,
+    TypeGuard,
+)
+
+from sqlalchemy import ColumnElement, Select, select
+
+from sanctumlabs_dbkit.exceptions import ModelNotFoundError
+from sanctumlabs_dbkit.exceptions import UnsupportedModelOperationError
+from sanctumlabs_dbkit.sql.session.async_session import AsyncSession
+from sanctumlabs_dbkit.sql.models import AbstractBaseModel
+from sanctumlabs_dbkit.sql.repository.types import T
+
+
+class AsyncRepository(Generic[T]):
+    """
+    A base class for implementing an async Repository or DAO.
+
+    ```python
+    job_dao = AsyncRepository(model=Job, session=async_session)
+
+    job = job_dao.find("123")
+    """
+
+    def __init__(self, model: Type[T], session: AsyncSession) -> None:
+        """Creates an instance of the Repository"""
+        self.model = model
+        self.session = session
+
+    @staticmethod
+    def _supports_soft_deletion(model: Type[T]) -> TypeGuard[Type[AbstractBaseModel]]:
+        """
+        Indicates if the provided model supports soft deletion (has a 'deleted_at' column). This function
+        takes in an argument due to mypy typeguarding requirements, and is thus static.
+        """
+        return issubclass(model, AbstractBaseModel)
+
+    def create(self, refresh: bool = False, **kwargs: Any) -> T:
+        """Creates a new entity
+
+        Args:
+            refresh (bool, optional): whether to refresh the model with the data in the return. Defaults to False.
+
+        Returns:
+            T: The created model instance
+        """
+        model_instance = self.model(**kwargs)
+        self.session.add(model_instance)
+
+        if refresh:
+            self.session.flush()
+            self.session.refresh(model_instance)
+
+        return cast(T, model_instance)
+
+    def query(self, include_deleted: bool = False) -> Select:
+        """Returns a select query with the model including deleted records if the include_deleted is set to True"""
+        selectable = select(self.model)
+
+        if not include_deleted and self._supports_soft_deletion(self.model):
+            selectable = selectable.where(
+                self.model.deleted_at == self.model.not_deleted_value()
+            )
+
+        return selectable
+
+    async def find(self, pk: Any, include_deleted: bool = False) -> Optional[T]:
+        """Retrieve a given model given its primary key"""
+        pk_column = cast(ColumnElement, getattr(self.model, self.model.pk))
+
+        statement = self.query(include_deleted).where(pk_column == pk).limit(1)
+        scalars = await self.session.scalars(statement)
+
+        return scalars.first()
+
+    async def find_or_raise(self, pk: Any, include_deleted: bool = False) -> T:
+        """Finds the given entity or raises an exception if the entity can not be found"""
+        entity = await self.find(pk, include_deleted)
+
+        if not entity:
+            raise ModelNotFoundError(
+                f"The model {self.model.__name__} {pk} does not exist"
+            )
+
+        return entity
+
+    async def all(self, include_deleted: bool = False) -> Sequence[T]:
+        """Retrieves all records for the given model"""
+        statement = self.query(include_deleted)
+        scalars = await self.session.scalars(statement)
+
+        return scalars.all()
+
+    async def delete(self, pk: Any) -> None:
+        """Deletes a given record with the given primary key"""
+        if not self._supports_soft_deletion(self.model):
+            raise UnsupportedModelOperationError(
+                f"The model {self.model.__name__} {pk} does not support soft deletion."
+            )
+
+        # Cast here as mypy type narrowing doesn't infer the type of entity
+        # correctly
+        entity = cast(AbstractBaseModel, await self.find(pk))
+
+        if entity:
+            entity.deleted_at = datetime.now(UTC)
+
+    async def list(
+        self, limit: int = 20, offset: int = 0, include_deleted: bool = False
+    ) -> Sequence[T]:
+        """Returns a list of records for the given database record"""
+        statement = (
+            self.query(include_deleted)
+            .order_by(self.model.created_at.desc())
+            .limit(limit)
+            .offset(offset)
+        )
+        scalars = await self.session.scalars(statement)
+
+        return scalars.all()

sanctumlabs_dbkit/sql/repository.py → sanctumlabs_dbkit/sql/repository/repository.py

@@ -9,21 +9,17 @@
     Optional,
     Sequence,
     Type,
-    TypeVar,
     cast,
     TypeGuard,
-    Union,
 )
+
 from sqlalchemy import ColumnElement, Select, select
 
 from sanctumlabs_dbkit.exceptions import ModelNotFoundError
-from sanctumlabs_dbkit.sql.models import AbstractBaseModel, BaseOutboxEvent
-from sanctumlabs_dbkit.sql.session import Session
 from sanctumlabs_dbkit.exceptions import UnsupportedModelOperationError
-
-RepositoryBaseModel = Union[AbstractBaseModel, BaseOutboxEvent]
-
-T = TypeVar("T", bound=RepositoryBaseModel)
+from sanctumlabs_dbkit.sql.models import AbstractBaseModel
+from sanctumlabs_dbkit.sql.repository.types import T
+from sanctumlabs_dbkit.sql.session import Session
 
 
 class Repository(Generic[T]):

sanctumlabs_dbkit/sql/repository/types.py

@@ -0,0 +1,10 @@
+from typing import (
+    TypeVar,
+    Union,
+)
+
+from sanctumlabs_dbkit.sql.models import AbstractBaseModel, BaseOutboxEvent
+
+RepositoryBaseModel = Union[AbstractBaseModel, BaseOutboxEvent]
+
+T = TypeVar("T", bound=RepositoryBaseModel)

sanctumlabs_dbkit/sql/session/__init__.py

@@ -0,0 +1,15 @@
+from sanctumlabs_dbkit.sql.session.async_session import (
+    AsyncSession,
+    AsyncSessionLocal,
+    async_transaction,
+)
+from sanctumlabs_dbkit.sql.session.session import Session, SessionLocal, transaction
+
+__all__ = [
+    "Session",
+    "AsyncSession",
+    "SessionLocal",
+    "transaction",
+    "AsyncSessionLocal",
+    "async_transaction",
+]

sanctumlabs_dbkit/sql/session/async_session.py

@@ -0,0 +1,103 @@
+"""
+Async Session module contains implementation logic for a database session
+"""
+
+import functools
+from typing import Any
+
+from sqlalchemy.ext.asyncio import (
+    AsyncSession as BaseAsyncSession,
+    AsyncSessionTransaction,
+    async_sessionmaker,
+)
+
+from sanctumlabs_dbkit.sql.session.types import FuncT
+
+
+class AsyncSession(BaseAsyncSession):
+    """
+    Session that subclasses SQLAlchemy Base Session class adding more functionality around a database session
+    """
+
+    def begin(self, nested: bool = False) -> AsyncSessionTransaction:
+        """Begins an async session transaction"""
+        if nested:
+            return super().begin_nested()
+        return super().begin()
+
+    def transaction(self, func: FuncT) -> FuncT:
+        """
+        A decorator to wrap a function within a transaction.
+
+        If we are already within a transaction, a nested transaction will be started.
+
+        Example:
+
+        ```python
+        from sanctumlabs_dbkit.sql.session import AsyncSessionLocal
+
+        session = SessionLocal()
+
+        @session.transaction
+        def create_user(payload) -> User:
+            user = User(**payload)
+            session.add(user)
+
+            return user
+
+        create_user({"first_name": "Bob"})
+        """
+
+        @functools.wraps(func)
+        async def wrapper(*args: Any, **kwargs: Any) -> Any:
+            async with self.begin():
+                return func(*args, **kwargs)
+
+        return wrapper
+
+
+async def async_transaction(func: FuncT) -> FuncT:
+    """
+    A decorator to wrap an instance method within a transaction.
+
+    If we are already within a transaction, a nested transaction will be started.
+
+    Example:
+
+    ```python
+    from sanctumlabs_dbkit.sql.session import AsyncSessionLocal
+    from sanctumlabs_dbkit.sql.sesison import async_transaction
+
+    class UserService():
+        def __init__(session: AsyncSession):
+            self.session = session
+
+        @async_transaction
+        def create(payload) -> User:
+            user = User(**payload)
+            self.session.add(user)
+
+            return user
+
+    session = AsyncSessionLocal()
+
+    user_service = UserService(session)
+    user_service.create({"first_name": "Bob"})
+    """
+
+    @functools.wraps(func)
+    async def wrapper(self: Any, *args: Any, **kwargs: Any) -> Any:
+        if not self.session or not isinstance(self.session, AsyncSession):
+            # pylint: disable=broad-exception-raised
+            raise Exception(
+                "The @transaction decorator requires that an instance variable `session` be set to an instance of a "
+                "`Session`."
+            )
+
+        async with self.session.begin():
+            return func(self, *args, **kwargs)
+
+    return wrapper
+
+
+AsyncSessionLocal = async_sessionmaker(class_=AsyncSession)

sanctumlabs_dbkit/sql/session.py → sanctumlabs_dbkit/sql/session/session.py

@@ -3,11 +3,11 @@
 """
 
 import functools
-from typing import Any, Callable, TypeVar, cast
+from typing import Any, cast
 
 from sqlalchemy.orm import SessionTransaction, Session as BaseSession, sessionmaker
 
-FuncT = TypeVar("FuncT", bound=Callable[..., Any])
+from sanctumlabs_dbkit.sql.session.types import FuncT
 
 
 class Session(BaseSession):

sanctumlabs_dbkit/sql/session/types.py

@@ -0,0 +1,3 @@
+from typing import Any, Callable, TypeVar
+
+FuncT = TypeVar("FuncT", bound=Callable[..., Any])

sanctumlabs_dbkit/sql/types.py

@@ -27,8 +27,10 @@
 from wrapt import ObjectProxy
 
 from sanctumlabs_dbkit.sql.session import Session
+from sanctumlabs_dbkit.sql.session.async_session import AsyncSession
 
 CommitCallback = Callable[[Session], None]
+CommitCallbackAsync = Callable[[AsyncSession], None]
 
 _T = TypeVar("_T", bound=BaseModel)