database.py

import json
import os
from contextlib import contextmanager

import psycopg2
from psycopg2.extras import RealDictCursor

DATABASE_URL = os.environ['DATABASE_URL']

TABLE_ANNOTATORS = 'annotators_pp1'
TABLE_COMBINATIONS = 'combinations_pp1'
TABLE_ASSIGNMENTS = 'combo_assignments_pp1'
TABLE_ANNOTATIONS = 'annotations_pp1'


@contextmanager
def get_conn():
    conn = psycopg2.connect(DATABASE_URL)
    try:
        yield conn
        conn.commit()
    except Exception:
        conn.rollback()
        raise
    finally:
        conn.close()


def create_tables() -> None:
    with get_conn() as conn:
        with conn.cursor() as cur:
            cur.execute(
                f'''
                CREATE EXTENSION IF NOT EXISTS pgcrypto;

                CREATE TABLE IF NOT EXISTS {TABLE_ANNOTATORS} (
                    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
                    code INT UNIQUE NOT NULL CHECK (code BETWEEN 1 AND 8),
                    created_at TIMESTAMPTZ DEFAULT now()
                );

                CREATE TABLE IF NOT EXISTS {TABLE_COMBINATIONS} (
                    id TEXT PRIMARY KEY,
                    base_id TEXT NOT NULL,
                    seed_id TEXT,
                    level INT,
                    class_tag TEXT,
                    ratio_type TEXT,
                    matchup_key TEXT NOT NULL,
                    model_left TEXT NOT NULL,
                    model_right TEXT NOT NULL,
                    n_total INT,
                    n_humans INT,
                    n_objects INT,
                    prompt TEXT,
                    prompt_zh TEXT,
                    humans JSONB,
                    objects JSONB
                );

                CREATE TABLE IF NOT EXISTS {TABLE_ASSIGNMENTS} (
                    combo_id TEXT PRIMARY KEY REFERENCES {TABLE_COMBINATIONS}(id) ON DELETE CASCADE,
                    group_key TEXT,
                    annotator_a_code INT NOT NULL CHECK (annotator_a_code BETWEEN 1 AND 8),
                    annotator_b_code INT NOT NULL CHECK (annotator_b_code BETWEEN 1 AND 8),
                    created_at TIMESTAMPTZ DEFAULT now()
                );

                CREATE TABLE IF NOT EXISTS {TABLE_ANNOTATIONS} (
                    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
                    annotator_id UUID NOT NULL REFERENCES {TABLE_ANNOTATORS}(id),
                    combo_id TEXT NOT NULL REFERENCES {TABLE_COMBINATIONS}(id),
                    model_a TEXT NOT NULL,
                    model_b TEXT NOT NULL,
                    a_existence INT CHECK (a_existence IN (0, 1)),
                    a_appearance INT CHECK (a_appearance IN (0, 1)),
                    a_interaction INT CHECK (a_interaction IN (0, 1)),
                    b_existence INT CHECK (b_existence IN (0, 1)),
                    b_appearance INT CHECK (b_appearance IN (0, 1)),
                    b_interaction INT CHECK (b_interaction IN (0, 1)),
                    preference TEXT CHECK (preference IN ('A', 'B')),
                    prompt_ilogical BOOLEAN NOT NULL DEFAULT FALSE,
                    created_at TIMESTAMPTZ DEFAULT now(),
                    UNIQUE (annotator_id, combo_id)
                );
                '''
            )
            for code in range(1, 9):
                cur.execute(
                    f"INSERT INTO {TABLE_ANNOTATORS} (code) VALUES (%s) ON CONFLICT (code) DO NOTHING",
                    (code,),
                )


def clear_round_tables() -> None:
    with get_conn() as conn:
        with conn.cursor() as cur:
            cur.execute(f'TRUNCATE TABLE {TABLE_ANNOTATIONS} CASCADE')
            cur.execute(f'TRUNCATE TABLE {TABLE_ASSIGNMENTS} CASCADE')
            cur.execute(f'TRUNCATE TABLE {TABLE_COMBINATIONS} CASCADE')


