-
Notifications
You must be signed in to change notification settings - Fork 6
/
extension.ts
83 lines (73 loc) · 2.19 KB
/
extension.ts
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
import { Prisma } from '@prisma/client/extension.js'
import { ConfigureReplicaCallback, ReplicaManager } from './ReplicaManager'
export type ReplicasOptions = {
url: string | string[]
}
const readOperations = [
'findFirst',
'findFirstOrThrow',
'findMany',
'findUnique',
'findUniqueOrThrow',
'groupBy',
'aggregate',
'count',
'queryRaw',
'queryRawUnsafe',
'findRaw',
'aggregateRaw',
]
export const readReplicas = (options: ReplicasOptions, configureReplicaClient?: ConfigureReplicaCallback) =>
Prisma.defineExtension((client) => {
const PrismaClient = Object.getPrototypeOf(client).constructor
const datasourceName = Object.keys(options).find((key) => !key.startsWith('$'))
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`)
}
const replicaManager = new ReplicaManager({
replicaUrls,
clientConstructor: PrismaClient,
configureCallback: configureReplicaClient,
})
return client.$extends({
client: {
$primary<T>(this: T): Omit<T, '$primary'> {
return client as unknown as Omit<T, '$primary'>
},
async $connect() {
await Promise.all([(client as any).$connect(), replicaManager.connectAll()])
},
async $disconnect() {
await Promise.all([(client as any).$disconnect(), replicaManager.disconnectAll()])
},
},
query: {
$allOperations({
args,
model,
operation,
query,
// @ts-expect-error
__internalParams: { transaction },
}) {
if (transaction) {
return query(args)
}
if (readOperations.includes(operation)) {
const replica = replicaManager.pickReplica()
if (model) {
return replica[model][operation](args)
}
return replica[operation](args)
}
return query(args)
},
},
})
})