F-r-ldraledigplannerare/backend/app/db/connection.py

"""Database connection management.

Production: the schema is owned by Alembic (`alembic upgrade head`). The
`run_migrations()` helper below calls Alembic programmatically from the app
lifespan so the first boot auto-migrates on SQLite (dev convenience).

Tests: `ensure_schema()` applies the full DDL directly to an in-memory
connection so test fixtures stay fast and do not depend on Alembic.
"""

from __future__ import annotations

import asyncio
import weakref
from collections.abc import AsyncIterator
from contextlib import asynccontextmanager
from contextvars import ContextVar

import aiosqlite

# Full schema used by test fixtures directly (avoids a full Alembic run in tests).
# This MUST stay in sync with the migration chain (migrations/versions/*.py),
# i.e. it always reflects the state after the latest migration, not just 001.
_FULL_SCHEMA = """
CREATE TABLE IF NOT EXISTS workspace (
    id TEXT PRIMARY KEY,
    name TEXT NOT NULL,
    created_at TIMESTAMP NOT NULL,
    take_pl_on_red_days BOOLEAN NOT NULL DEFAULT 0
);

CREATE TABLE IF NOT EXISTS personal_data (
    uuid TEXT PRIMARY KEY,
    workspace_id TEXT NOT NULL,
    name TEXT NOT NULL,
    salary_type TEXT NOT NULL DEFAULT 'monthly'
        CHECK (salary_type IN ('monthly','hourly')),
    employment_status TEXT NOT NULL DEFAULT 'employed'
        CHECK (employment_status IN ('employed','jobseeking','none')),
    min_vacation_days_per_year INTEGER NOT NULL DEFAULT 20
        CHECK (min_vacation_days_per_year >= 0 AND min_vacation_days_per_year <= 60),
    personnummer_enc TEXT,
    created_at TIMESTAMP NOT NULL,
    deleted_at TIMESTAMP,
    FOREIGN KEY (workspace_id) REFERENCES workspace(id) ON DELETE CASCADE
);

CREATE TABLE IF NOT EXISTS "user" (
    id TEXT PRIMARY KEY,
    personal_data_uuid TEXT NOT NULL UNIQUE,
    email TEXT NOT NULL UNIQUE,
    hashed_password TEXT NOT NULL,
    created_at TIMESTAMP NOT NULL,
    is_active INTEGER NOT NULL DEFAULT 1,
    FOREIGN KEY (personal_data_uuid) REFERENCES personal_data(uuid) ON DELETE CASCADE
);

CREATE TABLE IF NOT EXISTS workspace_invitation (
    id TEXT PRIMARY KEY,
    workspace_id TEXT NOT NULL,
    token_hash TEXT NOT NULL UNIQUE,
    invited_by_user_id TEXT NOT NULL,
    status TEXT NOT NULL CHECK (status IN ('pending','accepted','revoked','expired')),
    expires_at TIMESTAMP NOT NULL,
    created_at TIMESTAMP NOT NULL,
    accepted_at TIMESTAMP,
    FOREIGN KEY (workspace_id) REFERENCES workspace(id) ON DELETE CASCADE,
    FOREIGN KEY (invited_by_user_id) REFERENCES "user"(id) ON DELETE CASCADE
);

CREATE TABLE IF NOT EXISTS child (
    id TEXT PRIMARY KEY,
    workspace_id TEXT NOT NULL,
    name TEXT NOT NULL,
    birth_date DATE NOT NULL,
    preschool_start_date DATE,
    preschool_planning_days_per_year INTEGER NOT NULL DEFAULT 6
        CHECK (preschool_planning_days_per_year >= 0 AND preschool_planning_days_per_year <= 30),
    created_at TIMESTAMP NOT NULL,
    FOREIGN KEY (workspace_id) REFERENCES workspace(id) ON DELETE CASCADE
);

CREATE TABLE IF NOT EXISTS leave_wish (
    id TEXT PRIMARY KEY,
    workspace_id TEXT NOT NULL,
    child_id TEXT,
    name TEXT NOT NULL,
    name_message_key TEXT,
    window_start_month INTEGER NOT NULL DEFAULT 1
        CHECK (window_start_month >= 1 AND window_start_month <= 12),
    window_start_day INTEGER NOT NULL DEFAULT 1
        CHECK (window_start_day >= 1 AND window_start_day <= 31),
    window_end_month INTEGER NOT NULL DEFAULT 1
        CHECK (window_end_month >= 1 AND window_end_month <= 12),
    window_end_day INTEGER NOT NULL DEFAULT 1
        CHECK (window_end_day >= 1 AND window_end_day <= 31),
    target_weeks INTEGER NOT NULL,
    target_days INTEGER,
    first_year INTEGER NOT NULL DEFAULT 2026,
    last_year INTEGER,
    priority INTEGER NOT NULL DEFAULT 1 CHECK (priority >= 1),
    parent_mode TEXT NOT NULL DEFAULT 'either'
        CHECK (parent_mode IN ('either','both','specific')),
    parent_personal_data_uuid TEXT,
    window_mode TEXT NOT NULL DEFAULT 'fixed'
        CHECK (window_mode IN ('fixed','flexible')),
    series_id TEXT,
    wants_dubbeldag INTEGER NOT NULL DEFAULT 0,
    created_at TIMESTAMP NOT NULL,
    updated_at TIMESTAMP NOT NULL,
    FOREIGN KEY (workspace_id) REFERENCES workspace(id) ON DELETE CASCADE,
    FOREIGN KEY (child_id) REFERENCES child(id) ON DELETE CASCADE,
    FOREIGN KEY (parent_personal_data_uuid) REFERENCES personal_data(uuid) ON DELETE CASCADE
);

CREATE TABLE IF NOT EXISTS parental_leave_fetch (
    id TEXT PRIMARY KEY,
    personal_data_uuid TEXT NOT NULL,
    workspace_id TEXT NOT NULL,
    child_id TEXT NOT NULL,
    source TEXT NOT NULL CHECK (source IN ('scraped','manual','deduced')),
    fetched_at TIMESTAMP NOT NULL,
    total_days INTEGER NOT NULL,
    days_used INTEGER NOT NULL,
    days_left INTEGER NOT NULL,
    reserved_days_total INTEGER NOT NULL DEFAULT 90,
    reserved_days_left INTEGER NOT NULL,
    transferable_days_left INTEGER NOT NULL,
    dubbeldagar_left INTEGER NOT NULL,
    tier_data TEXT NOT NULL,
    created_at TIMESTAMP NOT NULL,
    updated_at TIMESTAMP NOT NULL,
    FOREIGN KEY (personal_data_uuid) REFERENCES personal_data(uuid) ON DELETE CASCADE,
    FOREIGN KEY (workspace_id) REFERENCES workspace(id) ON DELETE CASCADE,
    FOREIGN KEY (child_id) REFERENCES child(id) ON DELETE CASCADE
);

CREATE TABLE IF NOT EXISTS vacation_days (
    id TEXT PRIMARY KEY,
    personal_data_uuid TEXT NOT NULL,
    workspace_id TEXT NOT NULL,
    year INTEGER NOT NULL,
    total_days INTEGER NOT NULL,
    saved_days_in INTEGER NOT NULL DEFAULT 0,
    created_at TIMESTAMP NOT NULL,
    updated_at TIMESTAMP NOT NULL,
    UNIQUE(personal_data_uuid, year),
    FOREIGN KEY (personal_data_uuid) REFERENCES personal_data(uuid) ON DELETE CASCADE,
    FOREIGN KEY (workspace_id) REFERENCES workspace(id) ON DELETE CASCADE
);

CREATE TABLE IF NOT EXISTS planned_leave_day (
    id TEXT PRIMARY KEY,
    workspace_id TEXT NOT NULL,
    child_id TEXT NOT NULL,
    parent_personal_data_uuid TEXT NOT NULL,
    date DATE NOT NULL,
    day_type TEXT NOT NULL CHECK (day_type IN ('sjukpenning','lagsta','vacation')),
    is_dubbeldag INTEGER NOT NULL DEFAULT 0,
    wish_id TEXT,
    is_pinned INTEGER NOT NULL DEFAULT 0,
    created_at TIMESTAMP NOT NULL,
    updated_at TIMESTAMP NOT NULL,
    UNIQUE(workspace_id, child_id, date, parent_personal_data_uuid),
    FOREIGN KEY (workspace_id) REFERENCES workspace(id) ON DELETE CASCADE,
    FOREIGN KEY (child_id) REFERENCES child(id) ON DELETE CASCADE,
    FOREIGN KEY (parent_personal_data_uuid) REFERENCES personal_data(uuid) ON DELETE CASCADE,
    FOREIGN KEY (wish_id) REFERENCES leave_wish(id) ON DELETE SET NULL
);

CREATE TABLE IF NOT EXISTS day_override (
    id TEXT PRIMARY KEY,
    workspace_id TEXT NOT NULL,
    date DATE NOT NULL,
    is_working_day INTEGER NOT NULL,
    note TEXT,
    created_at TIMESTAMP NOT NULL,
    updated_at TIMESTAMP NOT NULL,
    UNIQUE(workspace_id, date),
    FOREIGN KEY (workspace_id) REFERENCES workspace(id) ON DELETE CASCADE
);

CREATE TABLE IF NOT EXISTS tip_dismissal (
    id TEXT PRIMARY KEY,
    workspace_id TEXT NOT NULL,
    tip_id TEXT NOT NULL,
    data_hash TEXT NOT NULL,
    dismissed_at TIMESTAMP NOT NULL,
    UNIQUE(workspace_id, tip_id),
    FOREIGN KEY (workspace_id) REFERENCES workspace(id) ON DELETE CASCADE
);

-- Indexes. Kept in step (same columns, same order, no partial indexes) with
-- migrations/versions/001_initial_schema.py so
-- tests/integration/test_schema_parity.py, which compares index sets between
-- the Alembic chain and this hand-maintained DDL, stays green. See the
-- backend-perf audit §4 for what each index serves.
CREATE INDEX IF NOT EXISTS idx_pld_parent_child_type_date
    ON planned_leave_day (parent_personal_data_uuid, child_id, day_type, date);
CREATE INDEX IF NOT EXISTS idx_pld_child_dubbeldag_date
    ON planned_leave_day (child_id, is_dubbeldag, date);
CREATE INDEX IF NOT EXISTS idx_pld_wish
    ON planned_leave_day (wish_id);
CREATE INDEX IF NOT EXISTS idx_plf_parent_child_fetched
    ON parental_leave_fetch (personal_data_uuid, child_id, fetched_at);
CREATE INDEX IF NOT EXISTS idx_plf_workspace
    ON parental_leave_fetch (workspace_id);
CREATE INDEX IF NOT EXISTS idx_plf_child
    ON parental_leave_fetch (child_id);
CREATE INDEX IF NOT EXISTS idx_leave_wish_workspace
    ON leave_wish (workspace_id);
CREATE INDEX IF NOT EXISTS idx_leave_wish_child
    ON leave_wish (child_id);
CREATE INDEX IF NOT EXISTS idx_leave_wish_parent
    ON leave_wish (parent_personal_data_uuid);
CREATE INDEX IF NOT EXISTS idx_leave_wish_series
    ON leave_wish (series_id);
CREATE INDEX IF NOT EXISTS idx_child_workspace_birth
    ON child (workspace_id, birth_date);
CREATE INDEX IF NOT EXISTS idx_personal_data_workspace_deleted
    ON personal_data (workspace_id, deleted_at);
CREATE INDEX IF NOT EXISTS idx_invitation_workspace_status_created
    ON workspace_invitation (workspace_id, status, created_at);
CREATE INDEX IF NOT EXISTS idx_invitation_invited_by
    ON workspace_invitation (invited_by_user_id);
CREATE INDEX IF NOT EXISTS idx_vacation_days_workspace
    ON vacation_days (workspace_id);
"""


