scripts/update_database.py

"""
update_database.py

This script updates the database schema by adding new columns to the sessions table
and populating them with appropriate values. It handles SQLite's limitations
with ALTER TABLE operations by checking if columns exist before attempting to add them.

The following columns are added:
- last_accessed (DateTime): Set to created_at for existing records
- is_important (Boolean): Set to False for existing records
- message_count (Integer): Calculated from the number of messages in chat_messages table

Usage:
    python update_database.py
"""

import sqlite3
import os
from datetime import datetime
from sqlalchemy import create_engine, inspect, text
from database import DATABASE_URL, SessionLocal, Base

def check_column_exists(engine, table_name, column_name):
    """Check if a column exists in a table."""
    inspector = inspect(engine)
    columns = inspector.get_columns(table_name)
    return any(col['name'] == column_name for col in columns)

def add_column_sqlite(db_path, table_name, column_name, column_type, default_value=None):
    """
    Add a column to a SQLite table by creating a new table, copying data, and renaming.
    This is necessary because SQLite has limited ALTER TABLE support.
    """
    conn = sqlite3.connect(db_path)
    cursor = conn.cursor()
    
    # Get current table info
    cursor.execute(f"PRAGMA table_info({table_name})")
    columns = cursor.fetchall()
    column_names = [col[1] for col in columns]
    
    # Create new table with additional column
    new_table_name = f"{table_name}_new"
    
    # Build new column list
    new_columns = []
    for col in columns:
        new_columns.append(f"{col[1]} {col[2]}")
    
    # Add the new column
    new_column_def = f"{column_name} {column_type}"
    if default_value is not None:
        new_column_def += f" DEFAULT {default_value}"
    new_columns.append(new_column_def)
    
    # Create new table
    columns_sql = ", ".join(new_columns)
    create_sql = f"CREATE TABLE {new_table_name} ({columns_sql})"
    cursor.execute(create_sql)
    
    # Copy data from old table to new table
    column_names_str = ", ".join(column_names)
    insert_sql = f"INSERT INTO {new_table_name} ({column_names_str}) SELECT {column_names_str} FROM {table_name}"
    cursor.execute(insert_sql)
    
    # Drop old table and rename new table
    cursor.execute(f"DROP TABLE {table_name}")
    cursor.execute(f"ALTER TABLE {new_table_name} RENAME TO {table_name}")
    
    conn.commit()
    conn.close()

def update_database():
    """Update the database schema and populate new columns."""
    # Create engine from DATABASE_URL
    engine = create_engine(DATABASE_URL)
    
    # Extract database path from DATABASE_URL for SQLite
    db_path = None
    if "sqlite" in DATABASE_URL:
        db_path = DATABASE_URL.replace("sqlite:///", "")
        # Handle relative paths
        if not os.path.isabs(db_path):
            db_path = os.path.join(os.path.dirname(__file__), db_path)
    
    print(f"Updating database at: {DATABASE_URL}")
    
    # Start a transaction
    db = SessionLocal()
    try:
        # Add last_accessed column if it doesn't exist
        if not check_column_exists(engine, 'sessions', 'last_accessed'):
            print("Adding last_accessed column...")
            if db_path:  # SQLite
                add_column_sqlite(db_path, 'sessions', 'last_accessed', 'DATETIME')
            else:  # Other databases
                with engine.connect() as conn:
                    conn.execute(text("ALTER TABLE sessions ADD COLUMN last_accessed DATETIME"))
                    conn.commit()
        
        # Add is_important column if it doesn't exist
        if not check_column_exists(engine, 'sessions', 'is_important'):
            print("Adding is_important column...")
            if db_path:  # SQLite
                add_column_sqlite(db_path, 'sessions', 'is_important', 'BOOLEAN', '0')
            else:  # Other databases
                with engine.connect() as conn:
                    conn.execute(text("ALTER TABLE sessions ADD COLUMN is_important BOOLEAN DEFAULT FALSE"))
                    conn.commit()
        
        # Add message_count column if it doesn't exist
        if not check_column_exists(engine, 'sessions', 'message_count'):
            print("Adding message_count column...")
            if db_path:  # SQLite
                add_column_sqlite(db_path, 'sessions', 'message_count', 'INTEGER', '0')
            else:  # Other databases
                with engine.connect() as conn:
                    conn.execute(text("ALTER TABLE sessions ADD COLUMN message_count INTEGER DEFAULT 0"))
                    conn.commit()
        
        # Populate last_accessed with created_at for existing records where last_accessed is NULL
        print("Populating last_accessed column...")
        with engine.connect() as conn:
            conn.execute(text("""
                UPDATE sessions 
                SET last_accessed = created_at 
                WHERE last_accessed IS NULL
            """))
            conn.commit()
        
        # Populate is_important with FALSE for existing records where is_important is NULL
        print("Populating is_important column...")
        with engine.connect() as conn:
            conn.execute(text("""
                UPDATE sessions 
                SET is_important = 0 
                WHERE is_important IS NULL
            """))
            conn.commit()
        
        # Calculate and populate message_count from chat_messages table
        print("Calculating and populating message_count column...")
        with engine.connect() as conn:
            # First, set all message_count to 0
            conn.execute(text("UPDATE sessions SET message_count = 0"))
            
            # Then, count messages for each session and update
            conn.execute(text("""
                UPDATE sessions 
                SET message_count = (
                    SELECT COUNT(*) 
                    FROM chat_messages 
                    WHERE chat_messages.session_id = sessions.id
                )
            """))
            conn.commit()
        
        print("Database update completed successfully!")
        
    except Exception as e:
        print(f"Error updating database: {e}")
        db.rollback()
        raise
    finally:
        db.close()

if __name__ == "__main__":
    update_database()