diff --git a/chatApp/config/database.py b/chatApp/config/database.py index 26181f6..2729bd9 100644 --- a/chatApp/config/database.py +++ b/chatApp/config/database.py @@ -6,6 +6,8 @@ AsyncIOMotorCollection, AsyncIOMotorDatabase, ) +from pymongo import ASCENDING, DESCENDING, IndexModel +from pymongo.errors import CollectionInvalid from .config import get_settings @@ -19,7 +21,8 @@ def __init__(self) -> None: self.db: AsyncIOMotorDatabase | None = None self.users_collection: AsyncIOMotorCollection | None = None self.messages_collection: AsyncIOMotorCollection | None = None - self.rooms_collection: AsyncIOMotorCollection | None = None + self.public_rooms_collection: AsyncIOMotorCollection | None = None + self.private_rooms_collection: AsyncIOMotorCollection | None = None async def connect_to_mongodb(self) -> None: try: @@ -33,15 +36,8 @@ async def connect_to_mongodb(self) -> None: self.db = self.db_client[settings.database_name] assert self.db is not None - # Initialize collections - self.users_collection = self.db.get_collection("users") - self.messages_collection = self.db.get_collection("messages") - self.public_rooms_collection = self.db.get_collection( - "public_rooms" - ) - self.private_rooms_collection = self.db.get_collection( - "private_rooms" - ) + # Define collections and schema validations + await self.create_collections() # Ping the server to validate the connection await self.db_client.admin.command("ismaster") @@ -50,14 +46,174 @@ async def connect_to_mongodb(self) -> None: logger.error(f"Could not connect to MongoDB: {e}") raise + async def create_collections(self) -> None: + # Define schema validation for each collection + user_schema = { + "$jsonSchema": { + "bsonType": "object", + "required": ["username", "email", "hashed_password"], + "properties": { + "username": {"bsonType": "string"}, + "email": {"bsonType": "string"}, + "hashed_password": {"bsonType": "string"}, + "is_active": {"bsonType": "bool"}, + "is_admin": {"bsonType": "bool"}, + "created_at": {"bsonType": "date"}, + "updated_at": {"bsonType": "date"}, + "last_login": {"bsonType": "date"}, + }, + } + } + + message_schema = { + "$jsonSchema": { + "bsonType": "object", + "required": ["user_id", "room_id", "room_type"], + "properties": { + "user_id": {"bsonType": "objectId"}, + "room_id": {"bsonType": "objectId"}, + "room_type": {"bsonType": "string"}, + "content": {"bsonType": "string"}, + "media": {"bsonType": "string"}, + "created_at": {"bsonType": "date"}, + }, + } + } + + public_room_schema = { + "$jsonSchema": { + "bsonType": "object", + "required": ["owner", "name"], + "properties": { + "owner": {"bsonType": "objectId"}, + "name": {"bsonType": "string"}, + "description": {"bsonType": "string"}, + "max_members": {"bsonType": "int"}, + "welcome_message": {"bsonType": "string"}, + "rules": {"bsonType": "string"}, + "allow_file_sharing": {"bsonType": "bool"}, + "members": { + "bsonType": "array", + "items": {"bsonType": "objectId"}, + }, + "ban_list": { + "bsonType": "array", + "items": {"bsonType": "objectId"}, + }, + "moderators": { + "bsonType": "array", + "items": {"bsonType": "objectId"}, + }, + "allow_users_access_message_history": {"bsonType": "bool"}, + "max_latest_messages_access": {"bsonType": "int"}, + "created_at": {"bsonType": "date"}, + }, + } + } + + private_room_schema = { + "$jsonSchema": { + "bsonType": "object", + "required": ["member1", "member2"], + "properties": { + "member1": {"bsonType": "objectId"}, + "member2": {"bsonType": "objectId"}, + "created_at": {"bsonType": "date"}, + }, + } + } + + if self.db is not None: + await self.create_or_update_collection("users", user_schema) + await self.create_or_update_collection("messages", message_schema) + await self.create_or_update_collection( + "public_rooms", public_room_schema + ) + await self.create_or_update_collection( + "private_rooms", private_room_schema + ) + + await self.create_indexes() + + async def create_or_update_collection( + self, name: str, validator: dict + ) -> None: + if self.db is not None: + try: + await self.db.create_collection(name, validator=validator) + except CollectionInvalid: + logger.info(f"Collection '{name}' already exists.") + except Exception as e: + logger.error(f"Could not create collection '{name}': {e}") + raise + else: + logger.error("MongoDB client is not initialized.") + + async def create_indexes(self) -> None: + if self.db is not None: + if self.users_collection is None: + self.users_collection = self.db["users"] + if self.messages_collection is None: + self.messages_collection = self.db["messages"] + if self.public_rooms_collection is None: + self.public_rooms_collection = self.db["public_rooms"] + if self.private_rooms_collection is None: + self.private_rooms_collection = self.db["private_rooms"] + + await self.users_collection.create_indexes( + [ + IndexModel([("username", ASCENDING)], unique=True), + IndexModel([("email", ASCENDING)], unique=True), + ] + ) + + await self.messages_collection.create_indexes( + [ + IndexModel( + [("room_id", ASCENDING), ("created_at", DESCENDING)] + ), + IndexModel( + [("room_type", ASCENDING), ("created_at", DESCENDING)] + ), + ] + ) + + await self.public_rooms_collection.create_indexes( + [IndexModel([("name", ASCENDING)], unique=True)] + ) + + await self.private_rooms_collection.create_indexes( + [ + IndexModel( + [("member1", ASCENDING), ("member2", ASCENDING)], + unique=True, + ) + ] + ) + async def close_mongodb_connection(self) -> None: if self.db_client: self.db_client.close() logger.info("Closed MongoDB connection") -# Create a global instance of MongoDB -mongo_db = MongoDB() +mongo_db = None + + +async def init_mongo_db(): + global mongo_db + mongo_db = MongoDB() + await mongo_db.connect_to_mongodb() + return mongo_db + + +async def shutdown_mongo_db(): + """ + Close the MongoDB connection. + """ + global mongo_db + if mongo_db is not None: + await mongo_db.close_mongodb_connection() @lru_cache @@ -68,10 +224,11 @@ def get_users_collection() -> AsyncIOMotorCollection: :return: The users collection instance. :raises RuntimeError: If the users collection is not initialized. """ - users_collection = mongo_db.users_collection - if users_collection is None: + if mongo_db is None: + raise RuntimeError("MongoDB instance is not initialized.") + if mongo_db.users_collection is None: raise RuntimeError("Users collection is not initialized.") - return users_collection + return mongo_db.users_collection @lru_cache @@ -82,10 +239,11 @@ def get_messages_collection() -> AsyncIOMotorCollection: :return: The messages collection instance. :raises RuntimeError: If the messages collection is not initialized. """ - messages_collection = mongo_db.messages_collection - if messages_collection is None: - raise RuntimeError("messages collection is not initialized.") - return messages_collection + if mongo_db is None: + raise RuntimeError("MongoDB instance is not initialized.") + if mongo_db.messages_collection is None: + raise RuntimeError("Messages collection is not initialized.") + return mongo_db.messages_collection @lru_cache @@ -96,10 +254,11 @@ def get_public_rooms_collection() -> AsyncIOMotorCollection: :return: The rooms collection instance. :raises RuntimeError: If the rooms collection is not initialized. """ - rooms_collection = mongo_db.public_rooms_collection - if rooms_collection is None: - raise RuntimeError("public rooms collection is not initialized.") - return rooms_collection + if mongo_db is None: + raise RuntimeError("MongoDB instance is not initialized.") + if mongo_db.public_rooms_collection is None: + raise RuntimeError("Public rooms collection is not initialized.") + return mongo_db.public_rooms_collection @lru_cache @@ -110,7 +269,8 @@ def get_private_rooms_collection() -> AsyncIOMotorCollection: :return: The rooms collection instance. :raises RuntimeError: If the rooms collection is not initialized. """ - rooms_collection = mongo_db.private_rooms_collection - if rooms_collection is None: - raise RuntimeError("private rooms collection is not initialized.") - return rooms_collection + if mongo_db is None: + raise RuntimeError("MongoDB instance is not initialized.") + if mongo_db.private_rooms_collection is None: + raise RuntimeError("Private rooms collection is not initialized.") + return mongo_db.private_rooms_collection diff --git a/chatApp/main.py b/chatApp/main.py index 99bef5d..94dbd62 100644 --- a/chatApp/main.py +++ b/chatApp/main.py @@ -4,7 +4,7 @@ from fastapi.middleware.trustedhost import TrustedHostMiddleware from chatApp.config.config import get_settings -from chatApp.config.database import mongo_db +from chatApp.config.database import init_mongo_db, shutdown_mongo_db from chatApp.middlewares.request_limit import RequestLimitMiddleware from chatApp.routes import auth, chat, user from chatApp.sockets import sio_app @@ -13,26 +13,16 @@ settings = get_settings() -# Define startup and shutdown event handlers -async def startup_event(): - await mongo_db.connect_to_mongodb() - - -async def shutdown_event(): - await mongo_db.close_mongodb_connection() - - # Create a FastAPI app instance app = FastAPI( title="FastAPI Chat App", description="A chat application built with FastAPI and socket.io", version="1.0.0", - on_startup=[startup_event], - on_shutdown=[shutdown_event], + on_startup=[init_mongo_db], + on_shutdown=[shutdown_mongo_db], ) ### Add middlewares ### - # Configure CORS using settings app.add_middleware( CORSMiddleware, diff --git a/chatApp/models/user.py b/chatApp/models/user.py index 29f85aa..0f6c826 100644 --- a/chatApp/models/user.py +++ b/chatApp/models/user.py @@ -14,7 +14,7 @@ class User(BaseModel): is_admin: bool = False created_at: datetime = Field(default_factory=lambda: datetime.now()) updated_at: datetime = Field(default_factory=lambda: datetime.now()) - last_login: datetime | None = None + last_login: datetime = Field(default_factory=lambda: datetime.now()) class UserInDB(User):