db.py

import os
from datetime import datetime
from typing import Dict, Any, List

from sqlalchemy import create_engine, text
from sqlalchemy.engine import Connection

from config import logger, LOG_LEVEL


def get_db_connection() -> Connection:
    """Establishes a connection to the PostgreSQL database using SQLAlchemy."""
    DB_URL = 'postgresql+psycopg2://{user}:{password}@{host}:{port}/{name}'.format(
        user=os.environ.get("DB_USER"),
        password=os.environ.get("DB_PASSWORD"),
        host=os.environ.get("DB_HOST"),
        port=os.environ.get("DB_PORT", "5432"),
        name=os.environ.get("DB_NAME")
    )

    # TCP keepalives stop the connection being dropped as idle while a long
    # query runs and no data is flowing (e.g. by NAT on the CI runner).
    connect_args = {
        "keepalives": 1,
        "keepalives_idle": 30,
        "keepalives_interval": 10,
        "keepalives_count": 5,
    }

    try:
        engine = create_engine(DB_URL, connect_args=connect_args)
        conn = engine.connect()
        return conn
    except Exception as e:
        logger.error(f"Could not connect to the database using SQLAlchemy: {e}")
        raise


def fetch_data(conn: Connection, sql_limit: int, eval_function_name: str, grade_params_json: str, seed: float, excluded_ids: List[str] = [], excluded_grade_param_values: Dict[str, List[str]] = None, included_ids: List[str] = None) -> List[Dict[str, Any]]:
    """Fetches data using the provided complex query with SQLAlchemy.
    When included_ids is set, only those submissions are fetched and the
    createdAt cut-off is skipped, so any listed submission can be replayed."""
    limit = max(1, sql_limit)

    where_clauses = ["EF.name = :name_param"]
    query_params = {
        "name_param": eval_function_name,
        "limit_param": limit,
    }

    if included_ids:
        cast_placeholders = ", ".join(f"CAST(:incl_{i} AS UUID)" for i in range(len(included_ids)))
        where_clauses.append(f"S.id IN ({cast_placeholders})")
        for i, inc_id in enumerate(included_ids):
            query_params[f"incl_{i}"] = inc_id
    else:
        where_clauses.append("S.\"createdAt\" >= :since_date_param")
        query_params["since_date_param"] = datetime(datetime.now().year - 1, 6, 1)

    if grade_params_json:
        where_clauses.append("RA.\"gradeParams\"::jsonb = (:params_param)::jsonb")
        query_params["params_param"] = grade_params_json

    if excluded_ids:
        cast_placeholders = ", ".join(f"CAST(:excl_{i} AS UUID)" for i in range(len(excluded_ids)))
        where_clauses.append(f"S.id NOT IN ({cast_placeholders})")
        for i, exc_id in enumerate(excluded_ids):
            query_params[f"excl_{i}"] = exc_id

    if excluded_grade_param_values:
        for i, (key, values) in enumerate(excluded_grade_param_values.items()):
            key_p = f"gpv_key_{i}"
            val_ps = [f"gpv_val_{i}_{j}" for j in range(len(values))]
            placeholders = ", ".join(f":{p}" for p in val_ps)
            where_clauses.append(
                f'(RA."gradeParams"->>(:{key_p}) IS NULL OR RA."gradeParams"->>(:{key_p}) NOT IN ({placeholders}))'
            )
            query_params[key_p] = key
            for p, val in zip(val_ps, values):
                query_params[p] = val

    where_sql = " AND ".join(where_clauses)

    debug_column = ', S."rawResponse"::json as raw_response' if LOG_LEVEL == 'DEBUG' else ''

    # Sample the submission ids first, so the JSON building below only runs on
    # the sampled rows rather than on every matching submission.
    sql_query_template = f"""
            WITH sampled AS (
                SELECT S.id, RA.id AS response_area_id, RANDOM() AS sort_key
                FROM "Submission" S
                    INNER JOIN "ResponseArea" RA ON S."responseAreaId" = RA.id
                    INNER JOIN "EvaluationFunction" EF ON RA."evaluationFunctionId" = EF.id
                    INNER JOIN "Part" P ON RA."partId" = P.id
                    INNER JOIN "QuestionVersion" QV ON P."questionVersionId" = QV.id
                    INNER JOIN "Question" Q ON QV."questionId" = Q.id AND QV.id = Q."publishedVersionId"
                WHERE
                    {where_sql}
                ORDER BY sort_key
                LIMIT :limit_param
            )
            SELECT
                S.id as submission_id, S.submission, S.answer::jsonb, S.grade::int::boolean, S.feedback,
                RA."gradeParams"::json as grade_params,
                (
                    SELECT json_agg(json_build_object(
                        'answer',     RAC.answer::jsonb,
                        'params',     RAC.params,
                        'feedback',   RAC.feedback,
                        'mark',       RAC.mark::int
                    ))
                    FROM "ResponseAreaCase" RAC
                    WHERE RAC."responseAreaId" = RA.id
                ) AS cases,
                (
                    SELECT json_object_agg(ISYM.code, json_build_object(
                        'latex',   ISYM.symbol,
                        'aliases', ISYM.aliases
                    ))
                    FROM "InputSymbol" ISYM
                    WHERE ISYM."responseAreaId" = RA.id AND ISYM.code IS NOT NULL
                ) AS symbols,
                S."rawResponse"::json -> 'error' ->> 'message' AS historical_error_message,
                S."rawResponse"::json -> 'error' ->> 'detail'  AS historical_error_detail
                {debug_column}
            FROM sampled
                INNER JOIN "Submission" S ON S.id = sampled.id
                INNER JOIN "ResponseArea" RA ON RA.id = sampled.response_area_id
            ORDER BY sampled.sort_key;
        """

    try:
        conn.execute(text("SELECT setseed(:seed_param)"), {"seed_param": seed})
        result = conn.execute(text(sql_query_template), query_params)
        data_records = [dict(row) for row in result.mappings()]

    except Exception as e:
        logger.error(f"Error fetching data with query: {e}")
        raise

    logger.info(f"Successfully fetched {len(data_records)} records.")
    if included_ids and len(data_records) < len(included_ids):
        logger.warning(f"{len(included_ids) - len(data_records)} of {len(included_ids)} requested submission ID(s) "
                       f"were not found for '{eval_function_name}' (or were filtered out by gradeParams).")
    return data_records