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.
This commit is contained in:
@@ -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,7 +92,9 @@ def _parse_args() -> argparse.Namespace:
|
||||
return p.parse_args()
|
||||
|
||||
|
||||
def _ensure_rank_only_table(conn) -> None:
|
||||
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(
|
||||
"""
|
||||
@@ -104,7 +105,6 @@ def _ensure_rank_only_table(conn) -> None:
|
||||
"""
|
||||
)
|
||||
)
|
||||
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(
|
||||
write.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()},
|
||||
{
|
||||
"tid": ticker_id,
|
||||
"start": start.isoformat(),
|
||||
"end": end.isoformat(),
|
||||
},
|
||||
)
|
||||
now = datetime.utcnow()
|
||||
session.bulk_insert_mappings(
|
||||
OHLCVRecord,
|
||||
now = datetime.now(timezone.utc).replace(tzinfo=None)
|
||||
write.execute(
|
||||
insert_ohlcv,
|
||||
[
|
||||
{
|
||||
"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,
|
||||
"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
|
||||
],
|
||||
)
|
||||
|
||||
# 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)"
|
||||
),
|
||||
{"tid": ticker.id, "sym": 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}")
|
||||
|
||||
Reference in New Issue
Block a user