# ---------------------------------------------------------------------------
# Unit of work — see backend-perf audit §3 for the full design rationale.
#
# The app runs on ONE shared aiosqlite connection (app.state.db_connection),
# so a multi-statement write needs (a) mutual exclusion for the whole unit of
# work and (b) the repos' own per-row `commit()` calls suppressed while one is
# in progress — otherwise a concurrent request's commit would durably commit
# this one's half-finished writes. `_in_uow` and `_locks` below exist to
# provide exactly that on top of the existing single-connection, per-repo-
# commit design, without touching any repo's read/write SQL.
# ---------------------------------------------------------------------------

# True while the current asyncio task is inside a `unit_of_work()` block, on
# whichever connection started it. Checked by `commit()` (skip the per-row
# commit; the unit of work commits once, at the end) and by `unit_of_work()`
# itself (a nested call is a no-op passthrough — see its docstring).
_in_uow: ContextVar[bool] = ContextVar("db_in_uow", default=False)

# One asyncio.Lock per connection, created lazily on first use. A module-level
# lock would break under pytest: each test's `connection` fixture opens a new
# in-memory connection under a fresh event loop, and an `asyncio.Lock` created
# against one loop raises "bound to a different event loop" if awaited from
# another. Keying by connection (and using a WeakKeyDictionary so a closed
# test connection's lock is garbage-collected instead of leaking) sidesteps
# that entirely: production has exactly one connection => exactly one lock.
_locks: "weakref.WeakKeyDictionary[aiosqlite.Connection, asyncio.Lock]" = (
    weakref.WeakKeyDictionary()
)


