Update database.py

database.py

@@ -1,163 +1,161 @@
-# database.py
-import sqlite3
-import hashlib
+import os
+from datetime import datetime
+from typing import Optional, Dict, Any, List
+from pymongo import MongoClient, ReturnDocument
+from pymongo.errors import DuplicateKeyError, ConnectionFailure
+import bcrypt
+from dotenv import load_dotenv
 import logging
-from pathlib import Path
-from typing import Optional, Dict, Any
 
-# Set up logging
+# Configure logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
+load_dotenv()
+
 class Database:
-    def __init__(self, db_path: str = "users.db"):
-        self.db_path = db_path
-        self.init_database()
-    
-    def init_database(self):
-        """Initialize the database and create tables"""
+    def __init__(self):
+        """Initialize database connection and ensure indexes"""
         try:
-            # Ensure directory exists
-            Path(self.db_path).parent.mkdir(parents=True, exist_ok=True)
-            
-            conn = sqlite3.connect(self.db_path)
-            cursor = conn.cursor()
-            
-            # Create users table with proper constraints
-            cursor.execute('''
-                CREATE TABLE IF NOT EXISTS users (
-                    id INTEGER PRIMARY KEY AUTOINCREMENT,
-                    username TEXT UNIQUE NOT NULL,
-                    password_hash TEXT NOT NULL,
-                    role TEXT NOT NULL,
-                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
-                )
-            ''')
-            
-            conn.commit()
-            conn.close()
-            logger.info("Database initialized successfully")
+            mongo_uri = os.getenv("MONGO_URI")
+            if not mongo_uri:
+                raise ValueError("MONGO_URI environment variable is not set")
+                
+            self.client = MongoClient(
+                mongo_uri,
+                serverSelectionTimeoutMS=5000,  # 5 second timeout
+                connectTimeoutMS=30000,        # 30 second connection timeout
+                socketTimeoutMS=45000,          # 45 second socket timeout
+                connect=False                   # Lazy connection
+            )
+            
+            # Test the connection
+            self.client.admin.command('ping')
+            
+            self.db = self.client[os.getenv("DB_NAME", "rag_system")]
+            self.users = self.db["users"]
+            self._create_indexes()
+            logger.info("Successfully connected to MongoDB")
             
         except Exception as e:
-            logger.error(f"Error initializing database: {e}")
+            logger.error(f"Failed to connect to MongoDB: {str(e)}")
             raise
     
-    def hash_password(self, password: str) -> str:
-        """Hash a password using SHA-256"""
-        return hashlib.sha256(password.encode()).hexdigest()
-    
-    def create_user(self, username: str, password: str, role: str) -> bool:
-        """Create a new user"""
+    def _create_indexes(self):
+        """Create necessary database indexes"""
         try:
-            conn = sqlite3.connect(self.db_path)
-            cursor = conn.cursor()
-            
-            password_hash = self.hash_password(password)
-            
-            cursor.execute('''
-                INSERT INTO users (username, password_hash, role)
-                VALUES (?, ?, ?)
-            ''', (username, password_hash, role))
+            # Create unique index on username
+            self.users.create_index("username", unique=True)
+            logger.info("Created database indexes")
+        except Exception as e:
+            logger.error(f"Error creating indexes: {str(e)}")
+            raise
+    
+    def add_user(self, username: str, password: str, role: str) -> bool:
+        """Add a new user to the database"""
+        if not username or not password or not role:
+            logger.warning("Missing required fields for user creation")
+            return False
             
-            conn.commit()
-            conn.close()
-            logger.info(f"User '{username}' created successfully with role '{role}'")
-            return True
+        try:
+            hashed = bcrypt.hashpw(password.encode('utf-8'), bcrypt.gensalt())
+            user_data = {
+                "username": username,
+                "password": hashed.decode('utf-8'),
+                "role": role.lower(),
+                "created_at": datetime.utcnow(),
+                "last_login": None
+            }
+            
+            result = self.users.insert_one(user_data)
+            if result.inserted_id:
+                logger.info(f"Created new user: {username}")
+                return True
+            return False
             
-        except sqlite3.IntegrityError as e:
-            if "UNIQUE constraint failed" in str(e):
-                logger.info(f"User '{username}' already exists, skipping...")
-                return False  # User already exists, not really an error
-            else:
-                logger.error(f"Integrity error creating user '{username}': {e}")
-                return False
+        except DuplicateKeyError:
+            logger.warning(f"Username already exists: {username}")
+            return False
         except Exception as e:
-            logger.error(f"Error creating user '{username}': {e}")
+            logger.error(f"Error adding user {username}: {str(e)}")
             return False
     
     def verify_user(self, username: str, password: str) -> Optional[Dict[str, Any]]:
-        """Verify user credentials and return user info"""
+        """Verify user credentials"""
         try:
-            conn = sqlite3.connect(self.db_path)
-            cursor = conn.cursor()
-            
-            password_hash = self.hash_password(password)
-            
-            cursor.execute('''
-                SELECT username, role FROM users
-                WHERE username = ? AND password_hash = ?
-            ''', (username, password_hash))
-            
-            result = cursor.fetchone()
-            conn.close()
-            
-            if result:
+            user = self.users.find_one({"username": username})
+            if not user:
+                logger.warning(f"Login attempt for non-existent user: {username}")
+                return None
+                
+            if bcrypt.checkpw(password.encode('utf-8'), user["password"].encode('utf-8')):
+                # Update last login time
+                self.users.update_one(
+                    {"_id": user["_id"]},
+                    {"$set": {"last_login": datetime.utcnow()}}
+                )
+                logger.info(f"Successful login for user: {username}")
                 return {
-                    "username": result[0],
-                    "role": result[1]
+                    "username": user["username"],
+                    "role": user["role"],
+                    "last_login": user.get("last_login")
                 }
+            
+            logger.warning(f"Failed login attempt for user: {username}")
             return None
             
         except Exception as e:
-            logger.error(f"Error verifying user '{username}': {e}")
+            logger.error(f"Error verifying user {username}: {str(e)}")
             return None
     
-    def list_users(self) -> list:
-        """List all users (for debugging)"""
+    def get_user(self, username: str) -> Optional[Dict[str, Any]]:
+        """Get user by username (without sensitive data)"""
         try:
-            conn = sqlite3.connect(self.db_path)
-            cursor = conn.cursor()
-            
-            cursor.execute('SELECT username, role FROM users')
-            users = cursor.fetchall()
-            conn.close()
-            
-            return [{"username": user[0], "role": user[1]} for user in users]
-            
+            user = self.users.find_one(
+                {"username": username},
+                {"password": 0}  # Exclude password from results
+            )
+            return user
         except Exception as e:
-            logger.error(f"Error listing users: {e}")
-            return []
+            logger.error(f"Error fetching user {username}: {str(e)}")
+            return None
 
-# Create global database instance
+# Initialize database connection
 db = Database()
 
 def initialize_users():
-    """Initialize default users"""
+    """
+    Initialize default users if they don't exist.
+    Returns tuple of (success_count, total_users, errors)
+    """
+    from datetime import datetime
+    
     default_users = [
-        ("Tony", "password123", "engineering"),
-        ("Bruce", "securepass", "marketing"),
-        ("Sam", "financepass", "finance"),
-        ("Natasha", "hrpass123", "hr"),
+        {"username": "Tony", "password": "password123", "role": "engineering"},
+        {"username": "Bruce", "password": "securepass", "role": "marketing"},
+        {"username": "Sam", "password": "financepass", "role": "finance"},
+        {"username": "Peter", "password": "pete123", "role": "engineering"},
+        {"username": "Sid", "password": "sidpass123", "role": "marketing"},
+        {"username": "Natasha", "password": "hrpass123", "role": "hr"}
     ]
     
     success_count = 0
-    error_count = 0
+    errors = []
     
-    logger.info("Initializing default users...")
-    
-    for username, password, role in default_users:
+    for user in default_users:
         try:
-            if db.create_user(username, password, role):
+            if db.add_user(user["username"], user["password"], user["role"]):
                 success_count += 1
+                logger.info(f"Initialized user: {user['username']}")
             else:
-                # User already exists, count as success but log as info
-                success_count += 1
-                # The error_count was being incremented for existing users
-                # which is why you were seeing "6 errors" - let's not count this as an error
+                errors.append(f"Failed to add user: {user['username']}")
         except Exception as e:
-            error_count += 1
-            logger.error(f"Failed to create user {username}: {e}")
-    
-    if error_count == 0:
-        logger.info(f"User initialization successful: {success_count} users ready")
-    else:
-        logger.info(f"User initialization completed: {success_count} successful, {error_count} errors")
+            error_msg = f"Error initializing user {user['username']}: {str(e)}"
+            logger.error(error_msg)
+            errors.append(error_msg)
     
-    # List all users for verification
-    users = db.list_users()
-    if users:
-        logger.info(f"Available users: {', '.join([f\"{u['username']}({u['role']})\" for u in users])}")
-    else:
-        logger.warning("No users found in database!")
+    logger.info(f"User initialization complete. Success: {success_count}/{len(default_users)}")
+    if errors:
+        logger.warning(f"Encountered {len(errors)} errors during user initialization")
     
-    return success_count, error_count
+    return success_count, len(default_users), errors