src/database.py

"""Database class for the evaluation."""

import json
import os
import subprocess
import threading
from contextlib import contextmanager
from datetime import datetime, timezone
from typing import Optional

import psycopg2
import psycopg2.extras
from huggingface_hub import CommitOperationAdd, HfApi

REQUIRED_TABLES = ('tasks', 'assignments', 'task_config', 'answers', 'allocation_history')


def _get_time_now() -> str:
    return datetime.now(timezone.utc).isoformat()


class DB:
    """Database for the evaluation."""

    def __init__(self, db_path: str, schema_path: Optional[str] = None, verify_only: bool = True):
        """Initialize the database.

        Note: db_path/schema_path kept for compatibility; Postgres uses env vars.
        Required env vars: PGHOST, PGPORT, PGUSER, PGPASSWORD, PGDATABASE
        """
        self.db_path = db_path  # kept for compatibility
        self._lock = threading.Lock()

        # Validate environment variables early
        self._pg_dsn = dict(
            host=os.environ.get('PGHOST'),
            port=int(os.environ.get('PGPORT')) if os.environ.get('PGPORT') else None,
            user=os.environ.get('PGUSER'),
            password=os.environ.get('PGPASSWORD'),
            database=os.environ.get('PGDATABASE'),
        )
        if not all([self._pg_dsn['host'], self._pg_dsn['user'], self._pg_dsn['password'], self._pg_dsn['database']]):
            raise OSError('Missing PostgreSQL env vars: PGHOST, PGUSER, PGPASSWORD, PGDATABASE')

        # Verify required tables exist
        if verify_only:
            with self._connect() as conn:
                missing = []
                for t in REQUIRED_TABLES:
                    cur = conn.execute('SELECT to_regclass(%s)', (f'public.{t}',))
                    exists = cur.fetchone()[0]
                    if exists is None:
                        missing.append(t)
                if missing:
                    raise FileNotFoundError(f'Missing required tables: {", ".join(missing)}')

    @contextmanager
    def _connect(self):
        """Context manager yielding a lightweight wrapper with conn.execute(...).

        - Uses psycopg2 with autocommit by default (emulates sqlite autocommit usage).
        - Supports BEGIN/COMMIT/ROLLBACK issued via execute for backward compatibility.
        """

        class _CursorWrapper:
            def __init__(self, cursor):
                self._cursor = cursor

            def fetchone(self):
                return self._cursor.fetchone()

            def fetchall(self):
                return self._cursor.fetchall()

            @property
            def rowcount(self):
                return self._cursor.rowcount

        class _ConnWrapper:
            def __init__(self, real_conn):
                self._conn = real_conn

            def execute(self, sql: str, params: tuple | list | None = None):
                sql_strip = (sql or '').strip().upper()
                # Map transaction control statements
                if sql_strip.startswith('BEGIN'):
                    self._conn.autocommit = False
                    cur = self._conn.cursor()
                    try:
                        cur.execute('BEGIN;')
                    finally:
                        cur.close()
                    # Return a dummy cursor wrapper

                    class _DummyCursor:
                        def fetchone(self):
                            return None

                        def fetchall(self):
                            return []

                        @property
                        def rowcount(self):
                            return 0

                    return _DummyCursor()

                if sql_strip == 'COMMIT':
                    try:
                        self._conn.commit()
                    finally:
                        self._conn.autocommit = True

                    class _DummyCursor:
                        def fetchone(self):
                            return None

                        def fetchall(self):
                            return []

                        @property
                        def rowcount(self):
                            return 0

                    return _DummyCursor()

                if sql_strip == 'ROLLBACK':
                    try:
                        self._conn.rollback()
                    finally:
                        self._conn.autocommit = True

                    class _DummyCursor:
                        def fetchone(self):
                            return None

                        def fetchall(self):
                            return []

                        @property
                        def rowcount(self):
                            return 0

                    return _DummyCursor()

                # Regular statement
                with self._conn.cursor() as c:
                    c.execute(sql, params or ())
                    # For SELECTs, we need a cursor able to fetch after context ends.
                    # Use named cursor copy via RealDictCursor? Simpler: re-execute on a new cursor kept open.
                    # Instead, materialize results if it's a query returning rows.
                    if c.description is not None:
                        rows = c.fetchall()

                        # Create a lightweight object mimicking cursor
                        class _Mat:
                            def __init__(self, rows):
                                self._rows = rows

                            def fetchone(self):
                                return self._rows[0] if self._rows else None

                            def fetchall(self):
                                return list(self._rows)

                            @property
                            def rowcount(self):
                                return len(self._rows)

                        return _Mat(rows)
                    else:
                        # Non-SELECT: return wrapper with rowcount captured
                        rc = c.rowcount

                        class _Affect:
                            def __init__(self, rc):
                                self._rc = rc

                            def fetchone(self):
                                return None

                            def fetchall(self):
                                return []

                            @property
                            def rowcount(self):
                                return self._rc

                        return _Affect(rc)

            def __enter__(self):
                return self

            def __exit__(self, exc_type, exc, tb):
                # If a transaction is open and no explicit commit/rollback occurred, rollback on error
                if not self._conn.autocommit:
                    if exc is None:
                        self._conn.commit()
                    else:
                        self._conn.rollback()
                self._conn.close()

        # Open a new psycopg2 connection per context
        conn = psycopg2.connect(**{k: v for k, v in self._pg_dsn.items() if v is not None})
        conn.autocommit = True
        wrapper = _ConnWrapper(conn)
        try:
            yield wrapper
        finally:
            # _ConnWrapper.__exit__ handles close when used with "with". If not, ensure close here.
            try:
                if not conn.closed:
                    conn.close()
            except Exception:
                pass

    # ========== Answer Recording ==========

    def record_answer(
        self,
        round_id: str,
        user_id: str,
        item_id: str,
        label: str,
        image_path: str,
        score: int,
        words_not_present: list[str],
    ) -> None:
        """Record an answer into the database."""
        now = _get_time_now()
        with self._connect() as conn, self._lock:
            conn.execute(
                (
                    'INSERT INTO answers (round_id, user_id, item_id, label, image_path, '
                    'score, words_not_present, answered_at) '
                    'VALUES (%s, %s, %s, %s, %s, %s, %s, %s) '
                    'ON CONFLICT (round_id, user_id, item_id) DO UPDATE SET '
                    'score = EXCLUDED.score, '
                    'words_not_present = EXCLUDED.words_not_present, '
                    'answered_at = EXCLUDED.answered_at'
                ),
                (
                    round_id,
                    user_id,
                    item_id,
                    label,
                    image_path,
                    score,
                    json.dumps(words_not_present, ensure_ascii=False),
                    now,
                ),
            )

    def get_answered_item_ids(self, round_id: str, user_id: str) -> set[str]:
        """Get all the answered item ids of a user in a round."""
        with self._connect() as conn:
            cur = conn.execute(
                'SELECT item_id FROM answers WHERE round_id=%s AND user_id=%s',
                (round_id, user_id),
            )
            return {r[0] for r in cur.fetchall()}  # fetchall returns a list of tuples

    def get_answer(self, round_id: str, user_id: str, item_id: str) -> Optional[tuple[int, list[str]]]:
        """Get an answer of a user in a round for a specific item.

        Returns:
            tuple[int, list[str]]: score (int) and list of words not present in the image.
            None if the answer does not exist.

        """
        with self._connect() as conn:
            cur = conn.execute(
                'SELECT score, words_not_present FROM answers WHERE round_id=%s AND user_id=%s AND item_id=%s',
                (round_id, user_id, item_id),
            )
            row = cur.fetchone()
            if not row:
                return None

            score, words_not_present = row
            return int(score), json.loads(words_not_present)

    # ========== Lease Management (for in_progress tasks) ==========

    def cleanup_expired_leases(self, round_id: str) -> int:
        """Auto-recover expired in_progress tasks to pending status.

        Returns the number of tasks recovered.
        """
        now = _get_time_now()
        with self._connect() as conn, self._lock:
            cur = conn.execute(
                """
                UPDATE assignments
                SET status = 'pending', lease_until = NULL, started_at = NULL
                WHERE round_id = %s
                  AND status = 'in_progress'
                  AND lease_until < %s
                """,
                (round_id, now),
            )
            return cur.rowcount or 0

    # ========== Auto Commit (For HF Spaces) ==========
    def commit_and_push_db(self) -> None:
        """Commit and push the database to the repository."""
        space_id = os.getenv('SPACE_ID', None)
        if not space_id:
            return
        try:
            repo_root = subprocess.run(
                ['git', 'rev-parse', '--show-toplevel'],
                check=True,
                capture_output=True,
                text=True,
            ).stdout.strip()

            ts = _get_time_now()
            db_rel = os.path.relpath(self.db_path, repo_root)
            subprocess.run(['git', 'add', db_rel], check=True, cwd=repo_root)
            api = HfApi(token=os.getenv('HF_TOKEN', None))
            api.create_commit(
                repo_id='JaneDing2025/IconEval_Pilot',
                repo_type='dataset',
                operations=[CommitOperationAdd(path_in_repo=db_rel, path_or_fileobj=self.db_path)],
                commit_message=f'update db {ts}',
            )
        except Exception as e:
            raise RuntimeError(f'Failed to commit and push the database: {e}')