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.
124 lines
3.9 KiB
Python
124 lines
3.9 KiB
Python
"""
|
|
Universe snapshot backfill script.
|
|
|
|
Usage:
|
|
python3 build_snapshots_batched.py [start_date] [end_date]
|
|
|
|
Defaults:
|
|
start_date: 2015-01-01
|
|
end_date: today (first day of current month)
|
|
|
|
Runs inside the stock_oracle_api container via:
|
|
docker exec -d stock_oracle_api bash -c "python3 /app/scripts/build_snapshots_batched.py > /tmp/build.log 2>&1"
|
|
|
|
Features:
|
|
- Skips tickers already fully covered in the date range (no redundant HTTP calls)
|
|
- Processes registry tickers in batches of BATCH_SIZE to limit memory usage
|
|
- Logs progress with batch index, snapshots created/failed, elapsed time
|
|
"""
|
|
import asyncio
|
|
import logging
|
|
import sys
|
|
import os
|
|
import time
|
|
from datetime import date, datetime, timezone
|
|
|
|
logging.basicConfig(
|
|
level=logging.INFO,
|
|
format="%(asctime)s %(levelname)s %(message)s",
|
|
datefmt="%H:%M:%S",
|
|
)
|
|
logger = logging.getLogger(__name__)
|
|
|
|
BATCH_SIZE = 30
|
|
|
|
def _resolve_dates():
|
|
today = date.today()
|
|
end = date(today.year, today.month, 1) # first day of current month
|
|
start = date(2015, 1, 1)
|
|
|
|
if len(sys.argv) >= 2:
|
|
start = date.fromisoformat(sys.argv[1])
|
|
if len(sys.argv) >= 3:
|
|
end = date.fromisoformat(sys.argv[2])
|
|
|
|
return start.isoformat(), end.isoformat()
|
|
|
|
|
|
async def get_all_tickers(engine):
|
|
from sqlalchemy import text
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
async with AsyncSession(engine) as session:
|
|
result = await session.execute(
|
|
text("SELECT ticker FROM universe_ticker_registry WHERE is_active = true ORDER BY ticker")
|
|
)
|
|
return [r[0] for r in result.fetchall()]
|
|
|
|
|
|
async def run_batch(svc, session_factory, tickers, batch_num, total_batches, start_date, end_date):
|
|
import gc
|
|
t0 = time.time()
|
|
try:
|
|
result = await svc.build_snapshots(
|
|
session_factory,
|
|
tickers=tickers,
|
|
start_date=start_date,
|
|
end_date=end_date,
|
|
force_rebuild=False,
|
|
)
|
|
elapsed = time.time() - t0
|
|
created = result.get("snapshots_created", "?")
|
|
failed = result.get("tickers_failed", "?")
|
|
processed = result.get("tickers_processed", len(tickers))
|
|
logger.info(
|
|
"[%d/%d] done in %.1fs — created=%s failed=%s processed=%d tickers=%s..%s",
|
|
batch_num, total_batches, elapsed, created, failed, processed,
|
|
tickers[0], tickers[-1],
|
|
)
|
|
except Exception as e:
|
|
logger.error("[%d/%d] FAILED (%s..%s): %s", batch_num, total_batches, tickers[0], tickers[-1], e)
|
|
finally:
|
|
gc.collect()
|
|
|
|
|
|
async def main():
|
|
sys.path.insert(0, "/app")
|
|
os.environ.setdefault("ENV", "production")
|
|
|
|
start_date, end_date = _resolve_dates()
|
|
|
|
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
|
|
from app.core.config import settings
|
|
from app.services.universe_service import UniverseService
|
|
|
|
engine = create_async_engine(
|
|
settings.DATABASE_URL,
|
|
echo=False,
|
|
pool_size=3,
|
|
max_overflow=0,
|
|
pool_timeout=60,
|
|
pool_pre_ping=True,
|
|
connect_args={"server_settings": {"jit": "off"}},
|
|
)
|
|
session_factory = async_sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)
|
|
|
|
all_tickers = await get_all_tickers(engine)
|
|
total = len(all_tickers)
|
|
batches = [all_tickers[i:i+BATCH_SIZE] for i in range(0, total, BATCH_SIZE)]
|
|
total_batches = len(batches)
|
|
|
|
logger.info("Backfill range: %s → %s", start_date, end_date)
|
|
logger.info("Registry tickers: %d | Batch size: %d | Total batches: %d", total, BATCH_SIZE, total_batches)
|
|
logger.info("Already-covered tickers will be skipped automatically per batch.")
|
|
|
|
svc = UniverseService()
|
|
for idx, batch in enumerate(batches, start=1):
|
|
await run_batch(svc, session_factory, batch, idx, total_batches, start_date, end_date)
|
|
|
|
await engine.dispose()
|
|
logger.info("All batches complete.")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(main())
|