Spaces:
Running
Running
File size: 2,911 Bytes
1e68437 76ccd90 1e68437 76ccd90 712805f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 | 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
)
|