From 1392fdeb6029e78ed3a73aa570eff0dc5da5a0b8 Mon Sep 17 00:00:00 2001 From: adithya-s-k Date: Wed, 28 Aug 2024 19:55:22 +0000 Subject: [PATCH] updated db connection method - Adithya S K --- backend/app/db.py | 136 +++++++++++++++------------------------------- 1 file changed, 45 insertions(+), 91 deletions(-) diff --git a/backend/app/db.py b/backend/app/db.py index 0df63dd..194e26f 100644 --- a/backend/app/db.py +++ b/backend/app/db.py @@ -1,94 +1,12 @@ -# import re -# import os -# import json -# from motor.motor_asyncio import AsyncIOMotorClient -# from dotenv import load_dotenv -# from app.models.user_model import User -# from app.core.security import get_password - -# load_dotenv() - -# MONGODB_URI = os.getenv("MONGODB_URI") -# MONGODB_NAME = os.getenv("MONGODB_NAME", "RAGSAAS") -# ADMIN_EMAIL = os.getenv("ADMIN_EMAIL") -# ADMIN_PASSWORD = os.getenv("ADMIN_PASSWORD") -# ADMIN_USERNAME = os.getenv("ADMIN_USERNAME", "admin") # Add this line - -# CONFIG_FILE = "rag_config.json" - - -# class MongoDB: -# client: AsyncIOMotorClient = None -# db = None - -# async def connect_to_database(self): -# self.client = AsyncIOMotorClient(MONGODB_URI) -# self.db = self.client[MONGODB_NAME] -# print("Connected to MongoDB") - -# async def close_database_connection(self): -# if self.client: -# self.client.close() -# print("Closed MongoDB connection") - -# async def database_init(self): -# users_collection = self.db.users -# config_collection = self.db.config - -# # Create admin user if not exists -# existing_user = await users_collection.find_one({"email": ADMIN_EMAIL}) -# if not existing_user: -# admin_user = User( -# first_name=ADMIN_USERNAME, -# last_name=ADMIN_USERNAME, -# email=ADMIN_EMAIL, -# hashed_password=get_password(ADMIN_PASSWORD), -# role="admin", -# ) -# await users_collection.insert_one(admin_user.to_mongo()) -# print(f"Admin user created with email: {ADMIN_EMAIL}") - -# # Check if config already exists -# existing_config = await config_collection.find_one({"_id": "app_config"}) -# if existing_config: -# print("Configuration already exists. Skipping initialization.") -# return - -# # Initialize system prompt and conversation starters -# system_prompt = os.getenv("SYSTEM_PROMPT", "") -# conversation_starters_raw = os.getenv("CONVERSATION_STARTERS", "") - -# # Split conversation starters by both newlines and commas -# conversation_starters = re.split(r"[,\n]+", conversation_starters_raw) -# # Clean up any extra whitespace and remove empty strings -# conversation_starters = [ -# starter.strip() for starter in conversation_starters if starter.strip() -# ] - -# initial_config = { -# "SYSTEM_PROMPT": system_prompt, -# "CONVERSATION_STARTERS": conversation_starters, -# } - -# # Insert new config in database -# await config_collection.insert_one({"_id": "app_config", **initial_config}) - -# # Save config to file -# with open(CONFIG_FILE, "w") as f: -# json.dump(initial_config, f, indent=2) - -# print("System prompt and conversation starters initialized") -# print(f"Conversation starters: {conversation_starters}") - - -# mongodb = MongoDB() - +import asyncio import re import os import json +import time from motor.motor_asyncio import AsyncIOMotorClient from pymongo import MongoClient from dotenv import load_dotenv +from pymongo.errors import ServerSelectionTimeoutError from app.models.user_model import User from app.core.security import get_password @@ -103,14 +21,35 @@ CONFIG_FILE = "rag_config.json" +MAX_RETRIES = 5 +RETRY_DELAY = 5 + + class AsyncMongoDB: client: AsyncIOMotorClient = None db = None async def connect_to_database(self): - self.client = AsyncIOMotorClient(MONGODB_URI) - self.db = self.client[MONGODB_NAME] - print("Connected to MongoDB (Async)") + for attempt in range(MAX_RETRIES): + try: + print( + f"Attempting to connect to MongoDB (Async) - Attempt {attempt + 1}" + ) + self.client = AsyncIOMotorClient(MONGODB_URI) + await ( + self.client.server_info() + ) # This will raise an exception if it can't connect + self.db = self.client[MONGODB_NAME] + print("Connected to MongoDB (Async)") + return + except ServerSelectionTimeoutError as e: + print(f"Connection attempt {attempt + 1} failed: {str(e)}") + if attempt < MAX_RETRIES - 1: + print(f"Retrying in {RETRY_DELAY} seconds...") + await asyncio.sleep(RETRY_DELAY) + else: + print("Max retries reached. Unable to connect to MongoDB.") + raise async def close_database_connection(self): if self.client: @@ -168,9 +107,24 @@ class SyncMongoDB: db = None def connect_to_database(self): - self.client = MongoClient(MONGODB_URI) - self.db = self.client[MONGODB_NAME] - print("Connected to MongoDB (Sync)") + for attempt in range(MAX_RETRIES): + try: + print( + f"Attempting to connect to MongoDB (Sync) - Attempt {attempt + 1}" + ) + self.client = MongoClient(MONGODB_URI) + self.client.server_info() # This will raise an exception if it can't connect + self.db = self.client[MONGODB_NAME] + print("Connected to MongoDB (Sync)") + return + except ServerSelectionTimeoutError as e: + print(f"Connection attempt {attempt + 1} failed: {str(e)}") + if attempt < MAX_RETRIES - 1: + print(f"Retrying in {RETRY_DELAY} seconds...") + time.sleep(RETRY_DELAY) + else: + print("Max retries reached. Unable to connect to MongoDB.") + raise def close_database_connection(self): if self.client: