Merge pull request #21 from casey-chow/expose-read-replicas

Expose replicas as $replica()

src/extension.ts

@@ -45,8 +45,26 @@ export const readReplicas = (options: ReplicasOptions, configureReplicaClient?:
 
     return client.$extends({
       client: {
-        $primary<T>(this: T): Omit<T, '$primary'> {
-          return client as unknown as Omit<T, '$primary'>
+        $primary<T extends object>(this: T): Omit<T, '$primary' | '$replica'> {
+          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<T, '$primary' | '$replica'>
+        },
+
+        $replica<T extends object>(this: T): Omit<T, '$primary' | '$replica'> {
+          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<T, '$primary' | '$replica'>
         },
 
         async $connect() {

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()