Update backend/database.py

backend/database.py

@@ -1,59 +1,143 @@
 # backend/database.py
 import sqlite3
-import hashlib
+import bcrypt
+import os
+import time
 
-DB_PATH = "code_agents_pro.db"
-PROJECT_ROOT = "./projects"
+DATABASE_PATH = 'data.db'
 
-def init_db():
-    with sqlite3.connect(DB_PATH) as conn:
-        cursor = conn.cursor()
-        cursor.executescript("""
-            CREATE TABLE IF NOT EXISTS users (id INTEGER PRIMARY KEY, username TEXT UNIQUE, password_hash TEXT);
-            CREATE TABLE IF NOT EXISTS projects (id INTEGER PRIMARY KEY, user_id INTEGER, title TEXT, description TEXT, status TEXT DEFAULT 'queued', zip_path TEXT, logs TEXT DEFAULT '', created_at DATETIME DEFAULT CURRENT_TIMESTAMP, FOREIGN KEY (user_id) REFERENCES users(id));
-            CREATE INDEX IF NOT EXISTS idx_user_status ON projects(user_id, status);
-        """)
-
-def _db_execute(query, params=(), fetchone=False, fetchall=False, commit=False):
-    try:
-        with sqlite3.connect(DB_PATH) as conn:
-            conn.row_factory = sqlite3.Row
-            cursor = conn.cursor()
-            cursor.execute(query, params)
-            if commit:
-                conn.commit()
-                return cursor.lastrowid
-            if fetchone:
-                return cursor.fetchone()
-            if fetchall:
-                return cursor.fetchall()
-    except sqlite3.Error as e:
-        print(f"Database error: {e}")
-        return None
-
-def hash_password(password):
-    return hashlib.sha256(password.encode()).hexdigest()
+def get_db_connection():
+    """Establishes a connection to the SQLite database."""
+    conn = sqlite3.connect(DATABASE_PATH)
+    return conn
 
-def verify_password(password, stored_hash):
-    return hash_password(password) == stored_hash
+def init_db():
+    """
+    Initializes the database by creating the necessary tables.
+    This function should be called once on application startup.
+    """
+    conn = get_db_connection()
+    c = conn.cursor()
+    
+    # Users table
+    c.execute("""
+        CREATE TABLE IF NOT EXISTS users (
+            id INTEGER PRIMARY KEY AUTOINCREMENT,
+            username TEXT UNIQUE NOT NULL,
+            password_hash TEXT NOT NULL
+        )
+    """)
+    
+    # Projects table
+    c.execute("""
+        CREATE TABLE IF NOT EXISTS projects (
+            id INTEGER PRIMARY KEY AUTOINCREMENT,
+            user_id INTEGER NOT NULL,
+            title TEXT NOT NULL,
+            prompt TEXT NOT NULL,
+            status TEXT NOT NULL,
+            created_at INTEGER NOT NULL,
+            logs TEXT,
+            zip_path TEXT,
+            FOREIGN KEY (user_id) REFERENCES users (id)
+        )
+    """)
+    
+    conn.commit()
+    conn.close()
 
 def create_user(username, password):
+    """
+    Creates a new user in the database.
+    Returns the user ID or None if the username is already taken.
+    """
+    conn = get_db_connection()
+    c = conn.cursor()
+    
+    hashed_password = bcrypt.hashpw(password.encode('utf-8'), bcrypt.gensalt()).decode('utf-8')
+    
     try:
-        return _db_execute("INSERT INTO users (username, password_hash) VALUES (?, ?)", (username, hash_password(password)), commit=True)
+        c.execute("INSERT INTO users (username, password_hash) VALUES (?, ?)", (username, hashed_password))
+        conn.commit()
+        return c.lastrowid
     except sqlite3.IntegrityError:
         return None
+    finally:
+        conn.close()
 
 def get_user_by_username(username):
-    return _db_execute("SELECT * FROM users WHERE username = ?", (username,), fetchone=True)
+    """Fetches a user by their username."""
+    conn = get_db_connection()
+    c = conn.cursor()
+    c.execute("SELECT * FROM users WHERE username = ?", (username,))
+    user = c.fetchone()
+    conn.close()
+    
+    if user:
+        columns = [desc[0] for desc in c.description]
+        return dict(zip(columns, user))
+    
+    return None
 
-def get_user_projects(user_id, limit=20):
-    return _db_execute("SELECT * FROM projects WHERE user_id = ? ORDER BY created_at DESC LIMIT ?", (user_id, limit), fetchall=True)
+def verify_password(password, password_hash):
+    """Verifies a password against its hash."""
+    return bcrypt.checkpw(password.encode('utf-8'), password_hash.encode('utf-8'))
 
-def create_project(user_id, title, description):
-    return _db_execute("INSERT INTO projects (user_id, title, description) VALUES (?, ?, ?)", (user_id, title, description), commit=True)
+def create_project(user_id, title, prompt):
+    """Creates a new project record and returns its ID."""
+    conn = get_db_connection()
+    c = conn.cursor()
+    c.execute("INSERT INTO projects (user_id, title, prompt, status, created_at) VALUES (?, ?, ?, ?, ?)",
+              (user_id, title, prompt, "queued", int(time.time())))
+    project_id = c.lastrowid
+    conn.commit()
+    conn.close()
+    return project_id
 
-def update_project_status(project_id, status, logs=None, zip_path=None):
-    _db_execute("UPDATE projects SET status=?, logs=COALESCE(?, logs), zip_path=COALESCE(?, zip_path) WHERE id=?", (status, logs, zip_path, project_id), commit=True)
+def get_user_projects(user_id):
+    """Retrieves all projects for a given user, ordered by creation date."""
+    conn = get_db_connection()
+    c = conn.cursor()
+    c.execute("SELECT * FROM projects WHERE user_id = ? ORDER BY created_at DESC", (user_id,))
+    projects = c.fetchall()
+    conn.close()
+    
+    project_list = []
+    if projects:
+        columns = [desc[0] for desc in c.description]
+        for project in projects:
+            project_list.append(dict(zip(columns, project)))
+    
+    return project_list
 
 def get_project(project_id):
-    return _db_execute("SELECT * FROM projects WHERE id = ?", (project_id,), fetchone=True)
+    """
+    Fetches a single project record by its unique ID.
+    This function is used by the front end for live log updates.
+    """
+    conn = get_db_connection()
+    c = conn.cursor()
+    c.execute("SELECT * FROM projects WHERE id = ?", (project_id,))
+    project = c.fetchone()
+    conn.close()
+    
+    if project:
+        columns = [desc[0] for desc in c.description]
+        return dict(zip(columns, project))
+    
+    return None
+
+def update_project_status(project_id, status, logs=None, zip_path=None):
+    """Updates the status and optional logs/zip path for a project."""
+    conn = get_db_connection()
+    c = conn.cursor()
+    
+    if logs is not None and zip_path is not None:
+        c.execute("UPDATE projects SET status = ?, logs = ?, zip_path = ? WHERE id = ?", (status, logs, zip_path, project_id))
+    elif logs is not None:
+        c.execute("UPDATE projects SET status = ?, logs = ? WHERE id = ?", (status, logs, project_id))
+    else:
+        c.execute("UPDATE projects SET status = ? WHERE id = ?", (status, project_id))
+        
+    conn.commit()
+    conn.close()