scripts/backfill_sectors.py

"""
Backfill missing sector on stocks via yfinance (batch, rate-limited).
Run: python scripts/backfill_sectors.py [--limit 500]
"""

import argparse
import asyncio
import os
import sys
import time
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent.parent))

from dotenv import load_dotenv

load_dotenv()

from app.services.database_service import DatabaseService

try:
    import yfinance as yf
except ImportError:
    print("pip install yfinance")
    sys.exit(1)


def yf_ticker(symbol: str, suffix: str) -> str:
    s = symbol.upper()
    if not suffix:
        return s
    suf = suffix if suffix.startswith(".") else f".{suffix}"
    return s if s.endswith(suf) else f"{s}{suf}"


async def main(limit: int) -> None:
    turso_url = os.getenv("TURSO_URL")
    turso_token = os.getenv("TURSO_TOKEN")
    if not turso_url or not turso_token:
        print("Set TURSO_URL and TURSO_TOKEN in backend/.env")
        sys.exit(1)

    db = DatabaseService(url=turso_url, token=turso_token)
    await db.initialize()

    rows = await db._execute(
        """
        SELECT symbol, country, yfinance_suffix, name
        FROM stocks
        WHERE sector IS NULL OR sector = ''
        ORDER BY market_cap DESC NULLS LAST
        LIMIT ?
        """,
        [limit],
    )

    print(f"Backfilling sector for {len(rows)} stocks...")
    updated = 0
    for i, row in enumerate(rows, 1):
        sym = yf_ticker(row["symbol"], row.get("yfinance_suffix") or "")
        try:
            info = yf.Ticker(sym).info or {}
            sector = info.get("sector") or info.get("industry")
            if sector:
                await db._execute(
                    "UPDATE stocks SET sector = ? WHERE symbol = ? AND country = ?",
                    [str(sector)[:120], row["symbol"], row["country"]],
                )
                updated += 1
                print(f"  [{i}/{len(rows)}] {sym} -> {sector}")
        except Exception as e:
            print(f"  [{i}/{len(rows)}] {sym} skip: {e}")
        time.sleep(0.25)

    print(f"Done. Updated {updated} rows.")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--limit", type=int, default=500)
    args = parser.parse_args()
    asyncio.run(main(args.limit))