apps/api/src/queue/postgres-store.ts

import { randomUUID } from "node:crypto";
import pg from "pg";

import {
  type DetectedMode,
  type ItemRow,
  type ItemStatus,
  type JobRow,
  type JobStatus,
  type SetItemStatusMeta,
  type Store,
  nowMs,
  truncateError,
} from "./store.js";

const { Pool } = pg;

export const POSTGRES_TAKE_NEXT_ITEM_SQL = `
WITH next_item AS (
  SELECT id
  FROM items
  WHERE job_id = $1 AND status = 'queued'
  ORDER BY ord ASC
  LIMIT 1
  FOR UPDATE SKIP LOCKED
)
UPDATE items
SET status = 'running',
    attempt_count = attempt_count + 1,
    error_preview = NULL,
    last_stage = 'running',
    updated_at = $2
FROM next_item
WHERE items.id = next_item.id
RETURNING items.*
`;

const SCHEMA_SQL = `
CREATE TABLE IF NOT EXISTS jobs (
  id TEXT PRIMARY KEY,
  status TEXT NOT NULL,
  options_snapshot TEXT NOT NULL,
  created_at BIGINT NOT NULL,
  updated_at BIGINT NOT NULL
);

CREATE TABLE IF NOT EXISTS items (
  id TEXT PRIMARY KEY,
  job_id TEXT NOT NULL REFERENCES jobs(id) ON DELETE CASCADE,
  ord INTEGER NOT NULL,
  filename TEXT NOT NULL,
  upload_path TEXT NOT NULL,
  detected_mode TEXT NOT NULL,
  status TEXT NOT NULL,
  evidence_path TEXT,
  task_package_path TEXT,
  judge_result_path TEXT,
  export_zip_path TEXT,
  error_preview TEXT,
  error_details_path TEXT,
  attempt_count INTEGER NOT NULL DEFAULT 0,
  last_stage TEXT,
  created_at BIGINT NOT NULL,
  updated_at BIGINT NOT NULL
);

ALTER TABLE items ADD COLUMN IF NOT EXISTS error_details_path TEXT;
ALTER TABLE items ADD COLUMN IF NOT EXISTS attempt_count INTEGER NOT NULL DEFAULT 0;
ALTER TABLE items ADD COLUMN IF NOT EXISTS last_stage TEXT;

CREATE INDEX IF NOT EXISTS idx_items_job_ord ON items(job_id, ord);
CREATE INDEX IF NOT EXISTS idx_items_status ON items(status);
CREATE INDEX IF NOT EXISTS idx_items_job_status_ord ON items(job_id, status, ord);
`;

function mapJob(row: any): JobRow {
  return {
    ...row,
    created_at: Number(row.created_at),
    updated_at: Number(row.updated_at),
  };
}

function mapItem(row: any): ItemRow {
  return {
    ...row,
    attempt_count: Number(row.attempt_count || 0),
    created_at: Number(row.created_at),
    updated_at: Number(row.updated_at),
    evidence_path: row.evidence_path ?? null,
    task_package_path: row.task_package_path ?? null,
    judge_result_path: row.judge_result_path ?? null,
    export_zip_path: row.export_zip_path ?? null,
    error_preview: row.error_preview ?? null,
    error_details_path: row.error_details_path ?? null,
    last_stage: row.last_stage ?? null,
  };
}

export class PostgresStore implements Store {
  private pool: pg.Pool;

  constructor(databaseUrl: string) {
    if (!databaseUrl) throw new Error("DATABASE_URL is required for Postgres store");
    this.pool = new Pool({ connectionString: databaseUrl });
  }

  async init(): Promise<void> {
    await this.pool.query(SCHEMA_SQL);
  }

  async close(): Promise<void> {
    await this.pool.end();
  }

  async createJob(optionsSnapshot: object): Promise<JobRow> {
    const id = randomUUID();
    const t = nowMs();
    const result = await this.pool.query(
      `INSERT INTO jobs (id, status, options_snapshot, created_at, updated_at)
       VALUES ($1, $2, $3, $4, $5)
       RETURNING *`,
      [id, "pending", JSON.stringify(optionsSnapshot), t, t],
    );
    return mapJob(result.rows[0]);
  }

  async getJob(id: string): Promise<JobRow | null> {
    const result = await this.pool.query("SELECT * FROM jobs WHERE id = $1", [id]);
    return result.rows[0] ? mapJob(result.rows[0]) : null;
  }

  async setJobStatus(id: string, status: JobStatus): Promise<void> {
    await this.pool.query("UPDATE jobs SET status = $1, updated_at = $2 WHERE id = $3", [
      status,
      nowMs(),
      id,
    ]);
  }

  async updateJobOptions(id: string, optionsSnapshot: object): Promise<void> {
    await this.pool.query(
      "UPDATE jobs SET options_snapshot = $1, updated_at = $2 WHERE id = $3",
      [JSON.stringify(optionsSnapshot), nowMs(), id],
    );
  }

