diff --git a/app/api/v1/endpoints/universe.py b/app/api/v1/endpoints/universe.py index 34280a7..ea9df5f 100644 --- a/app/api/v1/endpoints/universe.py +++ b/app/api/v1/endpoints/universe.py @@ -181,22 +181,42 @@ async def discover_tickers( async def _run_build_snapshots( tickers, start_date: str, end_date: str, force_rebuild: bool ): - """Background task wrapper for build_snapshots (needs its own DB session).""" + """Background task wrapper for build_snapshots. + + Spawns a completely separate thread with its own asyncio event loop so + the long-running build never touches the API's event loop or DB connection + pool, keeping the API responsive throughout. + """ + import asyncio as _asyncio + import threading from app.core.database import AsyncSessionLocal - async with AsyncSessionLocal() as db: - svc = UniverseService() + def _thread_main(): + loop = _asyncio.new_event_loop() + _asyncio.set_event_loop(loop) + + async def _build(): + svc = UniverseService() + try: + result = await svc.build_snapshots( + AsyncSessionLocal, + tickers=tickers, + start_date=start_date, + end_date=end_date, + force_rebuild=force_rebuild, + ) + logger.info(f"Universe background build complete: {result}") + except Exception as e: + logger.error(f"Universe background build failed: {e}") + try: - result = await svc.build_snapshots( - db, - tickers=tickers, - start_date=start_date, - end_date=end_date, - force_rebuild=force_rebuild, - ) - logger.info(f"Universe background build complete: {result}") - except Exception as e: - logger.error(f"Universe background build failed: {e}") + loop.run_until_complete(_build()) + finally: + loop.close() + + t = threading.Thread(target=_thread_main, daemon=True, name="universe-build") + t.start() + logger.info(f"Universe build started in isolated thread (daemon)") @router.post( @@ -231,7 +251,7 @@ async def build_snapshots( svc = UniverseService() try: result = await svc.build_snapshots( - db, + db, # existing session is fine for small quick jobs tickers=tickers, start_date=body.start_date, end_date=body.end_date, diff --git a/app/services/universe_service.py b/app/services/universe_service.py index fb45c63..8e1cb27 100644 --- a/app/services/universe_service.py +++ b/app/services/universe_service.py @@ -234,7 +234,7 @@ class UniverseService: async def build_snapshots( self, - db: AsyncSession, + db_or_factory, tickers: Optional[List[str]], start_date: str, end_date: str, @@ -243,65 +243,90 @@ class UniverseService: """ Build monthly market_cap snapshots for given tickers (or all registry). - For each ticker: - 1. Fetch shares_outstanding history from SEC EDGAR companyfacts - 2. Fetch monthly close prices via yfinance bulk download - 3. Compute market_cap = close × shares per month - 4. Upsert into universe_snapshot + db_or_factory: Either an AsyncSession (for small sync calls ≤20 tickers) + or an AsyncSessionLocal factory (for large background jobs). + Using a factory means each batch opens/closes its own DB + connection, preventing the long-running job from monopolising + the connection pool and blocking the API. Returns: { tickers_processed, tickers_failed, snapshots_created } """ - # Resolve ticker list - if tickers: - ticker_list = [t.upper() for t in tickers] - else: - result = await db.execute( - select(UniverseTickerRegistry.ticker) - .where(UniverseTickerRegistry.is_active == True) - .order_by(UniverseTickerRegistry.ticker) - ) - ticker_list = [r[0] for r in result.fetchall()] + from app.core.database import AsyncSessionLocal - if not ticker_list: - logger.warning("Universe: no tickers to process") - return {"tickers_processed": 0, "tickers_failed": 0, "snapshots_created": 0} - - logger.info( - f"Universe: building snapshots for {len(ticker_list)} tickers " - f"({start_date} → {end_date})" - ) + # Normalise: wrap a plain session in a no-op factory so the rest of the + # code can always call `async with _db_ctx() as db:` + from contextlib import asynccontextmanager - if force_rebuild: - await db.execute( - UniverseSnapshot.__table__.delete().where( - UniverseSnapshot.ticker.in_(ticker_list) + if callable(db_or_factory) and not isinstance(db_or_factory, AsyncSession): + # It's a factory (e.g. AsyncSessionLocal) + _db_ctx = db_or_factory + else: + # It's an existing session — wrap it so we never close it here + _session = db_or_factory + + @asynccontextmanager + async def _db_ctx(): + yield _session + + # ---- Resolve ticker list ---- + async with _db_ctx() as db: + if tickers: + ticker_list = [t.upper() for t in tickers] + else: + result = await db.execute( + select(UniverseTickerRegistry.ticker) + .where(UniverseTickerRegistry.is_active == True) + .order_by(UniverseTickerRegistry.ticker) ) - ) - await db.flush() + ticker_list = [r[0] for r in result.fetchall()] + + if not ticker_list: + logger.warning("Universe: no tickers to process") + return {"tickers_processed": 0, "tickers_failed": 0, "snapshots_created": 0} - # Load registry metadata (sector/industry/exchange/CIK) - reg_result = await db.execute( - select(UniverseTickerRegistry).where( - UniverseTickerRegistry.ticker.in_(ticker_list) + logger.info( + f"Universe: building snapshots for {len(ticker_list)} tickers " + f"({start_date} → {end_date})" ) - ) - registry_map: Dict[str, UniverseTickerRegistry] = { - r.ticker: r for r in reg_result.scalars().all() - } + + if force_rebuild: + # Delete in chunks to avoid huge IN clauses + for i in range(0, len(ticker_list), _CHUNK): + chunk = ticker_list[i:i + _CHUNK] + await db.execute( + UniverseSnapshot.__table__.delete().where( + UniverseSnapshot.ticker.in_(chunk) + ) + ) + await db.commit() + + # Load registry metadata (sector/industry/exchange/CIK) + registry_map: Dict[str, UniverseTickerRegistry] = {} + for i in range(0, len(ticker_list), _CHUNK): + chunk = ticker_list[i:i + _CHUNK] + reg_result = await db.execute( + select(UniverseTickerRegistry).where( + UniverseTickerRegistry.ticker.in_(chunk) + ) + ) + for r in reg_result.scalars().all(): + registry_map[r.ticker] = r total_snapshots = 0 total_failed = 0 + total_batches = (len(ticker_list) + _PRICE_BATCH - 1) // _PRICE_BATCH - # Process in batches of _PRICE_BATCH for yfinance bulk download + # ---- Process in batches — each batch uses its own DB session ---- for batch_start in range(0, len(ticker_list), _PRICE_BATCH): batch = ticker_list[batch_start:batch_start + _PRICE_BATCH] + batch_num = batch_start // _PRICE_BATCH + 1 - # ---- Fetch monthly prices (synchronous, in thread) ---- + # Fetch monthly prices (synchronous yfinance, runs in thread) price_data = await asyncio.to_thread( self._fetch_bulk_monthly_prices, batch, start_date, end_date ) - # ---- Fetch shares_outstanding from SEC EDGAR (concurrent) ---- + # Fetch shares_outstanding from SEC EDGAR (async, rate-limited) sem = asyncio.Semaphore(_SEC_CONCURRENCY) async def _fetch_one(tkr: str) -> Tuple[str, List]: @@ -321,7 +346,7 @@ class UniverseService: tkr, history = item shares_map[tkr] = history - # ---- Build snapshot rows ---- + # Build snapshot rows batch_rows = [] for ticker in batch: shares_history = shares_map.get(ticker, []) @@ -335,8 +360,6 @@ class UniverseService: # Use the LATEST available shares_outstanding (most recent SEC filing). # yfinance returns split-adjusted prices retroactively, so using the # post-split shares count gives correct market_cap across all periods. - # Small error (~10-20%) may occur for companies with large buyback programs, - # but this is acceptable for screening purposes. latest_shares = shares_history[-1][1] if shares_history else None for snap_date_key, close in ticker_prices.items(): @@ -351,14 +374,13 @@ class UniverseService: # Sanity checks: skip absurd values (data quality) if close > 1_000_000 or close <= 0: # max BRK-A ~$600K continue - if shares < 100_000: # too few shares for a real public co + if shares < 100_000: # too few shares for a public co continue market_cap = shares * close if math.isnan(market_cap) or market_cap <= 0 or market_cap > 5e12: continue - # Normalize to first of month snapshot_dt = datetime( snap_date.year, snap_date.month, 1, tzinfo=timezone.utc ) @@ -373,34 +395,33 @@ class UniverseService: "exchange": reg.exchange if reg else None, }) - # ---- Batch upsert ---- + # Upsert with a fresh short-lived DB session (releases connection immediately after) if batch_rows: - for i in range(0, len(batch_rows), _CHUNK): - chunk = batch_rows[i:i + _CHUNK] - stmt = pg_insert(UniverseSnapshot).values(chunk) - stmt = stmt.on_conflict_do_update( - constraint="uq_universe_snapshot", - set_={ - "close_price": stmt.excluded.close_price, - "shares_outstanding": stmt.excluded.shares_outstanding, - "market_cap": stmt.excluded.market_cap, - "sector": stmt.excluded.sector, - "industry": stmt.excluded.industry, - "exchange": stmt.excluded.exchange, - }, - ) - result = await db.execute(stmt) - total_snapshots += result.rowcount - await db.commit() + async with _db_ctx() as db: + for i in range(0, len(batch_rows), _CHUNK): + chunk = batch_rows[i:i + _CHUNK] + stmt = pg_insert(UniverseSnapshot).values(chunk) + stmt = stmt.on_conflict_do_update( + constraint="uq_universe_snapshot", + set_={ + "close_price": stmt.excluded.close_price, + "shares_outstanding": stmt.excluded.shares_outstanding, + "market_cap": stmt.excluded.market_cap, + "sector": stmt.excluded.sector, + "industry": stmt.excluded.industry, + "exchange": stmt.excluded.exchange, + }, + ) + result = await db.execute(stmt) + total_snapshots += result.rowcount + await db.commit() - batch_num = batch_start // _PRICE_BATCH + 1 - total_batches = (len(ticker_list) + _PRICE_BATCH - 1) // _PRICE_BATCH logger.info( f"Universe: batch {batch_num}/{total_batches} — " f"{len(batch)} tickers, {len(batch_rows)} snapshot rows" ) # Yield to event loop between batches to keep API responsive - await asyncio.sleep(1) + await asyncio.sleep(0) logger.info( f"Universe: build complete — {len(ticker_list)} tickers, "