Update db/migration.py

db/migration.py

@@ -1,22 +1,41 @@
 import os
 from dotenv import load_dotenv
 from sqlalchemy import create_engine, text
+import sys
 
-# Load .env variables
-load_dotenv()
+# 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__))))
 
-DB_NAME = os.getenv("DB_NAME")
-DB_USER = os.getenv("DB_USERNAME")
-DB_PASS = os.getenv("DB_PASSWORD")
-DB_HOST = os.getenv("DB_HOST", "localhost")
-DB_PORT = os.getenv("DB_PORT", "5432")
+# Import your existing database connection logic
+from db.connection import get_database_config
 
-DATABASE_URL = f"postgresql://{DB_USER}:{DB_PASS}@{DB_HOST}:{DB_PORT}/{DB_NAME}"
+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)
     
@@ -57,5 +76,199 @@ def add_generated_model_column():
             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__":
-    add_generated_model_column()
+    # Load environment variables
+    load_dotenv()
+    
+    # Run migrations
+    run_all_migrations()