def _lock_for(conn: aiosqlite.Connection) -> asyncio.Lock:
    lock = _locks.get(conn)
    if lock is None:
        lock = asyncio.Lock()
        _locks[conn] = lock
    return lock


async def commit(conn: aiosqlite.Connection) -> None:
    """Commit, unless a `unit_of_work()` is currently in progress on this task.

    Every sqlite repository's single-row write calls this instead of
    `conn.commit()` directly. Outside a unit of work this behaves exactly
    like today (commit immediately, once per write); inside one, it is a
    no-op, and the surrounding `unit_of_work()` block commits once, atomically,
    when it exits.
    """
    if not _in_uow.get():
        await conn.commit()


@asynccontextmanager
async def unit_of_work(conn: aiosqlite.Connection) -> AsyncIterator[None]:
    """Wrap the block in one atomic transaction on the shared connection.

    Acquires this connection's lock for the whole block (SQLite is a
    single-writer database anyway, so this costs no real throughput) and
    issues `BEGIN IMMEDIATE` up front — that claims the write lock
    immediately rather than at the first write statement, so a second,
    concurrent `unit_of_work()` call blocks here and only proceeds once this
    one has committed (or rolled back), then re-reads state as of that commit
    instead of racing it. Any exception rolls the whole block back; a normal
    exit commits it.

    Nested calls (e.g. `PlanningService.sync_plan` wrapping its own call to
    `apply_proposed_days`, which wraps its own writes) are a no-op passthrough
    when `_in_uow` is already set — the outer call owns the lock/transaction,
    and a naive re-entry would try to re-acquire the same (non-reentrant)
    `asyncio.Lock` from the same task and deadlock.
    """
    if _in_uow.get():
        yield
        return

    async with _lock_for(conn):
        token = _in_uow.set(True)
        try:
            await conn.execute("BEGIN IMMEDIATE")
            try:
                yield
            except BaseException:
                await conn.rollback()
                raise
            else:
                await conn.commit()
        finally:
            _in_uow.reset(token)


