From 30286111a8da09ad9c8aa5c45c4b4071d2399036 Mon Sep 17 00:00:00 2001 From: Dennis Thiessen Date: Sat, 18 Jul 2026 20:34:38 +0200 Subject: [PATCH] fix: per-symbol SQLite transactions in research snapshot extender Avoid inactive-transaction crashes from mixing connection.commit with ORM Session. Write path is raw SQL, one begin() block per symbol. --- scripts/extend_snapshot_universe.py | 182 ++++++++++++++++------------ 1 file changed, 105 insertions(+), 77 deletions(-) diff --git a/scripts/extend_snapshot_universe.py b/scripts/extend_snapshot_universe.py index 6b0f28b..9d5ab37 100644 --- a/scripts/extend_snapshot_universe.py +++ b/scripts/extend_snapshot_universe.py @@ -34,8 +34,7 @@ import time from datetime import date, datetime, timedelta, timezone from pathlib import Path -from sqlalchemy import create_engine, select, text -from sqlalchemy.orm import Session +from sqlalchemy import create_engine, text ROOT = Path(__file__).resolve().parents[1] if str(ROOT) not in sys.path: @@ -93,18 +92,19 @@ def _parse_args() -> argparse.Namespace: return p.parse_args() -def _ensure_rank_only_table(conn) -> None: - conn.execute( - text( - """ - CREATE TABLE IF NOT EXISTS research_rank_only ( - ticker_id INTEGER PRIMARY KEY, - symbol TEXT NOT NULL UNIQUE +def _ensure_rank_only_table(engine) -> None: + """DDL in its own connection/transaction (don't share with ORM Session).""" + with engine.begin() as conn: + conn.execute( + text( + """ + CREATE TABLE IF NOT EXISTS research_rank_only ( + ticker_id INTEGER PRIMARY KEY, + symbol TEXT NOT NULL UNIQUE + ) + """ ) - """ ) - ) - conn.commit() async def _resolve_pool() -> tuple[list[str], dict[str, str]]: @@ -213,8 +213,6 @@ async def _main() -> None: print(f"Updating existing research snapshot: {output}") from app.config import settings - from app.models.ohlcv import OHLCVRecord - from app.models.ticker import Ticker from app.providers.alpaca import AlpacaOHLCVProvider if not settings.alpaca_api_key or not settings.alpaca_api_secret: @@ -228,30 +226,32 @@ async def _main() -> None: pool, sources = await _resolve_pool() print(f"Pool size: {len(pool)} (sources={sources})") - # Sync sqlite via sqlalchemy core (simpler than async for bulk insert) - engine = create_engine(f"sqlite:///{output.resolve().as_posix()}") - with Session(engine) as session: - _ensure_rank_only_table(session.connection()) - existing = { - row.symbol: row - for row in session.execute(select(Ticker)).scalars().all() - } - prod_symbols = set(existing) + # Sync sqlite via raw SQL — one short transaction per symbol so a failed + # write never leaves the session in "transaction is inactive". + engine = create_engine( + f"sqlite:///{output.resolve().as_posix()}", + future=True, + ) + _ensure_rank_only_table(engine) + + with engine.connect() as conn: + existing_rows = conn.execute( + text("SELECT id, symbol FROM tickers") + ).fetchall() + existing_ids = {str(sym): int(tid) for tid, sym in existing_rows} + prod_symbols = set(existing_ids) - # Bar counts bar_counts: dict[str, int] = {} - for sym, ticker in existing.items(): - n = session.execute( + for sym, tid in existing_ids.items(): + n = conn.execute( text("SELECT COUNT(*) FROM ohlcv_records WHERE ticker_id = :tid"), - {"tid": ticker.id}, + {"tid": tid}, ).scalar_one() bar_counts[sym] = int(n) to_fetch: list[str] = [] for sym in pool: - if sym in existing and bar_counts.get(sym, 0) >= args.min_bars: - # Existing production or previously extended — keep rank_only - # only for *new* research names, not original prod universe. + if sym in existing_ids and bar_counts.get(sym, 0) >= args.min_bars: continue to_fetch.append(sym) @@ -262,6 +262,16 @@ async def _main() -> None: ok = 0 fail = 0 t0 = time.monotonic() + + insert_ohlcv = text( + """ + INSERT INTO ohlcv_records + (ticker_id, date, open, high, low, close, volume, created_at) + VALUES + (:ticker_id, :date, :open, :high, :low, :close, :volume, :created_at) + """ + ) + for index, sym in enumerate(to_fetch, 1): try: bars = await _fetch_symbol_bars( @@ -284,55 +294,71 @@ async def _main() -> None: print(f" [{index}/{len(to_fetch)}] {sym} empty") continue - ticker = existing.get(sym) - is_new = ticker is None - if ticker is None: - ticker = Ticker(symbol=sym, name=None, created_at=datetime.now(timezone.utc)) - session.add(ticker) - session.flush() - existing[sym] = ticker + try: + with engine.begin() as write: + ticker_id = existing_ids.get(sym) + is_new = ticker_id is None + if is_new: + write.execute( + text( + "INSERT INTO tickers (symbol, name, created_at) " + "VALUES (:sym, NULL, :created)" + ), + { + "sym": sym, + "created": datetime.now(timezone.utc).isoformat(), + }, + ) + ticker_id = int( + write.execute( + text("SELECT id FROM tickers WHERE symbol = :sym"), + {"sym": sym}, + ).scalar_one() + ) + existing_ids[sym] = ticker_id - # Upsert bars (delete+insert range for simplicity on research path) - session.execute( - text( - "DELETE FROM ohlcv_records WHERE ticker_id = :tid " - "AND date >= :start AND date <= :end" - ), - {"tid": ticker.id, "start": start.isoformat(), "end": end.isoformat()}, - ) - now = datetime.utcnow() - session.bulk_insert_mappings( - OHLCVRecord, - [ - { - "ticker_id": ticker.id, - "date": b.date, - "open": b.open, - "high": b.high, - "low": b.low, - "close": b.close, - "volume": b.volume, - "created_at": now, - } - for b in bars - ], - ) - - # rank_only only for names that were NOT in the original production - # snapshot at copy time (or are newly introduced to this research DB). - if is_new or sym not in prod_symbols: - # Re-evaluate: if source copy already had the symbol, don't flag. - # Only new inserts get rank_only. - if is_new: - session.execute( + write.execute( text( - "INSERT OR REPLACE INTO research_rank_only " - "(ticker_id, symbol) VALUES (:tid, :sym)" + "DELETE FROM ohlcv_records WHERE ticker_id = :tid " + "AND date >= :start AND date <= :end" ), - {"tid": ticker.id, "sym": sym}, + { + "tid": ticker_id, + "start": start.isoformat(), + "end": end.isoformat(), + }, ) + now = datetime.now(timezone.utc).replace(tzinfo=None) + write.execute( + insert_ohlcv, + [ + { + "ticker_id": ticker_id, + "date": b.date.isoformat(), + "open": float(b.open), + "high": float(b.high), + "low": float(b.low), + "close": float(b.close), + "volume": int(b.volume), + "created_at": now.isoformat(), + } + for b in bars + ], + ) + if is_new: + write.execute( + text( + "INSERT OR REPLACE INTO research_rank_only " + "(ticker_id, symbol) VALUES (:tid, :sym)" + ), + {"tid": ticker_id, "sym": sym}, + ) + except Exception as exc: + fail += 1 + if not args.quiet: + print(f" [{index}/{len(to_fetch)}] {sym} WRITE FAIL {exc}") + continue - session.commit() ok += 1 if not args.quiet and (index % 25 == 0 or index == len(to_fetch)): elapsed = time.monotonic() - t0 @@ -341,11 +367,13 @@ async def _main() -> None: f"elapsed={elapsed/60:.1f}m last={sym} bars={len(bars)}" ) - rank_only_n = session.execute( + rank_only_n = conn.execute( text("SELECT COUNT(*) FROM research_rank_only") ).scalar_one() - ticker_n = session.execute(text("SELECT COUNT(*) FROM tickers")).scalar_one() - ohlcv_n = session.execute(text("SELECT COUNT(*) FROM ohlcv_records")).scalar_one() + ticker_n = conn.execute(text("SELECT COUNT(*) FROM tickers")).scalar_one() + ohlcv_n = conn.execute( + text("SELECT COUNT(*) FROM ohlcv_records") + ).scalar_one() print("Done.") print(f" output: {output}")