scripts/init_database.py

"""
Database Initialization Script
Creates necessary tables and indexes in Turso database
"""

import asyncio
import os
import sys
from pathlib import Path

# Add parent directory to path
sys.path.insert(0, str(Path(__file__).parent.parent))

from app.services.database_service import DatabaseService
from dotenv import load_dotenv

load_dotenv()


async def run_migration(db: DatabaseService):
    """Run schema migration to add multi-country columns if missing"""
    print("\nRunning schema migration...")
    
    migration_queries = [
        ("exchange_name", "ALTER TABLE stocks ADD COLUMN exchange_name TEXT DEFAULT ''"),
        ("country", "ALTER TABLE stocks ADD COLUMN country TEXT DEFAULT 'us'"),
        ("yfinance_suffix", "ALTER TABLE stocks ADD COLUMN yfinance_suffix TEXT DEFAULT ''"),
    ]
    
    for col_name, alter_sql in migration_queries:
        try:
            await db._execute(alter_sql)
            print(f"  Added column: {col_name}")
        except Exception:
            # Column already exists - that's fine
            print(f"  Column already exists: {col_name}")

    # Create country index if not exists
    try:
        await db._execute("""
            CREATE INDEX IF NOT EXISTS idx_stocks_country
            ON stocks(country)
        """)
        print("  Created index: idx_stocks_country")
    except Exception as e:
        print(f"  Index already exists or could not be created: {e}")

    print("  Migration complete")


async def main():
    """Initialize database schema"""
    print("Initializing database...")
    
    # Get credentials from environment
    turso_url = os.getenv("TURSO_URL")
    turso_token = os.getenv("TURSO_TOKEN")
    
    if not turso_url or not turso_token:
        print("Error: TURSO_URL and TURSO_TOKEN must be set in .env file")
        sys.exit(1)
    
    # Initialize database service
    db = DatabaseService(url=turso_url, token=turso_token)
    
    try:
        await db.initialize()
        # Run migration to add multi-country columns
        await run_migration(db)
        
        print("\nDatabase initialized successfully!")
        print("\nCreated/updated tables:")
        print("  - forecasts (cached predictions)")
        print("  - stocks (stock metadata with multi-country support)")
        print("\nNew columns added:")
        print("  - exchange_name: Full exchange name")
        print("  - country: Country code (us, uk, ca, au, nz, ie)")
        print("  - yfinance_suffix: yfinance ticker suffix")
        print("\nNext steps:")
        print("  1. Run: python scripts/seed_all_stocks.py")
        print("     (This seeds stocks from all 6 countries)")
        print("  2. Start the API: uvicorn app.main:app --reload")
        
    except Exception as e:
        print(f"Initialization failed: {e}")
        sys.exit(1)
    finally:
        await db.close()


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