fix(engine): record per-parent backfill progress as one watermark

packages/sync-engine/src/state.ts

@@ -40,7 +40,21 @@ export interface BackfillState {
   status: 'pending' | 'running' | 'complete'
   /** Pagination cursor within the current listing. */
   after?: string | null
-  /** For N+1 streams: parents already fully listed. */
+  /**
+   * For N+1 streams: every parent whose id sorts at or before this one is
+   * fully listed. Parents are listed in id order, and WorkOS ids are
+   * time-ordered, so a parent created after a slice sorts past the watermark
+   * and the next slice lists it. One id however many parents the stream has:
+   * a list of every finished parent grew with the parent stream and was
+   * rewritten at every parent boundary, which a document store's size limit
+   * does not survive.
+   */
+  doneThrough?: string | null
+  /**
+   * For N+1 streams: parents fully listed by a release that recorded each one
+   * by id. Honored on resume and carried until `doneThrough` passes them;
+   * never written afresh.
+   */
   doneParents?: string[]
   /**
    * For N+1 streams: the parent whose listing `after` belongs to, when a page

packages/sync-engine/src/sync/backfill.ts

@@ -137,7 +137,16 @@ export async function backfillStream(
     }
   }
 
-  const doneParents = new Set(prior?.doneParents ?? [])
+  // Parents already fully listed: one watermark in id order, plus whatever an
+  // older release recorded parent by parent, each carried until the
+  // watermark passes it. See `BackfillState.doneThrough`.
+  let doneThrough = prior?.doneThrough ?? null
+  const recorded = new Set(prior?.doneParents ?? [])
+  const parentDone = (pid: string) => (doneThrough !== null && pid <= doneThrough) || recorded.has(pid)
+  const parentProgress = (): Pick<BackfillState, 'doneThrough' | 'doneParents'> => {
+    const carried = [...recorded].filter((pid) => doneThrough === null || pid > doneThrough)
+    return carried.length > 0 ? { doneThrough, doneParents: carried } : { doneThrough }
+  }
 
   const write = async (rows: Record<string, unknown>[]) => {
     // REST is the authoritative full object, so it replaces `_raw` rather than
@@ -182,13 +191,15 @@ export async function backfillStream(
 
   if (stream.list.parent) {
     const parentDef = streamByName(stream.list.parent.stream)
-    const parents = parentDef ? await destination.liveIds(parentDef, opts.env) : []
+    // Sorted here, not by the destination: the watermark compares ids the
+    // way this sort orders them, and a database collation need not.
+    const parents = (parentDef ? await destination.liveIds(parentDef, opts.env) : []).sort()
     opts.logger?.info(`${label} starting (per-parent)`, {
       stream: stream.name,
       parents: parents.length,
     })
     for (const pid of parents) {
-      if (doneParents.has(pid)) continue
+      if (parentDone(pid)) continue
       // The previous parent's checkpoint is already written, so stopping here
       // leaves nothing to redo.
       if (spent()) return paused({ parent: pid })
@@ -214,18 +225,20 @@ export async function backfillStream(
         await state.set(opts.env, stream.name, {
           kind: 'backfill',
           status: 'running',
-          doneParents: [...doneParents],
+          ...parentProgress(),
           parent: pid,
           after: page.after,
           count: stats.created + stats.updated,
         } satisfies BackfillState)
         if (spent()) return paused({ parent: pid, after: page.after })
       }
-      doneParents.add(pid)
+      // Every parent before this one in id order was either skipped as done
+      // or listed just now, so the watermark can move to it.
+      doneThrough = pid
       await state.set(opts.env, stream.name, {
         kind: 'backfill',
         status: 'running',
-        doneParents: [...doneParents],
+        ...parentProgress(),
         count: stats.created + stats.updated,
       } satisfies BackfillState)
     }

packages/sync-engine/test/backfill.test.ts

@@ -0,0 +1,159 @@
+/**
+ * The engine half of a per-parent backfill: what it records between slices
+ * and what a resumed slice lists. A destination is stubbed with an in-memory
+ * state store so the checkpoints themselves can be read; the same behavior
+ * against a real database is in sync-postgres/test/engine.test.ts.
+ */
+import { describe, it, expect, beforeEach } from 'vitest'
+import { backfillStream, pageBudget } from '../src/sync/backfill.js'
+import { StateStore, type BackfillState } from '../src/state.js'
+import { WorkOSClient } from '../src/workos/client.js'
+import { streamByName } from '../src/catalog.js'
+import type { Destination } from '../src/destination.js'
+import { emptyStats } from '../src/types.js'
+import { FakeWorkOS, resetIds } from './helpers/fake-workos.js'
+
+const memberships = streamByName('organization_memberships')!
+
+describe('per-parent backfill progress', () => {
+  let fake: FakeWorkOS
+  let api: WorkOSClient
+  let states: Map<string, unknown>
+  let parents: string[]
+  let destination: Destination
+  let state: StateStore
+
+  beforeEach(() => {
+    resetIds()
+    fake = new FakeWorkOS()
+    api = new WorkOSClient({ apiKey: 'sk_test_1', apiUrl: 'http://fake.workos.local', fetchImpl: fake.fetchImpl })
+    states = new Map()
+    parents = []
+    destination = {
+      init: async () => {},
+      close: async () => {},
+      upsert: async (_s, _e, rows) => ({ ...emptyStats(), created: rows.length }),
+      delete: async () => emptyStats(),
+      relink: async () => emptyStats(),
+      unlink: async () => emptyStats(),
+      liveIds: async () => [...parents],
+      rawRows: async () => [],
+      diffIds: async () => ({ matched: 0, localOnly: [], remoteOnly: [] }),
+      freshness: async (s) => ({
+        stream: s.name,
+        rows: 0,
+        deleted: 0,
+        lastSyncedAt: null,
+        stalenessSeconds: null,
+        lagSeconds: null,
+        lagSampleSize: 0,
+      }),
+      redact: async () => [],
+      getState: async <T>(env: string, stream: string) => (states.get(`${env}:${stream}`) as T | undefined) ?? null,
+      setState: async (env, stream, value) => {
+        states.set(`${env}:${stream}`, value)
+      },
+      setStateIfHeld: async () => true,
+      acquireLease: async (_e, holder) => ({ acquired: true, holder, expiresAt: null }),
+      releaseLease: async () => {},
+      getLease: async () => null,
+      reject: async () => {},
+      rejected: async () => [],
+      rejectedCount: async () => 0,
+      clearRejected: async () => 0,
+    }
+    state = new StateStore(destination)
+  })
+
+  /** An organization with `n` memberships, registered as a live parent. */
+  const org = (n = 1) => {
+    const o = fake.createOrganization()
+    const id = o['id'] as string
+    for (let i = 0; i < n; i++) fake.createMembership(fake.createUser()['id'] as string, id)
+    parents.push(id)
+    return id
+  }
+  const stored = () => states.get('prod:organization_memberships') as BackfillState
+  /** The parents listed so far, in request order. */
+  const listed = () =>
+    fake.requestLog
+      .filter((r) => r.startsWith('/user_management/organization_memberships'))
+      .map((r) => new URL(r, 'http://fake.workos.local').searchParams.get('organization_id'))
+  const run = (opts: { resume?: boolean; maxPages?: number } = {}) =>
+    backfillStream(destination, api, state, memberships, {
+      env: 'prod',
+      resume: opts.resume,
+      budget: pageBudget(opts.maxPages),
+    })
+
+  it('records the last parent finished, not every parent, and resumes after it', async () => {
+    const [a, b, c, d] = [org(), org(), org(), org()]
+
+    expect(await run({ maxPages: 2 })).toMatchObject({ pages: 2, complete: false })
+    expect(listed()).toEqual([a, b])
+    expect(stored()).toMatchObject({ status: 'running', doneThrough: b, count: 2 })
+    expect(stored()).not.toHaveProperty('doneParents')
+    expect(stored()).not.toHaveProperty('parent')
+
+    expect(await run({ resume: true })).toMatchObject({ pages: 2, complete: true })
+    expect(listed()).toEqual([a, b, c, d])
+    expect(stored()).toMatchObject({ status: 'complete' })
+  })
+
+  it('checkpoints inside a parent with the watermark of the parents before it', async () => {
+    const a = org()
+    const b = org(150)
+
+    // One page for a, then the first of b's two.
+    expect(await run({ maxPages: 2 })).toMatchObject({ pages: 2, complete: false })
+    expect(stored()).toMatchObject({ status: 'running', doneThrough: a, parent: b })
+    expect(stored().after).toEqual(expect.any(String))
+
+    // The resumed slice asks for b's second page and nothing else.
+    const before = fake.requestLog.length
+    expect(await run({ resume: true })).toMatchObject({ pages: 1, complete: true })
+    const since = fake.requestLog.slice(before)
+    expect(since).toHaveLength(1)
+    expect(since[0]).toContain(`organization_id=${b}`)
+    expect(since[0]).toContain('after=')
+  })
+
+  it('sorts the parents itself rather than trusting the order the destination returns', async () => {
+    const [a, b, c, d] = [org(), org(), org(), org()]
+    parents = [d, b, a, c]
+
+    expect(await run({ maxPages: 2 })).toMatchObject({ pages: 2, complete: false })
+    expect(listed()).toEqual([a, b])
+    expect(stored()).toMatchObject({ doneThrough: b })
+  })
+
+  it('lists a parent created after a slice, since its id sorts past the watermark', async () => {
+    const [a, b, c] = [org(), org(), org()]
+    await run({ maxPages: 2 })
+    expect(stored()).toMatchObject({ doneThrough: b })
+
+    const late = org()
+    expect(await run({ resume: true })).toMatchObject({ pages: 2, complete: true })
+    expect(listed()).toEqual([a, b, c, late])
+  })
+
+  it('honors a checkpoint that recorded each finished parent by id, and carries those past the watermark', async () => {
+    const [a, b, c, d] = [org(), org(), org(), org()]
+    // What a release before the watermark left behind after finishing a and c.
+    states.set('prod:organization_memberships', {
+      kind: 'backfill',
+      status: 'running',
+      doneParents: [c, a],
+      count: 2,
+    } satisfies BackfillState)
+
+    // b is the first parent not yet done; the one page budgeted goes to it.
+    expect(await run({ resume: true, maxPages: 1 })).toMatchObject({ pages: 1, complete: false })
+    expect(listed()).toEqual([b])
+    // a is behind the watermark now and dropped; c is still ahead of it and kept.
+    expect(stored()).toMatchObject({ doneThrough: b, doneParents: [c] })
+
+    expect(await run({ resume: true })).toMatchObject({ pages: 1, complete: true })
+    expect(listed()).toEqual([b, d])
+  })
+})