on-zero 0.16.19 → 0.16.20

package/src/createZeroClient.recovery.test.tsx

@@ -6,30 +6,65 @@ import { createRoot, type Root } from 'react-dom/client'
 import { afterEach, beforeEach, expect, test, vi } from 'vitest'
 
 const fakeZero = vi.hoisted(() => {
+  // the client group the next constructed instance joins: tests bump this to
+  // simulate newer code landing (a registry edit) between transitions.
+  const group = { current: 'g1' }
+
   class FakeZero {
     readonly context = {}
-    readonly connection = {
+    readonly connection!: {
       state: {
-        current: { name: 'closed' },
-        subscribe: () => () => {},
-      },
-      connect: vi.fn(),
+        current: { name: string }
+        subscribe: (listener: () => void) => () => void
+      }
+      connect: ReturnType<typeof vi.fn>
     }
     readonly delete = vi.fn(async () => ({ errors: [] }))
     readonly close = vi.fn()
     readonly run = vi.fn(async () => [])
+    readonly mutate = {
+      note: {
+        insert: vi.fn(() => ({
+          client: Promise.resolve({}),
+          server: Promise.resolve({}),
+        })),
+      },
+    }
     readonly preload = vi.fn(() => ({
       cleanup: () => {},
       complete: Promise.resolve(),
     }))
+    readonly clientGroupID: Promise<string>
 
     constructor(readonly options: Record<string, any>) {
       instances.push(this)
+      this.clientGroupID = Promise.resolve(group.current)
+      const listeners = new Set<() => void>()
+      this.connection = {
+        state: {
+          current: { name: 'closed' },
+          subscribe: (listener: () => void) => {
+            listeners.add(listener)
+            return () => {
+              listeners.delete(listener)
+            }
+          },
+        },
+        connect: vi.fn(),
+      }
+      ;(this as any).__connectionListeners = listeners
     }
   }
 
   const instances: FakeZero[] = []
-  return { FakeZero, instances }
+
+  function setConnectionState(instance: FakeZero, name: string) {
+    instance.connection.state.current.name = name
+    const listeners = (instance as any).__connectionListeners as Set<() => void>
+    for (const listener of [...listeners]) listener()
+  }
+
+  return { FakeZero, instances, group, setConnectionState }
 })
 
 vi.mock('@rocicorp/zero', async (importOriginal) => {
@@ -41,7 +76,12 @@ vi.mock('@rocicorp/zero', async (importOriginal) => {
 })
 
 import { createZeroClient } from './createZeroClient'
-import { resetRecoveryStateForTests } from './helpers/recoverZeroClient'
+import {
+  isGroupTransitionOpen,
+  resetRecoveryStateForTests,
+} from './helpers/recoverZeroClient'
+
+import type { ZeroEvent } from './types'
 
 declare global {
   // eslint-disable-next-line no-var
@@ -66,6 +106,7 @@ beforeEach(() => {
   window.sessionStorage.clear()
   resetRecoveryStateForTests()
   fakeZero.instances.length = 0
+  fakeZero.group.current = 'g1'
   container = document.createElement('div')
   root = null
 })
@@ -486,3 +527,325 @@ test('headless NewClientGroup recovery reconstructs its client without a provide
 
   await connection.close()
 })
+
+function makeTransitionClient(instanceName: string) {
+  return createZeroClient({
+    schema,
+    models: {},
+    groupedQueries: {},
+    instanceName,
+  })
+}
+
+function collectClientEvents(client: {
+  zeroEvents: { listen: (listener: (event: ZeroEvent | null) => void) => () => void }
+}) {
+  const events: ZeroEvent[] = []
+  const off = client.zeroEvents.listen((event) => {
+    if (event) events.push(event)
+  })
+  return { events, off }
+}
+
+async function mountTransitionClient(
+  client: ReturnType<typeof makeTransitionClient>,
+  userID: string
+) {
+  const el = document.createElement('div')
+  const mountedRoot = createRoot(el)
+  await act(async () => {
+    mountedRoot.render(
+      <client.ProvideZero cacheURL="http://127.0.0.1:7777/zero" userID={userID}>
+        <span>ok</span>
+      </client.ProvideZero>
+    )
+    await Promise.resolve()
+  })
+  return { el, root: mountedRoot, instance: fakeZero.instances.at(-1)! }
+}
+
+async function unmountTransitionClient(mounted: { root: Root }) {
+  await act(async () => {
+    mounted.root.unmount()
+  })
+}
+
+function withoutPageReload() {
+  const originalLocation = globalThis.location
+  Object.defineProperty(globalThis, 'location', {
+    configurable: true,
+    value: undefined,
+  })
+  return () => {
+    Object.defineProperty(globalThis, 'location', {
+      configurable: true,
+      value: originalLocation,
+    })
+  }
+}
+
+test('two providers ride one group transition and keep querying and mutating', async () => {
+  const restoreLocation = withoutPageReload()
+  const clientA = makeTransitionClient('transition-two-a')
+  const clientB = makeTransitionClient('transition-two-b')
+  const collectedA = collectClientEvents(clientA)
+  const collectedB = collectClientEvents(clientB)
+  // siblings of one user ride one transition.
+  const mountedA = await mountTransitionClient(clientA, 'transition-user')
+  const mountedB = await mountTransitionClient(clientB, 'transition-user')
+  const firstA = mountedA.instance
+  const firstB = mountedB.instance
+  const countBefore = fakeZero.instances.length
+  try {
+    // one registry change: both stale instances fire in the same tick.
+    firstA.options.onUpdateNeeded({ type: 'NewClientGroup' })
+    firstB.options.onUpdateNeeded({ type: 'NewClientGroup' })
+    await act(async () => {
+      await new Promise((resolve) => setTimeout(resolve, 0))
+    })
+
+    // each old instance closed and reminted exactly once.
+    expect(fakeZero.instances).toHaveLength(countBefore + 2)
+    const secondA = fakeZero.instances[countBefore]!
+    const secondB = fakeZero.instances[countBefore + 1]!
+    expect(secondA).not.toBe(firstA)
+    expect(secondB).not.toBe(firstB)
+    expect(firstA.close).toHaveBeenCalledTimes(1)
+    expect(firstB.close).toHaveBeenCalledTimes(1)
+
+    // the transition resolves only when every new group is connected.
+    expect(isGroupTransitionOpen()).toBe(true)
+    await act(async () => {
+      fakeZero.setConnectionState(secondA, 'connected')
+      fakeZero.setConnectionState(secondB, 'connected')
+      await Promise.resolve()
+    })
+    await vi.waitFor(() => expect(isGroupTransitionOpen()).toBe(false))
+
+    // one transition: exactly one recovering across both providers, no fatal.
+    const allEvents = [...collectedA.events, ...collectedB.events]
+    expect(allEvents.filter((event) => event.type === 'recovering')).toHaveLength(1)
+    expect(allEvents.filter((event) => event.type === 'fatal')).toHaveLength(0)
+    // NewClientGroup never deletes IndexedDB.
+    expect(firstA.delete).not.toHaveBeenCalled()
+    expect(firstB.delete).not.toHaveBeenCalled()
+
+    // both providers still query and mutate through their new instances.
+    await clientA.zero.run({} as never)
+    await clientB.zero.run({} as never)
+    expect(secondA.run).toHaveBeenCalled()
+    expect(secondB.run).toHaveBeenCalled()
+    // the fake instance carries an untyped mutate stub: the assertion is that
+    // the facade resolves it on the NEW instance at call time.
+    ;(clientA.zero.mutate as any).note.insert({ id: 'a1' })
+    ;(clientB.zero.mutate as any).note.insert({ id: 'b1' })
+    expect(secondA.mutate.note.insert).toHaveBeenCalledTimes(1)
+    expect(secondB.mutate.note.insert).toHaveBeenCalledTimes(1)
+  } finally {
+    collectedA.off()
+    collectedB.off()
+    await unmountTransitionClient(mountedA)
+    await unmountTransitionClient(mountedB)
+    restoreLocation()
+  }
+})
+
+test('a second registry change after recovery remints again without fatal', async () => {
+  const restoreLocation = withoutPageReload()
+  const clientA = makeTransitionClient('transition-second-a')
+  const clientB = makeTransitionClient('transition-second-b')
+  const collectedA = collectClientEvents(clientA)
+  const collectedB = collectClientEvents(clientB)
+  const mountedA = await mountTransitionClient(clientA, 'transition-second-user')
+  const mountedB = await mountTransitionClient(clientB, 'transition-second-user')
+  try {
+    const fireBoth = (instances: Array<{ options: Record<string, any> }>) => {
+      instances[0]!.options.onUpdateNeeded({ type: 'NewClientGroup' })
+      instances[1]!.options.onUpdateNeeded({ type: 'NewClientGroup' })
+    }
+    const settleRemints = async (countBefore: number) => {
+      await act(async () => {
+        await new Promise((resolve) => setTimeout(resolve, 0))
+      })
+      expect(fakeZero.instances).toHaveLength(countBefore + 2)
+      const next = fakeZero.instances.slice(countBefore)
+      await act(async () => {
+        for (const instance of next) fakeZero.setConnectionState(instance, 'connected')
+        await Promise.resolve()
+      })
+      await vi.waitFor(() => expect(isGroupTransitionOpen()).toBe(false))
+      return next
+    }
+
+    // first change.
+    const countBeforeFirst = fakeZero.instances.length
+    fireBoth(fakeZero.instances.slice(-2))
+    await settleRemints(countBeforeFirst)
+
+    // second change with newer code, well inside the old 12s remint cooldown.
+    fakeZero.group.current = 'g2'
+    const countBefore = fakeZero.instances.length
+    fireBoth(fakeZero.instances.slice(-2))
+    const third = await settleRemints(countBefore)
+
+    const allEvents = [...collectedA.events, ...collectedB.events]
+    expect(allEvents.filter((event) => event.type === 'recovering')).toHaveLength(2)
+    expect(allEvents.filter((event) => event.type === 'fatal')).toHaveLength(0)
+    // the facades resolve to the newest instances.
+    await clientA.zero.run({} as never)
+    await clientB.zero.run({} as never)
+    expect(third[0]!.run).toHaveBeenCalled()
+    expect(third[1]!.run).toHaveBeenCalled()
+  } finally {
+    collectedA.off()
+    collectedB.off()
+    await unmountTransitionClient(mountedA)
+    await unmountTransitionClient(mountedB)
+    restoreLocation()
+  }
+})
+
+test('negative control: an incompatible schema emits exactly one terminal error, never a loop', async () => {
+  const restoreLocation = withoutPageReload()
+  const consoleError = vi.spyOn(console, 'error').mockImplementation(() => {})
+  const clientA = makeTransitionClient('transition-negative-a')
+  const clientB = makeTransitionClient('transition-negative-b')
+  const collectedA = collectClientEvents(clientA)
+  const collectedB = collectClientEvents(clientB)
+  const mountedA = await mountTransitionClient(clientA, 'transition-negative-user')
+  const mountedB = await mountTransitionClient(clientB, 'transition-negative-user')
+  const firstA = mountedA.instance
+  const firstB = mountedB.instance
+  const countBefore = fakeZero.instances.length
+  try {
+    firstA.options.onUpdateNeeded({ type: 'SchemaVersionNotSupported' })
+    firstB.options.onUpdateNeeded({ type: 'SchemaVersionNotSupported' })
+    await act(async () => {
+      await new Promise((resolve) => setTimeout(resolve, 0))
+    })
+    // schema incompatibility drops the stale stores, then remints once each.
+    expect(firstA.delete).toHaveBeenCalledTimes(1)
+    expect(firstB.delete).toHaveBeenCalledTimes(1)
+    expect(fakeZero.instances).toHaveLength(countBefore + 2)
+    const secondA = fakeZero.instances[countBefore]!
+    const secondB = fakeZero.instances[countBefore + 1]!
+    expect(isGroupTransitionOpen()).toBe(true)
+
+    // the new code still cannot join the server group: same group fires again.
+    secondA.options.onUpdateNeeded({ type: 'SchemaVersionNotSupported' })
+    await act(async () => {
+      await new Promise((resolve) => setTimeout(resolve, 0))
+    })
+
+    const allEvents = () => [...collectedA.events, ...collectedB.events]
+    expect(allEvents().filter((event) => event.type === 'fatal')).toHaveLength(1)
+    const [fatal] = allEvents().filter((event) => event.type === 'fatal')
+    expect(fatal).toMatchObject({
+      reasonKey: 'SchemaVersionNotSupported',
+      reason: expect.stringContaining('g1'),
+    })
+    expect(consoleError).toHaveBeenCalledTimes(1)
+    expect(consoleError.mock.calls[0]![0]).toContain('g1')
+
+    // every later re-fire stays silent and remints nothing: never a loop.
+    secondB.options.onUpdateNeeded({ type: 'SchemaVersionNotSupported' })
+    secondA.options.onUpdateNeeded({ type: 'SchemaVersionNotSupported' })
+    await act(async () => {
+      await new Promise((resolve) => setTimeout(resolve, 0))
+    })
+    expect(allEvents().filter((event) => event.type === 'fatal')).toHaveLength(1)
+    expect(consoleError).toHaveBeenCalledTimes(1)
+    expect(fakeZero.instances).toHaveLength(countBefore + 2)
+  } finally {
+    collectedA.off()
+    collectedB.off()
+    await unmountTransitionClient(mountedA)
+    await unmountTransitionClient(mountedB)
+    restoreLocation()
+    consoleError.mockRestore()
+  }
+})
+
+test('concurrent update signals across providers schedule a single reload', async () => {
+  const reloadPage = vi.fn()
+  const originalLocation = globalThis.location
+  Object.defineProperty(globalThis, 'location', {
+    configurable: true,
+    value: { href: 'http://localhost/reload-test', reload: reloadPage },
+  })
+  const scheduleA = vi.fn()
+  const scheduleB = vi.fn()
+  const clientA = makeTransitionClient('transition-reload-a')
+  const clientB = makeTransitionClient('transition-reload-b')
+  const collectedA = collectClientEvents(clientA)
+  const collectedB = collectClientEvents(clientB)
+  const elA = document.createElement('div')
+  const elB = document.createElement('div')
+  const rootA = createRoot(elA)
+  const rootB = createRoot(elB)
+  try {
+    await act(async () => {
+      rootA.render(
+        <clientA.ProvideZero
+          cacheURL="http://127.0.0.1:7777/zero"
+          userID="transition-reload-user"
+          scheduleReload={scheduleA}
+        >
+          <span>ok</span>
+        </clientA.ProvideZero>
+      )
+      await Promise.resolve()
+    })
+    await act(async () => {
+      rootB.render(
+        <clientB.ProvideZero
+          cacheURL="http://127.0.0.1:7777/zero"
+          userID="transition-reload-user"
+          scheduleReload={scheduleB}
+        >
+          <span>ok</span>
+        </clientB.ProvideZero>
+      )
+      await Promise.resolve()
+    })
+    const firstA = fakeZero.instances.at(-2)!
+    const firstB = fakeZero.instances.at(-1)!
+
+    firstA.options.onUpdateNeeded({ type: 'NewClientGroup' })
+    firstB.options.onUpdateNeeded({ type: 'NewClientGroup' })
+    await act(async () => {
+      await new Promise((resolve) => setTimeout(resolve, 0))
+    })
+
+    // one transition drives one deferred reload, however many providers join.
+    expect(scheduleA.mock.calls.length + scheduleB.mock.calls.length).toBe(1)
+    expect(reloadPage).not.toHaveBeenCalled()
+    const allEvents = [...collectedA.events, ...collectedB.events]
+    expect(allEvents.filter((event) => event.type === 'fatal')).toHaveLength(0)
+
+    const ctx = (scheduleA.mock.calls[0] ?? scheduleB.mock.calls[0])![0]
+    await act(async () => {
+      await ctx.performReload()
+    })
+    expect(reloadPage).toHaveBeenCalledTimes(1)
+    expect(isGroupTransitionOpen()).toBe(false)
+    // the reload persisted the recovered-from groups for the next page load.
+    const marker = window.sessionStorage.getItem(
+      'on-zero-recover-http://localhost/reload-test-NewClientGroup'
+    )
+    expect(marker).toContain('g1')
+  } finally {
+    collectedA.off()
+    collectedB.off()
+    await act(async () => {
+      rootA.unmount()
+    })
+    await act(async () => {
+      rootB.unmount()
+    })
+    Object.defineProperty(globalThis, 'location', {
+      configurable: true,
+      value: originalLocation,
+    })
+  }
+})

package/src/createZeroClient.tsx

@@ -39,8 +39,10 @@ import { createMutationLifecycle } from './helpers/mutationLifecycle'
 import { IS_SERVER_RUNTIME } from './helpers/platform'
 import {
   composeRecoveryLogSink,
+  isGroupTransitionOpen,
   isRecoverableZeroStalePokeMessage,
   makeZeroRecovery,
+  primeGroupCache,
   type RecoveryGuardStorage,
   type ScheduleReloadContext,
   type ZeroLogPattern,
@@ -49,6 +51,7 @@ import {
 import { observeMutation, reportMutationInvocationError } from './helpers/useMutation'
 import { registerClientInstance } from './instanceRegistry'
 import { getAllMutationsPermissions, getMutationsPermissions } from './modelRegistry'
+import { assertQueryJson } from './queryJson'
 import { registerQuery } from './queryRegistry'
 import { resolveQuery, type PlainQueryFn } from './resolveQuery'
 import { setCustomQueries } from './run'
@@ -207,9 +210,11 @@ export function createZeroClientInternal<
     for (const [name, fn] of Object.entries(queries)) {
       registerQuery(fn, `${namespace}.${name}`)
       // wrap each plain function in defineQuery
-      wrappedNamespaces[namespace][name] = defineQuery(({ args }: { args: any }) =>
-        fn(args)
-      )
+      wrappedNamespaces[namespace][name] = defineQuery(({ args }: { args: any }) => {
+        const query = fn(args)
+        assertQueryJson(`${namespace}.${name}`, query.ast, 'ast')
+        return query
+      })
     }
   }
 
@@ -351,6 +356,7 @@ export function createZeroClientInternal<
     guardStorage?: RecoveryGuardStorage
     benignLogPatterns?: readonly ZeroLogPattern[]
     recoverInPlace?: () => Promise<boolean>
+    awaitReconnected?: () => Promise<boolean>
     reload?: () => void | boolean
   }
 
@@ -366,6 +372,7 @@ export function createZeroClientInternal<
     guardStorage,
     benignLogPatterns,
     recoverInPlace,
+    awaitReconnected,
     reload,
   }: ConstructZeroInstanceArgs): ZeroInstance {
     let installedTransport: Pick<HttpPullTransport, 'subscribeLifecycle'> | undefined
@@ -397,7 +404,13 @@ export function createZeroClientInternal<
       ],
       onRecovery: () => mutationLifecycle.fence(),
       recoverInPlace: recoverInPlace ?? (() => remint({ dropLocalState: false })),
