Add beanie ODM

README.md

@@ -8,7 +8,7 @@ This backend facilitates an Event Mangement System application, which is essenti
 
 ## Technology
 
-The backend uses the [FastAPI](https://fastapi.tiangolo.com/) framework with [Python](https://www.python.org/) 3.12 or higher. To facilitate MongoDB connection, I will be using the [motor](https://pypi.org/project/motor/) library. The project will also contain a Dockerfile and a docker-compose file, which will make starting it very easy and reproducible.
+The backend uses the [FastAPI](https://fastapi.tiangolo.com/) framework with [Python](https://www.python.org/) 3.12 or higher. To facilitate MongoDB connection, I will be using the [motor](https://pypi.org/project/motor/) library with [Beanie](https://beanie-odm.dev/) object-document mapper (ODM). The project will also contain a Dockerfile and a docker-compose file, which will make starting it very easy and reproducible.
 
 ## Installation
 

pyproject.toml

@@ -10,6 +10,8 @@ dependencies = [
     "python-decouple>=3.8",
     "coloredlogs>=15.0.1",
     "motor[srv]>=3.6.0",
+    "beanie>=1.27.0",
+    "pydantic[email]>=2.10.2",
 ]
 readme = "README.md"
 requires-python = ">= 3.12"

requirements-dev.lock

@@ -15,16 +15,22 @@ annotated-types==0.7.0
 anyio==4.6.2.post1
     # via starlette
 basedpyright==1.22.0
+beanie==1.27.0
+    # via event-management
 cfgv==3.4.0
     # via pre-commit
 click==8.1.7
+    # via beanie
     # via uvicorn
 coloredlogs==15.0.1
     # via event-management
 distlib==0.3.9
     # via virtualenv
 dnspython==2.7.0
+    # via email-validator
     # via pymongo
+email-validator==2.2.0
+    # via pydantic
 fastapi==0.115.5
     # via event-management
 filelock==3.16.1
@@ -37,7 +43,11 @@ identify==2.6.3
     # via pre-commit
 idna==3.10
     # via anyio
+    # via email-validator
+lazy-model==0.2.0
+    # via beanie
 motor==3.6.0
+    # via beanie
     # via event-management
 nodeenv==1.9.1
     # via pre-commit
@@ -51,7 +61,10 @@ poethepoet==0.31.1
     # via event-management
 pre-commit==4.0.1
 pydantic==2.10.2
+    # via beanie
+    # via event-management
     # via fastapi
+    # via lazy-model
 pydantic-core==2.27.1
     # via pydantic
 pymongo==4.9.2
@@ -66,7 +79,10 @@ sniffio==1.3.1
     # via anyio
 starlette==0.41.3
     # via fastapi
+toml==0.10.2
+    # via beanie
 typing-extensions==4.12.2
+    # via beanie
     # via fastapi
     # via pydantic
     # via pydantic-core

requirements.lock

@@ -14,12 +14,18 @@ annotated-types==0.7.0
     # via pydantic
 anyio==4.6.2.post1
     # via starlette
+beanie==1.27.0
+    # via event-management
 click==8.1.7
+    # via beanie
     # via uvicorn
 coloredlogs==15.0.1
     # via event-management
 dnspython==2.7.0
+    # via email-validator
     # via pymongo
+email-validator==2.2.0
+    # via pydantic
 fastapi==0.115.5
     # via event-management
 h11==0.14.0
@@ -28,14 +34,21 @@ humanfriendly==10.0
     # via coloredlogs
 idna==3.10
     # via anyio
+    # via email-validator
+lazy-model==0.2.0
+    # via beanie
 motor==3.6.0
+    # via beanie
     # via event-management
 pastel==0.2.1
     # via poethepoet
 poethepoet==0.31.1
     # via event-management
 pydantic==2.10.2
+    # via beanie
+    # via event-management
     # via fastapi
+    # via lazy-model
 pydantic-core==2.27.1
     # via pydantic
 pymongo==4.9.2
@@ -48,7 +61,10 @@ sniffio==1.3.1
     # via anyio
 starlette==0.41.3
     # via fastapi
+toml==0.10.2
+    # via beanie
 typing-extensions==4.12.2
+    # via beanie
     # via fastapi
     # via pydantic
     # via pydantic-core

src/api.py

@@ -7,9 +7,8 @@ from fastapi import FastAPI
 from fastapi.openapi.docs import get_swagger_ui_html
 from fastapi.requests import Request
 from fastapi.responses import HTMLResponse, PlainTextResponse
-from motor.motor_asyncio import AsyncIOMotorClient
 
-from src.settings import MONGODB_URI
+from src.db.init import initialize_db
 from src.utils.logging import get_logger
 
 log = get_logger(__name__)
@@ -20,15 +19,9 @@ __all__ = ["app"]
 @asynccontextmanager
 async def lifespan(app: FastAPI) -> AsyncIterator[None]:
     """Context manager wrapping the FastAPI app, to handle initialization and exit cleanup."""
-    log.info("Initialization complete.")
+    app.state.db_client, app.state.database = await initialize_db()
 
-    app.state.db_client = AsyncIOMotorClient(MONGODB_URI)
-    app.state.database = app.state.db_client.get_default_database()
-    ping_response = await app.state.database.command("ping")
-    if int(ping_response["ok"]) != 1:
-        raise ConnectionError("Database connection failed")
-    else:
-        log.info("Connected to MongoDB")
+    log.info("Initialization Complete")
 
     try:
         yield

src/db/init.py

@@ -0,0 +1,84 @@
+import importlib
+import pkgutil
+from inspect import getmembers, isclass
+from typing import Any, NoReturn
+
+from beanie import Document, init_beanie
+from motor.motor_asyncio import AsyncIOMotorClient, AsyncIOMotorDatabase
+
+from src.settings import MONGODB_URI
+from src.utils.logging import get_logger
+
+DOCUMENTS_PACKAGE_PATH = "src.db.models"
+
+log = get_logger(__name__)
+
+
+def gather_documents() -> list[type[Document]]:
+    """Collect all MongoDB document model classes defined in `models` submodule."""
+
+    def on_error(name: str) -> NoReturn:
+        """Handle an error encountered while walking packages."""
+        raise ImportError(name=name)
+
+    def ignore_module(module: pkgutil.ModuleInfo) -> bool:
+        """Return whether the module with name `name` should be ignored."""
+        return any(name.startswith("_") for name in module.name.split("."))
+
+    def is_document_class(cls: type) -> bool:
+        """Return whether given class is a document class."""
+        # We're not interested in classes that don't inherit from Document
+        if not issubclass(cls, Document):
+            return False
+
+        # We only want Document subclasses, not the Document base class itself
+        return cls is not Document
+
+    document_classes: set[type[Document]] = set()
+
+    log.debug(f"Loading database modules from {DOCUMENTS_PACKAGE_PATH}")
+    db_module = importlib.import_module(DOCUMENTS_PACKAGE_PATH)
+    for module_info in pkgutil.walk_packages(db_module.__path__, f"{db_module.__name__}.", onerror=on_error):
+        if ignore_module(module_info):
+            continue
+
+        log.debug(f"Loading database module: {module_info.name}")
+        module = importlib.import_module(module_info.name)
+
+        document_classes.update(cls for _, cls in getmembers(module, isclass) if is_document_class(cls))
+
+    log.debug(
+        f"Found {len(document_classes)} document (db) classes: "
+        f"{', '.join(cls.__qualname__ for cls in document_classes)}"
+    )
+
+    # Resolve any forwardref annotations, by this time, everything needs to be loaded fully
+    for document_model in document_classes:
+        if document_model.model_rebuild():
+            log.trace(f"Document model {document_model.__qualname__!r} rebuilt")
+
+    return list(document_classes)
+
+
+async def initialize_db() -> tuple[AsyncIOMotorClient[Any], AsyncIOMotorDatabase[Any]]:
+    """Initialize the MongoDB database connection."""
+    log.info("Connecting to MongoDB...")
+
+    db_client = AsyncIOMotorClient(MONGODB_URI)
+    database = db_client.get_default_database()
+
+    try:
+        ping_response = await database.command("ping")
+    except Exception as exc:
+        log.exception("Database connection failed")
+        raise ConnectionError("Database connection failed") from exc
+
+    if int(ping_response["ok"]) != 1:
+        log.critical(f"Database connection failed, ping response: {ping_response!r}")
+        raise ConnectionError("Database connection failed")
+
+    document_models = gather_documents()
+    await init_beanie(database=database, document_models=document_models)
+    log.info(f"Connected to {database.name!r} MongoDB database")
+
+    return db_client, database

src/db/models/category.py

@@ -0,0 +1,21 @@
+from datetime import UTC, datetime
+from typing import Annotated, ClassVar, final
+
+from beanie import Document, Indexed, Link
+from pydantic import Field
+
+from src.db.models.user import User
+
+
+@final
+class Category(Document):
+    """Category table."""
+
+    user: Annotated[User, Annotated[Link[User], Indexed()]]
+    name: str
+    color: str  # TODO: Consider using a complex rgb type
+    created_at: datetime = Field(default_factory=lambda: datetime.now(UTC))
+
+    @final
+    class Settings:
+        name: ClassVar = "category"

src/db/models/event.py

@@ -0,0 +1,29 @@
+from datetime import UTC, date, datetime, time
+from typing import Annotated, ClassVar, final
+
+from beanie import Document, Indexed, Link
+from pydantic import Field
+
+from src.db.models.category import Category
+from src.db.models.user import User
+
+
+@final
+class Event(Document):
+    """Event table."""
+
+    user: Annotated[User, Annotated[Link[User], Indexed()]]
+    title: str
+    description: str
+    category: Annotated[Category, Annotated[Link[Category], Indexed()]]
+    start_date: Annotated[date, Indexed()]
+    start_time: time
+    end_date: Annotated[date, Indexed()]
+    end_time: time
+    color: str  # TODO: Consider using a complex rgb type
+    attendees: Annotated[list[Link[User]], Indexed()]
+    created_at: datetime = Field(default_factory=lambda: datetime.now(UTC))
+
+    @final
+    class Settings:
+        name: ClassVar = "events"

src/db/models/invitation.py

@@ -0,0 +1,23 @@
+from datetime import UTC, datetime
+from typing import Annotated, ClassVar, Literal, final
+
+from beanie import Document, Indexed, Link
+from pydantic import Field
+
+from src.db.models.event import Event
+from src.db.models.user import User
+
+
+@final
+class Invitation(Document):
+    """Invitation table."""
+
+    event: Annotated[Event, Annotated[Link[Event], Indexed()]]
+    invitee: Annotated[User, Annotated[Link[User], Indexed()]]
+    status: Literal["accepted", "declined", "pending"] = "pending"
+    sent_at: datetime = Field(default_factory=lambda: datetime.now(UTC))
+    responded_at: datetime | None = None
+
+    @final
+    class Settings:
+        name: ClassVar = "invitations"

src/db/models/notificaton.py

@@ -0,0 +1,24 @@
+from datetime import UTC, datetime
+from typing import Annotated, Any, ClassVar, Literal, final
+
+from beanie import Document, Indexed, Link
+from pydantic import Field
+
+from src.db.models.user import User
+
+
+@final
+class Notification(Document):
+    """Notification table."""
+
+    user: Annotated[User, Annotated[Link[User], Indexed()]]
+    event_type: Annotated[Literal["reminder", "invitation"], Indexed()]
+    message: str
+    data: Any
+    read: bool
+    created_at: datetime = Field(default_factory=lambda: datetime.now(UTC))
+    read_at: datetime | None = None
+
+    @final
+    class Settings:
+        name: ClassVar = "notifications"

src/db/models/token.py

@@ -0,0 +1,31 @@
+from datetime import UTC, datetime
+from typing import Annotated, ClassVar, Literal, final
+
+from beanie import Document, Indexed, Link
+from pydantic import Field
+from pymongo import ASCENDING
+
+from src.db.models.user import User
+
+
+@final
+class Token(Document):
+    """Token table."""
+
+    user: Annotated[User, Annotated[Link[User], Indexed()]]
+    tok_type: Literal["access", "refresh"]
+    parent_token: Annotated["Token | None", Annotated[Link["Token"] | None, Indexed()]] = None
+    revoked: bool = False
+    issued_at: datetime = Field(default_factory=lambda: datetime.now(UTC))
+    expires_at: datetime
+
+    @final
+    class Settings:
+        name = "tokens"
+        indexes: ClassVar = [
+            # Compound index for finding all tokens of given type from a specific user
+            [
+                ("user", ASCENDING),
+                ("token_type", ASCENDING),
+            ],
+        ]

src/db/models/user.py

@@ -0,0 +1,19 @@
+from datetime import UTC, datetime
+from typing import Annotated, final
+
+from beanie import Document, Indexed
+from pydantic import EmailStr, Field
+
+
+@final
+class User(Document):
+    """User table."""
+
+    username: Annotated[str, Indexed(unique=True)]
+    password_hash: bytes
+    email: Annotated[EmailStr, Indexed(unique=True)]
+    created_at: datetime = Field(default_factory=lambda: datetime.now(UTC))
+
+    @final
+    class Settings:
+        name = "users"