diff --git a/src/extension.ts b/src/extension.ts index ff994e0..f74dc02 100644 --- a/src/extension.ts +++ b/src/extension.ts @@ -45,8 +45,26 @@ export const readReplicas = (options: ReplicasOptions, configureReplicaClient?: return client.$extends({ client: { - $primary(this: T): Omit { - return client as unknown as Omit + $primary(this: T): Omit { + const context = Prisma.getExtensionContext(this) + // If we're in a transaction, the current client is connected to the + // primary. + if (!('$transaction' in context && typeof context.$transaction === 'function')) { + return context + } + + return client as unknown as Omit + }, + + $replica(this: T): Omit { + const context = Prisma.getExtensionContext(this) + // If we're in a transaction, the current client is connected to the + // primary. + if (!('$transaction' in context && typeof context.$transaction === 'function')) { + throw new Error(`Cannot use $replica inside of a transaction`) + } + + return replicaManager.pickReplica() as unknown as Omit }, async $connect() { diff --git a/tests/extension.test.ts b/tests/extension.test.ts index 34a6c2b..a256c8e 100644 --- a/tests/extension.test.ts +++ b/tests/extension.test.ts @@ -77,6 +77,26 @@ test('write query is executed against primary', async () => { expect(logs).toEqual([{ server: 'primary', operation: 'updateMany' }]) }) +test('$executeRaw and $executeRawUnsafe are executed against primary', async () => { + await prisma.$executeRaw`SELECT 1;` + await prisma.$executeRawUnsafe('SELECT $1 as id;', 1) + + expect(logs).toEqual([ + { server: 'primary', operation: '$executeRaw' }, + { server: 'primary', operation: '$executeRawUnsafe' }, + ]) +}) + +test('$executeRaw and $executeRawUnsafe are executed against replica if $replica() is used', async () => { + await prisma.$replica().$executeRaw`SELECT 1;` + await prisma.$replica().$executeRawUnsafe('SELECT $1 as id;', 1) + + expect(logs).toEqual([ + { server: 'replica', operation: '$executeRaw' }, + { server: 'replica', operation: '$executeRawUnsafe' }, + ]) +}) + test('read query is executed against primary if $primary() is used', async () => { await prisma.$primary().user.findMany()