feat: Add DB async connection (#37)

* feat: Add DB async connection

* style: use built-in list for type hints

Author: gastonva

.env.example

@@ -1,8 +1,9 @@
-# The name of the project will be used to build the final docker image and has to 
+# The name of the project will be used to build the final docker image and has to
 # be exactly the same as project.name as defined in pyproject.toml
 PROJECT_NAME="python-template"
 
 DATABASE_URL=postgresql://dev:dev@postgres:5432/dev
+ASYNC_DATABASE_URL=postgresql+asyncpg://dev:dev@postgres:5432/dev
 LOG_LEVEL=DEBUG
 SERVER_URL=example.com
 ACCESS_TOKEN_EXPIRE_MINUTES=15

poetry.lock

@@ -6,7 +6,7 @@ description = "Xmartlabs' Python project template"
 authors = [{ name = "Xmartlabs", email = "getintouch@xmartlabs.com" }]
 readme = "README.md"
 requires-python = "^3.13"
-dependencies = []
+dependencies = ["asyncpg (>=0.30.0,<0.31.0)"]
 
 [tool.poetry]
 # TODO(remer): this can be removed when the source files are moved to project name folder within src

pyproject.toml

@@ -1,8 +1,13 @@
-from typing import Iterator
+from typing import AsyncIterator, Iterator
 
 from fastapi import Depends, Request
 
-from src.core.database import Session, SessionLocal
+from src.core.database import (
+    AsyncSession,
+    Session,
+    SessionLocal,
+    async_session_generator,
+)
 from src.core.security import AuthManager
 from src.models import User
 
@@ -15,6 +20,18 @@ def db_session() -> Iterator[Session]:
         db.close()
 
 
+async def async_db_session() -> AsyncIterator[AsyncSession]:
+    try:
+        async_session = async_session_generator()
+        async with async_session() as session:
+            yield session
+    except Exception:
+        await session.rollback()
+        raise
+    finally:
+        await session.close()
+
+
 def get_user(request: Request, session: Session = Depends(db_session)) -> User:
     manager = AuthManager()
     return manager(request=request, session=session)

src/api/dependencies.py

@@ -4,10 +4,10 @@
 from fastapi_pagination import Page
 from fastapi_pagination.ext.sqlalchemy import paginate
 
-from src.api.dependencies import db_session, get_user
-from src.api.v1.schemas import Item, ItemCreate
+from src.api.dependencies import async_db_session, db_session, get_user
+from src.api.v1.schemas import BulkItemCreate, Item, ItemCreate
 from src.controllers import ItemController
-from src.core.database import Session
+from src.core.database import AsyncSession, Session
 from src.models import User
 
 router = APIRouter()
@@ -27,3 +27,24 @@ def create_item(
     session: Session = Depends(db_session),
 ) -> Any:
     return ItemController.create(item_data=item_data, owner_id=user.id, session=session)
+
+
+@router.get("/async", response_model=Page[Item])
+async def get_items_async(
+    user: User = Depends(get_user),
+    async_session: AsyncSession = Depends(async_db_session),
+) -> Any:
+    """Get items asynchronously."""
+    return await paginate(async_session, user.get_items())
+
+
+@router.post("/async", response_model=list[Item], status_code=201)
+async def create_item_async(
+    item_data: BulkItemCreate,
+    user: User = Depends(get_user),
+    async_session: AsyncSession = Depends(async_db_session),
+) -> Any:
+    """Create items asynchronously."""
+    return await ItemController.bulk_create(
+        items_data=item_data.items, owner_id=user.id, async_session=async_session
+    )

src/api/v1/routers/item.py

@@ -1,3 +1,3 @@
-from .item import Item, ItemCreate
+from .item import BulkItemCreate, Item, ItemCreate
 from .token import Token, TokenPayload
 from .user import User, UserCreate

src/api/v1/schemas/__init__.py

@@ -12,3 +12,7 @@ class ItemCreate(BaseModel):
 class Item(ItemCreate):
     owner_id: UUID
     model_config = ConfigDict(from_attributes=True)
+
+
+class BulkItemCreate(BaseModel):
+    items: list[ItemCreate]

src/api/v1/schemas/item.py

@@ -1,8 +1,9 @@
+from typing import Sequence
 from uuid import UUID
 
 from src import models
 from src.api.v1 import schemas
-from src.core.database import Session
+from src.core.database import AsyncSession, Session
 
 
 class ItemController:
@@ -13,3 +14,21 @@ def create(
         item_data = schemas.Item(owner_id=owner_id, **item_data.model_dump())
         item = models.Item.objects(session).create(item_data.model_dump())
         return item
+
+    @staticmethod
+    async def bulk_create(
+        items_data: Sequence[schemas.ItemCreate],
+        owner_id: UUID,
+        async_session: AsyncSession,
+    ) -> Sequence[models.Item]:
+        items_data = [
+            schemas.Item(owner_id=owner_id, **item_data.model_dump())
+            for item_data in items_data
+        ]
+        items = await models.Item.async_objects(async_session).bulk_create(
+            [item_data.model_dump() for item_data in items_data]
+        )
+        for item in items:
+            await async_session.refresh(item)
+
+        return items

src/controllers/item.py

@@ -14,6 +14,7 @@ class LogLevel(str, Enum):
 
 class Settings(BaseSettings):
     database_url: PostgresDsn
+    async_database_url: PostgresDsn
     log_level: LogLevel = LogLevel.debug
     server_url: str
 

src/core/config.py

@@ -4,6 +4,12 @@
 
 from fastapi import HTTPException
 from sqlalchemy import create_engine, func, select
+from sqlalchemy.ext.asyncio import (
+    AsyncEngine,
+    AsyncSession,
+    async_sessionmaker,
+    create_async_engine,
+)
 from sqlalchemy.orm import (
     DeclarativeBase,
     Mapped,
@@ -19,9 +25,19 @@
 from src.helpers.casing import snakecase
 from src.helpers.sql import random_uuid, utcnow
 
+# Sync engine and session
 engine = create_engine(str(settings.database_url))
 SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
 
+# Async engine and session
+async_engine: AsyncEngine = create_async_engine(str(settings.async_database_url))
+
+
+def async_session_generator() -> async_sessionmaker[AsyncSession]:
+    return async_sessionmaker(
+        autocommit=False, autoflush=False, bind=async_engine, class_=AsyncSession
+    )
+
 
 class SQLBase(DeclarativeBase):
     __abstract__ = True
@@ -34,6 +50,12 @@ def __tablename__(cls) -> str:
     def objects(cls: Type["_Model"], session: Session) -> "Objects[_Model]":
         return Objects(cls, session)
 
+    @classmethod
+    def async_objects(
+        cls: Type["_Model"], session: AsyncSession
+    ) -> "AsyncObjects[_Model]":
+        return AsyncObjects(cls, session)
+
 
 _Model = TypeVar("_Model", bound=SQLBase)
 
@@ -94,6 +116,72 @@ def create(self, data: Dict[str, Any]) -> _Model:
         return obj
 
 
+class AsyncObjects(Generic[_Model]):
+    cls: Type[_Model]
+    session: AsyncSession
+    base_statement: Select
+    queryset_filters: Any = None
+
+    def __init__(
+        self,
+        cls: Type[_Model],
+        session: AsyncSession,
+        *queryset_filters: Any,
+    ) -> None:
+        self.cls = cls
+        self.session = session
+        base_statement = select(cls)
+        if queryset_filters:
+            self.queryset_filters = queryset_filters
+            base_statement = base_statement.where(*queryset_filters)
+        self.base_statement = base_statement
+
+    async def all(self) -> Sequence[_Model]:
+        result = await self.session.execute(self.base_statement)
+        return result.scalars().unique().all()
+
+    async def get(self, *where_clause: Any) -> _Model | None:
+        statement = self.base_statement.where(*where_clause)
+        result = await self.session.execute(statement)
+        return result.unique().scalar_one_or_none()
+
+    async def get_or_404(self, *where_clause: Any) -> _Model:
+        obj = await self.get(*where_clause)
+        if obj is None:
+            raise HTTPException(
+                status_code=404, detail=f"{self.cls.__name__} not found"
+            )
+        return obj
+
+    async def get_all(self, *where_clause: Any) -> Sequence[_Model]:
+        statement = self.base_statement.where(*where_clause)
+        result = await self.session.execute(statement)
+        return result.scalars().unique().all()
+
+    async def count(self, *where_clause: Any) -> int:
+        statement = select(func.count()).select_from(self.cls)
+        if self.queryset_filters:
+            statement = statement.where(*self.queryset_filters)
+        if where_clause:
+            statement = statement.where(*where_clause)
+
+        result = await self.session.execute(statement)
+        return result.scalar_one()
+
+    async def create(self, data: Dict[str, Any]) -> _Model:
+        obj = self.cls(**data)
+        self.session.add(obj)
+        await self.session.commit()
+        await self.session.refresh(obj)
+        return obj
+
+    async def bulk_create(self, data: Sequence[Dict[str, Any]]) -> Sequence[_Model]:
+        objs = [self.cls(**item) for item in data]
+        self.session.add_all(objs)
+        await self.session.commit()
+        return objs
+
+
 @declarative_mixin
 class TableIdMixin:
     id: Mapped[uuid.UUID] = mapped_column(