async def connect(path: str) -> aiosqlite.Connection:
    """Open a connection and enable row-by-name access + FK enforcement.

    Schema is applied via Alembic in the app lifespan (`run_migrations`).
    For test fixtures, call `ensure_schema` directly.

    PRAGMAs: `journal_mode=WAL` + `synchronous=NORMAL` cut single-row commit
    latency roughly 4x (measured 88.7ms -> 21.7ms per 300 sequential commits)
    and let Alembic/`seed_demo.py` co-access the same file wait instead of
    hitting "database is locked" outright; `busy_timeout` backs that wait with
    an actual retry budget. Durability trade-off of NORMAL: at most the last
    transaction(s) are lost on an OS crash/power loss, never on an app crash —
    acceptable here (see backend-perf audit §3).
    """
    conn = await aiosqlite.connect(path)
    conn.row_factory = aiosqlite.Row
    await conn.execute("PRAGMA foreign_keys = ON;")
    await conn.execute("PRAGMA journal_mode = WAL;")
    await conn.execute("PRAGMA synchronous = NORMAL;")
    await conn.execute("PRAGMA busy_timeout = 5000;")
    return conn


async def ensure_schema(conn: aiosqlite.Connection) -> None:
    """Apply the full DDL directly. Used by test fixtures only."""
    await conn.executescript(_FULL_SCHEMA)
    await conn.commit()


def run_migrations(database_url: str | None = None) -> None:
    """Run any pending Alembic migrations (called from app lifespan).

    `database_url` overrides alembic.ini's sqlalchemy.url, allowing the app
    config's `sqlite_path` to take precedence over the ini default.
    """
    from alembic import command
    from alembic.config import Config

    cfg = Config("alembic.ini")
    if database_url:
        cfg.set_main_option("sqlalchemy.url", database_url)
    command.upgrade(cfg, "head")