Skip to content

Commit 8829912

Browse files
igalklebanovigalklebanov
committed
onReserveConnection.
1 parent 3594038 commit 8829912

File tree

5 files changed

+208
-155
lines changed

5 files changed

+208
-155
lines changed

‎src/dialect-config.mts‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,14 @@
1+
import type { DatabaseConnection } from 'kysely'
2+
13
export interface PostgresJSDialectConfig {
4+
/**
5+
* Called every time a connection is acquired from the pool.
6+
*/
7+
onReserveConnection?: (connection: DatabaseConnection) => Promise<void>
8+
9+
/**
10+
* An instance, or a factory returning an instance, of `postgres`'s `Sql` (returned by `postgres(...)`) or Bun's `SQL` class.
11+
*/
212
readonly postgres:
313
| PostgresJSSql
414
| (() => PostgresJSSql | Promise<PostgresJSSql>)

‎src/driver.mts‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,11 @@ export class PostgresJSDriver extends PostgresDriver {
2626
// biome-ignore lint/style/noNonNullAssertion: `init` ran at this point.
2727
const reservedConnection = await this.#postgres!.reserve()
2828

29-
return new PostgresJSConnection(reservedConnection)
29+
const connection = new PostgresJSConnection(reservedConnection)
30+
31+
await this.#config.onReserveConnection?.(connection)
32+
33+
return connection
3034
}
3135

3236
override async destroy(): Promise<void> {

‎tests/crud.test.mts‎

Lines changed: 131 additions & 134 deletions
Original file line numberDiff line numberDiff line change
@@ -8,27 +8,25 @@ import {
88
} from './test-setup.mjs'
99

1010
for (const dialect of SUPPORTED_DIALECTS) {
11-
describe.skipIf(dialect === 'bun' && !isBun)(
12-
`postgres.js: ${dialect}`,
13-
() => {
14-
let ctx: TestContext
11+
describe.skipIf(dialect === 'bun' && !isBun)(dialect, () => {
12+
let ctx: TestContext
1513

16-
beforeAll(async () => {
17-
ctx = await initTest(dialect)
18-
})
14+
beforeAll(async () => {
15+
ctx = await initTest(dialect)
16+
})
1917

20-
beforeEach(async () => {
21-
await resetState()
22-
})
18+
beforeEach(async () => {
19+
await resetState()
20+
})
2321

24-
afterAll(async () => {
25-
await ctx.db.destroy()
26-
})
22+
afterAll(async () => {
23+
await ctx.db.destroy()
24+
})
2725

28-
it('should execute select queries', async () => {
29-
const result = await ctx.db.selectFrom('person').selectAll().execute()
26+
it('should execute select queries', async () => {
27+
const result = await ctx.db.selectFrom('person').selectAll().execute()
3028

31-
expect(result).toMatchInlineSnapshot(`
29+
expect(result).toMatchInlineSnapshot(`
3230
[
3331
{
3432
"id": "48856ed4-9f1f-4111-ba7f-6092a1be96eb",
@@ -48,160 +46,160 @@ for (const dialect of SUPPORTED_DIALECTS) {
4846
},
4947
]
5048
`)
51-
})
49+
})
5250

53-
it('should execute insert queries - no returning', async () => {
54-
const result = await ctx.db
55-
.insertInto('person')
56-
.values({ name: 'johnny' })
57-
.executeTakeFirstOrThrow()
51+
it('should execute insert queries - no returning', async () => {
52+
const result = await ctx.db
53+
.insertInto('person')
54+
.values({ name: 'johnny' })
55+
.executeTakeFirstOrThrow()
5856

59-
expect(result).toMatchInlineSnapshot(`
57+
expect(result).toMatchInlineSnapshot(`
6058
InsertResult {
6159
"insertId": undefined,
6260
"numInsertedOrUpdatedRows": 1n,
6361
}
6462
`)
65-
})
66-
67-
it('should execute insert queries - with returning', async () => {
68-
const result = await ctx.db
69-
.insertInto('person')
70-
.values({
71-
id: 'c260ee08-5c8a-4ed8-8415-575af63f22da',
72-
name: 'duncan',
73-
})
74-
.returning('id')
75-
.executeTakeFirstOrThrow()
76-
77-
expect(result).toMatchInlineSnapshot(`
63+
})
64+
65+
it('should execute insert queries - with returning', async () => {
66+
const result = await ctx.db
67+
.insertInto('person')
68+
.values({
69+
id: 'c260ee08-5c8a-4ed8-8415-575af63f22da',
70+
name: 'duncan',
71+
})
72+
.returning('id')
73+
.executeTakeFirstOrThrow()
74+
75+
expect(result).toMatchInlineSnapshot(`
7876
{
7977
"id": "c260ee08-5c8a-4ed8-8415-575af63f22da",
8078
}
8179
`)
82-
})
80+
})
8381

84-
it('should execute update queries - no returning', async () => {
85-
const result = await ctx.db
86-
.updateTable('person')
87-
.set('name', 'alex')
88-
.where('name', '=', 'moshe')
89-
.executeTakeFirstOrThrow()
82+
it('should execute update queries - no returning', async () => {
83+
const result = await ctx.db
84+
.updateTable('person')
85+
.set('name', 'alex')
86+
.where('name', '=', 'moshe')
87+
.executeTakeFirstOrThrow()
9088

91-
expect(result).toMatchInlineSnapshot(`
89+
expect(result).toMatchInlineSnapshot(`
9290
UpdateResult {
9391
"numChangedRows": undefined,
9492
"numUpdatedRows": 1n,
9593
}
9694
`)
97-
})
95+
})
9896

99-
it('should execute update queries - with returning', async () => {
100-
const result = await ctx.db
101-
.updateTable('person')
102-
.set('name', 'alex')
103-
.where('name', '=', 'moshe')
104-
.returning('id')
105-
.executeTakeFirstOrThrow()
97+
it('should execute update queries - with returning', async () => {
98+
const result = await ctx.db
99+
.updateTable('person')
100+
.set('name', 'alex')
101+
.where('name', '=', 'moshe')
102+
.returning('id')
103+
.executeTakeFirstOrThrow()
106104

107-
expect(result).toMatchInlineSnapshot(`
105+
expect(result).toMatchInlineSnapshot(`
108106
{
109107
"id": "48856ed4-9f1f-4111-ba7f-6092a1be96eb",
110108
}
111109
`)
112-
})
110+
})
113111

114-
it('should execute delete queries - no returning', async () => {
115-
const result = await ctx.db
116-
.deleteFrom('person')
117-
.where('id', '=', '48856ed4-9f1f-4111-ba7f-6092a1be96eb')
118-
.executeTakeFirstOrThrow()
112+
it('should execute delete queries - no returning', async () => {
113+
const result = await ctx.db
114+
.deleteFrom('person')
115+
.where('id', '=', '48856ed4-9f1f-4111-ba7f-6092a1be96eb')
116+
.executeTakeFirstOrThrow()
119117

120-
expect(result).toMatchInlineSnapshot(`
118+
expect(result).toMatchInlineSnapshot(`
121119
DeleteResult {
122120
"numDeletedRows": 1n,
123121
}
124122
`)
125-
})
123+
})
126124

127-
it('should execute delete queries - with returning', async () => {
128-
const result = await ctx.db
129-
.deleteFrom('person')
130-
.where('id', '=', '48856ed4-9f1f-4111-ba7f-6092a1be96eb')
131-
.returning('name')
132-
.executeTakeFirstOrThrow()
125+
it('should execute delete queries - with returning', async () => {
126+
const result = await ctx.db
127+
.deleteFrom('person')
128+
.where('id', '=', '48856ed4-9f1f-4111-ba7f-6092a1be96eb')
129+
.returning('name')
130+
.executeTakeFirstOrThrow()
133131

134-
expect(result).toMatchInlineSnapshot(`
132+
expect(result).toMatchInlineSnapshot(`
135133
{
136134
"name": "moshe",
137135
}
138136
`)
139-
})
140-
141-
it('should execute merge queries - no returning', async () => {
142-
const result = await ctx.db
143-
.mergeInto('person as target')
144-
.using('person as source', (jb) =>
145-
jb
146-
.onRef('source.name', '!=', 'target.name')
147-
.on((eb) =>
148-
eb(
149-
eb.fn('left', ['source.name', eb.val(1)]),
150-
'=',
151-
eb.fn('left', ['target.name', eb.val(1)]),
152-
),
137+
})
138+
139+
it('should execute merge queries - no returning', async () => {
140+
const result = await ctx.db
141+
.mergeInto('person as target')
142+
.using('person as source', (jb) =>
143+
jb
144+
.onRef('source.name', '!=', 'target.name')
145+
.on((eb) =>
146+
eb(
147+
eb.fn('left', ['source.name', eb.val(1)]),
148+
'=',
149+
eb.fn('left', ['target.name', eb.val(1)]),
153150
),
154-
)
155-
.whenMatched()
156-
.thenDelete()
157-
.executeTakeFirstOrThrow()
151+
),
152+
)
153+
.whenMatched()
154+
.thenDelete()
155+
.executeTakeFirstOrThrow()
158156

159-
expect(result).toMatchInlineSnapshot(`
157+
expect(result).toMatchInlineSnapshot(`
160158
MergeResult {
161159
"numChangedRows": 2n,
162160
}
163161
`)
164-
})
165-
166-
it('should execute merge queries - with returning', async () => {
167-
const result = await ctx.db
168-
.mergeInto('person as target')
169-
.using('person as source', (jb) =>
170-
jb
171-
.onRef('source.name', '!=', 'target.name')
172-
.on((eb) =>
173-
eb(
174-
eb.fn('left', ['source.name', eb.val(1)]),
175-
'=',
176-
eb.fn('left', ['target.name', eb.val(1)]),
177-
),
162+
})
163+
164+
it('should execute merge queries - with returning', async () => {
165+
const result = await ctx.db
166+
.mergeInto('person as target')
167+
.using('person as source', (jb) =>
168+
jb
169+
.onRef('source.name', '!=', 'target.name')
170+
.on((eb) =>
171+
eb(
172+
eb.fn('left', ['source.name', eb.val(1)]),
173+
'=',
174+
eb.fn('left', ['target.name', eb.val(1)]),
178175
),
179-
)
180-
.whenMatched()
181-
.thenDelete()
182-
.returningAll()
183-
.executeTakeFirstOrThrow()
184-
185-
expect(result).toMatchInlineSnapshot(`
176+
),
177+
)
178+
.whenMatched()
179+
.thenDelete()
180+
.returningAll()
181+
.executeTakeFirstOrThrow()
182+
183+
expect(result).toMatchInlineSnapshot(`
186184
{
187185
"id": "28175ebc-02ec-4c87-9a84-b3d25193fefa",
188186
"name": "haim",
189187
}
190188
`)
191-
})
189+
})
192190

193-
it.skipIf(dialect === 'bun')(
194-
'should stream select queries: exhaust',
195-
async () => {
196-
const items = []
191+
it.skipIf(dialect === 'bun')(
192+
'should stream select queries: exhaust',
193+
async () => {
194+
const items = []
197195

198-
const iterator = ctx.db.selectFrom('person').selectAll().stream()
196+
const iterator = ctx.db.selectFrom('person').selectAll().stream()
199197

200-
for await (const item of iterator) {
201-
items.push(item)
202-
}
198+
for await (const item of iterator) {
199+
items.push(item)
200+
}
203201

204-
expect(items).toMatchInlineSnapshot(`
202+
expect(items).toMatchInlineSnapshot(`
205203
[
206204
{
207205
"id": "48856ed4-9f1f-4111-ba7f-6092a1be96eb",
@@ -221,32 +219,31 @@ for (const dialect of SUPPORTED_DIALECTS) {
221219
},
222220
]
223221
`)
224-
},
225-
)
222+
},
223+
)
226224

227-
it.skipIf(dialect === 'bun')(
228-
'should stream select queries: break',
229-
async () => {
230-
const items = []
225+
it.skipIf(dialect === 'bun')(
226+
'should stream select queries: break',
227+
async () => {
228+
const items = []
231229

232-
const iterator = ctx.db.selectFrom('person').selectAll().stream()
230+
const iterator = ctx.db.selectFrom('person').selectAll().stream()
233231

234-
for await (const item of iterator) {
235-
items.push(item)
232+
for await (const item of iterator) {
233+
items.push(item)
236234

237-
break
238-
}
235+
break
236+
}
239237

240-
expect(items).toMatchInlineSnapshot(`
238+
expect(items).toMatchInlineSnapshot(`
241239
[
242240
{
243241
"id": "48856ed4-9f1f-4111-ba7f-6092a1be96eb",
244242
"name": "moshe",
245243
},
246244
]
247245
`)
248-
},
249-
)
250-
},
251-
)
246+
},
247+
)
248+
})
252249
}

0 commit comments

Comments
 (0)