Fix Postgres-only rank-tracking & site-audit workflow failures (#317)

src/server/features/rank-tracking/repositories/RankTrackingRepository.ts

@@ -202,8 +202,23 @@ async function insertSnapshots(
     Omit<InferInsertModel<typeof rankSnapshots>, "id" | "checkedAt">
   >,
 ) {
+  // Target the (run, keyword, device) unique index explicitly. An UNtargeted
+  // ON CONFLICT DO NOTHING also swallows a primary-key collision, which would
+  // silently drop every row if the `id` serial sequence ever drifts behind
+  // max(id) (e.g. after a data import that copied explicit ids). Scoping the
+  // clause to the intended dedupe index keeps re-runs idempotent while letting
+  // a pk collision surface as a loud duplicate-key error instead of data loss.
   await executeInBatches(snapshots, (tx, snapshot) =>
-    tx.insert(rankSnapshots).values(snapshot).onConflictDoNothing(),
+    tx
+      .insert(rankSnapshots)
+      .values(snapshot)
+      .onConflictDoNothing({
+        target: [
+          rankSnapshots.runId,
+          rankSnapshots.trackingKeywordId,
+          rankSnapshots.device,
+        ],
+      }),
   );
 }
 

src/server/workflows/RankCheckWorkflow.ts

@@ -13,6 +13,7 @@ import {
   runQueuedCheck,
   type QueuedCheckStats,
 } from "@/server/workflows/rankCheckPaths";
+import { pgStep } from "@/server/workflows/pgStep";
 import { createDataforseoClient } from "@/server/lib/dataforseo";
 import { captureServerEvent } from "@/server/lib/posthog";
 import { AppError } from "@/server/lib/errors";
@@ -277,7 +278,8 @@ export class RankCheckWorkflow extends WorkflowEntrypoint<
     const client = createDataforseoClient(billingCustomer);
 
     // Guard: skip if config was archived after the workflow was triggered
-    const configCheck = await step.do(
+    const configCheck = await pgStep(
+      step,
       "check-active",
       { retries: { limit: 0, delay: "1 second" } },
       async () => {
@@ -298,7 +300,8 @@ export class RankCheckWorkflow extends WorkflowEntrypoint<
         `[rank-check] ${runId} starting (trigger=${trigger}, devices=${devices})`,
       );
 
-      const prepareResult = await step.do(
+      const prepareResult = await pgStep(
+        step,
         "prepare",
         { retries: { limit: 0, delay: "1 second" } },
         async () =>
@@ -345,7 +348,7 @@ export class RankCheckWorkflow extends WorkflowEntrypoint<
         console.warn(`[rank-check] ${runId} partial failure: ${batchError}`);
       }
 
-      await step.do("finalize", SINGLE_ATTEMPT_STEP_CONFIG, async () =>
+      await pgStep(step, "finalize", SINGLE_ATTEMPT_STEP_CONFIG, async () =>
         finalizeRankCheckRun({
           runId,
           configId,
@@ -358,7 +361,7 @@ export class RankCheckWorkflow extends WorkflowEntrypoint<
       );
     } catch (error) {
       console.error(`Rank check ${runId} failed:`, error);
-      await step.do("mark-failed", SINGLE_ATTEMPT_STEP_CONFIG, async () =>
+      await pgStep(step, "mark-failed", SINGLE_ATTEMPT_STEP_CONFIG, async () =>
         markRankCheckRunFailed({
           runId,
           configId,

src/server/workflows/SiteAuditWorkflow.ts

@@ -15,6 +15,7 @@ import { AuditRepository } from "@/server/features/audit/repositories/AuditRepos
 import type { AuditConfig } from "@/server/lib/audit/types";
 import { captureServerEvent } from "@/server/lib/posthog";
 import { runAuditPhases } from "@/server/workflows/siteAuditWorkflowPhases";
+import { pgStep } from "@/server/workflows/pgStep";
 
 interface AuditParams {
   auditId: string;
@@ -63,7 +64,7 @@ export class SiteAuditWorkflow extends WorkflowEntrypoint<Env, AuditParams> {
       });
     } catch (error) {
       console.error(`Audit ${auditId} failed:`, error);
-      await step.do("mark-failed", async () => {
+      await pgStep(step, "mark-failed", undefined, async () => {
         await AuditRepository.failAudit(auditId, event.instanceId);
 
         const latestAudit = await AuditRepository.getAuditForWorkflow(

src/server/workflows/pgStep.ts

@@ -0,0 +1,29 @@
+import type { WorkflowStep, WorkflowStepConfig } from "cloudflare:workers";
+import { withPgClient } from "@/db";
+
+/**
+ * `step.do` with a request-scoped Postgres client active inside the step body.
+ *
+ * Cloudflare Workflows invoke each step callback in its own execution context —
+ * steps are independently persisted and can resume in a fresh invocation, so the
+ * `AsyncLocalStorage` scope opened by `withPgClient` around `run()` does NOT
+ * propagate into a step. Each DB-touching step must therefore open its own
+ * client. In D1 mode `withPgClient` is a no-op, so this is just a plain
+ * `step.do`. The client is lazy (postgres-js only connects on first query), so
+ * wrapping a step that happens not to touch the DB costs nothing.
+ *
+ * `T` mirrors `step.do`'s own `Rpc.Serializable<T>` bound so step results stay
+ * serializable (the workflow engine persists and replays them). Pass `undefined`
+ * for `config` to use the engine's default step config (matches the 2-arg
+ * `step.do(name, fn)` form).
+ */
+export function pgStep<T extends Rpc.Serializable<T>>(
+  step: WorkflowStep,
+  name: string,
+  config: WorkflowStepConfig | undefined,
+  fn: () => Promise<T>,
+): Promise<T> {
+  return config
+    ? step.do(name, config, () => withPgClient(fn))
+    : step.do(name, () => withPgClient(fn));
+}

src/server/workflows/rankCheckPaths.ts

@@ -12,6 +12,7 @@ import type {
 } from "@/server/lib/dataforseo";
 import type { RankTrackingConfig } from "@/types/schemas/rank-tracking";
 import { KEYWORDS_PER_BATCH } from "@/shared/rank-tracking";
+import { pgStep } from "@/server/workflows/pgStep";
 
 const SINGLE_ATTEMPT_STEP_CONFIG = {
   retries: { limit: 0, delay: "1 second" as const },
@@ -131,7 +132,8 @@ export async function runLiveCheck(
     const batchIndex = Math.floor(i / KEYWORDS_PER_BATCH);
     const keywordsChecked = i + keywordBatch.length;
 
-    await step.do(
+    await pgStep(
+      step,
       `live-batch-${batchIndex}`,
       SINGLE_ATTEMPT_STEP_CONFIG,
       async () => {
@@ -283,7 +285,8 @@ export async function runQueuedCheck(
     const postIndex = Math.floor(i / MAX_TASKS_PER_POST);
     let posted: PostedRankCheckTask[];
     try {
-      posted = await step.do(
+      posted = await pgStep(
+        step,
         `post-tasks-${postIndex}`,
         SINGLE_ATTEMPT_STEP_CONFIG,
         async () =>
@@ -339,8 +342,11 @@ export async function runQueuedCheck(
 
     let outcome: CollectRoundOutcome;
     try {
-      outcome = await step.do(`collect-${round}`, COLLECT_STEP_CONFIG, () =>
+      outcome = await pgStep(
-        collectQueuedRound(ctx, batch),
+        step,
+        `collect-${round}`,
+        COLLECT_STEP_CONFIG,
+        () => collectQueuedRound(ctx, batch),
       );
     } catch (error) {
       console.warn(`[rank-check] ${ctx.runId} collect-${round} failed:`, error);
@@ -369,7 +375,8 @@ export async function runQueuedCheck(
     const batch = stragglers.slice(i, i + KEYWORDS_PER_BATCH);
     const batchIndex = Math.floor(i / KEYWORDS_PER_BATCH);
 
-    stats.fallbackChecked += await step.do(
+    stats.fallbackChecked += await pgStep(
+      step,
       `fallback-batch-${batchIndex}`,
       SINGLE_ATTEMPT_STEP_CONFIG,
       () => checkBatchLive(ctx, batch),

src/server/workflows/siteAuditWorkflowCrawl.ts

@@ -5,6 +5,7 @@ import { isSameOrigin, normalizeUrl } from "@/server/lib/audit/url-utils";
 import { AuditRepository } from "@/server/features/audit/repositories/AuditRepository";
 import { AuditProgressKV } from "@/server/lib/audit/progress-kv";
 import { crawlPage } from "@/server/workflows/site-audit-workflow-helpers";
+import { pgStep } from "@/server/workflows/pgStep";
 
 const CRAWL_CONCURRENCY = 25;
 
@@ -236,10 +237,15 @@ async function persistCrawlProgress(params: {
     );
   });
 
-  await step.do(`progress-batch-${crawlBatchIndex}`, async () => {
+  await pgStep(
+    step,
+    `progress-batch-${crawlBatchIndex}`,
+    undefined,
+    async () => {
       await AuditRepository.updateAuditProgress(auditId, workflowInstanceId, {
         pagesCrawled,
         pagesTotal: Math.min(visitedCount + queueLength, maxPages),
       });
-  });
+    },
+  );
 }

src/server/workflows/siteAuditWorkflowPhases.ts

@@ -15,6 +15,7 @@ import type {
 } from "@/server/lib/audit/types";
 import { captureServerEvent } from "@/server/lib/posthog";
 import { runCrawlPhase } from "@/server/workflows/siteAuditWorkflowCrawl";
+import { pgStep } from "@/server/workflows/pgStep";
 
 const LIGHTHOUSE_URL_BATCH_SIZE = 10;
 
@@ -103,7 +104,7 @@ async function runDiscoveryPhase(
   origin: string,
   maxPages: number,
 ) {
-  return step.do("discover-urls", async () => {
+  return pgStep(step, "discover-urls", undefined, async () => {
     const result = await discoverUrls(origin, maxPages);
     await AuditRepository.updateAuditProgress(auditId, workflowInstanceId, {
       pagesTotal: Math.min(result.urls.length + 1, maxPages),
@@ -168,8 +169,10 @@ async function runLighthousePhase(
     const counts = countLighthouseBatchResults(lighthouseBatchResults);
     failedChecks += counts.failed;
     completedChecks += counts.completed;
-    await step.do(
+    await pgStep(
+      step,
       `lighthouse-progress-batch-${lighthouseBatchIndex}`,
+      undefined,
       async () => {
         await AuditRepository.updateAuditProgress(auditId, workflowInstanceId, {
           lighthouseCompleted: completedChecks,
@@ -192,7 +195,7 @@ async function selectLighthousePages(params: {
 }) {
   const { step, auditId, workflowInstanceId, allPages, startUrl, strategy } =
     params;
-  return step.do("select-lighthouse-sample", async () => {
+  return pgStep(step, "select-lighthouse-sample", undefined, async () => {
     const sample = selectLighthouseSample(allPages, startUrl, strategy);
     const selectedUrls = new Set(sample);
 
@@ -274,7 +277,7 @@ async function finalizeAudit(args: {
     lighthouseResults,
   } = args;
 
-  await step.do("finalize", async () => {
+  await pgStep(step, "finalize", undefined, async () => {
     await AuditRepository.updateAuditProgress(auditId, workflowInstanceId, {
       currentPhase: "finalizing",
     });