import os from dotenv import load_dotenv from sqlalchemy import create_engine, text import sys # Add the parent directory to Python path to import your connection module sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) # Import your existing database connection logic from db.connection import get_database_config def get_database_url(): """Get database URL using the same logic as your main app""" db_config = get_database_config() if not db_config: print("āŒ No database configuration found") return None # Construct connection URL DATABASE_URL = f"postgresql://{db_config['user']}:{db_config['password']}@{db_config['host']}:{db_config['port']}/{db_config['dbname']}" # Add SSL mode for Supabase if db_config['host'].endswith('supabase.co'): DATABASE_URL += "?sslmode=require" print(f"šŸ”— Using database: {db_config['host']}:{db_config['port']}") return DATABASE_URL def add_generated_model_column(): """Add generated_model column to content_history table""" print("šŸ”„ Running database migration: adding generated_model column...") DATABASE_URL = get_database_url() if not DATABASE_URL: print("āŒ Cannot proceed without database configuration") return # Create engine engine = create_engine(DATABASE_URL) with engine.connect() as conn: try: # Check if column already exists result = conn.execute(text(""" SELECT column_name FROM information_schema.columns WHERE table_name='content_history' AND column_name='generated_model' """)) if result.fetchone() is None: print("šŸ“ Adding generated_model column...") # Add the column conn.execute(text(""" ALTER TABLE content_history ADD COLUMN generated_model VARCHAR DEFAULT 'groq' """)) print("āœ… Added generated_model column to content_history table") # Update existing records (set all existing content to groq) conn.execute(text(""" UPDATE content_history SET generated_model = 'groq' WHERE generated_model IS NULL """)) print("āœ… Set existing records to 'groq'") conn.commit() print("šŸŽ‰ Migration completed successfully!") else: print("āœ… generated_model column already exists") except Exception as e: print(f"āŒ Migration failed: {e}") conn.rollback() raise def add_auth_columns_to_users(): """Add authentication columns to users table - FIXED for existing data""" print("šŸ”„ Running database migration: adding authentication columns to users table...") DATABASE_URL = get_database_url() if not DATABASE_URL: print("āŒ Cannot proceed without database configuration") return # Create engine engine = create_engine(DATABASE_URL) with engine.connect() as conn: try: # Check if users table exists first result = conn.execute(text(""" SELECT table_name FROM information_schema.tables WHERE table_name='users' """)) if result.fetchone() is None: print("āŒ Users table doesn't exist. Please run the initial database setup first.") return # List of columns to add with their definitions columns_to_add = [ { 'name': 'fullname', 'definition': 'VARCHAR(200) NOT NULL DEFAULT \'\'', 'update_existing': "UPDATE users SET fullname = 'User' WHERE fullname IS NULL OR fullname = ''" }, { 'name': 'username', 'definition': 'VARCHAR(50) NOT NULL DEFAULT \'\'', # Remove UNIQUE constraint initially 'update_existing': "UPDATE users SET username = 'user_' || id::text WHERE username IS NULL OR username = ''" }, { 'name': 'password_hash', 'definition': 'VARCHAR(128) NOT NULL DEFAULT \'\'', 'update_existing': None }, { 'name': 'salt', 'definition': 'VARCHAR(32) NOT NULL DEFAULT \'\'', 'update_existing': None }, { 'name': 'last_login', 'definition': 'TIMESTAMP', 'update_existing': None } ] added_columns = [] for column in columns_to_add: # Check if column already exists result = conn.execute(text(f""" SELECT column_name FROM information_schema.columns WHERE table_name='users' AND column_name='{column['name']}' """)) if result.fetchone() is None: print(f"šŸ“ Adding {column['name']} column...") try: # Add the column without UNIQUE constraint initially conn.execute(text(f""" ALTER TABLE users ADD COLUMN {column['name']} {column['definition']} """)) # Update existing records if needed if column['update_existing']: conn.execute(text(column['update_existing'])) print(f"āœ… Updated existing records for {column['name']}") added_columns.append(column['name']) print(f"āœ… Added {column['name']} column to users table") except Exception as column_error: print(f"āš ļø Failed to add {column['name']} column: {column_error}") continue else: print(f"āœ… {column['name']} column already exists") # NOW add the UNIQUE constraint after all usernames are populated if 'username' in added_columns: print("šŸ”§ Adding UNIQUE constraint to username column...") try: conn.execute(text(""" ALTER TABLE users ADD CONSTRAINT users_username_unique UNIQUE (username) """)) print("āœ… Added UNIQUE constraint to username column") except Exception as e: print(f"āš ļø Could not add UNIQUE constraint: {e}") # This might happen if there are duplicate usernames, but we generated unique ones if added_columns: conn.commit() print(f"šŸŽ‰ Migration completed successfully! Added columns: {', '.join(added_columns)}") else: print("āœ… All authentication columns already exist") # Create index on username for faster lookups print("šŸ”§ Creating index on username column...") conn.execute(text(""" CREATE INDEX IF NOT EXISTS idx_users_username ON users(username) """)) print("āœ… Username index created") conn.commit() except Exception as e: print(f"āŒ Migration failed: {e}") conn.rollback() raise def migrate_existing_users(): """Migrate existing users to have proper usernames and fullnames""" print("šŸ”„ Migrating existing users to new authentication system...") DATABASE_URL = get_database_url() if not DATABASE_URL: print("āŒ Cannot proceed without database configuration") return engine = create_engine(DATABASE_URL) with engine.connect() as conn: try: # Check if we have users without proper usernames result = conn.execute(text(""" SELECT COUNT(*) as count FROM users WHERE username = '' OR username IS NULL """)) users_without_username = result.fetchone()[0] if users_without_username > 0: print(f"šŸ“ Found {users_without_username} users without usernames, migrating...") # Generate unique usernames for existing users conn.execute(text(""" UPDATE users SET username = 'user_' || id::text, fullname = 'User ' || id::text WHERE username = '' OR username IS NULL """)) print(f"āœ… Migrated {users_without_username} users to new authentication system") else: print("āœ… All users already have usernames") conn.commit() except Exception as e: print(f"āŒ User migration failed: {e}") conn.rollback() raise def run_all_migrations(): """Run all database migrations""" print("šŸš€ Starting database migrations...") try: # 1. Add generated_model column to content_history add_generated_model_column() print("\n" + "="*50 + "\n") # 2. Add authentication columns to users table add_auth_columns_to_users() print("\n" + "="*50 + "\n") # 3. Migrate existing users migrate_existing_users() print("\nšŸŽ‰ All migrations completed successfully!") except Exception as e: print(f"āŒ Migration process failed: {e}") raise if __name__ == "__main__": # Load environment variables load_dotenv() # Run migrations run_all_migrations()