Skip to content

Commit

Permalink
Merge pull request #10 from sinasezza/sinasezza
Browse files Browse the repository at this point in the history 
Sinasezza
  • Loading branch information
@sinasezza
sinasezza authored Aug 4, 2024
2 parents 22736ee + 6bc448c commit aa57323
Show file tree
Hide file tree
Showing 7 changed files with 325 additions and 150 deletions.
214 changes: 187 additions & 27 deletions chatApp/config/database.py
View file
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
View file
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

0 comments on commit aa57323

Please sign in to comment.