  async listResumableJobs(): Promise<JobRow[]> {
    const result = await this.pool.query(
      "SELECT * FROM jobs WHERE status = 'running' ORDER BY created_at",
    );
    return result.rows.map(mapJob);
  }

  async recoverInterruptedJob(id: string): Promise<void> {
    const t = nowMs();
    const client = await this.pool.connect();
    try {
      await client.query("BEGIN");
      await client.query(
        `UPDATE items
         SET status = 'queued',
             error_preview = NULL,
             last_stage = 'recovered',
             updated_at = $1
         WHERE job_id = $2 AND status IN ('running', 'judging')`,
        [t, id],
      );
      await client.query("UPDATE jobs SET status = 'running', updated_at = $1 WHERE id = $2", [
        t,
        id,
      ]);
      await client.query("COMMIT");
    } catch (err) {
      await client.query("ROLLBACK");
      throw err;
    } finally {
      client.release();
    }
  }

  async createItem(input: {
    jobId: string;
    ord: number;
    filename: string;
    uploadPath: string;
    detectedMode: DetectedMode;
  }): Promise<ItemRow> {
    const id = randomUUID();
    const t = nowMs();
    const result = await this.pool.query(
      `INSERT INTO items
       (id, job_id, ord, filename, upload_path, detected_mode, status,
        created_at, updated_at)
       VALUES ($1, $2, $3, $4, $5, $6, 'queued', $7, $8)
       RETURNING *`,
      [id, input.jobId, input.ord, input.filename, input.uploadPath, input.detectedMode, t, t],
    );
    return mapItem(result.rows[0]);
  }

  async getItem(id: string): Promise<ItemRow | null> {
    const result = await this.pool.query("SELECT * FROM items WHERE id = $1", [id]);
    return result.rows[0] ? mapItem(result.rows[0]) : null;
  }

  async listItemsByJob(jobId: string): Promise<ItemRow[]> {
    const result = await this.pool.query(
      "SELECT * FROM items WHERE job_id = $1 ORDER BY ord ASC",
      [jobId],
    );
    return result.rows.map(mapItem);
  }

  async setItemStatus(
    id: string,
    status: ItemStatus,
    errorPreview?: string,
    meta: SetItemStatusMeta = {},
  ): Promise<void> {
    await this.pool.query(
      `UPDATE items
       SET status = $1,
           error_preview = $2,
           last_stage = COALESCE($3, last_stage),
           error_details_path = COALESCE($4, error_details_path),
           updated_at = $5
       WHERE id = $6`,
      [
        status,
        truncateError(errorPreview),
        meta.lastStage ?? null,
        meta.errorDetailsPath ?? null,
        nowMs(),
        id,
      ],
    );
  }

  async setItemPaths(
    id: string,
    paths: Partial<{
      evidencePath: string;
      taskPackagePath: string;
      judgeResultPath: string;
      exportZipPath: string;
    }>,
  ): Promise<void> {
    const fields: string[] = [];
    const values: unknown[] = [];
    const add = (column: string, value: unknown) => {
      values.push(value);
      fields.push(`${column} = $${values.length}`);
    };
    if (paths.evidencePath !== undefined) add("evidence_path", paths.evidencePath);
    if (paths.taskPackagePath !== undefined) add("task_package_path", paths.taskPackagePath);
    if (paths.judgeResultPath !== undefined) add("judge_result_path", paths.judgeResultPath);
    if (paths.exportZipPath !== undefined) add("export_zip_path", paths.exportZipPath);
    if (!fields.length) return;
    values.push(nowMs());
    fields.push(`updated_at = $${values.length}`);
    values.push(id);
    await this.pool.query(`UPDATE items SET ${fields.join(", ")} WHERE id = $${values.length}`, values);
  }

  async takeNextQueuedItem(jobId: string): Promise<ItemRow | null> {
    const client = await this.pool.connect();
    try {
      await client.query("BEGIN");
      const result = await client.query(POSTGRES_TAKE_NEXT_ITEM_SQL, [jobId, nowMs()]);
      await client.query("COMMIT");
      return result.rows[0] ? mapItem(result.rows[0]) : null;
    } catch (err) {
      await client.query("ROLLBACK");
      throw err;
    } finally {
      client.release();
    }
  }

  async resetItemForRetry(id: string): Promise<void> {
    await this.pool.query(
      `UPDATE items
       SET status = 'queued',
           error_preview = NULL,
           error_details_path = NULL,
           attempt_count = 0,
           last_stage = NULL,
           evidence_path = NULL,
           task_package_path = NULL,
           judge_result_path = NULL,
           export_zip_path = NULL,
           updated_at = $1
       WHERE id = $2`,
      [nowMs(), id],
    );
  }
}