def get_annotator_id_by_code(code: int) -> str:
    with get_conn() as conn:
        with conn.cursor() as cur:
            cur.execute(f'SELECT id FROM {TABLE_ANNOTATORS} WHERE code = %s', (code,))
            row = cur.fetchone()
            if not row:
                raise ValueError(f'annotator code not found: {code}')
            return str(row[0])


def get_candidate_combos(annotator_code: int, limit: int = 500) -> list[dict]:
    with get_conn() as conn:
        with conn.cursor(cursor_factory=RealDictCursor) as cur:
            cur.execute(
                f'''
                SELECT c.*
                FROM {TABLE_COMBINATIONS} c
                JOIN {TABLE_ASSIGNMENTS} a ON a.combo_id = c.id
                JOIN {TABLE_ANNOTATORS} an ON an.code = %s
                WHERE (a.annotator_a_code = %s OR a.annotator_b_code = %s)
                  AND NOT EXISTS (
                    SELECT 1 FROM {TABLE_ANNOTATIONS} ann
                    WHERE ann.combo_id = c.id AND ann.annotator_id = an.id
                  )
                ORDER BY c.seed_id, c.matchup_key, c.level DESC, c.id
                LIMIT %s
                ''',
                (annotator_code, annotator_code, annotator_code, limit),
            )
            return [dict(r) for r in cur.fetchall()]


def count_remaining(annotator_code: int) -> int:
    with get_conn() as conn:
        with conn.cursor() as cur:
            cur.execute(
                f'''
                SELECT COUNT(*)
                FROM {TABLE_COMBINATIONS} c
                JOIN {TABLE_ASSIGNMENTS} a ON a.combo_id = c.id
                JOIN {TABLE_ANNOTATORS} an ON an.code = %s
                WHERE (a.annotator_a_code = %s OR a.annotator_b_code = %s)
                  AND NOT EXISTS (
                    SELECT 1 FROM {TABLE_ANNOTATIONS} ann
                    WHERE ann.combo_id = c.id AND ann.annotator_id = an.id
                  )
                ''',
                (annotator_code, annotator_code, annotator_code),
            )
            return int(cur.fetchone()[0])


def save_annotation_by_code(data: dict) -> None:
    annotator_id = get_annotator_id_by_code(int(data['annotator_code']))
    payload = {
        'annotator_id': annotator_id,
        'combo_id': data['combo_id'],
        'model_a': data['model_a'],
        'model_b': data['model_b'],
        'a_existence': data['a_existence'],
        'a_appearance': data['a_appearance'],
        'a_interaction': data['a_interaction'],
        'b_existence': data['b_existence'],
        'b_appearance': data['b_appearance'],
        'b_interaction': data['b_interaction'],
        'preference': data['preference'],
        'prompt_ilogical': bool(data.get('prompt_ilogical', False)),
    }
    with get_conn() as conn:
        with conn.cursor() as cur:
            cur.execute(
                f'''
                INSERT INTO {TABLE_ANNOTATIONS}
                (annotator_id, combo_id, model_a, model_b,
                 a_existence, a_appearance, a_interaction,
                 b_existence, b_appearance, b_interaction,
                 preference, prompt_ilogical)
                VALUES
                (%(annotator_id)s, %(combo_id)s, %(model_a)s, %(model_b)s,
                 %(a_existence)s, %(a_appearance)s, %(a_interaction)s,
                 %(b_existence)s, %(b_appearance)s, %(b_interaction)s,
                 %(preference)s, %(prompt_ilogical)s)
                ON CONFLICT (annotator_id, combo_id) DO UPDATE SET
                    model_a = EXCLUDED.model_a,
                    model_b = EXCLUDED.model_b,
                    a_existence = EXCLUDED.a_existence,
                    a_appearance = EXCLUDED.a_appearance,
                    a_interaction = EXCLUDED.a_interaction,
                    b_existence = EXCLUDED.b_existence,
                    b_appearance = EXCLUDED.b_appearance,
                    b_interaction = EXCLUDED.b_interaction,
                    preference = EXCLUDED.preference,
                    prompt_ilogical = EXCLUDED.prompt_ilogical
                ''',
                payload,
            )