+      awaitReconnected: awaitReconnected ?? awaitReconnectedAfterRemint,
       reload,
+      // sibling instances of one user ride one group transition.
+      clientIdentity: typeof options.userID === 'string' ? options.userID : undefined,
+      getGroupID: () =>
+        instanceRef.current?.clientGroupID ??
+        Promise.reject(new Error('[on-zero] no zero instance')),
     }
     const recovery = makeZeroRecovery(recoveryDeps)
     const createdInstance = new ZeroClient<Schema, ZeroMutators>({
@@ -421,6 +434,7 @@ export function createZeroClientInternal<
     })
     instanceRef.current = createdInstance
     if (installedTransport) installedTransports.set(createdInstance, installedTransport)
+    primeGroupCache(recoveryDeps)
     return createdInstance
   }
 
@@ -689,6 +703,67 @@ export function createZeroClientInternal<
   // revives that instance instead of constructing a twin beside it.
   const retiredZero = new Map<string, CachedZeroEntry>()
 
+  // resolve when this instance's connection reports connected. never rejects:
+  // a transition participant that never connects simply never resolves its
+  // transition, and a still-stale instance re-fires into a new one.
+  function waitForInstanceConnected(instanceToWatch: ZeroInstance): Promise<boolean> {
+    return new Promise((resolve) => {
+      let unsubscribe: (() => void) | undefined
+      let done = false
+      const check = () => {
+        if (done) return
+        let name: string | undefined
+        try {
+          name = instanceToWatch.connection?.state?.current?.name
+        } catch {
+          name = undefined
+        }
+        if (name === 'connected') {
+          done = true
+          try {
+            unsubscribe?.()
+          } catch {}
+          resolve(true)
+        }
+      }
+      check()
+      if (done) return
+      try {
+        unsubscribe = instanceToWatch.connection.state.subscribe(check)
+      } catch {
+        done = true
+        resolve(false)
+      }
+    })
+  }
+
+  // the group transition's connected gate for the provider path: remint()
+  // bumps the mounted provider, the rotate publishes a replacement instance,
+  // and the transition resolves only once that replacement is connected.
+  async function awaitReconnectedAfterRemint(): Promise<boolean> {
+    // the pre-remint instance, captured synchronously when remint() ran:
+    // reading zeroRuntime.zero here would race the rotate (an already-new
+    // instance must not wait for a second rotation that never comes, and a
+    // drop-remint's transient null must wait through instead of resolving).
+    const previous = remintPreviousInstance
+    let current = zeroRuntime.zero
+    if (current === previous || current === null) {
+      // the rotate publishes asynchronously after the bump; the re-check runs
+      // synchronously inside the executor so a publish racing this wait still
+      // resolves instead of dangling the waiter.
+      current = await new Promise<ZeroInstance | null>((resolve) => {
+        const latest = zeroRuntime.zero
+        if (latest !== previous && latest !== null) {
+          resolve(latest)
+          return
+        }
+        zeroRuntime.readyWaiters.add(resolve as (instance: ZeroInstance) => void)
+      })
+    }
+    if (!current) return false
+    return waitForInstanceConnected(current)
+  }
+
   // in-place re-mint: drop the current instance's local state then reconstruct a
   // fresh client WITHOUT a page reload — the native-safe recovery path (a reload
   // may never land on prod native, wedging the module latch). the mounted
