Added database ORM support (Async SQLAlchemy)

nb-programmer · nb-programmer · commit 0cbcb3032ffe · 2024-11-03T12:53:23.000+05:30
* Updated .env.example file with sample config
* Added database setup, session and dependency functions
* Added utils file for common utility functions
* Added User management resource to demonstrate database queries
diff --git a/.env.example b/.env.example
@@ -1,2 +1,5 @@
 SECRET_KEY = "SecretSecretSecret!"
 UVICORN_PORT = 8000
+
+DATABASE_URL = "sqlite+aiosqlite:///dev/app.db"
+ECHO_SQL = false
diff --git a/sample_fastapi/app/app_factory.py b/sample_fastapi/app/app_factory.py
@@ -5,7 +5,7 @@
 from fastapi import FastAPI
 from tqdm.contrib.logging import logging_redirect_tqdm
 
-from . import middleware, resources
+from . import database, middleware, resources
 
 LOG = logging.getLogger(__name__)
 
@@ -18,16 +18,22 @@ async def app_lifespan_handler(app: FastAPI):
     """
 
     async with AsyncExitStack() as stack:
-        LOG.info('Adding lifespan handler for "resources"...')
+        LOG.info('Executing lifespan handler for "resources"...')
         await stack.enter_async_context(resources.lifespan(app))
-        LOG.info('Adding lifespan handler for "tqdm_logging"...')
+        LOG.info('Executing lifespan handler for "database"...')
+        await stack.enter_async_context(database.lifespan(app))
+        LOG.info('Executing lifespan handler for "tqdm_logging"...')
         stack.enter_context(logging_redirect_tqdm())
         yield
 
 
 def init_app():
+    db_settings = database.get_db_settings()
+
     app = FastAPI(lifespan=app_lifespan_handler)
+
     resources.init_app(app)
     middleware.init_middleware(app)
+    database.init_app_db(app, db_settings)
 
     return app
diff --git a/sample_fastapi/app/database/__init__.py b/sample_fastapi/app/database/__init__.py
@@ -0,0 +1,16 @@
+from .base import BaseDBModel
+from .config import DBSettings
+from .depends import get_app_db_session_manager, get_db_session, get_db_settings
+from .engine import init_app_db, lifespan
+from .session import AsyncSession
+
+__all__ = [
+    "DBSettings",
+    "get_db_settings",
+    "get_app_db_session_manager",
+    "get_db_session",
+    "lifespan",
+    "init_app_db",
+    "BaseDBModel",
+    "AsyncSession",
+]
diff --git a/sample_fastapi/app/database/base.py b/sample_fastapi/app/database/base.py
@@ -0,0 +1,5 @@
+from sqlalchemy.orm import DeclarativeBase
+
+
+class BaseDBModel(DeclarativeBase):
+    pass
diff --git a/sample_fastapi/app/database/config.py b/sample_fastapi/app/database/config.py
@@ -0,0 +1,19 @@
+import logging
+
+from pydantic import Field
+
+from ..config import BaseAppSettings
+
+LOG = logging.getLogger(__name__)
+
+APP_DB_SESSION_MANAGER_KEY = "_dbsm"  # Lol
+
+
+def _default_db_dsn():
+    LOG.warning("Using SQLite In-memory database as `database_url` config was not specified. Content WILL be lost after shutting down the app!")
+    return "sqlite+aiosqlite:///:memory:"
+
+
+class DBSettings(BaseAppSettings):
+    database_url: str = Field(default_factory=_default_db_dsn)
+    echo_sql: bool = Field(default=True)
diff --git a/sample_fastapi/app/database/depends.py b/sample_fastapi/app/database/depends.py
@@ -0,0 +1,29 @@
+from functools import lru_cache
+from typing import Annotated
+
+from fastapi import Depends, Request
+
+from .config import APP_DB_SESSION_MANAGER_KEY, DBSettings
+from .session import AsyncDatabaseSessionManager
+
+
+@lru_cache()
+def get_db_settings():
+    return DBSettings()
+
+
+def get_app_db_session_manager(req: Request) -> AsyncDatabaseSessionManager:
+    sessionmanager: AsyncDatabaseSessionManager | None = req.app.extra.get(APP_DB_SESSION_MANAGER_KEY)
+    if sessionmanager is None:
+        raise RuntimeError(
+            "Couldn't get the app's DB session manager - DB Session Manager with key `%s` was not initialized"
+            % APP_DB_SESSION_MANAGER_KEY
+        )
+    return sessionmanager
+
+
+async def get_db_session(sessionmanager: Annotated[AsyncDatabaseSessionManager, Depends(get_app_db_session_manager)]):
+    """Dependency to get a session object. Use this in request handlers to get a database session."""
+
+    async with sessionmanager.session() as session:
+        yield session
diff --git a/sample_fastapi/app/database/engine.py b/sample_fastapi/app/database/engine.py
@@ -0,0 +1,42 @@
+import logging
+from contextlib import asynccontextmanager
+from typing import Annotated
+
+from fastapi import FastAPI
+
+from .. import utils
+from .config import APP_DB_SESSION_MANAGER_KEY, DBSettings
+from .session import AsyncDatabaseSessionManager
+
+LOG = logging.getLogger(__name__)
+
+
+async def create_schema_models(sessionmanager: AsyncDatabaseSessionManager):
+    """Create schema tables from models under `BaseDBModel`."""
+    async with sessionmanager._engine.connect() as connection:
+        from .base import BaseDBModel
+
+        LOG.info("Creating schema (if not already exists)...")
+        LOG.info("DB models registered: [%s]", ",".join(map(utils.str_quote, BaseDBModel.metadata.tables.keys())))
+
+        await connection.run_sync(BaseDBModel.metadata.create_all)
+        await connection.commit()
+
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    """Session manager, manage database connections for this app"""
+    sessionmanager: AsyncDatabaseSessionManager = app.extra[APP_DB_SESSION_MANAGER_KEY]
+    async with sessionmanager:
+        await create_schema_models(sessionmanager)
+        yield
+
+
+def init_app_db(app: FastAPI, db_settings: Annotated[DBSettings, DBSettings()]):
+    """Initialize database session manager and bind it with the FastAPI application"""
+    app.extra[APP_DB_SESSION_MANAGER_KEY] = AsyncDatabaseSessionManager(
+        db_settings.database_url,
+        {
+            "echo": db_settings.echo_sql,
+        },
+    )
diff --git a/sample_fastapi/app/database/session.py b/sample_fastapi/app/database/session.py
@@ -0,0 +1,62 @@
+import contextlib
+import logging
+from typing import Any, AsyncIterator
+
+from sqlalchemy.ext.asyncio import (
+    AsyncConnection,
+    AsyncEngine,
+    AsyncSession,
+    async_sessionmaker,
+    create_async_engine,
+)
+
+__all__ = [
+    "AsyncConnection",
+    "AsyncSession",
+    "AsyncDatabaseSessionManager",
+]
+
+
+LOG = logging.getLogger(__name__)
+
+
+# Heavily inspired by https://praciano.com.br/fastapi-and-async-sqlalchemy-20-with-pytest-done-right.html
+
+
+class AsyncDatabaseSessionManager:
+    def __init__(self, host: str, engine_kwargs: dict[str, Any] = {}, session_kwargs: dict[str, Any] = {}):
+        session_kwargs.setdefault("autocommit", False)
+        session_kwargs.setdefault("expire_on_commit", False)
+
+        self._engine: AsyncEngine = create_async_engine(host, **engine_kwargs)
+        self._sessionmaker: async_sessionmaker[AsyncSession] = async_sessionmaker(bind=self._engine, **session_kwargs)
+
+    async def close(self):
+        await self._engine.dispose()
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        await self.close()
+
+    @contextlib.asynccontextmanager
+    async def connect(self) -> AsyncIterator[AsyncConnection]:
+        async with self._engine.begin() as connection:
+            try:
+                yield connection
+            except Exception:
+                LOG.exception("DB connection was rolled back due to an exception:")
+                await connection.rollback()
+                raise
+
+    @contextlib.asynccontextmanager
+    async def session(self) -> AsyncIterator[AsyncSession]:
+        async with self._sessionmaker() as sess:
+            yield sess
+            # try:
+            # except Exception:
+            #     LOG.exception("DB transaction was rolled back due to an exception:")
+            #     # TODO: Not sure if this is needed. Test if it does auto-rollback
+            #     await sess.rollback()
+            #     raise
diff --git a/sample_fastapi/app/resources/routes.py b/sample_fastapi/app/resources/routes.py
@@ -3,6 +3,7 @@
 
 def init_app(app: FastAPI):
     """Initialize routes, middleware, sub-apps, etc. to the given Application"""
-    from . import calculator, hello
+    from . import calculator, hello, user
     hello.init_app(app)
     calculator.init_app(app)
+    user.init_app(app)
diff --git a/sample_fastapi/app/resources/user/__init__.py b/sample_fastapi/app/resources/user/__init__.py
@@ -0,0 +1,3 @@
+from .routes import init_app, init_routes
+
+__all__ = ["init_app", "init_routes"]
diff --git a/sample_fastapi/app/resources/user/db_models.py b/sample_fastapi/app/resources/user/db_models.py
@@ -0,0 +1,31 @@
+import datetime
+import uuid
+
+from sqlalchemy.sql.schema import Column, ForeignKey
+from sqlalchemy.sql.sqltypes import DateTime, Integer, String, Uuid
+
+from ...database.base import BaseDBModel
+
+
+class User(BaseDBModel):
+    __tablename__ = "users"
+
+    id: Column[uuid.UUID] = Column(Uuid, primary_key=True, default=uuid.uuid4)
+    username = Column(String, unique=True)
+
+    first_name = Column(String, nullable=False)
+    middle_name = Column(String, nullable=True)
+    last_name = Column(String, nullable=True)
+
+    created_ts = Column(DateTime, default=datetime.datetime.now)
+    updated_ts = Column(DateTime, nullable=True)
+
+
+# TODO:
+# class UserPassword(BaseDBModel):
+#     __tablename__ = "user_passwords"
+#     id = Column(Integer, primary_key=True, autoincrement=True)
+#     user_id = Column(Uuid, ForeignKey(User.id))
+#     password_hash = Column(String, nullable=False)
+#     created_ts = Column(DateTime, default=datetime.datetime.now)
+#     expiry_ts = Column(DateTime, nullable=True)
diff --git a/sample_fastapi/app/resources/user/db_queries.py b/sample_fastapi/app/resources/user/db_queries.py
@@ -0,0 +1,8 @@
+from sqlalchemy import bindparam, insert, select
+
+from .db_models import User
+
+QUERY_ALL_USERS = select(User)
+
+QUERY_USER_BY_ID = QUERY_ALL_USERS.where(User.id == bindparam("id"))
+QUERY_USER_BY_USERNAME = QUERY_ALL_USERS.where(User.username == bindparam("username"))
diff --git a/sample_fastapi/app/resources/user/models.py b/sample_fastapi/app/resources/user/models.py
@@ -0,0 +1,24 @@
+from uuid import UUID
+
+from pydantic import BaseModel, ConfigDict, Field
+
+
+class UserCreatedResponse(BaseModel):
+    model_config = ConfigDict(from_attributes=True)
+
+    user_id: UUID = Field(validation_alias="id")
+    username: str
+
+
+class UserInfoResponse(BaseModel):
+    model_config = ConfigDict(from_attributes=True)
+
+    user_id: UUID = Field(validation_alias="id")
+    username: str
+    first_name: str
+    middle_name: str | None
+    last_name: str | None
+
+
+class UserCreateError(ValueError):
+    pass
diff --git a/sample_fastapi/app/resources/user/routes.py b/sample_fastapi/app/resources/user/routes.py
@@ -0,0 +1,20 @@
+from fastapi import APIRouter, FastAPI
+
+from .views import user_create, user_data, user_list
+
+
+def init_routes():
+    router = APIRouter(prefix="/user", tags=["user"])
+
+    router.add_api_route("/register", user_create, methods={"POST"})
+    router.add_api_route("/list", user_list, methods={"GET"})
+    router.add_api_route("/{user_id}", user_data, methods={"GET"})
+
+    return router
+
+
+def init_app(app: FastAPI):
+    app.include_router(init_routes())
+
+    # Bind models
+    from . import db_models  # noqa
diff --git a/sample_fastapi/app/resources/user/services.py b/sample_fastapi/app/resources/user/services.py
@@ -0,0 +1,67 @@
+from typing import Annotated
+from uuid import UUID
+
+import sqlalchemy
+from fastapi import Depends
+
+from ...database import AsyncSession, get_db_session
+from . import db_models, db_queries
+from .models import UserCreatedResponse, UserCreateError, UserInfoResponse
+
+
+class UserService:
+    """User management service"""
+
+    def __init__(self, db: Annotated[AsyncSession, Depends(get_db_session)]):
+        self._db = db
+
+    async def create_user(
+        self,
+        username: str,
+        first_name: str,
+        middle_name: str | None = None,
+        last_name: str | None = None,
+    ) -> UserCreatedResponse:
+        """Create a new user (if `username` doesn't exist) with the
+        given details and save it in the database.
+
+        :param str username: Username of the new user
+        :param str first_name: First name
+        :param str or None middle_name: Middle name
+        :param str or None last_name: Last name
+        :return UserCreatedResponse: Contains the ID of the newly created user
+
+        :raises UserCreateError: If an error occurres while creating the user (Eg. already existing username)"""
+        try:
+            new_user = db_models.User(
+                username=username,
+                first_name=first_name,
+                middle_name=middle_name,
+                last_name=last_name,
+            )
+
+            self._db.add(new_user)
+            await self._db.commit()
+
+            return UserCreatedResponse.model_validate(new_user)
+        except sqlalchemy.exc.IntegrityError as ex:
+            raise UserCreateError("Username `%s` already exists." % username) from ex
+
+    async def user_query_all(self) -> list[UserInfoResponse]:
+        """Returns a list of users that have registered.
+
+        :return list[UserInfoResponse]: A list of `UserInfoResponse` items containing various information about each user
+        """
+        results = await self._db.scalars(db_queries.QUERY_ALL_USERS)
+        user_list = map(UserInfoResponse.model_validate, results.all())
+        return list(user_list)
+
+    async def user_query_id(self, user_id: UUID) -> UserInfoResponse | None:
+        result = await self._db.scalar(db_queries.QUERY_USER_BY_ID, {"id": user_id})
+        if result:
+            return UserInfoResponse.model_validate(result)
+
+    async def user_query_username(self, username: str) -> UserInfoResponse | None:
+        result = await self._db.scalar(db_queries.QUERY_USER_BY_USERNAME, {"username": username})
+        if result:
+            return UserInfoResponse.model_validate(result)
diff --git a/sample_fastapi/app/resources/user/views.py b/sample_fastapi/app/resources/user/views.py
diff --git a/sample_fastapi/app/utils.py b/sample_fastapi/app/utils.py

-Original file line number
+Diff line change
@@ @@ -0,0 +1,5 @@ @@
 +from sqlalchemy.orm import DeclarativeBase
++
++
 +class BaseDBModel(DeclarativeBase):
 +    pass
Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,3 @@`
	`1`	`+from .routes import init_app, init_routes`
	`2`	`+`
	`3`	`+__all__ = ["init_app", "init_routes"]`