// SPDX-License-Identifier: AGPL-3.0-or-later import {BatchBuilder, deleteOneOrMany, fetchMany, fetchOne, upsertOne} from '../database/CassandraQueryExecution'; import {Db} from '../database/CassandraTypes'; import type {JobActiveRow, JobByDayBucketRow, JobByIdRow, JobStatus} from '../database/types/JobLedgerTypes'; import {JobsActive, JobsByDayBucket, JobsById} from '../Tables'; import { type CreateJobInput, IJobLedgerRepository, type ListJobsCursor, type ListJobsFilters, type ListJobsResult, } from './IJobLedgerRepository'; const FETCH_JOB_BY_ID_QUERY = JobsById.select({ where: JobsById.where.eq('job_id'), }); const FETCH_CANCEL_REQUESTED_QUERY = JobsById.select({ where: JobsById.where.eq('job_id'), }); const ACTIVE_JOBS_QUERY = JobsActive.select(); function bucketDayFor(d: Date): string { return d.toISOString().slice(0, 10); } export class JobLedgerRepository extends IJobLedgerRepository { async createJob(input: CreateJobInput): Promise { const now = new Date(); const status: JobStatus = 'queued'; const idRow: JobByIdRow = { job_id: input.jobId, task_type: input.taskType, status, progress_current: null, progress_total: null, progress_message: null, payload: JSON.stringify(input.payload), result: null, error_message: null, created_at: now, started_at: null, completed_at: null, requested_by_user_id: input.requestedByUserId, audit_log_reason: input.auditLogReason, jet_stream_seq: input.jetStreamSeq, jet_stream_lane: input.jetStreamLane, attempts: 0, max_attempts: input.maxAttempts, run_at: input.runAt, cancel_requested: false, context_link: null, }; const bucketRow: JobByDayBucketRow = { bucket_day: bucketDayFor(now), created_at: now, job_id: input.jobId, task_type: input.taskType, status, requested_by_user_id: input.requestedByUserId, }; const activeRow: JobActiveRow = { job_id: input.jobId, task_type: input.taskType, status, requested_by_user_id: input.requestedByUserId, created_at: now, started_at: null, }; const batch = new BatchBuilder(); batch.addPrepared(JobsById.insert(idRow)); batch.addPrepared(JobsByDayBucket.insert(bucketRow)); batch.addPrepared(JobsActive.insert(activeRow)); await batch.executeChunked(10, true); } async getJob(jobId: bigint): Promise { return fetchOne(FETCH_JOB_BY_ID_QUERY.bind({job_id: jobId})); } async markRunning(jobId: bigint, lane: string): Promise { const startedAt = new Date(); const status: JobStatus = 'running'; await upsertOne( JobsById.patchByPk( {job_id: jobId}, {status: Db.set(status), started_at: Db.set(startedAt), jet_stream_lane: Db.set(lane)}, ), ); await upsertOne(JobsActive.patchByPk({job_id: jobId}, {status: Db.set(status), started_at: Db.set(startedAt)})); } async markSucceeded(jobId: bigint, result: Record | null): Promise { const completedAt = new Date(); const status: JobStatus = 'succeeded'; await upsertOne( JobsById.patchByPk( {job_id: jobId}, { status: Db.set(status), completed_at: Db.set(completedAt), result: result === null ? Db.clear() : Db.set(JSON.stringify(result)), }, ), ); await deleteOneOrMany(JobsActive.deleteByPk({job_id: jobId})); } async markFailed(jobId: bigint, errorMessage: string): Promise { const completedAt = new Date(); const status: JobStatus = 'failed'; await upsertOne( JobsById.patchByPk( {job_id: jobId}, {status: Db.set(status), completed_at: Db.set(completedAt), error_message: Db.set(errorMessage)}, ), ); await deleteOneOrMany(JobsActive.deleteByPk({job_id: jobId})); } async markCancelled(jobId: bigint): Promise { const completedAt = new Date(); const status: JobStatus = 'cancelled'; await upsertOne(JobsById.patchByPk({job_id: jobId}, {status: Db.set(status), completed_at: Db.set(completedAt)})); await deleteOneOrMany(JobsActive.deleteByPk({job_id: jobId})); } async markDeadletter(jobId: bigint, errorMessage: string): Promise { const completedAt = new Date(); const status: JobStatus = 'deadletter'; await upsertOne( JobsById.patchByPk( {job_id: jobId}, {status: Db.set(status), completed_at: Db.set(completedAt), error_message: Db.set(errorMessage)}, ), ); await deleteOneOrMany(JobsActive.deleteByPk({job_id: jobId})); } async reportProgress(jobId: bigint, current: number, total: number | null, message: string | null): Promise { await upsertOne( JobsById.patchByPk( {job_id: jobId}, { progress_current: Db.set(BigInt(current)), progress_total: total === null ? Db.clear() : Db.set(BigInt(total)), progress_message: message === null ? Db.clear() : Db.set(message), }, ), ); } async setContextLink(jobId: bigint, link: string): Promise { await upsertOne(JobsById.patchByPk({job_id: jobId}, {context_link: Db.set(link)})); } async requestCancel(jobId: bigint): Promise { await upsertOne(JobsById.patchByPk({job_id: jobId}, {cancel_requested: Db.set(true)})); } async isCancelRequested(jobId: bigint): Promise { const row = await fetchOne<{ cancel_requested: boolean | null; }>(FETCH_CANCEL_REQUESTED_QUERY.bind({job_id: jobId})); return row?.cancel_requested === true; } async incrementAttempts(jobId: bigint): Promise { const row = await this.getJob(jobId); if (!row) return; await upsertOne(JobsById.patchByPk({job_id: jobId}, {attempts: Db.set(row.attempts + 1)})); } async listJobs(opts: { limit: number; cursor: ListJobsCursor | null; filters: ListJobsFilters; maxLookbackDays: number; }): Promise { const {limit, cursor, filters, maxLookbackDays} = opts; const startBucket = cursor ? new Date(`${cursor.bucketDay}T00:00:00Z`) : new Date(); const collected: Array = []; let nextCursor: ListJobsCursor | null = null; for (let dayOffset = 0; dayOffset <= maxLookbackDays && collected.length < limit; dayOffset++) { const bucketDate = new Date(startBucket); bucketDate.setUTCDate(bucketDate.getUTCDate() - dayOffset); const bucketDay = bucketDayFor(bucketDate); const remaining = limit - collected.length + 1; const useCursor = dayOffset === 0 && cursor !== null; let bucketRows: Array; if (useCursor && cursor) { const query = JobsByDayBucket.select({ where: [JobsByDayBucket.where.eq('bucket_day'), JobsByDayBucket.where.lt('created_at')], limit: remaining, }); bucketRows = await fetchMany( query.bind({bucket_day: bucketDay, created_at: cursor.createdAt}), ); } else { const query = JobsByDayBucket.select({ where: JobsByDayBucket.where.eq('bucket_day'), limit: remaining, }); bucketRows = await fetchMany(query.bind({bucket_day: bucketDay})); } for (const r of bucketRows) { if (filters.status && r.status !== filters.status) continue; if (filters.taskType && r.task_type !== filters.taskType) continue; if (filters.requestedByUserId !== undefined && filters.requestedByUserId !== null) { if (r.requested_by_user_id !== filters.requestedByUserId) continue; } if (collected.length >= limit) { nextCursor = {bucketDay, createdAt: r.created_at, jobId: r.job_id}; break; } const fullRow = await this.getJob(r.job_id); if (fullRow) collected.push(fullRow); } if (nextCursor) break; } return {jobs: collected, nextCursor}; } async listActiveJobs(): Promise> { const activeRows = await fetchMany(ACTIVE_JOBS_QUERY.bind({})); const fullRows = await Promise.all(activeRows.map((r) => this.getJob(r.job_id))); return fullRows.filter((r): r is JobByIdRow => r !== null); } async listActiveJobsByTaskType(taskType: string): Promise> { const activeRows = await fetchMany(ACTIVE_JOBS_QUERY.bind({})); const matching = activeRows.filter((r) => r.task_type === taskType); const fullRows = await Promise.all(matching.map((r) => this.getJob(r.job_id))); return fullRows.filter((r): r is JobByIdRow => r !== null); } }