Spaces:
Sleeping
Sleeping
File size: 6,610 Bytes
f825419 d69ace5 f825419 d69ace5 f825419 248d767 f825419 248d767 f825419 d69ace5 4b09d2d f825419 4b09d2d f825419 4b09d2d 8cc4299 4b09d2d f825419 4b09d2d f825419 4b09d2d f825419 4b09d2d | 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 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 | """Database setup - SQLite (local dev) or PostgreSQL (deployments).
Backend chosen by DATABASE_URL:
* unset -> SQLite file under CCR_DATA_DIR (local dev; zero setup),
* postgres URL -> PostgreSQL (Supabase free tier recommended: persistent,
backed up, and already the auth provider - one vendor).
The models use no backend-specific features, so this is a connection-string
change, not a rewrite. On the ephemeral-disk hosts (HF Spaces free), SQLite is
wiped on every restart; Postgres is what makes accounts and data survive.
DATA_DIR still holds corpora/results/cache files locally; for durable FILE
storage on ephemeral hosts, additionally set CCR_STORAGE=s3 (storage.py).
"""
import os
from pathlib import Path
from sqlalchemy import create_engine, event
from sqlalchemy.orm import DeclarativeBase, sessionmaker
DATA_DIR = Path(
os.environ.get("CCR_DATA_DIR", Path(__file__).resolve().parent.parent / "data")
)
DATA_DIR.mkdir(parents=True, exist_ok=True)
(DATA_DIR / "corpora").mkdir(exist_ok=True)
(DATA_DIR / "results").mkdir(exist_ok=True)
def _normalize_pg_url(url: str) -> str:
"""Force the psycopg (v3) driver; accept the bare postgres:// URL that
dashboards (Supabase) hand out."""
if url.startswith("postgres://"):
url = "postgresql://" + url[len("postgres://"):]
if url.startswith("postgresql://"):
url = "postgresql+psycopg://" + url[len("postgresql://"):]
return url
_DATABASE_URL = os.environ.get("DATABASE_URL", "").strip()
IS_POSTGRES = _DATABASE_URL.startswith(("postgres://", "postgresql://"))
if IS_POSTGRES:
engine = create_engine(
_normalize_pg_url(_DATABASE_URL),
pool_pre_ping=True, # survive Supabase idle-connection drops
pool_recycle=1800,
pool_size=5,
max_overflow=5,
)
else:
DB_PATH = DATA_DIR / "ccr.db"
engine = create_engine(
f"sqlite:///{DB_PATH}",
connect_args={"check_same_thread": False}, # FastAPI threadpool access
)
@event.listens_for(engine, "connect")
def _sqlite_pragmas(dbapi_conn, _record):
"""WAL lets the API read while the job worker writes; busy_timeout
absorbs brief lock contention instead of raising immediately.
foreign_keys is OFF by default in SQLite, which silently makes the
dev/test backend more permissive than the deployed one: an ordering
bug that Postgres rejects with a ForeignKeyViolation passes locally
and in CI. Turning it on keeps both backends honest about the same
constraints.
"""
cur = dbapi_conn.cursor()
cur.execute("PRAGMA journal_mode=WAL")
cur.execute("PRAGMA busy_timeout=5000")
cur.execute("PRAGMA synchronous=NORMAL")
cur.execute("PRAGMA foreign_keys=ON")
cur.close()
SessionLocal = sessionmaker(bind=engine, autoflush=False, expire_on_commit=False)
class Base(DeclarativeBase):
pass
def get_db():
"""FastAPI dependency yielding a request-scoped session."""
db = SessionLocal()
try:
yield db
finally:
db.close()
def _default_literal(value, dialect_name: str) -> str:
if isinstance(value, bool):
if dialect_name == "postgresql":
return "TRUE" if value else "FALSE"
return "1" if value else "0"
if isinstance(value, (int, float)):
return str(value)
return "'" + str(value).replace("'", "''") + "'"
def lock_down_public_schema(target_engine, metadata) -> list[str]:
"""Enable Row-Level Security on every app table (Postgres only).
Supabase auto-exposes the public schema through its REST API (PostgREST):
any table WITHOUT RLS is readable AND writable by anyone holding the
project URL + anon key - which for this app would mean users (password
hashes), corpora, jobs, everything. This app never uses that REST API:
the backend talks to Postgres directly as the table owner, and owners
bypass RLS. So RLS-with-no-policies cleanly closes the public door
without touching app behavior (Supabase linter: rls_disabled_in_public).
Runs at every startup AFTER create_all, so tables added later are locked
down the day they appear, not when someone remembers. Idempotent.
"""
import logging
from sqlalchemy import text
if target_engine.dialect.name != "postgresql":
return [] # SQLite has no exposed REST surface (and no RLS)
locked = []
with target_engine.begin() as conn:
for table in metadata.sorted_tables:
conn.execute(text(f'ALTER TABLE "{table.name}" ENABLE ROW LEVEL SECURITY'))
locked.append(table.name)
logging.getLogger("ccr.db").info("RLS enabled on: %s", ", ".join(locked))
return locked
def auto_migrate_sqlite(target_engine, metadata) -> list[str]:
"""Add ORM columns missing from existing tables (additive only).
Named for history; runs on both backends. create_all() creates missing
tables but never alters existing ones, so a DB from last week 500s on this
week's new column. This closes that gap for additive changes; anything
non-additive (renames, drops, type changes) waits for Alembic. Columns with
scalar defaults get that default; callable defaults (uuid/now) are added
nullable and filled by the ORM on new rows. A brand-new Postgres database
needs none of this (create_all already made every current column).
"""
import logging
from sqlalchemy import inspect, text
added: list[str] = []
dialect = target_engine.dialect.name
inspector = inspect(target_engine)
with target_engine.begin() as conn:
for table in metadata.sorted_tables:
if table.name not in inspector.get_table_names():
continue # create_all handles brand-new tables
existing = {c["name"] for c in inspector.get_columns(table.name)}
for column in table.columns:
if column.name in existing:
continue
col_type = column.type.compile(target_engine.dialect)
ddl = f'ALTER TABLE {table.name} ADD COLUMN "{column.name}" {col_type}'
default = getattr(column.default, "arg", None)
if default is not None and not callable(default):
ddl += f" DEFAULT {_default_literal(default, dialect)}"
conn.execute(text(ddl))
added.append(f"{table.name}.{column.name}")
if added:
logging.getLogger("ccr.db").warning("auto-migrated columns: %s", ", ".join(added))
return added
|