import os from pathlib import Path from sqlalchemy import event, text from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession, async_sessionmaker from database.models import Base # Setup database directory dynamically relative to project root BASE_DIR = Path(__file__).resolve().parent.parent DATA_DIR = BASE_DIR / "data" DATA_DIR.mkdir(parents=True, exist_ok=True) DATABASE_URL = f"sqlite+aiosqlite:///{DATA_DIR}/quantforge_v2.db" # Create async database engine engine = create_async_engine( DATABASE_URL, echo=False, connect_args={"timeout": 30} ) # Listen to connect event to set WAL mode for SQLite @event.listens_for(engine.sync_engine, "connect") def set_sqlite_pragma(dbapi_connection, connection_record): """ Sets SQLite pragmas on connection to enable WAL (Write-Ahead Logging) mode and normal synchronization for high concurrency and performance. """ cursor = dbapi_connection.cursor() try: cursor.execute("PRAGMA journal_mode=WAL") cursor.execute("PRAGMA synchronous=NORMAL") except Exception as e: # Fallback print or log if needed, though typically silent pass finally: cursor.close() # Async session maker AsyncSessionLocal = async_sessionmaker( bind=engine, class_=AsyncSession, expire_on_commit=False, autocommit=False, autoflush=False ) async def init_db() -> None: """ Asynchronously creates all tables defined in models.py if they do not exist. """ async with engine.begin() as conn: await conn.run_sync(Base.metadata.create_all) columns = { row[1] for row in (await conn.execute(text("PRAGMA table_info(alphas)"))).fetchall() } if "raw_result_json" not in columns: await conn.execute(text("ALTER TABLE alphas ADD COLUMN raw_result_json TEXT")) # --------------------------------------------------------- # Synchronous Database Components for Backward Compatibility # --------------------------------------------------------- from sqlalchemy import create_engine from sqlalchemy.orm import sessionmaker SYNC_DATABASE_URL = f"sqlite:///{DATA_DIR}/quantforge_v2.db" # Create sync database engine (pointing to the exact same file) sync_engine = create_engine( SYNC_DATABASE_URL, echo=False, connect_args={"timeout": 30} ) # Listen to connect event to set WAL mode for SQLite on sync engine @event.listens_for(sync_engine, "connect") def set_sync_sqlite_pragma(dbapi_connection, connection_record): cursor = dbapi_connection.cursor() try: cursor.execute("PRAGMA journal_mode=WAL") cursor.execute("PRAGMA synchronous=NORMAL") except Exception: pass finally: cursor.close() # Synchronous session maker SessionLocal = sessionmaker( bind=sync_engine, autocommit=False, autoflush=False, expire_on_commit=False )