atelier/sidecar/db/generations/queries.ts

import { and, asc, desc, eq, inArray } from 'drizzle-orm';
import { db } from '../db';
import { generations } from './schema';
import type { GenerationStatus } from './schema';

export type Generation = typeof generations.$inferSelect;
export type NewGeneration = typeof generations.$inferInsert;

export async function createGeneration(values: NewGeneration): Promise<Generation> {
  const [row] = await db.insert(generations).values(values).returning();
  return row!;
}

export async function getGeneration(id: string): Promise<Generation | undefined> {
  const [row] = await db.select().from(generations).where(eq(generations.id, id)).limit(1);
  return row;
}

export async function listGenerations(limit = 50): Promise<Generation[]> {
  return db.select().from(generations).orderBy(desc(generations.createdAt)).limit(limit);
}

/**
 * Delete generations outright, returning the rows that went — the caller needs `image_path` to unlink
 * the files, and after the delete there is nowhere left to read it from.
 *
 * `collection_items` rows go with them by foreign key, so an image deleted while filed in three
 * collections leaves none of the three holding a dangling membership.
 */
export async function deleteGenerations(ids: string[]): Promise<Generation[]> {
  if (ids.length === 0) return [];
  return db.delete(generations).where(inArray(generations.id, ids)).returning();
}

/** Every image of one request, in the order it was asked for. This is what polling a job reads. */
export async function getRequest(requestId: string): Promise<Generation[]> {
  return db.select().from(generations).where(eq(generations.requestId, requestId)).orderBy(asc(generations.index));
}

export async function createRequest(rows: NewGeneration[]): Promise<Generation[]> {
  return db.insert(generations).values(rows).returning();
}

/** The backlog, oldest first — what the queue picks up on a tick and after a restart. */
export async function claimableGenerations(limit: number): Promise<Generation[]> {
  return db
    .select()
    .from(generations)
    .where(eq(generations.status, 'queued'))
    .orderBy(asc(generations.createdAt))
    .limit(limit);
}

/**
 * Move a row to `running`, but only from `queued`.
 *
 * The `and(... eq(status, 'queued'))` is the whole point: two workers reaching for the same row both
 * update, and only the one that got there first matches. The loser sees no returned row and moves on,
 * so a generation is never paid for twice.
 */
export async function markRunning(id: string): Promise<Generation | undefined> {
  const [row] = await db
    .update(generations)
    .set({ status: 'running', startedAt: new Date() })
    .where(and(eq(generations.id, id), eq(generations.status, 'queued')))
    .returning();
  return row;
}

export async function markSucceeded(
  id: string,
  values: {
    imagePath: string;
    width?: number | null;
    height?: number | null;
    hasAlpha?: boolean | null;
    opaque?: boolean | null;
    costUsd?: string | null;
    usage?: Record<string, unknown> | null;
    result?: Record<string, unknown> | null;
    durationMs: number;
  },
): Promise<void> {
  await db
    .update(generations)
    .set({ ...values, status: 'succeeded', finishedAt: new Date() })
    .where(eq(generations.id, id));
}

export async function markFailed(id: string, error: string, durationMs?: number): Promise<void> {
  await db
    .update(generations)
    .set({ status: 'failed', error, durationMs: durationMs ?? null, finishedAt: new Date() })
    .where(eq(generations.id, id));
}

/**
 * Fail whatever was mid-flight when the process died.
 *
 * A row left `running` has no worker behind it — the only thing that would have finished it went away
 * with the last process. Without this it stays running forever, and a caller polling it waits on
 * something that will never answer.
 */
export async function failInterrupted(): Promise<number> {
  const rows = await db
    .update(generations)
    .set({ status: 'failed', error: 'Interrupted by a server restart', finishedAt: new Date() })
    .where(inArray(generations.status, ['running' satisfies GenerationStatus]))
    .returning({ id: generations.id });
  return rows.length;
}