Merge pull request #30 from yicrotkd/add-replica-client-options

Add options for replica clients

README.md

@@ -67,6 +67,27 @@ const prisma = new PrismaClient().$extends(
 
 In this case, a replica for each read query will be selected randomly.
 
+### Pre-configured clients
+
+If you want to supply additional options to replica client, you can also pass pre-configured read clients instead of urls:
+
+```ts
+const replicaClient = new PrismaClient({
+  datasourceUrl: 'postgresql://replica.example.com:5432/db'
+  log: [{ level: 'query', emit: 'event' }]
+})
+
+replicaClient.$on('query', (event) => console.log('Replica event', event))
+
+const prisma = new PrismaClient().$extends(
+  readReplicas({
+    replicas: [
+      replicaClient
+    ],
+  }),
+)
+```
+
 ### Bypassing the replicas
 
 If you want to execute a read query against the primary server, you can use the `$primary()` method on your extended client:

src/ReplicaManager.ts

@@ -9,16 +9,26 @@ interface PrismaClientConstructor {
   new (options?: PrismaConstructorOptions): PrismaClient
 }
 
-type ReplicaManagerOptions = {
-  clientConstructor: PrismaClientConstructor
-  replicaUrls: string[]
-  configureCallback: ConfigureReplicaCallback | undefined
-}
+export type ReplicaManagerOptions =
+  | {
+      clientConstructor: PrismaClientConstructor
+      replicaUrls: string[]
+      configureCallback: ConfigureReplicaCallback | undefined
+    }
+  | {
+      replicas: PrismaClient[]
+    }
 
 export class ReplicaManager {
   private _replicaClients: PrismaClient[]
 
-  constructor({ replicaUrls, clientConstructor, configureCallback }: ReplicaManagerOptions) {
+  constructor(options: ReplicaManagerOptions) {
+    if ('replicas' in options) {
+      this._replicaClients = options.replicas
+      return
+    }
+
+    const { replicaUrls, clientConstructor, configureCallback } = options
     this._replicaClients = replicaUrls.map((datasourceUrl) => {
       const client = new clientConstructor({
         datasourceUrl,

src/extension.ts

@@ -1,10 +1,13 @@
-import { Prisma } from '@prisma/client/extension.js'
+import { Prisma, PrismaClient } from '@prisma/client/extension.js'
 
-import { ConfigureReplicaCallback, ReplicaManager } from './ReplicaManager'
+import { ConfigureReplicaCallback, ReplicaManager, type ReplicaManagerOptions } from './ReplicaManager'
 
-export type ReplicasOptions = {
-  url: string | string[]
-}
+export type ReplicasOptions =
+  | {
+      url: string | string[]
+      replicas?: undefined
+    }
+  | { url?: undefined; replicas: PrismaClient[] }
 
 const readOperations = [
   'findFirst',
@@ -26,20 +29,43 @@ export const readReplicas = (options: ReplicasOptions, configureReplicaClient?:
     if (!datasourceName) {
       throw new Error(`Read replicas options must specify a datasource`)
     }
-    let replicaUrls = options.url
-    if (typeof replicaUrls === 'string') {
-      replicaUrls = [replicaUrls]
-    } else if (!Array.isArray(replicaUrls)) {
-      throw new Error(`Replica URLs must be a string or list of strings`)
-    } else if (replicaUrls.length === 0) {
-      throw new Error(`At least one replica URL must be specified`)
+
+    if ('url' in options && 'replicas' in options) {
+      throw new Error(`Only one of 'url' or 'replicas' can be specified`)
     }
 
-    const replicaManager = new ReplicaManager({
-      replicaUrls,
-      clientConstructor: PrismaClient,
-      configureCallback: configureReplicaClient,
-    })
+    let replicaManagerOptions: ReplicaManagerOptions
+
+    if (options.url) {
+      let replicaUrls = options.url
+
+      if (typeof replicaUrls === 'string') {
+        replicaUrls = [replicaUrls]
+      } else if (replicaUrls && !Array.isArray(replicaUrls)) {
+        throw new Error(`Replica URLs must be a string or list of strings`)
+      }
+
+      if (replicaUrls?.length === 0) {
+        throw new Error(`At least one replica URL must be specified`)
+      }
+
+      replicaManagerOptions = {
+        replicaUrls: replicaUrls,
+        clientConstructor: PrismaClient,
+        configureCallback: configureReplicaClient,
+      }
+    } else if (options.replicas) {
+      if (options.replicas.length === 0) {
+        throw new Error(`At least one replica must be specified`)
+      }
+      replicaManagerOptions = {
+        replicas: options.replicas,
+      }
+    } else {
+      throw new Error(`Invalid read replicas options`)
+    }
+
+    const replicaManager = new ReplicaManager(replicaManagerOptions)
 
     return client.$extends({
       client: {

tests/extension.test.ts

@@ -54,7 +54,7 @@ beforeEach(async () => {
   logs = []
 })
 
-test('client throws an error when given an empty read replica list', async () => {
+test('client throws an error when given an empty read replica url list', async () => {
   const createInstance = () =>
     basePrisma.$extends(
       readReplicas({
@@ -65,6 +65,47 @@ test('client throws an error when given an empty read replica list', async () =>
   expect(createInstance).toThrowError('At least one replica URL must be specified')
 })
 
+test('client throws an error when given an empty read replica list', async () => {
+  const createInstance = () =>
+    basePrisma.$extends(
+      readReplicas({
+        replicas: [],
+      }),
+    )
+
+  expect(createInstance).toThrowError('At least one replica must be specified')
+})
+
+test('client throws an error when given both a URL and a list of replicas', async () => {
+  // eslint-disable-next-line @typescript-eslint/no-var-requires
+  const clientModule = require('./client')
+  const replicaPrisma = new clientModule.PrismaClient({
+    datasourceUrl: process.env.REPLICA_URL!,
+    log: [{ emit: 'event', level: 'query' }],
+  })
+  const createInstance = () =>
+    basePrisma.$extends(
+      readReplicas({
+        url: process.env.REPLICA_URL!,
+        replicas: [replicaPrisma],
+      }),
+    )
+
+  expect(createInstance).toThrowError(`Only one of 'url' or 'replicas' can be specified`)
+})
+
+test('client throws an error when given an invalid read replicas options', async () => {
+  const createInstance = () =>
+    basePrisma.$extends(
+      readReplicas({
+        // @ts-expect-error
+        foo: 'bar',
+      }),
+    )
+
+  expect(createInstance).toThrowError('Invalid read replicas options')
+})
+
 test('read query is executed against replica', async () => {
   await prisma.user.findMany()
 
@@ -124,6 +165,39 @@ test('transactional queries are executed against primary (itx)', async () => {
   ])
 })
 
+test('replica client with options', async () => {
+  let resolve: (value: unknown) => void
+  const promise = new Promise((res) => {
+    resolve = res
+  })
+  // eslint-disable-next-line @typescript-eslint/no-var-requires
+  const clientModule = require('./client')
+  const replicaPrisma = new clientModule.PrismaClient({
+    datasourceUrl: process.env.REPLICA_URL!,
+    log: [{ emit: 'event', level: 'query' }],
+  })
+
+  replicaPrisma.$on('query', () => {
+    logs.push({ server: 'replica', operation: 'replica logger' })
+    resolve('logged')
+  })
+
+  const prisma = basePrisma.$extends(
+    readReplicas({
+      replicas: [replicaPrisma],
+    }),
+  )
+
+  await prisma.user.findMany()
+  await promise
+  expect(logs).toEqual([
+    {
+      operation: `replica logger`,
+      server: 'replica',
+    },
+  ])
+})
+
 afterAll(async () => {
   await prisma.$disconnect()
 })