fix(drizzle): ensure getPrimaryDb is used for all write operations (#16240)

r1tsuu · web-flow · commit aa446490e4c6 · 2026-04-10T14:32:05.000Z
PR #16083 added `getPrimaryDb` logic to fix the broken behaviour with the `readReplica` config option but didn't include it to some other operations like with versions or jobs. This PR ensures that all the write operations use it.
diff --git a/packages/db-postgres/src/types.ts b/packages/db-postgres/src/types.ts
@@ -9,9 +9,10 @@ import type {
 import type { DrizzleConfig } from 'drizzle-orm'
 import type { NodePgDatabase } from 'drizzle-orm/node-postgres'
 import type { PgSchema, PgTableFn, PgTransactionConfig, PgWithReplicas } from 'drizzle-orm/pg-core'
+import type pg from 'pg'
 import type { Pool, PoolConfig } from 'pg'
 
-type PgDependency = typeof import('pg').default
+type PgDependency = typeof pg
 
 export type Args = {
   /**
diff --git a/packages/drizzle/src/createGlobal.ts b/packages/drizzle/src/createGlobal.ts
@@ -5,6 +5,7 @@ import toSnakeCase from 'to-snake-case'
 import type { DrizzleAdapter } from './types.js'
 
 import { upsertRow } from './upsertRow/index.js'
+import { getPrimaryDb } from './utilities/getPrimaryDb.js'
 import { getTransaction } from './utilities/getTransaction.js'
 
 export async function createGlobal<T extends Record<string, unknown>>(
@@ -17,7 +18,7 @@ export async function createGlobal<T extends Record<string, unknown>>(
 
   data.createdAt = new Date().toISOString()
 
-  const db = await getTransaction(this, req)
+  const db = getPrimaryDb(this, await getTransaction(this, req))
 
   const result = await upsertRow<{ globalType: string } & T>({
     adapter: this,
diff --git a/packages/drizzle/src/createGlobalVersion.ts b/packages/drizzle/src/createGlobalVersion.ts
@@ -8,6 +8,7 @@ import toSnakeCase from 'to-snake-case'
 import type { DrizzleAdapter } from './types.js'
 
 import { upsertRow } from './upsertRow/index.js'
+import { getPrimaryDb } from './utilities/getPrimaryDb.js'
 import { getTransaction } from './utilities/getTransaction.js'
 
 export async function createGlobalVersion<T extends JsonObject = JsonObject>(
@@ -29,7 +30,7 @@ export async function createGlobalVersion<T extends JsonObject = JsonObject>(
 
   const tableName = this.tableNameMap.get(`_${toSnakeCase(global.slug)}${this.versionsSuffix}`)
 
-  const db = await getTransaction(this, req)
+  const db = getPrimaryDb(this, await getTransaction(this, req))
 
   const result = await upsertRow<TypeWithVersion<T>>({
     adapter: this,
diff --git a/packages/drizzle/src/createVersion.ts b/packages/drizzle/src/createVersion.ts
@@ -8,6 +8,7 @@ import toSnakeCase from 'to-snake-case'
 import type { DrizzleAdapter } from './types.js'
 
 import { upsertRow } from './upsertRow/index.js'
+import { getPrimaryDb } from './utilities/getPrimaryDb.js'
 import { getTransaction } from './utilities/getTransaction.js'
 
 export async function createVersion<T extends JsonObject = JsonObject>(
@@ -52,7 +53,7 @@ export async function createVersion<T extends JsonObject = JsonObject>(
     version,
   }
 
-  const db = await getTransaction(this, req)
+  const db = getPrimaryDb(this, await getTransaction(this, req))
 
   const result = await upsertRow<TypeWithVersion<T>>({
     adapter: this,
diff --git a/packages/drizzle/src/deleteMany.ts b/packages/drizzle/src/deleteMany.ts
@@ -55,6 +55,9 @@ export const deleteMany: DeleteMany = async function deleteMany(
     )
   }
 
+  // No getPrimaryDb needed: db is only used for deleteWhere (a write, always routed to primary
+  // by drizzle's withReplicas). findMany resolves its own db via getTransaction, which returns
+  // the transaction db (always primary) or falls back to shouldReadFromPrimary.
   const db = await getTransaction(this, req)
 
   await this.deleteWhere({
diff --git a/packages/drizzle/src/deleteOne.ts b/packages/drizzle/src/deleteOne.ts
@@ -9,6 +9,7 @@ import { buildFindManyArgs } from './find/buildFindManyArgs.js'
 import { buildQuery } from './queries/buildQuery.js'
 import { selectDistinct } from './queries/selectDistinct.js'
 import { transform } from './transform/read/index.js'
+import { getPrimaryDb } from './utilities/getPrimaryDb.js'
 import { getTransaction } from './utilities/getTransaction.js'
 import { markWrite } from './utilities/readAfterWrite.js'
 
@@ -30,7 +31,7 @@ export const deleteOne: DeleteOne = async function deleteOne(
     where: whereArg,
   })
 
-  const db = await getTransaction(this, req)
+  const db = getPrimaryDb(this, await getTransaction(this, req))
 
   const selectDistinctResult = await selectDistinct({
     adapter: this,
diff --git a/packages/drizzle/src/deleteVersions.ts b/packages/drizzle/src/deleteVersions.ts
@@ -52,6 +52,9 @@ export const deleteVersions: DeleteVersions = async function deleteVersion(
   })
 
   if (ids.length > 0) {
+    // No getPrimaryDb needed: db is only used for deleteWhere (a write, always routed to primary
+    // by drizzle's withReplicas). findMany resolves its own db via getTransaction, which returns
+    // the transaction db (always primary) or falls back to shouldReadFromPrimary.
     const db = await getTransaction(this, req)
 
     await this.deleteWhere({
diff --git a/packages/drizzle/src/updateGlobal.ts b/packages/drizzle/src/updateGlobal.ts
@@ -5,6 +5,7 @@ import toSnakeCase from 'to-snake-case'
 import type { DrizzleAdapter } from './types.js'
 
 import { upsertRow } from './upsertRow/index.js'
+import { getPrimaryDb } from './utilities/getPrimaryDb.js'
 import { getTransaction } from './utilities/getTransaction.js'
 
 export async function updateGlobal<T extends Record<string, unknown>>(
@@ -14,7 +15,7 @@ export async function updateGlobal<T extends Record<string, unknown>>(
   const globalConfig = this.payload.globals.config.find((config) => config.slug === slug)
   const tableName = this.tableNameMap.get(toSnakeCase(globalConfig.slug))
 
-  const db = await getTransaction(this, req)
+  const db = getPrimaryDb(this, await getTransaction(this, req))
   const existingGlobal = await db.query[tableName].findFirst({})
 
   const result = await upsertRow<{ globalType: string } & T>({
diff --git a/packages/drizzle/src/updateGlobalVersion.ts b/packages/drizzle/src/updateGlobalVersion.ts
@@ -12,6 +12,7 @@ import type { DrizzleAdapter } from './types.js'
 
 import { buildQuery } from './queries/buildQuery.js'
 import { upsertRow } from './upsertRow/index.js'
+import { getPrimaryDb } from './utilities/getPrimaryDb.js'
 import { getTransaction } from './utilities/getTransaction.js'
 
 export async function updateGlobalVersion<T extends JsonObject = JsonObject>(
@@ -46,7 +47,7 @@ export async function updateGlobalVersion<T extends JsonObject = JsonObject>(
     where: whereToUse,
   })
 
-  const db = await getTransaction(this, req)
+  const db = getPrimaryDb(this, await getTransaction(this, req))
 
   const result = await upsertRow<TypeWithVersion<T>>({
     id,
diff --git a/packages/drizzle/src/updateJobs.ts b/packages/drizzle/src/updateJobs.ts
@@ -67,7 +67,7 @@ export const updateJobs: UpdateJobs = async function updateMany(
     return []
   }
 
-  const db = await getTransaction(this, req)
+  const db = getPrimaryDb(this, await getTransaction(this, req))
 
   const results = []
 
diff --git a/packages/drizzle/src/updateVersion.ts b/packages/drizzle/src/updateVersion.ts
@@ -12,6 +12,7 @@ import type { DrizzleAdapter } from './types.js'
 
 import { buildQuery } from './queries/buildQuery.js'
 import { upsertRow } from './upsertRow/index.js'
+import { getPrimaryDb } from './utilities/getPrimaryDb.js'
 import { getTransaction } from './utilities/getTransaction.js'
 
 export async function updateVersion<T extends JsonObject = JsonObject>(
@@ -43,7 +44,7 @@ export async function updateVersion<T extends JsonObject = JsonObject>(
     where: whereToUse,
   })
 
-  const db = await getTransaction(this, req)
+  const db = getPrimaryDb(this, await getTransaction(this, req))
 
   const result = await upsertRow<TypeWithVersion<T>>({
     id,
diff --git a/test/database/pg-replica/int.spec.ts b/test/database/pg-replica/int.spec.ts
@@ -40,6 +40,30 @@ describeReplica('postgres read replicas', () => {
               type: 'text',
             },
           ],
+          versions: {
+            drafts: true,
+          },
+        },
+      ],
+      globals: [
+        {
+          slug: 'settings',
+          fields: [
+            {
+              name: 'siteTitle',
+              type: 'text',
+            },
+          ],
+        },
+        {
+          slug: 'nav',
+          fields: [
+            {
+              name: 'label',
+              type: 'text',
+            },
+          ],
+          versions: true,
         },
       ],
     })
@@ -178,4 +202,146 @@ describeReplica('postgres read replicas', () => {
       expect(found.title).toBe('old-write')
     })
   })
+
+  describe('version operations use primary for read-back', () => {
+    it('should create a document with a version without errors', async () => {
+      // This exercises createVersion which previously lacked getPrimaryDb,
+      // causing the read-back after insert to hit a stale replica
+      const doc = await (payload as any).create({
+        collection: 'posts',
+        data: { title: 'versioned-doc', _status: 'draft' },
+        draft: true,
+      })
+
+      expect(doc).toBeDefined()
+      expect(doc.title).toBe('versioned-doc')
+
+      const versions = await (payload as any).findVersions({
+        collection: 'posts',
+        where: { parent: { equals: doc.id } },
+      })
+
+      expect(versions.docs.length).toBeGreaterThanOrEqual(1)
+    })
+
+    it('should update a draft and create a new version without errors', async () => {
+      const doc = await (payload as any).create({
+        collection: 'posts',
+        data: { title: 'draft-original', _status: 'draft' },
+        draft: true,
+      })
+
+      // This triggers updateOne (has getPrimaryDb) + createVersion (now fixed)
+      const updated = await (payload as any).update({
+        collection: 'posts',
+        id: doc.id,
+        data: { title: 'draft-updated' },
+        draft: true,
+      })
+
+      expect(updated.title).toBe('draft-updated')
+
+      const versions = await (payload as any).findVersions({
+        collection: 'posts',
+        where: { parent: { equals: doc.id } },
+      })
+
+      expect(versions.docs.length).toBeGreaterThanOrEqual(2)
+    })
+
+    it('should restore a version without errors', async () => {
+      const doc = await (payload as any).create({
+        collection: 'posts',
+        data: { title: 'restore-v1', _status: 'draft' },
+        draft: true,
+      })
+
+      await (payload as any).update({
+        collection: 'posts',
+        id: doc.id,
+        data: { title: 'restore-v2' },
+        draft: true,
+      })
+
+      const versions = await (payload as any).findVersions({
+        collection: 'posts',
+        where: { parent: { equals: doc.id } },
+        sort: '-updatedAt',
+      })
+
+      const firstVersion = versions.docs[versions.docs.length - 1]!
+
+      // restoreVersion triggers updateVersion (now fixed)
+      const restored = await (payload as any).restoreVersion({
+        collection: 'posts',
+        id: firstVersion.id,
+      })
+
+      expect(restored.title).toBe('restore-v1')
+    })
+  })
+
+  describe('global operations use primary for read-back', () => {
+    it('should create and update a global without errors', async () => {
+      // First update creates the global row (createGlobal, now fixed)
+      const result = await (payload as any).updateGlobal({
+        slug: 'settings',
+        data: { siteTitle: 'My Site' },
+      })
+
+      expect(result).toBeDefined()
+      expect(result.siteTitle).toBe('My Site')
+
+      // Second update uses updateGlobal path (also fixed)
+      const updated = await (payload as any).updateGlobal({
+        slug: 'settings',
+        data: { siteTitle: 'Updated Site' },
+      })
+
+      expect(updated.siteTitle).toBe('Updated Site')
+    })
+
+    it('should create and update a versioned global without errors', async () => {
+      // This exercises createGlobalVersion (now fixed)
+      const result = await (payload as any).updateGlobal({
+        slug: 'nav',
+        data: { label: 'Home' },
+      })
+
+      expect(result).toBeDefined()
+      expect(result.label).toBe('Home')
+
+      const versions = await (payload as any).findGlobalVersions({
+        slug: 'nav',
+      })
+
+      expect(versions.docs.length).toBeGreaterThanOrEqual(1)
+
+      // Update again to exercise updateGlobalVersion path
+      const updated = await (payload as any).updateGlobal({
+        slug: 'nav',
+        data: { label: 'Updated Home' },
+      })
+
+      expect(updated.label).toBe('Updated Home')
+    })
+  })
+
+  describe('delete operations use primary for pre-delete read', () => {
+    it('should delete a document and return the deleted data without errors', async () => {
+      const doc = await payload.create({
+        collection: 'posts',
+        data: { title: 'to-delete-replica-test' },
+      })
+
+      // deleteOne reads the doc before deleting (now uses getPrimaryDb)
+      const deleted = await payload.delete({
+        collection: 'posts',
+        id: doc.id,
+      })
+
+      expect(deleted).toBeDefined()
+      expect(deleted.title).toBe('to-delete-replica-test')
+    })
+  })
 })

Original file line number	Diff line number	Diff line change
`@@ -55,6 +55,9 @@ export const deleteMany: DeleteMany = async function deleteMany(`
`55`	`55`	`)`
`56`	`56`	`}`
`57`	`57`
	`58`	`+ // No getPrimaryDb needed: db is only used for deleteWhere (a write, always routed to primary`
	`59`	`+ // by drizzle's withReplicas). findMany resolves its own db via getTransaction, which returns`
	`60`	`+ // the transaction db (always primary) or falls back to shouldReadFromPrimary.`
`58`	`61`	`const db = await getTransaction(this, req)`
`59`	`62`
`60`	`63`	`await this.deleteWhere({`
Original file line number	Diff line number	Diff line change
`@@ -67,7 +67,7 @@ export const updateJobs: UpdateJobs = async function updateMany(`
`67`	`67`	`return []`
`68`	`68`	`}`
`69`	`69`
`70`		`- const db = await getTransaction(this, req)`
	`70`	`+ const db = getPrimaryDb(this, await getTransaction(this, req))`
`71`	`71`
`72`	`72`	`const results = []`
`73`	`73`