grokbot_ty
85 lines
2.1 KiB
Python
85 lines
2.1 KiB
Python
"""SQLite engine and session management.
|
|
|
|
Uses SQLAlchemy async engine with aiosqlite driver.
|
|
|
|
For tests, pass `test_db_url` via the settings or use a test fixture
|
|
to point at an in-memory SQLite database.
|
|
"""
|
|
|
|
import logging
|
|
|
|
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
|
|
|
|
from app.db.models import Base
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Default DB path for production
|
|
DEFAULT_DB_URL = "sqlite+aiosqlite:////data/slab.db"
|
|
|
|
_engine: object = None
|
|
_session_factory: object = None
|
|
|
|
_test_url: str | None = None
|
|
|
|
|
|
def set_test_url(url: str) -> None:
|
|
"""Override the database URL for testing. Must be called before create_app."""
|
|
global _test_url
|
|
_test_url = url
|
|
|
|
|
|
def get_db_url() -> str:
|
|
"""Return the database URL, preferring test URL if set."""
|
|
if _test_url:
|
|
return _test_url
|
|
return DEFAULT_DB_URL
|
|
|
|
|
|
async def init_db() -> None:
|
|
"""Initialize the async engine, session factory, and create tables.
|
|
|
|
For in-memory SQLite tests, creates tables directly.
|
|
For file-based SQLite, attempts alembic migrations first, falling
|
|
back to direct table creation.
|
|
"""
|
|
global _engine, _session_factory
|
|
|
|
db_url = get_db_url()
|
|
logger.info("Connecting to %s", db_url)
|
|
|
|
_engine = create_async_engine(
|
|
db_url,
|
|
echo=False,
|
|
connect_args={"check_same_thread": False},
|
|
)
|
|
|
|
_session_factory = async_sessionmaker(
|
|
_engine,
|
|
class_=AsyncSession,
|
|
expire_on_commit=False,
|
|
autocommit=False,
|
|
autoflush=False,
|
|
)
|
|
|
|
# Create all tables (works for both in-memory and file-based)
|
|
async with _engine.begin() as conn:
|
|
await conn.run_sync(Base.metadata.create_all)
|
|
|
|
logger.info("Database initialized")
|
|
|
|
|
|
async def close_db_connection() -> None:
|
|
"""Dispose of the engine (closes all connections)."""
|
|
global _engine
|
|
if _engine is not None:
|
|
await _engine.dispose()
|
|
_engine = None
|
|
|
|
|
|
async def get_session() -> AsyncSession:
|
|
"""Get a database session."""
|
|
if _session_factory is None:
|
|
raise RuntimeError("Database not initialized. Call init_db() first.")
|
|
return _session_factory()
|