def get_progress(annotator_code: int) -> dict:
    annotator_id = get_annotator_id_by_code(annotator_code)
    with get_conn() as conn:
        with conn.cursor() as cur:
            cur.execute(f'SELECT COUNT(*) FROM {TABLE_ANNOTATIONS} WHERE annotator_id = %s', (annotator_id,))
            completed = int(cur.fetchone()[0])
            cur.execute(
                f'''
                SELECT COUNT(*)
                FROM {TABLE_COMBINATIONS} c
                JOIN {TABLE_ASSIGNMENTS} a ON a.combo_id = c.id
                WHERE (a.annotator_a_code = %s OR a.annotator_b_code = %s)
                ''',
                (annotator_code, annotator_code),
            )
            quota = int(cur.fetchone()[0])
            cur.execute(
                f'''
                SELECT COUNT(*) FROM (
                    SELECT combo_id
                    FROM {TABLE_ANNOTATIONS}
                    GROUP BY combo_id
                    HAVING COUNT(DISTINCT annotator_id) >= 2
                ) s
                '''
            )
            overall_completed = int(cur.fetchone()[0])
            cur.execute(f'SELECT COUNT(*) FROM {TABLE_COMBINATIONS}')
            overall_total = int(cur.fetchone()[0])
            return {
                'completed': completed,
                'quota': quota,
                'overall_completed': overall_completed,
                'overall_total': overall_total,
            }


def get_progress_board() -> dict:
    board = []
    overall_completed = 0
    overall_total = 0
    for code in range(1, 9):
        p = get_progress(code)
        board.append({'code': code, 'completed': p['completed'], 'quota': p['quota']})
        overall_completed = p['overall_completed']
        overall_total = p['overall_total']
    return {'annotators': board, 'overall_completed': overall_completed, 'overall_total': overall_total}


def upsert_combinations(combos: list[dict]) -> None:
    with get_conn() as conn:
        with conn.cursor() as cur:
            for c in combos:
                cur.execute(
                    f'''
                    INSERT INTO {TABLE_COMBINATIONS}
                    (id, base_id, seed_id, level, class_tag, ratio_type, matchup_key, model_left, model_right,
                     n_total, n_humans, n_objects, prompt, prompt_zh, humans, objects)
                    VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
                    ON CONFLICT (id) DO UPDATE SET
                        base_id = EXCLUDED.base_id,
                        seed_id = EXCLUDED.seed_id,
                        level = EXCLUDED.level,
                        class_tag = EXCLUDED.class_tag,
                        ratio_type = EXCLUDED.ratio_type,
                        matchup_key = EXCLUDED.matchup_key,
                        model_left = EXCLUDED.model_left,
                        model_right = EXCLUDED.model_right,
                        n_total = EXCLUDED.n_total,
                        n_humans = EXCLUDED.n_humans,
                        n_objects = EXCLUDED.n_objects,
                        prompt = EXCLUDED.prompt,
                        prompt_zh = EXCLUDED.prompt_zh,
                        humans = EXCLUDED.humans,
                        objects = EXCLUDED.objects
                    ''',
                    (
                        c['id'], c['base_id'], c.get('seed_id', ''), int(c.get('level', 0)), c.get('class_tag', ''),
                        c.get('ratio_type', ''), c['matchup_key'], c['model_left'], c['model_right'],
                        int(c.get('n_total', 0)), int(c.get('n_humans', 0)), int(c.get('n_objects', 0)),
                        c.get('prompt', ''), c.get('prompt_zh', ''), json.dumps(c.get('humans', [])), json.dumps(c.get('objects', []))
                    ),
                )


def upsert_assignments(assignments: list[dict]) -> None:
    with get_conn() as conn:
        with conn.cursor() as cur:
            for a in assignments:
                cur.execute(
                    f'''
                    INSERT INTO {TABLE_ASSIGNMENTS}
                    (combo_id, group_key, annotator_a_code, annotator_b_code)
                    VALUES (%s, %s, %s, %s)
                    ON CONFLICT (combo_id) DO UPDATE SET
                        group_key = EXCLUDED.group_key,
                        annotator_a_code = EXCLUDED.annotator_a_code,
                        annotator_b_code = EXCLUDED.annotator_b_code
                    ''',
                    (a['combo_id'], a['group_key'], int(a['annotator_a_code']), int(a['annotator_b_code'])),
                )