You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
94 lines
3.1 KiB
Python
94 lines
3.1 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Backfill sector/industry/exchange/country for all active tickers in universe_ticker_registry
|
|
that currently have sector=NULL.
|
|
|
|
Usage:
|
|
python scripts/backfill_registry_sector.py [--batch 50] [--dry-run]
|
|
|
|
Expected runtime: ~30 min for 9376 tickers at semaphore=5.
|
|
"""
|
|
|
|
import asyncio
|
|
import argparse
|
|
import logging
|
|
import sys
|
|
import os
|
|
|
|
sys.path.insert(0, os.path.dirname(os.path.dirname(__file__)))
|
|
|
|
from sqlalchemy import select
|
|
from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine, async_sessionmaker
|
|
|
|
from app.core.config import settings
|
|
from app.models.universe_snapshot import UniverseTickerRegistry
|
|
from app.services import company_metadata_service as cms
|
|
|
|
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
async def backfill(batch_size: int, dry_run: bool) -> None:
|
|
engine = create_async_engine(settings.DATABASE_URL, echo=False)
|
|
Session = async_sessionmaker(engine, expire_on_commit=False)
|
|
|
|
async with Session() as db:
|
|
result = await db.execute(
|
|
select(UniverseTickerRegistry.ticker)
|
|
.where(
|
|
UniverseTickerRegistry.is_active == True,
|
|
UniverseTickerRegistry.sector == None,
|
|
)
|
|
.order_by(UniverseTickerRegistry.ticker)
|
|
)
|
|
tickers = [row[0] for row in result.fetchall()]
|
|
|
|
logger.info("Found %d tickers with NULL sector", len(tickers))
|
|
if dry_run:
|
|
logger.info("Dry-run mode — no changes will be written.")
|
|
return
|
|
|
|
success = 0
|
|
failed = 0
|
|
|
|
for i in range(0, len(tickers), batch_size):
|
|
batch = tickers[i:i + batch_size]
|
|
logger.info("Batch %d/%d — processing %d tickers", i // batch_size + 1,
|
|
(len(tickers) + batch_size - 1) // batch_size, len(batch))
|
|
|
|
async def _enrich(ticker: str) -> bool:
|
|
async with Session() as db:
|
|
try:
|
|
meta = await cms.get_metadata(db, ticker)
|
|
sector = meta.get("sector")
|
|
logger.info(" %s → sector=%s", ticker, sector or "NULL")
|
|
return bool(sector)
|
|
except ValueError:
|
|
logger.debug(" %s → invalid ticker, skipping", ticker)
|
|
return False
|
|
except Exception as e:
|
|
logger.warning(" %s → error: %s", ticker, e)
|
|
return False
|
|
|
|
results = await asyncio.gather(*[_enrich(t) for t in batch])
|
|
success += sum(results)
|
|
failed += sum(1 for r in results if not r)
|
|
|
|
# Brief pause between batches to be kind to yfinance
|
|
await asyncio.sleep(1)
|
|
|
|
logger.info("Done. success=%d failed/skipped=%d total=%d", success, failed, len(tickers))
|
|
await engine.dispose()
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--batch", type=int, default=50)
|
|
parser.add_argument("--dry-run", action="store_true")
|
|
args = parser.parse_args()
|
|
asyncio.run(backfill(args.batch, args.dry_run))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|