update_db.py

from sqlalchemy import create_engine, text
import os
from dotenv import load_dotenv

load_dotenv()

# Database URL
DB_URL = os.getenv("DATABASE_URL", "postgresql://curator:curator_secret@127.0.0.1:5432/reddit_curator")

def migrate():
    print(f"Connecting to database at {DB_URL}...")
    try:
        engine = create_engine(DB_URL)
        print("Engine created.")
        
        with engine.connect() as conn:
            print("Connected successfully.")
            with conn.begin() as trans:
                # Set lock timeout to 5 seconds
                conn.execute(text("SET LOCAL lock_timeout = '5s';"))
                
                # 1. Add columns to tagged_results
                print("Migrating tagged_results...")
                conn.execute(text("ALTER TABLE tagged_results ADD COLUMN IF NOT EXISTS url VARCHAR(500);"))
                conn.execute(text("ALTER TABLE tagged_results ADD COLUMN IF NOT EXISTS permalink VARCHAR(500);"))
                
                # 2. Add columns to search_sessions
                print("Migrating search_sessions...")
                conn.execute(text("ALTER TABLE search_sessions ADD COLUMN IF NOT EXISTS results JSON;"))
                
                print("Migration completed successfully.")
    except Exception as e:
        print(f"Migration failed: {e}")
            
if __name__ == "__main__":
    migrate()