fithia2/scripts/backfill_finra_short_volume_cdn.py

"""Backfill short_sale_daily from FINRA CDN daily files.

FINRA publishes daily short-volume files at:
  https://cdn.finra.org/equity/regsho/daily/CNMSshvol{yyyymmdd}.txt

This script downloads each file for the requested date range and bulk-inserts
into the short_sale_daily table, filtering to the midlarge universe only.

Usage:
  .venv/bin/python3 scripts/backfill_finra_short_volume_cdn.py \
      --start-date 2025-07-01 --end-date 2026-04-22

Idempotent: uses ON CONFLICT DO NOTHING.
"""
from __future__ import annotations

import argparse
import asyncio
import datetime as dt
import os
import sys
import time
from pathlib import Path

import asyncpg
import requests
import yaml

sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))

from libs.common.config import get_settings
from libs.common.time_utils import trading_days_between

FINRA_CDN_URL = "https://cdn.finra.org/equity/regsho/daily/CNMSshvol{yyyymmdd}.txt"
UNIVERSE_FILE = "configs/symbols_midlarge_snapshot_exact.yaml"
BATCH_SIZE = 1000


def load_universe() -> set[str]:
    with open(UNIVERSE_FILE) as f:
        udata = yaml.safe_load(f)
    syms = udata.get("symbols", udata) if isinstance(udata, dict) else udata
    return {s.upper() for s in syms}


def fetch_day(date: dt.date, universe: set[str]) -> list[dict]:
    url = FINRA_CDN_URL.format(yyyymmdd=date.strftime("%Y%m%d"))
    try:
        resp = requests.get(url, timeout=15)
    except requests.RequestException as e:
        print(f"  {date}: request error: {e}")
        return []

    if resp.status_code != 200:
        return []

    lines = [l.strip() for l in resp.text.splitlines() if l.strip()]
    if not lines:
        return []

    header = [p.strip() for p in lines[0].split("|")]
    try:
        sym_idx = header.index("Symbol")
        short_idx = header.index("ShortVolume")
        total_idx = header.index("TotalVolume")
    except ValueError:
        return []

    rows = []
    for line in lines[1:]:
        parts = [p.strip() for p in line.split("|")]
        if len(parts) <= max(sym_idx, short_idx, total_idx):
            continue
        sym = parts[sym_idx].upper()
        if sym not in universe:
            continue
        try:
            short_vol = int(float(parts[short_idx]))
            total_vol = int(float(parts[total_idx]))
        except (ValueError, IndexError):
            continue
        if total_vol <= 0:
            continue
        rows.append({
            "ticker_raw": sym,
            "trade_date": date,
            "short_volume": short_vol,
            "short_exempt_volume": None,
            "total_volume": total_vol,
        })
    return rows


async def insert_batch(conn: asyncpg.Connection, rows: list[dict]) -> int:
    if not rows:
        return 0
    now = dt.datetime.now(tz=dt.timezone.utc)
    inserted = await conn.executemany(
        """
        INSERT INTO short_sale_daily
            (ticker_raw, trade_date, short_volume, short_exempt_volume, total_volume, source_name, created_at_utc)
        VALUES ($1, $2, $3, $4, $5, 'finra_cdn', $6)
        ON CONFLICT ON CONSTRAINT uq_short_sale_ticker_date_source DO NOTHING
        """,
        [
            (r["ticker_raw"], r["trade_date"], r["short_volume"],
             r["short_exempt_volume"], r["total_volume"], now)
            for r in rows
        ],
    )
    return len(rows)


async def main(start_date: dt.date, end_date: dt.date) -> None:
    universe = load_universe()
    print(f"Universe: {len(universe)} tickers")

    # Get trading days in range
    trading_days = [
        d for d in trading_days_between(start_date, end_date)
        if start_date <= d <= end_date
    ]
    print(f"Trading days: {len(trading_days)} ({trading_days[0]} → {trading_days[-1]})")

    dsn = str(get_settings().postgres_dsn).replace("+asyncpg", "")
    conn = await asyncpg.connect(dsn)

    total_rows = 0
    total_days = 0

    try:
        pending: list[dict] = []
        for i, day in enumerate(trading_days):
            rows = fetch_day(day, universe)
            if rows:
                pending.extend(rows)
                total_days += 1
                if (i + 1) % 20 == 0:
                    print(f"  [{i+1}/{len(trading_days)}] day={day}, pending={len(pending)}")
            else:
                if (i + 1) % 20 == 0:
                    print(f"  [{i+1}/{len(trading_days)}] day={day}, skip (no data)")

            if len(pending) >= BATCH_SIZE:
                n = await insert_batch(conn, pending)
                total_rows += n
                pending = []
                print(f"    → inserted batch, total={total_rows}")

            time.sleep(0.05)  # gentle rate limiting

        if pending:
            n = await insert_batch(conn, pending)
            total_rows += n

    finally:
        await conn.close()

    print(f"\nDone. {total_days} trading days fetched, {total_rows} rows inserted.")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--start-date", default="2025-07-01")
    parser.add_argument("--end-date", default=dt.date.today().isoformat())
    args = parser.parse_args()

    asyncio.run(main(
        dt.date.fromisoformat(args.start_date),
        dt.date.fromisoformat(args.end_date),
    ))