Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Sinasezza #10

Merged
merged 2 commits into from
Aug 4, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
214 changes: 187 additions & 27 deletions chatApp/config/database.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,8 @@
AsyncIOMotorCollection,
AsyncIOMotorDatabase,
)
from pymongo import ASCENDING, DESCENDING, IndexModel
from pymongo.errors import CollectionInvalid

from .config import get_settings

Expand All @@ -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:
Expand All @@ -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")
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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
16 changes: 3 additions & 13 deletions chatApp/main.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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,
Expand Down
Loading
Loading