diff --git a/apps/api/src/app/api/automations/dispatch/[id]/route.test.ts b/apps/api/src/app/api/automations/dispatch/[id]/route.test.ts new file mode 100644 index 00000000000..529b0d04a48 --- /dev/null +++ b/apps/api/src/app/api/automations/dispatch/[id]/route.test.ts @@ -0,0 +1,405 @@ +import { beforeEach, describe, expect, mock, test } from "bun:test"; + +const scheduledFor = new Date("2026-08-06T18:46:00.000Z"); +const previousUpdatedAt = new Date("2026-08-01T00:00:00.123Z"); +const previousUpdatedAtToken = "2026-08-01T00:00:00.123456Z"; +const terminalDispatchToken = new Date(scheduledFor.getTime() + 1); +const userPausedUpdatedAt = new Date("2026-08-02T00:00:00.000Z"); +const automationId = "75c82d06-77af-454c-9f0c-e6c617ea702b"; + +type DispatchOutcome = + | { status: "dispatched"; runId: string } + | { status: "conflict" }; + +let dispatchOutcome: DispatchOutcome = { + status: "dispatched", + runId: "3166c37c-add6-4382-ad07-44c816edb03e", +}; +const dispatchAutomation = mock(async () => dispatchOutcome); +const updateValues: unknown[] = []; +const updateWhereValues: unknown[] = []; +let claimSucceeds = true; +let claimResults: boolean[] | undefined; + +let automation = { + id: automationId, + enabled: true, + nextRunAt: scheduledFor, + updatedAt: previousUpdatedAt, +}; + +mock.module("@/env", () => ({ + env: { + NEXT_PUBLIC_API_URL: "http://localhost:3001", + RELAY_URL: "https://relay.example.com", + QSTASH_CURRENT_SIGNING_KEY: "current-key", + QSTASH_NEXT_SIGNING_KEY: "next-key", + }, +})); + +mock.module("@/lib/relay-url", () => ({ + getRelayUrl: mock(async () => "https://relay.example.com"), +})); + +mock.module("@upstash/qstash", () => ({ + Receiver: class { + verify = mock(async () => true); + }, +})); + +mock.module("@superset/db/schema", () => ({ + automations: { + id: "id", + enabled: "enabled", + nextRunAt: "nextRunAt", + updatedAt: "updatedAt", + }, +})); + +mock.module("@superset/db/client", () => ({ + dbWs: { + select: () => ({ + from: () => ({ + where: () => ({ + limit: async () => [automation], + }), + }), + }), + update: () => ({ + set: (values: unknown) => { + updateValues.push(values); + return { + where: (condition: unknown) => { + updateWhereValues.push(condition); + return { + returning: async () => + (claimResults?.shift() ?? claimSucceeds) + ? [{ id: automationId }] + : [], + }; + }, + }; + }, + }), + }, +})); + +mock.module("@superset/trpc/automation-dispatch", () => ({ + dispatchAutomation, +})); + +mock.module("drizzle-orm", () => ({ + and: (...conditions: unknown[]) => ({ type: "and", conditions }), + eq: (field: unknown, value: unknown) => ({ type: "eq", field, value }), + sql: (strings: TemplateStringsArray, ...values: unknown[]) => ({ + type: "sql", + strings: [...strings], + values, + }), +})); + +const { POST } = await import("./route"); + +function request(payload: Record): Request { + return new Request( + `http://localhost:3001/api/automations/dispatch/${automationId}`, + { + method: "POST", + headers: { "upstash-signature": "valid-signature" }, + body: JSON.stringify(payload), + }, + ); +} + +const params = Promise.resolve({ id: automationId }); + +describe("automations dispatch route", () => { + beforeEach(() => { + dispatchAutomation.mockClear(); + dispatchOutcome = { + status: "dispatched", + runId: "3166c37c-add6-4382-ad07-44c816edb03e", + }; + automation = { + id: automationId, + enabled: true, + nextRunAt: scheduledFor, + updatedAt: previousUpdatedAt, + }; + updateValues.length = 0; + updateWhereValues.length = 0; + claimSucceeds = true; + claimResults = undefined; + }); + + test("claims an enabled terminal occurrence before dispatching it", async () => { + const response = await POST( + request({ + automationId, + scheduledFor: scheduledFor.toISOString(), + terminal: true, + terminalDispatchToken: terminalDispatchToken.toISOString(), + terminalPreviousUpdatedAt: previousUpdatedAtToken, + }), + { params }, + ); + + expect(response.status).toBe(200); + expect(await response.json()).toEqual({ + ok: true, + outcome: { + status: "dispatched", + runId: "3166c37c-add6-4382-ad07-44c816edb03e", + }, + }); + expect(dispatchAutomation).toHaveBeenCalledTimes(1); + expect(updateValues).toEqual([ + { enabled: false, updatedAt: terminalDispatchToken }, + ]); + expect(updateWhereValues[0]).toEqual({ + type: "and", + conditions: [ + { type: "eq", field: "id", value: automationId }, + { type: "eq", field: "enabled", value: true }, + { type: "eq", field: "nextRunAt", value: scheduledFor }, + { + type: "sql", + strings: ["", " = ", "::timestamptz"], + values: ["updatedAt", previousUpdatedAtToken], + }, + ], + }); + }); + + test("returns non-2xx for a terminal dispatch conflict", async () => { + automation.enabled = false; + automation.updatedAt = terminalDispatchToken; + dispatchOutcome = { status: "conflict" }; + + const response = await POST( + request({ + automationId, + scheduledFor: scheduledFor.toISOString(), + terminal: true, + terminalDispatchToken: terminalDispatchToken.toISOString(), + terminalPreviousUpdatedAt: previousUpdatedAtToken, + }), + { params }, + ); + + expect(response.status).toBe(409); + expect(await response.json()).toEqual({ + ok: false, + error: "Terminal automation dispatch is already in progress", + }); + expect(dispatchAutomation).toHaveBeenCalledTimes(1); + }); + + test("does not dispatch an automation intentionally disabled by the user", async () => { + automation.enabled = false; + automation.updatedAt = userPausedUpdatedAt; + + const response = await POST( + request({ + automationId, + scheduledFor: scheduledFor.toISOString(), + }), + { params }, + ); + + expect(response.status).toBe(200); + expect(await response.json()).toEqual({ + ok: true, + skipped: "disabled", + }); + expect(dispatchAutomation).not.toHaveBeenCalled(); + expect(updateValues).toEqual([]); + }); + + test("dispatches a terminal occurrence reserved by evaluate", async () => { + automation.enabled = false; + automation.updatedAt = terminalDispatchToken; + + const response = await POST( + request({ + automationId, + scheduledFor: scheduledFor.toISOString(), + terminal: true, + terminalDispatchToken: terminalDispatchToken.toISOString(), + terminalPreviousUpdatedAt: previousUpdatedAtToken, + }), + { params }, + ); + + expect(response.status).toBe(200); + expect(await response.json()).toEqual({ + ok: true, + outcome: { + status: "dispatched", + runId: "3166c37c-add6-4382-ad07-44c816edb03e", + }, + }); + expect(dispatchAutomation).toHaveBeenCalledTimes(1); + expect(updateValues).toEqual([{ updatedAt: terminalDispatchToken }]); + }); + + test("does not dispatch a terminal message after a user pause", async () => { + automation.enabled = false; + automation.updatedAt = userPausedUpdatedAt; + + const response = await POST( + request({ + automationId, + scheduledFor: scheduledFor.toISOString(), + terminal: true, + terminalDispatchToken: terminalDispatchToken.toISOString(), + terminalPreviousUpdatedAt: previousUpdatedAtToken, + }), + { params }, + ); + + expect(await response.json()).toEqual({ + ok: true, + skipped: "disabled", + }); + expect(dispatchAutomation).not.toHaveBeenCalled(); + expect(updateValues).toEqual([]); + }); + + test("skips a stale terminal occurrence", async () => { + const response = await POST( + request({ + automationId, + scheduledFor: new Date(scheduledFor.getTime() - 60_000).toISOString(), + terminal: true, + terminalDispatchToken: terminalDispatchToken.toISOString(), + terminalPreviousUpdatedAt: previousUpdatedAtToken, + }), + { params }, + ); + + expect(response.status).toBe(200); + expect(await response.json()).toEqual({ + ok: true, + skipped: "stale", + }); + expect(dispatchAutomation).not.toHaveBeenCalled(); + expect(updateValues).toEqual([]); + }); + + test("does not dispatch when the terminal claim loses a concurrent update", async () => { + claimSucceeds = false; + + const response = await POST( + request({ + automationId, + scheduledFor: scheduledFor.toISOString(), + terminal: true, + terminalDispatchToken: terminalDispatchToken.toISOString(), + terminalPreviousUpdatedAt: previousUpdatedAtToken, + }), + { params }, + ); + + expect(await response.json()).toEqual({ + ok: true, + skipped: "stale", + }); + expect(dispatchAutomation).not.toHaveBeenCalled(); + }); + + test("claims the evaluator reservation if it wins after the initial select", async () => { + claimResults = [false, true]; + + const response = await POST( + request({ + automationId, + scheduledFor: scheduledFor.toISOString(), + terminal: true, + terminalDispatchToken: terminalDispatchToken.toISOString(), + terminalPreviousUpdatedAt: previousUpdatedAtToken, + }), + { params }, + ); + + expect(response.status).toBe(200); + expect(dispatchAutomation).toHaveBeenCalledTimes(1); + expect(updateValues).toEqual([ + { enabled: false, updatedAt: terminalDispatchToken }, + { updatedAt: terminalDispatchToken }, + ]); + expect(updateWhereValues[1]).toEqual({ + type: "and", + conditions: [ + { type: "eq", field: "id", value: automationId }, + { type: "eq", field: "enabled", value: false }, + { type: "eq", field: "nextRunAt", value: scheduledFor }, + { + type: "sql", + strings: ["", " = ", "::timestamptz"], + values: ["updatedAt", terminalDispatchToken.toISOString()], + }, + ], + }); + }); + + test("skips a terminal message for a different reservation", async () => { + automation.enabled = false; + automation.updatedAt = new Date("2026-08-03T00:00:00.000Z"); + + const response = await POST( + request({ + automationId, + scheduledFor: scheduledFor.toISOString(), + terminal: true, + terminalDispatchToken: terminalDispatchToken.toISOString(), + terminalPreviousUpdatedAt: previousUpdatedAtToken, + }), + { params }, + ); + + expect(await response.json()).toEqual({ + ok: true, + skipped: "disabled", + }); + expect(dispatchAutomation).not.toHaveBeenCalled(); + }); + + test("skips a reserved terminal message after its occurrence changes", async () => { + automation.enabled = false; + automation.updatedAt = terminalDispatchToken; + automation.nextRunAt = new Date(scheduledFor.getTime() + 60_000); + + const response = await POST( + request({ + automationId, + scheduledFor: scheduledFor.toISOString(), + terminal: true, + terminalDispatchToken: terminalDispatchToken.toISOString(), + terminalPreviousUpdatedAt: previousUpdatedAtToken, + }), + { params }, + ); + + expect(await response.json()).toEqual({ + ok: true, + skipped: "stale", + }); + expect(dispatchAutomation).not.toHaveBeenCalled(); + expect(updateValues).toEqual([]); + }); + + test("keeps dispatching enabled non-terminal occurrences", async () => { + const response = await POST( + request({ + automationId, + scheduledFor: scheduledFor.toISOString(), + }), + { params }, + ); + + expect(response.status).toBe(200); + expect(dispatchAutomation).toHaveBeenCalledTimes(1); + expect(updateValues).toEqual([]); + }); +}); diff --git a/apps/api/src/app/api/automations/dispatch/[id]/route.ts b/apps/api/src/app/api/automations/dispatch/[id]/route.ts index 0197fe8e5c5..f37a7190d6f 100644 --- a/apps/api/src/app/api/automations/dispatch/[id]/route.ts +++ b/apps/api/src/app/api/automations/dispatch/[id]/route.ts @@ -1,12 +1,16 @@ import { dbWs } from "@superset/db/client"; -import { automations } from "@superset/db/schema"; +import { automations, type SelectAutomation } from "@superset/db/schema"; import { dispatchAutomation } from "@superset/trpc/automation-dispatch"; import { Receiver } from "@upstash/qstash"; -import { eq } from "drizzle-orm"; +import { and, eq, sql } from "drizzle-orm"; import { z } from "zod"; import { env } from "@/env"; import { getRelayUrl } from "@/lib/relay-url"; +import { + matchesTerminalOccurrence, + matchesTerminalReservation, +} from "../../terminal-occurrence"; export const maxDuration = 60; export const dynamic = "force-dynamic"; @@ -19,6 +23,13 @@ const receiver = new Receiver({ const payloadSchema = z.object({ automationId: z.string().uuid(), scheduledFor: z.string().datetime(), + // The token identifies the evaluator's terminal reservation. The previous + // value lets dispatch claim an enabled row only if no edit won the race. + terminal: z.boolean().default(false), + terminalDispatchToken: z.string().datetime().optional(), + terminalPreviousUpdatedAt: z.string().datetime().optional(), + // Accept messages created before the updatedAt reservation was introduced. + terminalPendingNextRunAt: z.string().datetime().optional(), }); export async function POST( @@ -56,17 +67,140 @@ export async function POST( if (!automation) { return Response.json({ ok: true, skipped: "deleted" }); } - if (!automation.enabled) { + + const scheduledFor = new Date(parsed.data.scheduledFor); + if (parsed.data.terminal) { + const terminalReservation = matchesTerminalReservation({ + updatedAt: automation.updatedAt, + terminalDispatchToken: parsed.data.terminalDispatchToken, + }); + if (!automation.enabled && !terminalReservation) { + return Response.json({ ok: true, skipped: "disabled" }); + } + if ( + !matchesTerminalOccurrence({ + nextRunAt: automation.nextRunAt, + scheduledFor, + legacyPendingNextRunAt: parsed.data.terminalPendingNextRunAt, + }) + ) { + return Response.json({ ok: true, skipped: "stale" }); + } + + const claimed = await claimTerminalOccurrence({ + automation, + automationId: parsed.data.automationId, + terminalDispatchToken: parsed.data.terminalDispatchToken, + terminalPreviousUpdatedAt: parsed.data.terminalPreviousUpdatedAt, + }); + if (!claimed) { + return Response.json({ + ok: true, + skipped: automation.enabled ? "stale" : "disabled", + }); + } + } else if (!automation.enabled) { return Response.json({ ok: true, skipped: "disabled" }); } const outcome = await dispatchAutomation({ automation, - scheduledFor: new Date(parsed.data.scheduledFor), + scheduledFor, // The owner's host may be on an overridden relay (relay-url-override); // env.RELAY_URL alone reaches only hosts still on the default relay. relayUrl: await getRelayUrl(automation.ownerUserId), }); + if (parsed.data.terminal && outcome.status === "conflict") { + // A conflict can mean that another worker owns a still-dispatching run. + // Do not acknowledge the message or close the recurrence: QStash retries, + // then its failure callback records dispatch_failed if the owner never + // reaches a terminal run state. + return Response.json( + { + ok: false, + error: "Terminal automation dispatch is already in progress", + }, + { status: 409 }, + ); + } + return Response.json({ ok: true, outcome }); } + +async function claimTerminalOccurrence({ + automation, + automationId, + terminalDispatchToken, + terminalPreviousUpdatedAt, +}: { + automation: SelectAutomation; + automationId: string; + terminalDispatchToken?: string; + terminalPreviousUpdatedAt?: string; +}): Promise { + if (terminalDispatchToken !== undefined) { + if (terminalPreviousUpdatedAt === undefined) { + return false; + } + + const token = new Date(terminalDispatchToken); + if (automation.enabled) { + const [claimed] = await dbWs + .update(automations) + .set({ enabled: false, updatedAt: token }) + .where( + and( + eq(automations.id, automationId), + eq(automations.enabled, true), + eq(automations.nextRunAt, automation.nextRunAt), + // Compare the original database value without converting it to a + // millisecond-only JavaScript Date. + sql`${automations.updatedAt} = ${terminalPreviousUpdatedAt}::timestamptz`, + ), + ) + .returning({ id: automations.id }); + + if (claimed !== undefined) return true; + } + + // Evaluation may have completed the reservation after the SELECT above + // but before this claim. Claim that exact reservation instead of + // acknowledging the queued occurrence without a run. + const [reserved] = await dbWs + .update(automations) + .set({ updatedAt: token }) + .where( + and( + eq(automations.id, automationId), + eq(automations.enabled, false), + eq(automations.nextRunAt, automation.nextRunAt), + sql`${automations.updatedAt} = ${terminalDispatchToken}::timestamptz`, + ), + ) + .returning({ id: automations.id }); + + return reserved !== undefined; + } + + // Rolling-deployment compatibility for messages that carry only the old + // nextRunAt reservation. They may claim an enabled row, but never bypass a + // user-disabled row. + if (!automation.enabled) { + return false; + } + + const [claimed] = await dbWs + .update(automations) + .set({ enabled: false }) + .where( + and( + eq(automations.id, automationId), + eq(automations.enabled, true), + eq(automations.nextRunAt, automation.nextRunAt), + ), + ) + .returning({ id: automations.id }); + + return claimed !== undefined; +} diff --git a/apps/api/src/app/api/automations/evaluate/route.test.ts b/apps/api/src/app/api/automations/evaluate/route.test.ts new file mode 100644 index 00000000000..6f90f7c3fed --- /dev/null +++ b/apps/api/src/app/api/automations/evaluate/route.test.ts @@ -0,0 +1,198 @@ +import { describe, expect, mock, test } from "bun:test"; +import { nextOccurrenceAfter } from "@superset/shared/rrule"; + +const terminalOccurrence = new Date("2026-08-06T18:46:30.000Z"); +const nonTerminalOccurrence = new Date("2026-08-05T18:46:30.000Z"); +const terminalUpdatedAt = new Date("2026-08-01T00:00:00.123Z"); +const terminalUpdatedAtToken = "2026-08-01T00:00:00.123456Z"; +const nonTerminalUpdatedAt = new Date("2026-08-01T00:00:00.000Z"); +const nonTerminalUpdatedAtToken = "2026-08-01T00:00:00.000000Z"; +const automationId = "75c82d06-77af-454c-9f0c-e6c617ea702b"; +const terminalDispatchToken = new Date(terminalOccurrence.getTime() + 1); + +let dueAutomations: Array<{ + id: string; + nextRunAt: Date; + rrule: string; + dtstart: Date; + timezone: string; + updatedAt: Date; + updatedAtToken: string; +}> = []; +const updateValues: unknown[] = []; +const updateWhereValues: unknown[] = []; +const batchJSON = mock(async (_messages: unknown[]) => undefined); + +mock.module("@/env", () => ({ + env: { + NEXT_PUBLIC_API_URL: "http://localhost:3001", + QSTASH_TOKEN: "test-token", + QSTASH_URL: "https://qstash.example.com", + QSTASH_CURRENT_SIGNING_KEY: "current-key", + QSTASH_NEXT_SIGNING_KEY: "next-key", + }, +})); + +mock.module("@upstash/qstash", () => ({ + Client: class { + batchJSON = batchJSON; + }, + Receiver: class { + verify = mock(async () => true); + }, +})); + +mock.module("@superset/db/schema", () => ({ + automations: { + id: "id", + enabled: "enabled", + nextRunAt: "nextRunAt", + updatedAt: "updatedAt", + }, + automationRuns: { + automationId: "automationId", + organizationId: "organizationId", + scheduledFor: "scheduledFor", + }, +})); + +mock.module("@superset/db/client", () => ({ + dbWs: { + select: () => ({ + from: () => ({ + where: () => ({ + orderBy: () => ({ + limit: async () => dueAutomations, + }), + }), + }), + }), + update: () => ({ + set: (values: unknown) => { + updateValues.push(values); + return { + where: async (condition: unknown) => { + updateWhereValues.push(condition); + return undefined; + }, + }; + }, + }), + }, +})); + +mock.module("drizzle-orm", () => ({ + and: (...conditions: unknown[]) => ({ type: "and", conditions }), + eq: (field: unknown, value: unknown) => ({ type: "eq", field, value }), + getTableColumns: () => ({}), + lte: (field: unknown, value: unknown) => ({ type: "lte", field, value }), + sql: (strings: TemplateStringsArray, ...values: unknown[]) => ({ + type: "sql", + strings: [...strings], + values, + }), +})); + +const { POST } = await import("./route"); + +type EnqueuedMessage = { + body: { + automationId: string; + scheduledFor: string; + terminal: boolean; + }; +}; + +function request(): Request { + return new Request("http://localhost:3001/api/automations/evaluate", { + method: "POST", + headers: { "upstash-signature": "valid-signature" }, + body: "{}", + }); +} + +describe("automations evaluate route", () => { + test("marks a COUNT=1 occurrence as terminal in the queued payload", async () => { + dueAutomations = [ + { + id: automationId, + nextRunAt: terminalOccurrence, + rrule: "FREQ=DAILY;COUNT=1", + dtstart: terminalOccurrence, + timezone: "UTC", + updatedAt: terminalUpdatedAt, + updatedAtToken: terminalUpdatedAtToken, + }, + ]; + updateValues.length = 0; + updateWhereValues.length = 0; + batchJSON.mockClear(); + + const response = await POST(request()); + + expect(response.status).toBe(200); + expect(batchJSON).toHaveBeenCalledTimes(1); + const messages = batchJSON.mock.calls[0]?.[0] as EnqueuedMessage[]; + expect(messages[0]?.body).toEqual({ + automationId, + scheduledFor: new Date( + Math.floor(terminalOccurrence.getTime() / 60_000) * 60_000, + ).toISOString(), + terminal: true, + terminalDispatchToken: terminalDispatchToken.toISOString(), + terminalPreviousUpdatedAt: terminalUpdatedAtToken, + }); + expect(updateValues).toEqual([ + { + enabled: false, + updatedAt: terminalDispatchToken, + }, + ]); + expect(updateWhereValues[0]).toEqual({ + type: "and", + conditions: [ + { type: "eq", field: "id", value: automationId }, + { type: "eq", field: "enabled", value: true }, + { type: "eq", field: "nextRunAt", value: terminalOccurrence }, + { + type: "sql", + strings: ["", " = ", "::timestamptz"], + values: ["updatedAt", terminalUpdatedAtToken], + }, + ], + }); + }); + + test("keeps the existing non-terminal advance path", async () => { + dueAutomations = [ + { + id: automationId, + nextRunAt: nonTerminalOccurrence, + rrule: "FREQ=DAILY", + dtstart: nonTerminalOccurrence, + timezone: "UTC", + updatedAt: nonTerminalUpdatedAt, + updatedAtToken: nonTerminalUpdatedAtToken, + }, + ]; + updateValues.length = 0; + updateWhereValues.length = 0; + batchJSON.mockClear(); + + const response = await POST(request()); + + expect(response.status).toBe(200); + const messages = batchJSON.mock.calls[0]?.[0] as EnqueuedMessage[]; + expect(messages[0]?.body.terminal).toBe(false); + expect(updateValues).toEqual([ + { + nextRunAt: nextOccurrenceAfter({ + rrule: "FREQ=DAILY", + dtstart: nonTerminalOccurrence, + timezone: "UTC", + after: nonTerminalOccurrence, + }), + }, + ]); + }); +}); diff --git a/apps/api/src/app/api/automations/evaluate/route.ts b/apps/api/src/app/api/automations/evaluate/route.ts index f0c8671c087..235cc137da1 100644 --- a/apps/api/src/app/api/automations/evaluate/route.ts +++ b/apps/api/src/app/api/automations/evaluate/route.ts @@ -1,8 +1,8 @@ import { dbWs } from "@superset/db/client"; import { automations } from "@superset/db/schema"; -import { nextOccurrenceAfter } from "@superset/shared/rrule"; +import { bucketToMinute, nextOccurrenceAfter } from "@superset/shared/rrule"; import { Client, Receiver } from "@upstash/qstash"; -import { and, eq, lte } from "drizzle-orm"; +import { and, eq, getTableColumns, lte, sql } from "drizzle-orm"; import { env } from "@/env"; @@ -19,12 +19,6 @@ const receiver = new Receiver({ const BATCH_SIZE = 2000; -function bucketToMinute(d: Date): Date { - const copy = new Date(d.getTime()); - copy.setUTCSeconds(0, 0); - return copy; -} - export async function POST(request: Request): Promise { const body = await request.text(); const signature = request.headers.get("upstash-signature"); @@ -43,7 +37,11 @@ export async function POST(request: Request): Promise { const now = new Date(); const due = await dbWs - .select() + .select({ + ...getTableColumns(automations), + // Keep the database's full timestamp precision in a signed ISO token. + updatedAtToken: sql`to_char(${automations.updatedAt} AT TIME ZONE 'UTC', 'YYYY-MM-DD"T"HH24:MI:SS.US"Z"')`, + }) .from(automations) .where(and(eq(automations.enabled, true), lte(automations.nextRunAt, now))) .orderBy(automations.nextRunAt) @@ -53,14 +51,36 @@ export async function POST(request: Request): Promise { return Response.json({ enqueued: 0 }); } + const dueWithNext = due.map((automation) => { + const next = nextOccurrenceAfter({ + rrule: automation.rrule, + dtstart: automation.dtstart, + timezone: automation.timezone, + after: automation.nextRunAt, + }); + // Derive the token from the occurrence so concurrent evaluators produce + // the same signed payload and idempotent reservation value. + const terminalDispatchToken = + next === null ? new Date(automation.nextRunAt.getTime() + 1) : undefined; + return { + automation, + next, + terminalDispatchToken, + }; + }); + await qstash.batchJSON( - due.map((automation) => { + dueWithNext.map(({ automation, next, terminalDispatchToken }) => { const scheduledFor = bucketToMinute(automation.nextRunAt); return { url: `${env.NEXT_PUBLIC_API_URL}/api/automations/dispatch/${automation.id}`, body: { automationId: automation.id, scheduledFor: scheduledFor.toISOString(), + terminal: next === null, + terminalDispatchToken: terminalDispatchToken?.toISOString(), + terminalPreviousUpdatedAt: + next === null ? automation.updatedAtToken : undefined, }, deduplicationId: `${automation.id}_${scheduledFor.getTime()}`, retries: 2, @@ -70,16 +90,31 @@ export async function POST(request: Request): Promise { ); const advanceResults = await Promise.allSettled( - due.map((automation) => { - const next = nextOccurrenceAfter({ - rrule: automation.rrule, - dtstart: automation.dtstart, - timezone: automation.timezone, - after: automation.nextRunAt, - }); + dueWithNext.map(({ automation, next, terminalDispatchToken }) => { + if (!next) { + if (!terminalDispatchToken) { + throw new Error("Missing terminal dispatch token"); + } + + // Keep nextRunAt as the real occurrence for API/UI consumers. The + // updatedAt token records that this disabled row belongs to the + // evaluator's queued terminal occurrence. + return dbWs + .update(automations) + .set({ enabled: false, updatedAt: terminalDispatchToken }) + .where( + and( + eq(automations.id, automation.id), + eq(automations.enabled, true), + eq(automations.nextRunAt, automation.nextRunAt), + sql`${automations.updatedAt} = ${automation.updatedAtToken}::timestamptz`, + ), + ); + } + return dbWs .update(automations) - .set(next ? { nextRunAt: next } : { enabled: false }) + .set({ nextRunAt: next }) .where(eq(automations.id, automation.id)); }), ); diff --git a/apps/api/src/app/api/automations/run-failed/route.test.ts b/apps/api/src/app/api/automations/run-failed/route.test.ts new file mode 100644 index 00000000000..2f75741f66b --- /dev/null +++ b/apps/api/src/app/api/automations/run-failed/route.test.ts @@ -0,0 +1,195 @@ +import { beforeEach, expect, mock, test } from "bun:test"; + +const automationId = "75c82d06-77af-454c-9f0c-e6c617ea702b"; +const organizationId = "3ee200f3-c54c-46b1-b8b8-24a7f27348f3"; +const terminalOccurrence = new Date("2026-08-06T18:46:30.000Z"); +const scheduledFor = new Date("2026-08-06T18:46:00.000Z"); +const previousUpdatedAt = new Date("2026-08-01T00:00:00.123Z"); +const previousUpdatedAtToken = "2026-08-01T00:00:00.123456Z"; +const terminalDispatchToken = new Date(terminalOccurrence.getTime() + 1); +const insertValues: unknown[] = []; +const updateValues: unknown[] = []; +const updateWhereValues: unknown[] = []; + +let automation = { + organizationId, + name: "Nightly automation", + enabled: true, + nextRunAt: terminalOccurrence, + updatedAt: previousUpdatedAt, +}; + +beforeEach(() => { + automation = { + organizationId, + name: "Nightly automation", + enabled: true, + nextRunAt: terminalOccurrence, + updatedAt: previousUpdatedAt, + }; + insertValues.length = 0; + updateValues.length = 0; + updateWhereValues.length = 0; +}); + +mock.module("@/env", () => ({ + env: { + NEXT_PUBLIC_API_URL: "http://localhost:3001", + QSTASH_CURRENT_SIGNING_KEY: "current-key", + QSTASH_NEXT_SIGNING_KEY: "next-key", + }, +})); + +mock.module("@sentry/nextjs", () => ({ + captureException: mock(() => undefined), +})); + +mock.module("@upstash/qstash", () => ({ + Receiver: class { + verify = mock(async () => true); + }, +})); + +mock.module("@superset/db/schema", () => ({ + automations: { + id: "id", + organizationId: "organizationId", + name: "name", + enabled: "enabled", + nextRunAt: "nextRunAt", + updatedAt: "updatedAt", + }, + automationRuns: { + automationId: "automationId", + organizationId: "organizationId", + scheduledFor: "scheduledFor", + status: "status", + }, +})); + +mock.module("@superset/db/client", () => ({ + dbWs: { + select: () => ({ + from: () => ({ + where: () => ({ + limit: async () => [automation], + }), + }), + }), + insert: () => ({ + values: (values: unknown) => { + insertValues.push(values); + return { + onConflictDoUpdate: async () => undefined, + }; + }, + }), + update: () => ({ + set: (values: unknown) => { + updateValues.push(values); + return { + where: async (condition: unknown) => { + updateWhereValues.push(condition); + return undefined; + }, + }; + }, + }), + }, +})); + +mock.module("drizzle-orm", () => ({ + and: (...conditions: unknown[]) => ({ type: "and", conditions }), + eq: (field: unknown, value: unknown) => ({ type: "eq", field, value }), + sql: (strings: TemplateStringsArray, ...values: unknown[]) => ({ + type: "sql", + strings: [...strings], + values, + }), +})); + +const { POST } = await import("./route"); + +function requestFor({ + scheduledFor: sourceScheduledFor = scheduledFor, + terminal = true, +}: { + scheduledFor?: Date; + terminal?: boolean; +} = {}): Request { + const sourceBody = Buffer.from( + JSON.stringify({ + automationId, + scheduledFor: sourceScheduledFor.toISOString(), + terminal, + terminalDispatchToken: terminalDispatchToken.toISOString(), + terminalPreviousUpdatedAt: previousUpdatedAtToken, + }), + ).toString("base64"); + + return new Request("http://localhost:3001/api/automations/run-failed", { + method: "POST", + headers: { "upstash-signature": "valid-signature" }, + body: JSON.stringify({ + sourceMessageId: "qstash-message-id", + sourceBody, + status: 500, + error: "delivery failed", + }), + }); +} + +test("records a terminal delivery failure and closes the recurrence", async () => { + const response = await POST(requestFor()); + + expect(response.status).toBe(200); + expect(insertValues).toEqual([ + expect.objectContaining({ + automationId, + status: "dispatch_failed", + scheduledFor, + }), + ]); + expect(updateValues).toEqual([ + { enabled: false, updatedAt: terminalDispatchToken }, + ]); + expect(updateWhereValues[0]).toEqual({ + type: "and", + conditions: [ + { type: "eq", field: "id", value: automationId }, + { type: "eq", field: "enabled", value: true }, + { type: "eq", field: "nextRunAt", value: terminalOccurrence }, + { + type: "sql", + strings: ["", " = ", "::timestamptz"], + values: ["updatedAt", previousUpdatedAtToken], + }, + ], + }); +}); + +test("does not close a changed terminal occurrence", async () => { + const response = await POST( + requestFor({ + scheduledFor: new Date(scheduledFor.getTime() - 60_000), + }), + ); + + expect(response.status).toBe(200); + expect(insertValues).toHaveLength(1); + expect(updateValues).toEqual([]); + expect(updateWhereValues).toEqual([]); +}); + +test("does not close a user-paused automation", async () => { + automation.enabled = false; + // Keep the original version so this assertion isolates the enabled guard. + automation.updatedAt = previousUpdatedAt; + + const response = await POST(requestFor()); + + expect(response.status).toBe(200); + expect(insertValues).toHaveLength(1); + expect(updateValues).toEqual([]); + expect(updateWhereValues).toEqual([]); +}); diff --git a/apps/api/src/app/api/automations/run-failed/route.ts b/apps/api/src/app/api/automations/run-failed/route.ts index e782e49a21e..f7d750b62ee 100644 --- a/apps/api/src/app/api/automations/run-failed/route.ts +++ b/apps/api/src/app/api/automations/run-failed/route.ts @@ -2,10 +2,11 @@ import * as Sentry from "@sentry/nextjs"; import { dbWs } from "@superset/db/client"; import { automationRuns, automations } from "@superset/db/schema"; import { Receiver } from "@upstash/qstash"; -import { eq } from "drizzle-orm"; +import { and, eq, sql } from "drizzle-orm"; import { z } from "zod"; import { env } from "@/env"; +import { matchesTerminalOccurrence } from "../terminal-occurrence"; export const dynamic = "force-dynamic"; @@ -25,6 +26,11 @@ const failurePayloadSchema = z.object({ const sourceBodySchema = z.object({ automationId: z.string().uuid(), scheduledFor: z.string().datetime(), + terminal: z.boolean().default(false), + terminalDispatchToken: z.string().datetime().optional(), + terminalPreviousUpdatedAt: z.string().datetime().optional(), + // Accept messages created before the updatedAt reservation was introduced. + terminalPendingNextRunAt: z.string().datetime().optional(), }); export async function POST(request: Request): Promise { @@ -78,6 +84,8 @@ export async function POST(request: Request): Promise { .select({ organizationId: automations.organizationId, name: automations.name, + enabled: automations.enabled, + nextRunAt: automations.nextRunAt, }) .from(automations) .where(eq(automations.id, automationId)) @@ -102,8 +110,51 @@ export async function POST(request: Request): Promise { .onConflictDoUpdate({ target: [automationRuns.automationId, automationRuns.scheduledFor], set: { status: "dispatch_failed", error: errorText }, + // A retry can collide with an already completed or offline run; preserve + // that observable outcome. + setWhere: eq(automationRuns.status, "dispatching"), }); + const terminalOccurrenceMatches = + source.data.terminal && + matchesTerminalOccurrence({ + nextRunAt: automation.nextRunAt, + scheduledFor: new Date(scheduledFor), + legacyPendingNextRunAt: source.data.terminalPendingNextRunAt, + }); + const canClaimUnreservedTerminal = + terminalOccurrenceMatches && + automation.enabled && + (source.data.terminalDispatchToken === undefined + ? true + : source.data.terminalPreviousUpdatedAt !== undefined); + + if (canClaimUnreservedTerminal) { + const terminalDispatchToken = source.data.terminalDispatchToken; + await dbWs + .update(automations) + .set( + terminalDispatchToken === undefined + ? { enabled: false } + : { + enabled: false, + updatedAt: new Date(terminalDispatchToken), + }, + ) + .where( + and( + eq(automations.id, automationId), + eq(automations.enabled, true), + eq(automations.nextRunAt, automation.nextRunAt), + ...(source.data.terminalPreviousUpdatedAt === undefined + ? [] + : [ + sql`${automations.updatedAt} = ${source.data.terminalPreviousUpdatedAt}::timestamptz`, + ]), + ), + ); + } + Sentry.captureException( new Error(`automation dispatch failed: ${automationId}`), { diff --git a/apps/api/src/app/api/automations/terminal-occurrence.ts b/apps/api/src/app/api/automations/terminal-occurrence.ts new file mode 100644 index 00000000000..9eb5b0bde34 --- /dev/null +++ b/apps/api/src/app/api/automations/terminal-occurrence.ts @@ -0,0 +1,41 @@ +import { bucketToMinute } from "@superset/shared/rrule"; + +/** + * Checks that a terminal message still refers to the occurrence stored on the + * automation. `legacyPendingNextRunAt` keeps already-queued messages from the + * previous reservation format deliverable during a rolling deployment. + */ +export function matchesTerminalOccurrence({ + nextRunAt, + scheduledFor, + legacyPendingNextRunAt, +}: { + nextRunAt: Date; + scheduledFor: Date; + legacyPendingNextRunAt?: string; +}): boolean { + if ( + legacyPendingNextRunAt !== undefined && + nextRunAt.getTime() === new Date(legacyPendingNextRunAt).getTime() + ) { + return true; + } + + return ( + bucketToMinute(nextRunAt).getTime() === + bucketToMinute(scheduledFor).getTime() + ); +} + +export function matchesTerminalReservation({ + updatedAt, + terminalDispatchToken, +}: { + updatedAt: Date; + terminalDispatchToken?: string; +}): boolean { + return ( + terminalDispatchToken !== undefined && + updatedAt.getTime() === new Date(terminalDispatchToken).getTime() + ); +} diff --git a/packages/shared/src/rrule.ts b/packages/shared/src/rrule.ts index d89628279a5..8dfa842919a 100644 --- a/packages/shared/src/rrule.ts +++ b/packages/shared/src/rrule.ts @@ -369,6 +369,13 @@ export function nextOccurrenceAfter(args: { return next ? rruleDateToUtc(next, args.timezone) : null; } +/** Floors a timestamp to the UTC minute used by automation run deduplication. */ +export function bucketToMinute(date: Date): Date { + const copy = new Date(date.getTime()); + copy.setUTCSeconds(0, 0); + return copy; +} + /** Parses + validates an RRule body, returning the next occurrence. */ export function parseRrule(args: { rrule: string; diff --git a/packages/trpc/src/router/automation/automation.ts b/packages/trpc/src/router/automation/automation.ts index 37f5b384c37..0496ee491f7 100644 --- a/packages/trpc/src/router/automation/automation.ts +++ b/packages/trpc/src/router/automation/automation.ts @@ -8,6 +8,7 @@ import { v2Workspaces, } from "@superset/db/schema"; import { + bucketToMinute, describeSchedule, nextOccurrences, parseRrule, @@ -634,16 +635,6 @@ export const automationRouter = { }), } satisfies TRPCRouterRecord; -/** - * Floors a Date down to the minute so two dispatches in the same minute bucket - * collide on the unique index. - */ -function bucketToMinute(date: Date): Date { - const copy = new Date(date.getTime()); - copy.setUTCSeconds(0, 0); - return copy; -} - function safeDescribeRrule(row: { rrule: string } | null | undefined): string { if (!row) return ""; try {