@@ -702,6 +777,7 @@ export function createZeroClientInternal<
   const remintControl: { bump: (() => void) | null } = { bump: null }
   let lastRemintAt = 0
   let remintAttempts = 0
+  let remintPreviousInstance: ZeroInstance | null = null
 
   function unpublishZeroInstance(instanceToInvalidate: ZeroInstance): boolean {
     if (zeroRuntime.zero !== instanceToInvalidate) return false
@@ -752,13 +828,22 @@ export function createZeroClientInternal<
     // no mounted provider to reconstruct through — bail BEFORE the guard so an
     // unmounted call doesn't burn an attempt or start the cooldown.
     if (!remintControl.bump) return false
-    const now = Date.now()
-    const sinceLast = now - lastRemintAt
-    if (lastRemintAt > 0 && sinceLast < REMINT_GUARD_MS) return false
-    if (sinceLast > REMINT_ATTEMPT_RESET_MS) remintAttempts = 0
-    if (remintAttempts >= REMINT_MAX_ATTEMPTS) return false
-    lastRemintAt = now
-    remintAttempts += 1
+    // capture the outgoing instance synchronously for the transition's
+    // connected gate (see awaitReconnectedAfterRemint).
+    remintPreviousInstance = zeroRuntime.zero
+    // a group transition in flight owns the retry policy (one remint per
+    // provider per transition, terminal-on-failure): its remints bypass the
+    // cooldown and the attempt budget untouched, so rapid legitimate
+    // transitions all land while lone storms stay guarded.
+    if (!isGroupTransitionOpen()) {
+      const now = Date.now()
+      const sinceLast = now - lastRemintAt
+      if (lastRemintAt > 0 && sinceLast < REMINT_GUARD_MS) return false
+      if (sinceLast > REMINT_ATTEMPT_RESET_MS) remintAttempts = 0
+      if (remintAttempts >= REMINT_MAX_ATTEMPTS) return false
+      lastRemintAt = now
+      remintAttempts += 1
+    }
 
     const { dropLocalState = true } = opts
     if (dropLocalState && zeroRuntime.zero) {
@@ -1234,6 +1319,11 @@ export function createZeroClientInternal<
         // a headless host has no provider state to bump and no page it owns to
         // reload. replace the disabled instance directly, then republish the
         // stable module facade and move the connection watcher with it.
+        // the transition resolves once the replacement connects.
+        awaitReconnected: async () => {
+          if (closed) return true
+          return waitForInstanceConnected(activeInstance)
+        },
         recoverInPlace: async () => {
           // an update can arrive after the host intentionally closed this
           // connection during preview turnover. there is no live client left

package/src/generate.test.ts

@@ -139,6 +139,37 @@ export const relations = defineRelations(schema, (r) => ({
     )
   })
 
+  test('preserves schemas owned by drizzle-zero configuration', async () => {
+    const dataDir = join(testDir, 'src/data')
+    const schemaPath = join(dataDir, 'generated/schema.ts')
+    const externalSchema = `import { createSchema, table, string } from '@rocicorp/zero'
+export const schema = createSchema({
+  tables: [table('post').columns({ id: string() }).primaryKey('id')],
+})
+`
+    writeFileSync(join(dataDir, 'post.ts'), `export const posts = () => zql.post`)
+    writeFileSync(
+      join(testDir, 'src/database/schema.ts'),
+      `export * from './schema-public'`
+    )
+    writeFileSync(
+      join(testDir, 'src/database/schema-public.ts'),
+      `
+import { pgTable, text } from 'drizzle-orm/pg-core'
+export const post = pgTable('post', { id: text('id').primaryKey() })
+`
+    )
+    writeFileSync(
+      join(testDir, 'src/database/drizzle-zero.config.ts'),
+      `export default {}`
+    )
+    writeFileSync(schemaPath, externalSchema)
+
+    await generate({ dir: dataDir, silent: true })
+
+    expect(readFileSync(schemaPath, 'utf8')).toBe(externalSchema)
+  })
+
   test('generates every crud slot for a default table registration', async () => {
     writeFileSync(
       join(testDir, 'post/mutations.ts'),