Spaces:
Sleeping
Sleeping
File size: 2,652 Bytes
fc5156b f65a6b3 fc5156b | 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 | """SQLAlchemy engine, session, and Base used across the app."""
from sqlalchemy import create_engine, event
from sqlalchemy.engine import Engine
from sqlalchemy.orm import declarative_base, sessionmaker
from .config import DATABASE_URL
# check_same_thread=False is required for SQLite used by FastAPI's thread pool.
engine = create_engine(
DATABASE_URL,
connect_args={"check_same_thread": False},
)
@event.listens_for(Engine, "connect")
def _enable_sqlite_foreign_keys(dbapi_connection, _connection_record):
"""SQLite ignores ON DELETE CASCADE unless foreign keys are enabled per
connection. Without this, the FK constraints on chunks/qa_history would be
dead weight and rely solely on ORM-level cascades."""
cursor = dbapi_connection.cursor()
cursor.execute("PRAGMA foreign_keys=ON")
cursor.close()
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
Base = declarative_base()
def ensure_schema() -> None:
"""Tiny, idempotent migration for an MVP that uses create_all (no Alembic).
create_all() creates missing TABLES but never adds a new COLUMN to a table
that already exists. So when a model gains a column (e.g. chunks.embedding for
hybrid retrieval), an older on-disk SQLite file would 'no such column' on the
next query. This adds any such columns in place. Safe to call repeatedly and
before the tables exist (it simply skips what isn't there yet), and never
raises — a migration hiccup must not brick startup.
"""
from sqlalchemy import inspect, text
# (table, column, sqlite_type) added after the initial schema.
additions = [("chunks", "embedding", "BLOB")]
try:
insp = inspect(engine)
tables = set(insp.get_table_names())
except Exception:
return
for table, column, col_type in additions:
if table not in tables:
continue # fresh DB: create_all() will make it with the column
try:
cols = {c["name"] for c in insp.get_columns(table)}
if column not in cols:
with engine.begin() as conn:
conn.execute(text(f"ALTER TABLE {table} ADD COLUMN {column} {col_type}"))
except Exception:
pass # leave the DB as-is rather than crash on import
# Run on import so every entrypoint (app, eval CLIs, scripts) gets a current
# schema before its first query, including processes that never call create_all().
ensure_schema()
def get_db():
"""FastAPI dependency that yields a DB session and always closes it."""
db = SessionLocal()
try:
yield db
finally:
db.close()
|