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