Update database.py

database.py

@@ -1,161 +1,163 @@
-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
-
-# Configure logging
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
-load_dotenv()
-
-class Database:
-    def __init__(self):
-        """Initialize database connection and ensure indexes"""
-        try:
-            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"Failed to connect to MongoDB: {str(e)}")
-            raise
-    
-    def _create_indexes(self):
-        """Create necessary database indexes"""
-        try:
-            # 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
-            
-        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 DuplicateKeyError:
-            logger.warning(f"Username already exists: {username}")
-            return False
-        except Exception as 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"""
-        try:
-            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": 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}: {str(e)}")
-            return None
-    
-    def get_user(self, username: str) -> Optional[Dict[str, Any]]:
-        """Get user by username (without sensitive data)"""
-        try:
-            user = self.users.find_one(
-                {"username": username},
-                {"password": 0}  # Exclude password from results
-            )
-            return user
-        except Exception as e:
-            logger.error(f"Error fetching user {username}: {str(e)}")
-            return None
-
-# Initialize database connection
-db = Database()
-
-def initialize_users():
-    """
-    Initialize default users if they don't exist.
-    Returns tuple of (success_count, total_users, errors)
-    """
-    from datetime import datetime
-    
-    default_users = [
-        {"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
-    errors = []
-    
-    for user in default_users:
-        try:
-            if db.add_user(user["username"], user["password"], user["role"]):
-                success_count += 1
-                logger.info(f"Initialized user: {user['username']}")
-            else:
-                errors.append(f"Failed to add user: {user['username']}")
-        except Exception as e:
-            error_msg = f"Error initializing user {user['username']}: {str(e)}"
-            logger.error(error_msg)
-            errors.append(error_msg)
-    
-    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, len(default_users), errors
+# database.py
+import sqlite3
+import hashlib
+import logging
+from pathlib import Path
+from typing import Optional, Dict, Any
+
+# Set up logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+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"""
+        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")
+            
+        except Exception as e:
+            logger.error(f"Error initializing database: {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"""
+        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))
+            
+            conn.commit()
+            conn.close()
+            logger.info(f"User '{username}' created successfully with role '{role}'")
+            return True
+            
+        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 Exception as e:
+            logger.error(f"Error creating user '{username}': {e}")
+            return False
+    
+    def verify_user(self, username: str, password: str) -> Optional[Dict[str, Any]]:
+        """Verify user credentials and return user info"""
+        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:
+                return {
+                    "username": result[0],
+                    "role": result[1]
+                }
+            return None
+            
+        except Exception as e:
+            logger.error(f"Error verifying user '{username}': {e}")
+            return None
+    
+    def list_users(self) -> list:
+        """List all users (for debugging)"""
+        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]
+            
+        except Exception as e:
+            logger.error(f"Error listing users: {e}")
+            return []
+
+# Create global database instance
+db = Database()
+
+def initialize_users():
+    """Initialize default users"""
+    default_users = [
+        ("Tony", "password123", "engineering"),
+        ("Bruce", "securepass", "marketing"),
+        ("Sam", "financepass", "finance"),
+        ("Natasha", "hrpass123", "hr"),
+    ]
+    
+    success_count = 0
+    error_count = 0
+    
+    logger.info("Initializing default users...")
+    
+    for username, password, role in default_users:
+        try:
+            if db.create_user(username, password, role):
+                success_count += 1
+            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
+        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")
+    
+    # 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!")
+    
+    return success_count, error_count