database/db.py

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
)