echoes-of-the-ash/scripts/migrate_fix_telegram_id.py

"""
Migration: Fix telegram_id to allow NULL for web users

Changes:
- Drop existing primary key constraint on telegram_id
- Make telegram_id nullable
- Ensure id column exists and is unique
- Add constraint that either telegram_id OR username must be NOT NULL

Run this migration to allow web-only users without telegram_id.
"""

import asyncio
import os
from sqlalchemy.ext.asyncio import create_async_engine
from sqlalchemy import text

DB_USER = os.getenv("POSTGRES_USER")
DB_PASS = os.getenv("POSTGRES_PASSWORD")
DB_NAME = os.getenv("POSTGRES_DB")
DB_HOST = os.getenv("POSTGRES_HOST", "localhost")
DB_PORT = os.getenv("POSTGRES_PORT", "5432")

DATABASE_URL = f"postgresql+psycopg://{DB_USER}:{DB_PASS}@{DB_HOST}:{DB_PORT}/{DB_NAME}"


async def migrate():
    engine = create_async_engine(DATABASE_URL)
    
    async with engine.begin() as conn:
        print("Starting migration: Fix telegram_id to allow NULL...")
        
        try:
            # First, check current state
            result = await conn.execute(text("""
                SELECT 
                    c.column_name,
                    c.is_nullable,
                    c.column_default
                FROM information_schema.columns c
                WHERE c.table_name = 'players' 
                    AND c.column_name IN ('telegram_id', 'id', 'username')
                ORDER BY c.ordinal_position
            """))
            print("\nCurrent columns:")
            for row in result:
                print(f"  - {row[0]}: nullable={row[1]}, default={row[2]}")
            
            # Store foreign keys that reference players.telegram_id
            print("\nFinding foreign keys that reference players(telegram_id)...")
            result = await conn.execute(text("""
                SELECT 
                    tc.constraint_name,
                    tc.table_name,
                    kcu.column_name
                FROM information_schema.table_constraints tc
                JOIN information_schema.key_column_usage kcu
                    ON tc.constraint_name = kcu.constraint_name
                JOIN information_schema.constraint_column_usage ccu
                    ON tc.constraint_name = ccu.constraint_name
                WHERE tc.constraint_type = 'FOREIGN KEY'
                    AND ccu.table_name = 'players'
                    AND ccu.column_name = 'telegram_id'
            """))
            fk_constraints = list(result)
            
            # Drop foreign key constraints temporarily
            for fk_name, table_name, column_name in fk_constraints:
                print(f"  Dropping FK: {table_name}.{fk_name}")
                await conn.execute(text(f"ALTER TABLE {table_name} DROP CONSTRAINT {fk_name}"))
            
            # Check for primary key constraint
            result = await conn.execute(text("""
                SELECT constraint_name, constraint_type
                FROM information_schema.table_constraints
                WHERE table_name = 'players' AND constraint_type = 'PRIMARY KEY'
            """))
            pk_constraints = list(result)
            
            if pk_constraints:
                pk_name = pk_constraints[0][0]
                print(f"\nDropping PRIMARY KEY constraint: {pk_name}")
                await conn.execute(text(f"ALTER TABLE players DROP CONSTRAINT {pk_name}"))
            else:
                print("\n✓ No PRIMARY KEY constraint found")
            
            # Make telegram_id nullable
            print("Making telegram_id nullable...")
            await conn.execute(text("""
                ALTER TABLE players 
                ALTER COLUMN telegram_id DROP NOT NULL
            """))
            
            # Make telegram_id unique if not already
            print("Ensuring telegram_id is unique...")
            try:
                await conn.execute(text("""
                    ALTER TABLE players 
                    ADD CONSTRAINT players_telegram_id_unique UNIQUE (telegram_id)
                """))
            except Exception as e:
                if "already exists" in str(e):
                    print("✓ telegram_id unique constraint already exists")
                else:
                    raise
            
            # Ensure id column exists and is unique
            result = await conn.execute(text("""
                SELECT column_name 
                FROM information_schema.columns 
                WHERE table_name='players' AND column_name='id'
            """))
            if not list(result):
                print("Adding id column...")
                await conn.execute(text("""
                    ALTER TABLE players 
                    ADD COLUMN id SERIAL UNIQUE
                """))
            else:
                print("✓ id column exists")
                
                # Make sure it's unique
                try:
                    await conn.execute(text("""
                        ALTER TABLE players 
                        ADD CONSTRAINT players_id_unique UNIQUE (id)
                    """))
                except Exception as e:
                    if "already exists" in str(e):
                        print("✓ id unique constraint already exists")
                    else:
                        raise
            
            # Add check constraint that either telegram_id OR username must be NOT NULL
            print("Adding constraint: either telegram_id OR username must be NOT NULL...")
            try:
                await conn.execute(text("""
                    ALTER TABLE players 
                    ADD CONSTRAINT players_id_check 
                    CHECK (telegram_id IS NOT NULL OR username IS NOT NULL)
                """))
            except Exception as e:
                if "already exists" in str(e):
                    print("✓ Check constraint already exists")
                else:
                    raise
            
            # Recreate foreign key constraints
            print("\nRecreating foreign key constraints...")
            for fk_name, table_name, column_name in fk_constraints:
                print(f"  Adding FK: {table_name}.{fk_name}")
                await conn.execute(text(f"""
                    ALTER TABLE {table_name} 
                    ADD CONSTRAINT {fk_name} 
                    FOREIGN KEY ({column_name}) 
                    REFERENCES players(telegram_id) 
                    ON DELETE CASCADE
                """))
            
            print("\n✅ Migration completed successfully!")
            print("\nNow players can be created with:")
            print("  - telegram_id (Telegram users)")
            print("  - username + password_hash (web users)")
            print("  - Both telegram_id is unique and id is unique")
            
        except Exception as e:
            print(f"\n❌ Migration failed: {e}")
            raise
    
    await engine.dispose()


if __name__ == "__main__":
    asyncio.run(migrate())