@dbos-inc/drizzle-datasource 4.10.15 → 4.11.4-preview
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- package/dist/index.d.ts +3 -3
- package/dist/index.d.ts.map +1 -1
- package/dist/index.js +23 -15
- package/dist/index.js.map +1 -1
- package/index.ts +24 -14
- package/package.json +1 -1
- package/tests/custom-pool.test.ts +125 -0
package/dist/index.d.ts
CHANGED
|
@@ -1,4 +1,4 @@
|
|
|
1
|
-
import { ClientConfig, PoolConfig } from 'pg';
|
|
1
|
+
import { Client, ClientConfig, Pool, PoolClient, PoolConfig } from 'pg';
|
|
2
2
|
import { FunctionName } from '@dbos-inc/dbos-sdk';
|
|
3
3
|
import { DBOSDataSource } from '@dbos-inc/dbos-sdk/datasource';
|
|
4
4
|
import { NodePgDatabase } from 'drizzle-orm/node-postgres';
|
|
@@ -21,8 +21,8 @@ export declare class DrizzleDataSource<CT = NodePgDatabase<{
|
|
|
21
21
|
[key: string]: object;
|
|
22
22
|
}>;
|
|
23
23
|
get client(): CT;
|
|
24
|
-
static initializeDBOSSchema(
|
|
25
|
-
constructor(name: string,
|
|
24
|
+
static initializeDBOSSchema(configOrClient: ClientConfig | Client | PoolClient, schemaName?: string): Promise<void>;
|
|
25
|
+
constructor(name: string, configOrPool: PoolConfig | Pool, entities?: {
|
|
26
26
|
[key: string]: object;
|
|
27
27
|
}, schemaName?: string);
|
|
28
28
|
runTransaction<T>(func: () => Promise<T>, config?: TransactionConfig): Promise<T>;
|
package/dist/index.d.ts.map
CHANGED
|
@@ -1 +1 @@
|
|
|
1
|
-
{"version":3,"file":"index.d.ts","sourceRoot":"","sources":["../index.ts"],"names":[],"mappings":"AAAA,OAAO,
|
|
1
|
+
{"version":3,"file":"index.d.ts","sourceRoot":"","sources":["../index.ts"],"names":[],"mappings":"AAAA,OAAO,EAAE,MAAM,EAAE,YAAY,EAAE,IAAI,EAAE,UAAU,EAAE,UAAU,EAAE,MAAM,IAAI,CAAC;AACxE,OAAO,EAAmC,YAAY,EAAE,MAAM,oBAAoB,CAAC;AACnF,OAAO,EAQL,cAAc,EAIf,MAAM,+BAA+B,CAAC;AACvC,OAAO,EAAW,cAAc,EAAE,MAAM,2BAA2B,CAAC;AAGpE,OAAO,EAAE,mBAAmB,EAAE,MAAM,qBAAqB,CAAC;AAQ1D,MAAM,MAAM,iBAAiB,GAAG,IAAI,CAAC,mBAAmB,EAAE,gBAAgB,GAAG,YAAY,CAAC,GAAG;IAAE,IAAI,CAAC,EAAE,MAAM,CAAA;CAAE,CAAC;AAI/G,MAAM,WAAW,sBAAsB;IACrC,WAAW,EAAE,MAAM,CAAC;IACpB,YAAY,EAAE,MAAM,CAAC;IACrB,MAAM,EAAE,MAAM,GAAG,IAAI,CAAC;IACtB,KAAK,EAAE,MAAM,GAAG,IAAI,CAAC;CACtB;AA8MD,qBAAa,iBAAiB,CAAC,EAAE,GAAG,cAAc,CAAC;IAAE,CAAC,GAAG,EAAE,MAAM,GAAG,MAAM,CAAA;CAAE,CAAC,CAC3E,YAAW,cAAc,CAAC,iBAAiB,CAAC;;IA+C1C,QAAQ,CAAC,IAAI,EAAE,MAAM;IA9BvB,MAAM,KAAK,MAAM;;OAEhB;IAED,IAAI,MAAM,OAET;WAEY,oBAAoB,CAC/B,cAAc,EAAE,YAAY,GAAG,MAAM,GAAG,UAAU,EAClD,UAAU,GAAE,MAAe,GAC1B,OAAO,CAAC,IAAI,CAAC;gBAmBL,IAAI,EAAE,MAAM,EACrB,YAAY,EAAE,UAAU,GAAG,IAAI,EAC/B,QAAQ,GAAE;QAAE,CAAC,GAAG,EAAE,MAAM,GAAG,MAAM,CAAA;KAAO,EACxC,UAAU,GAAE,MAAe;IAMvB,cAAc,CAAC,CAAC,EAAE,IAAI,EAAE,MAAM,OAAO,CAAC,CAAC,CAAC,EAAE,MAAM,CAAC,EAAE,iBAAiB;IAI1E,mBAAmB,CAAC,IAAI,EAAE,IAAI,SAAS,OAAO,EAAE,EAAE,MAAM,EACtD,IAAI,EAAE,CAAC,IAAI,EAAE,IAAI,EAAE,GAAG,IAAI,EAAE,IAAI,KAAK,OAAO,CAAC,MAAM,CAAC,EACpD,MAAM,CAAC,EAAE,iBAAiB,GAAG,YAAY,GACxC,CAAC,IAAI,EAAE,IAAI,EAAE,GAAG,IAAI,EAAE,IAAI,KAAK,OAAO,CAAC,MAAM,CAAC;IAKjD,WAAW,CAAC,MAAM,CAAC,EAAE,iBAAiB,kDAI1B,MAAM,eACD,WAAW,cACZ,wBAAwB,CAAC,IAAI,EAAE,IAAI,EAAE,GAAG,IAAI,EAAE,IAAI,KAAK,QAAQ,MAAM,CAAC,CAAC,oCAAxC,IAAI,WAAW,IAAI,KAAK,QAAQ,MAAM,CAAC;CAevF"}
|
package/dist/index.js
CHANGED
|
@@ -11,22 +11,24 @@ const drizzle_orm_1 = require("drizzle-orm");
|
|
|
11
11
|
const asyncLocalCtx = new async_hooks_1.AsyncLocalStorage();
|
|
12
12
|
class DrizzleTransactionHandler {
|
|
13
13
|
name;
|
|
14
|
-
|
|
14
|
+
configOrPool;
|
|
15
15
|
entities;
|
|
16
16
|
dsType = 'drizzle';
|
|
17
17
|
#connection;
|
|
18
18
|
schemaName;
|
|
19
|
-
|
|
19
|
+
#userProvidedPool;
|
|
20
|
+
constructor(name, configOrPool, entities = {}, schemaName = 'dbos') {
|
|
20
21
|
this.name = name;
|
|
21
|
-
this.
|
|
22
|
+
this.configOrPool = configOrPool;
|
|
22
23
|
this.entities = entities;
|
|
23
24
|
this.schemaName = schemaName;
|
|
25
|
+
this.#userProvidedPool = configOrPool instanceof pg_1.Pool;
|
|
24
26
|
}
|
|
25
27
|
async initialize() {
|
|
26
28
|
const conn = this.#connection;
|
|
27
|
-
const driver = new pg_1.Pool(this.
|
|
29
|
+
const driver = this.configOrPool instanceof pg_1.Pool ? this.configOrPool : new pg_1.Pool(this.configOrPool);
|
|
28
30
|
const db = (0, node_postgres_1.drizzle)(driver, { schema: this.entities });
|
|
29
|
-
this.#connection = { db, end: () => driver.end() };
|
|
31
|
+
this.#connection = { db, end: this.#userProvidedPool ? async () => { } : () => driver.end() };
|
|
30
32
|
await conn?.end();
|
|
31
33
|
let installed = false;
|
|
32
34
|
try {
|
|
@@ -183,21 +185,27 @@ class DrizzleDataSource {
|
|
|
183
185
|
get client() {
|
|
184
186
|
return DrizzleDataSource.#getClient(this.#provider);
|
|
185
187
|
}
|
|
186
|
-
static async initializeDBOSSchema(
|
|
187
|
-
|
|
188
|
-
|
|
189
|
-
await
|
|
190
|
-
await client.query((0, datasource_1.createTransactionCompletionSchemaPG)(schemaName));
|
|
191
|
-
await client.query((0, datasource_1.createTransactionCompletionTablePG)(schemaName));
|
|
188
|
+
static async initializeDBOSSchema(configOrClient, schemaName = 'dbos') {
|
|
189
|
+
if (typeof configOrClient === 'object' && 'query' in configOrClient) {
|
|
190
|
+
await configOrClient.query((0, datasource_1.createTransactionCompletionSchemaPG)(schemaName));
|
|
191
|
+
await configOrClient.query((0, datasource_1.createTransactionCompletionTablePG)(schemaName));
|
|
192
192
|
}
|
|
193
|
-
|
|
194
|
-
|
|
193
|
+
else {
|
|
194
|
+
const client = new pg_1.Client(configOrClient);
|
|
195
|
+
try {
|
|
196
|
+
await client.connect();
|
|
197
|
+
await client.query((0, datasource_1.createTransactionCompletionSchemaPG)(schemaName));
|
|
198
|
+
await client.query((0, datasource_1.createTransactionCompletionTablePG)(schemaName));
|
|
199
|
+
}
|
|
200
|
+
finally {
|
|
201
|
+
await client.end();
|
|
202
|
+
}
|
|
195
203
|
}
|
|
196
204
|
}
|
|
197
205
|
#provider;
|
|
198
|
-
constructor(name,
|
|
206
|
+
constructor(name, configOrPool, entities = {}, schemaName = 'dbos') {
|
|
199
207
|
this.name = name;
|
|
200
|
-
this.#provider = new DrizzleTransactionHandler(name,
|
|
208
|
+
this.#provider = new DrizzleTransactionHandler(name, configOrPool, entities, schemaName);
|
|
201
209
|
(0, datasource_1.registerDataSource)(this.#provider);
|
|
202
210
|
}
|
|
203
211
|
async runTransaction(func, config) {
|
package/dist/index.js.map
CHANGED
|
@@ -1 +1 @@
|
|
|
1
|
-
{"version":3,"file":"index.js","sourceRoot":"","sources":["../index.ts"],"names":[],"mappings":";;;AAAA,
|
|
1
|
+
{"version":3,"file":"index.js","sourceRoot":"","sources":["../index.ts"],"names":[],"mappings":";;;AAAA,2BAAwE;AACxE,iDAAmF;AACnF,8DAYuC;AACvC,6DAAoE;AACpE,6CAAgD;AAChD,yCAAsC;AAEtC,6CAAkC;AASlC,MAAM,aAAa,GAAG,IAAI,+BAAiB,EAAmB,CAAC;AAc/D,MAAM,yBAAyB;IAOlB;IACQ;IACA;IARV,MAAM,GAAG,SAAS,CAAC;IAC5B,WAAW,CAAgC;IAClC,UAAU,CAAS;IACnB,iBAAiB,CAAU;IAEpC,YACW,IAAY,EACJ,YAA+B,EAC/B,WAAsC,EAAE,EACzD,aAAqB,MAAM;QAHlB,SAAI,GAAJ,IAAI,CAAQ;QACJ,iBAAY,GAAZ,YAAY,CAAmB;QAC/B,aAAQ,GAAR,QAAQ,CAAgC;QAGzD,IAAI,CAAC,UAAU,GAAG,UAAU,CAAC;QAC7B,IAAI,CAAC,iBAAiB,GAAG,YAAY,YAAY,SAAI,CAAC;IACxD,CAAC;IAED,KAAK,CAAC,UAAU;QACd,MAAM,IAAI,GAAG,IAAI,CAAC,WAAW,CAAC;QAE9B,MAAM,MAAM,GAAG,IAAI,CAAC,YAAY,YAAY,SAAI,CAAC,CAAC,CAAC,IAAI,CAAC,YAAY,CAAC,CAAC,CAAC,IAAI,SAAI,CAAC,IAAI,CAAC,YAAY,CAAC,CAAC;QACnG,MAAM,EAAE,GAAG,IAAA,uBAAO,EAAC,MAAM,EAAE,EAAE,MAAM,EAAE,IAAI,CAAC,QAAQ,EAAE,CAAC,CAAC;QACtD,IAAI,CAAC,WAAW,GAAG,EAAE,EAAE,EAAE,GAAG,EAAE,IAAI,CAAC,iBAAiB,CAAC,CAAC,CAAC,KAAK,IAAI,EAAE,GAAE,CAAC,CAAC,CAAC,CAAC,GAAG,EAAE,CAAC,MAAM,CAAC,GAAG,EAAE,EAAE,CAAC;QAC7F,MAAM,IAAI,EAAE,GAAG,EAAE,CAAC;QAElB,IAAI,SAAS,GAAG,KAAK,CAAC;QACtB,IAAI,CAAC;YACH,MAAM,GAAG,GAAG,CAAC,MAAM,EAAE,CAAC,OAAO,CAAC,iBAAG,CAAC,GAAG,CAAC,IAAA,sCAAyB,EAAC,IAAI,CAAC,UAAU,CAAC,CAAC,CAAC,CAAY,CAAC;YAC/F,MAAM,GAAG,GAAG,KAAK,CAAC,OAAO,CAAC,GAAG,CAAC;gBAC5B,CAAC,CAAE,GAAuC,CAAC,CAAC,CAAC;gBAC7C,CAAC,CAAC,CAAE,GAAkD,CAAC,IAAI,EAAE,CAAC,CAAC,CAAC;oBAC7D,GAAuC,CAAC,CAAC,CAAC,CAAC,CAAC;YACjD,SAAS,GAAG,CAAC,CAAC,GAAG,EAAE,aAAa,IAAI,CAAC,CAAC,GAAG,EAAE,YAAY,CAAC;QAC1D,CAAC;QAAC,OAAO,CAAC,EAAE,CAAC;YACX,MAAM,IAAI,KAAK,CACb,4CAA4C,IAAI,CAAC,IAAI,oCAAqC,CAAW,CAAC,OAAO,EAAE,CAChH,CAAC;QACJ,CAAC;QAED,IAAI,CAAC,SAAS,EAAE,CAAC;YACf,IAAI,CAAC;gBACH,MAAM,EAAE,CAAC,OAAO,CAAC,iBAAG,CAAC,GAAG,CAAC,IAAA,gDAAmC,EAAC,IAAI,CAAC,UAAU,CAAC,CAAC,CAAC,CAAC;gBAChF,MAAM,EAAE,CAAC,OAAO,CAAC,iBAAG,CAAC,GAAG,CAAC,IAAA,+CAAkC,EAAC,IAAI,CAAC,UAAU,CAAC,CAAC,CAAC,CAAC;YACjF,CAAC;YAAC,OAAO,GAAG,EAAE,CAAC;gBACb,MAAM,IAAI,KAAK,CACb,4CAA4C,IAAI,CAAC,IAAI,UAAU,IAAI,CAAC,UAAU;0GACkB,CACjG,CAAC;YACJ,CAAC;QACH,CAAC;IACH,CAAC;IAED,KAAK,CAAC,OAAO;QACX,MAAM,IAAI,GAAG,IAAI,CAAC,WAAW,CAAC;QAE9B,IAAI,CAAC,WAAW,GAAG,SAAS,CAAC;QAE7B,MAAM,IAAI,EAAE,GAAG,EAAE,CAAC;IACpB,CAAC;IAED,IAAI,QAAQ;QACV,IAAI,CAAC,IAAI,CAAC,WAAW,EAAE,CAAC;YACtB,MAAM,IAAI,KAAK,CAAC,cAAc,IAAI,CAAC,IAAI,sBAAsB,CAAC,CAAC;QACjE,CAAC;QACD,OAAO,IAAI,CAAC,WAAW,CAAC,EAAE,CAAC;IAC7B,CAAC;IAED,KAAK,CAAC,eAAe,CACnB,UAAkB,EAClB,MAAc;QAId,MAAM,SAAS,GAAG,IAAA,iBAAG,EAAA;oCACW,iBAAG,CAAC,UAAU,CAAC,IAAI,CAAC,UAAU,CAAC;8BACrC,UAAU,uBAAuB,MAAM,EAAE,CAAC;QACpE,MAAM,MAAM,GAAG,MAAM,IAAI,CAAC,QAAQ,CAAC,OAAO,CAAS,SAAS,CAAC,CAAC;QAE9D,IAAI,MAAM,CAAC,IAAI,CAAC,MAAM,KAAK,CAAC,EAAE,CAAC;YAC7B,OAAO,SAAS,CAAC;QACnB,CAAC;QAED,MAAM,EAAE,MAAM,EAAE,KAAK,EAAE,GAAG,MAAM,CAAC,IAAI,CAAC,CAAC,CAAC,CAAC;QACzC,OAAO,KAAK,KAAK,IAAI,CAAC,CAAC,CAAC,EAAE,KAAK,EAAE,CAAC,CAAC,CAAC,EAAE,MAAM,EAAE,CAAC;IACjD,CAAC;IAED,MAAM,CAAC,KAAK,CAAC,aAAa,CACxB,MAAiD,EACjD,UAAkB,EAClB,MAAc,EACd,MAAc,EACd,UAAkB;QAElB,IAAI,CAAC;YACH,MAAM,SAAS,GAAG,IAAA,iBAAG,EAAA;sBACL,iBAAG,CAAC,UAAU,CAAC,UAAU,CAAC;kBAC9B,UAAU,KAAK,MAAM,KAAK,MAAM,GAAG,CAAC;YAChD,MAAM,MAAM,CAAC,OAAO,CAAC,SAAS,CAAC,CAAC;QAClC,CAAC;QAAC,OAAO,KAAK,EAAE,CAAC;YACf,IAAI,IAAA,iCAAoB,EAAC,KAAK,CAAC,EAAE,CAAC;gBAChC,MAAM,IAAI,oCAAyB,CAAC,UAAU,CAAC,CAAC;YAClD,CAAC;iBAAM,CAAC;gBACN,MAAM,KAAK,CAAC;YACd,CAAC;QACH,CAAC;IACH,CAAC;IAED,KAAK,CAAC,YAAY,CAAC,UAAkB,EAAE,MAAc,EAAE,KAAa;QAClE,IAAI,CAAC;YACH,MAAM,SAAS,GAAG,IAAA,iBAAG,EAAA;sBACL,iBAAG,CAAC,UAAU,CAAC,IAAI,CAAC,UAAU,CAAC;kBACnC,UAAU,KAAK,MAAM,KAAK,KAAK,GAAG,CAAC;YAC/C,MAAM,IAAI,CAAC,QAAQ,CAAC,OAAO,CAAC,SAAS,CAAC,CAAC;QACzC,CAAC;QAAC,OAAO,KAAK,EAAE,CAAC;YACf,IAAI,IAAA,iCAAoB,EAAC,KAAK,CAAC,EAAE,CAAC;gBAChC,MAAM,IAAI,oCAAyB,CAAC,UAAU,CAAC,CAAC;YAClD,CAAC;iBAAM,CAAC;gBACN,MAAM,KAAK,CAAC;YACd,CAAC;QACH,CAAC;IACH,CAAC;IAED,4DAA4D;IAC5D,KAAK,CAAC,yBAAyB,CAC7B,MAAqC,EACrC,MAAY,EACZ,IAAoD,EACpD,GAAG,IAAU;QAEb,MAAM,UAAU,GAAG,eAAI,CAAC,UAAU,CAAC;QACnC,MAAM,MAAM,GAAG,eAAI,CAAC,MAAM,CAAC;QAC3B,IAAI,UAAU,KAAK,SAAS,IAAI,MAAM,KAAK,SAAS,EAAE,CAAC;YACrD,MAAM,IAAI,KAAK,CAAC,6CAA6C,CAAC,CAAC;QACjE,CAAC;QAED,MAAM,QAAQ,GAAG,MAAM,EAAE,UAAU,KAAK,WAAW,CAAC,CAAC,CAAC,IAAI,CAAC,CAAC,CAAC,KAAK,CAAC;QACnE,MAAM,WAAW,GAAG,CAAC,QAAQ,IAAI,UAAU,KAAK,SAAS,CAAC;QAE1D,4BAA4B;QAC5B,IAAI,WAAW,GAAG,CAAC,CAAC;QACpB,MAAM,aAAa,GAAG,GAAG,CAAC;QAC1B,MAAM,cAAc,GAAG,IAAI,CAAC,CAAC,0BAA0B;QAEvD,OAAO,IAAI,EAAE,CAAC;YACZ,oDAAoD;YACpD,MAAM,cAAc,GAAG,WAAW,CAAC,CAAC,CAAC,MAAM,IAAI,CAAC,eAAe,CAAC,UAAU,EAAE,MAAO,CAAC,CAAC,CAAC,CAAC,SAAS,CAAC;YACjG,IAAI,cAAc,EAAE,CAAC;gBACnB,eAAI,CAAC,IAAI,EAAE,YAAY,CAAC,QAAQ,EAAE,IAAI,CAAC,CAAC;gBAExC,IAAI,OAAO,IAAI,cAAc,EAAE,CAAC;oBAC9B,MAAM,qBAAS,CAAC,KAAK,CAAC,cAAc,CAAC,KAAK,CAAC,CAAC;gBAC9C,CAAC;gBACD,OAAO,CAAC,cAAc,CAAC,MAAM,CAAC,CAAC,CAAC,qBAAS,CAAC,KAAK,CAAC,cAAc,CAAC,MAAM,CAAC,CAAC,CAAC,CAAC,IAAI,CAAW,CAAC;YAC3F,CAAC;YAED,IAAI,CAAC;gBACH,MAAM,MAAM,GAAG,MAAM,IAAI,CAAC,QAAQ,CAAC,WAAW,CAC5C,KAAK,EAAE,MAAM,EAAE,EAAE;oBACf,sCAAsC;oBACtC,MAAM,MAAM,GAAG,MAAM,aAAa,CAAC,GAAG,CAAC,EAAE,MAAM,EAAE,KAAK,EAAE,IAAI,EAAE,EAAE,KAAK,IAAI,EAAE;wBACzE,OAAO,MAAM,IAAI,CAAC,IAAI,CAAC,MAAM,EAAE,GAAG,IAAI,CAAC,CAAC;oBAC1C,CAAC,CAAC,CAAC;oBAEH,6CAA6C;oBAC7C,IAAI,WAAW,EAAE,CAAC;wBAChB,MAAM,yBAAyB,CAAC,aAAa,CAC3C,MAAM,EACN,UAAU,EACV,MAAO,EACP,qBAAS,CAAC,SAAS,CAAC,MAAM,CAAC,EAC3B,IAAI,CAAC,UAAU,CAChB,CAAC;oBACJ,CAAC;oBAED,OAAO,MAAM,CAAC;gBAChB,CAAC,EACD,EAAE,UAAU,EAAE,MAAM,EAAE,UAAU,EAAE,cAAc,EAAE,MAAM,EAAE,cAAc,EAAE,CAC3E,CAAC;gBAEF,OAAO,MAAM,CAAC;YAChB,CAAC;YAAC,OAAO,KAAK,EAAE,CAAC;gBACf,IAAI,IAAA,0CAA6B,EAAC,KAAK,CAAC,EAAE,CAAC;oBACzC,eAAI,CAAC,IAAI,EAAE,QAAQ,CAAC,2BAA2B,EAAE,EAAE,eAAe,EAAE,WAAW,EAAE,EAAE,WAAW,CAAC,GAAG,EAAE,CAAC,CAAC;oBACtG,gCAAgC;oBAChC,MAAM,IAAI,OAAO,CAAC,CAAC,OAAO,EAAE,EAAE,CAAC,UAAU,CAAC,OAAO,EAAE,WAAW,CAAC,CAAC,CAAC;oBACjE,WAAW,GAAG,IAAI,CAAC,GAAG,CAAC,WAAW,GAAG,aAAa,EAAE,cAAc,CAAC,CAAC;oBACpE,SAAS;gBACX,CAAC;qBAAM,CAAC;oBACN,IAAI,WAAW,EAAE,CAAC;wBAChB,MAAM,OAAO,GAAG,qBAAS,CAAC,SAAS,CAAC,KAAK,CAAC,CAAC;wBAC3C,MAAM,IAAI,CAAC,YAAY,CAAC,UAAU,EAAE,MAAO,EAAE,OAAO,CAAC,CAAC;oBACxD,CAAC;oBAED,MAAM,KAAK,CAAC;gBACd,CAAC;YACH,CAAC;QACH,CAAC;IACH,CAAC;CACF;AAED,MAAa,iBAAiB;IAgDjB;IA7CX,mDAAmD;IACnD,MAAM,CAAC,UAAU,CAAC,CAA6B;QAC7C,IAAI,CAAC,eAAI,CAAC,eAAe,EAAE,EAAE,CAAC;YAC5B,MAAM,IAAI,KAAK,CAAC,uEAAuE,CAAC,CAAC;QAC3F,CAAC;QACD,MAAM,GAAG,GAAG,aAAa,CAAC,QAAQ,EAAE,CAAC;QACrC,IAAI,CAAC,GAAG,EAAE,CAAC;YACT,MAAM,IAAI,KAAK,CAAC,uEAAuE,CAAC,CAAC;QAC3F,CAAC;QACD,IAAI,CAAC,IAAI,CAAC,KAAK,GAAG,CAAC,KAAK,EAAE,CAAC;YACzB,MAAM,IAAI,KAAK,CAAC,2EAA2E,CAAC,CAAC;QAC/F,CAAC;QACD,OAAO,GAAG,CAAC,MAAM,CAAC;IACpB,CAAC;IAED,MAAM,KAAK,MAAM;QACf,OAAO,iBAAiB,CAAC,UAAU,CAAC,SAAS,CAAC,CAAC;IACjD,CAAC;IAED,IAAI,MAAM;QACR,OAAO,iBAAiB,CAAC,UAAU,CAAC,IAAI,CAAC,SAAS,CAAO,CAAC;IAC5D,CAAC;IAED,MAAM,CAAC,KAAK,CAAC,oBAAoB,CAC/B,cAAkD,EAClD,aAAqB,MAAM;QAE3B,IAAI,OAAO,cAAc,KAAK,QAAQ,IAAI,OAAO,IAAI,cAAc,EAAE,CAAC;YACpE,MAAM,cAAc,CAAC,KAAK,CAAC,IAAA,gDAAmC,EAAC,UAAU,CAAC,CAAC,CAAC;YAC5E,MAAM,cAAc,CAAC,KAAK,CAAC,IAAA,+CAAkC,EAAC,UAAU,CAAC,CAAC,CAAC;QAC7E,CAAC;aAAM,CAAC;YACN,MAAM,MAAM,GAAG,IAAI,WAAM,CAAC,cAAc,CAAC,CAAC;YAC1C,IAAI,CAAC;gBACH,MAAM,MAAM,CAAC,OAAO,EAAE,CAAC;gBACvB,MAAM,MAAM,CAAC,KAAK,CAAC,IAAA,gDAAmC,EAAC,UAAU,CAAC,CAAC,CAAC;gBACpE,MAAM,MAAM,CAAC,KAAK,CAAC,IAAA,+CAAkC,EAAC,UAAU,CAAC,CAAC,CAAC;YACrE,CAAC;oBAAS,CAAC;gBACT,MAAM,MAAM,CAAC,GAAG,EAAE,CAAC;YACrB,CAAC;QACH,CAAC;IACH,CAAC;IAED,SAAS,CAA4B;IAErC,YACW,IAAY,EACrB,YAA+B,EAC/B,WAAsC,EAAE,EACxC,aAAqB,MAAM;QAHlB,SAAI,GAAJ,IAAI,CAAQ;QAKrB,IAAI,CAAC,SAAS,GAAG,IAAI,yBAAyB,CAAC,IAAI,EAAE,YAAY,EAAE,QAAQ,EAAE,UAAU,CAAC,CAAC;QACzF,IAAA,+BAAkB,EAAC,IAAI,CAAC,SAAS,CAAC,CAAC;IACrC,CAAC;IAED,KAAK,CAAC,cAAc,CAAI,IAAsB,EAAE,MAA0B;QACxE,OAAO,MAAM,IAAA,2BAAc,EAAC,IAAI,EAAE,MAAM,EAAE,IAAI,IAAI,IAAI,CAAC,IAAI,EAAE,EAAE,MAAM,EAAE,IAAI,CAAC,IAAI,EAAE,MAAM,EAAE,CAAC,CAAC;IAC9F,CAAC;IAED,mBAAmB,CACjB,IAAoD,EACpD,MAAyC;QAEzC,OAAO,IAAA,gCAAmB,EAAC,IAAI,CAAC,IAAI,EAAE,IAAI,EAAE,MAAM,CAAC,CAAC;IACtD,CAAC;IAED,YAAY;IACZ,WAAW,CAAC,MAA0B;QACpC,4DAA4D;QAC5D,MAAM,EAAE,GAAG,IAAI,CAAC;QAChB,OAAO,SAAS,SAAS,CACvB,MAAc,EACd,WAAwB,EACxB,UAAmF;YAEnF,IAAI,CAAC,UAAU,CAAC,KAAK,EAAE,CAAC;gBACtB,MAAM,IAAI,KAAK,CAAC,oDAAoD,CAAC,CAAC;YACxE,CAAC;YAED,UAAU,CAAC,KAAK,GAAG,EAAE,CAAC,mBAAmB,CAAC,UAAU,CAAC,KAAK,EAAE;gBAC1D,GAAG,MAAM;gBACT,WAAW,EAAE,MAAM;gBACnB,IAAI,EAAE,MAAM,EAAE,IAAI,IAAI,MAAM,CAAC,WAAW,CAAC;aAC1C,CAAC,CAAC;YAEH,OAAO,UAAU,CAAC;QACpB,CAAC,CAAC;IACJ,CAAC;CACF;AA1FD,8CA0FC"}
|
package/index.ts
CHANGED
|
@@ -1,4 +1,4 @@
|
|
|
1
|
-
import { Client, ClientConfig, Pool, PoolConfig } from 'pg';
|
|
1
|
+
import { Client, ClientConfig, Pool, PoolClient, PoolConfig } from 'pg';
|
|
2
2
|
import { DBOS, DBOSWorkflowConflictError, FunctionName } from '@dbos-inc/dbos-sdk';
|
|
3
3
|
import {
|
|
4
4
|
type DataSourceTransactionHandler,
|
|
@@ -44,22 +44,24 @@ class DrizzleTransactionHandler implements DataSourceTransactionHandler {
|
|
|
44
44
|
readonly dsType = 'drizzle';
|
|
45
45
|
#connection: DrizzleConnection | undefined;
|
|
46
46
|
readonly schemaName: string;
|
|
47
|
+
readonly #userProvidedPool: boolean;
|
|
47
48
|
|
|
48
49
|
constructor(
|
|
49
50
|
readonly name: string,
|
|
50
|
-
private readonly
|
|
51
|
+
private readonly configOrPool: PoolConfig | Pool,
|
|
51
52
|
private readonly entities: { [key: string]: object } = {},
|
|
52
53
|
schemaName: string = 'dbos',
|
|
53
54
|
) {
|
|
54
55
|
this.schemaName = schemaName;
|
|
56
|
+
this.#userProvidedPool = configOrPool instanceof Pool;
|
|
55
57
|
}
|
|
56
58
|
|
|
57
59
|
async initialize(): Promise<void> {
|
|
58
60
|
const conn = this.#connection;
|
|
59
61
|
|
|
60
|
-
const driver = new Pool(this.
|
|
62
|
+
const driver = this.configOrPool instanceof Pool ? this.configOrPool : new Pool(this.configOrPool);
|
|
61
63
|
const db = drizzle(driver, { schema: this.entities });
|
|
62
|
-
this.#connection = { db, end: () => driver.end() };
|
|
64
|
+
this.#connection = { db, end: this.#userProvidedPool ? async () => {} : () => driver.end() };
|
|
63
65
|
await conn?.end();
|
|
64
66
|
|
|
65
67
|
let installed = false;
|
|
@@ -263,14 +265,22 @@ export class DrizzleDataSource<CT = NodePgDatabase<{ [key: string]: object }>>
|
|
|
263
265
|
return DrizzleDataSource.#getClient(this.#provider) as CT;
|
|
264
266
|
}
|
|
265
267
|
|
|
266
|
-
static async initializeDBOSSchema(
|
|
267
|
-
|
|
268
|
-
|
|
269
|
-
|
|
270
|
-
|
|
271
|
-
await
|
|
272
|
-
|
|
273
|
-
|
|
268
|
+
static async initializeDBOSSchema(
|
|
269
|
+
configOrClient: ClientConfig | Client | PoolClient,
|
|
270
|
+
schemaName: string = 'dbos',
|
|
271
|
+
): Promise<void> {
|
|
272
|
+
if (typeof configOrClient === 'object' && 'query' in configOrClient) {
|
|
273
|
+
await configOrClient.query(createTransactionCompletionSchemaPG(schemaName));
|
|
274
|
+
await configOrClient.query(createTransactionCompletionTablePG(schemaName));
|
|
275
|
+
} else {
|
|
276
|
+
const client = new Client(configOrClient);
|
|
277
|
+
try {
|
|
278
|
+
await client.connect();
|
|
279
|
+
await client.query(createTransactionCompletionSchemaPG(schemaName));
|
|
280
|
+
await client.query(createTransactionCompletionTablePG(schemaName));
|
|
281
|
+
} finally {
|
|
282
|
+
await client.end();
|
|
283
|
+
}
|
|
274
284
|
}
|
|
275
285
|
}
|
|
276
286
|
|
|
@@ -278,11 +288,11 @@ export class DrizzleDataSource<CT = NodePgDatabase<{ [key: string]: object }>>
|
|
|
278
288
|
|
|
279
289
|
constructor(
|
|
280
290
|
readonly name: string,
|
|
281
|
-
|
|
291
|
+
configOrPool: PoolConfig | Pool,
|
|
282
292
|
entities: { [key: string]: object } = {},
|
|
283
293
|
schemaName: string = 'dbos',
|
|
284
294
|
) {
|
|
285
|
-
this.#provider = new DrizzleTransactionHandler(name,
|
|
295
|
+
this.#provider = new DrizzleTransactionHandler(name, configOrPool, entities, schemaName);
|
|
286
296
|
registerDataSource(this.#provider);
|
|
287
297
|
}
|
|
288
298
|
|
package/package.json
CHANGED
|
@@ -0,0 +1,125 @@
|
|
|
1
|
+
import { DBOS } from '@dbos-inc/dbos-sdk';
|
|
2
|
+
import { Client, Pool } from 'pg';
|
|
3
|
+
import { DrizzleDataSource } from '..';
|
|
4
|
+
import { dropDB, ensureDB } from './test-helpers';
|
|
5
|
+
import { randomUUID } from 'crypto';
|
|
6
|
+
import SuperJSON from 'superjson';
|
|
7
|
+
import { pgTable, text, integer } from 'drizzle-orm/pg-core';
|
|
8
|
+
import { drizzle } from 'drizzle-orm/node-postgres';
|
|
9
|
+
import { pushSchema } from 'drizzle-kit/api';
|
|
10
|
+
import { sql } from 'drizzle-orm';
|
|
11
|
+
|
|
12
|
+
const config = { user: 'postgres', database: 'drizzle_pool_test_userdb' };
|
|
13
|
+
|
|
14
|
+
const greetingsTable = pgTable('greetings', {
|
|
15
|
+
name: text('name').primaryKey().notNull(),
|
|
16
|
+
greet_count: integer('greet_count').default(0),
|
|
17
|
+
});
|
|
18
|
+
|
|
19
|
+
const customPool = new Pool(config);
|
|
20
|
+
const poolDataSource = new DrizzleDataSource('pool-db', customPool);
|
|
21
|
+
|
|
22
|
+
async function insertFunction(user: string) {
|
|
23
|
+
const result = await poolDataSource.client
|
|
24
|
+
.insert(greetingsTable)
|
|
25
|
+
.values({ name: user, greet_count: 1 })
|
|
26
|
+
.onConflictDoUpdate({
|
|
27
|
+
target: greetingsTable.name,
|
|
28
|
+
set: {
|
|
29
|
+
greet_count: sql`${greetingsTable.greet_count} + 1`,
|
|
30
|
+
},
|
|
31
|
+
})
|
|
32
|
+
.returning({ greet_count: greetingsTable.greet_count });
|
|
33
|
+
const row = result.length > 0 ? result[0] : undefined;
|
|
34
|
+
return { user, greet_count: row?.greet_count };
|
|
35
|
+
}
|
|
36
|
+
|
|
37
|
+
const regInsertFunction = poolDataSource.registerTransaction(insertFunction);
|
|
38
|
+
|
|
39
|
+
async function insertWorkflow(user: string) {
|
|
40
|
+
return await regInsertFunction(user);
|
|
41
|
+
}
|
|
42
|
+
|
|
43
|
+
const regInsertWorkflow = DBOS.registerWorkflow(insertWorkflow);
|
|
44
|
+
|
|
45
|
+
interface transaction_completion {
|
|
46
|
+
workflow_id: string;
|
|
47
|
+
function_num: number;
|
|
48
|
+
output: string | null;
|
|
49
|
+
error: string | null;
|
|
50
|
+
}
|
|
51
|
+
|
|
52
|
+
describe('DrizzleDataSource with custom Pool', () => {
|
|
53
|
+
beforeAll(async () => {
|
|
54
|
+
const client = new Client({ ...config, database: 'postgres' });
|
|
55
|
+
try {
|
|
56
|
+
await client.connect();
|
|
57
|
+
await dropDB(client, 'drizzle_pool_test', true);
|
|
58
|
+
await dropDB(client, 'drizzle_pool_test_dbos_sys', true);
|
|
59
|
+
await dropDB(client, config.database, true);
|
|
60
|
+
await ensureDB(client, config.database);
|
|
61
|
+
} finally {
|
|
62
|
+
await client.end();
|
|
63
|
+
}
|
|
64
|
+
|
|
65
|
+
const poolClient = await customPool.connect();
|
|
66
|
+
try {
|
|
67
|
+
await DrizzleDataSource.initializeDBOSSchema(poolClient);
|
|
68
|
+
} finally {
|
|
69
|
+
poolClient.release();
|
|
70
|
+
}
|
|
71
|
+
|
|
72
|
+
const drizzlePool = new Pool(config);
|
|
73
|
+
const db = drizzle(drizzlePool);
|
|
74
|
+
try {
|
|
75
|
+
const res = await pushSchema({ greetingsTable }, db);
|
|
76
|
+
await res.apply();
|
|
77
|
+
} finally {
|
|
78
|
+
await drizzlePool.end();
|
|
79
|
+
}
|
|
80
|
+
});
|
|
81
|
+
|
|
82
|
+
afterAll(async () => {
|
|
83
|
+
await customPool.end();
|
|
84
|
+
});
|
|
85
|
+
|
|
86
|
+
beforeEach(async () => {
|
|
87
|
+
DBOS.setConfig({ name: 'drizzle-pool-test' });
|
|
88
|
+
await DBOS.launch();
|
|
89
|
+
});
|
|
90
|
+
|
|
91
|
+
afterEach(async () => {
|
|
92
|
+
await DBOS.shutdown();
|
|
93
|
+
});
|
|
94
|
+
|
|
95
|
+
test('insert with custom pool', async () => {
|
|
96
|
+
const user = 'poolTest1';
|
|
97
|
+
await customPool.query('DELETE FROM greetings WHERE name = $1', [user]);
|
|
98
|
+
const workflowID = randomUUID();
|
|
99
|
+
|
|
100
|
+
await expect(DBOS.withNextWorkflowID(workflowID, () => regInsertWorkflow(user))).resolves.toMatchObject({
|
|
101
|
+
user,
|
|
102
|
+
greet_count: 1,
|
|
103
|
+
});
|
|
104
|
+
|
|
105
|
+
const { rows } = await customPool.query<transaction_completion>(
|
|
106
|
+
'SELECT * FROM dbos.transaction_completion WHERE workflow_id = $1',
|
|
107
|
+
[workflowID],
|
|
108
|
+
);
|
|
109
|
+
expect(rows.length).toBe(1);
|
|
110
|
+
expect(rows[0].output).not.toBeNull();
|
|
111
|
+
expect(SuperJSON.parse(rows[0].output!)).toMatchObject({ user, greet_count: 1 });
|
|
112
|
+
});
|
|
113
|
+
|
|
114
|
+
test('custom pool is not closed after DBOS shutdown', async () => {
|
|
115
|
+
// Explicitly shut down DBOS to trigger destroy() on all datasources
|
|
116
|
+
await DBOS.shutdown();
|
|
117
|
+
// The custom pool should still be usable after DBOS.shutdown()
|
|
118
|
+
// because DrizzleDataSource should not close a user-provided pool
|
|
119
|
+
const { rows } = await customPool.query<{ val: number }>('SELECT 1 as val');
|
|
120
|
+
expect(rows[0].val).toBe(1);
|
|
121
|
+
// Re-launch so the afterEach shutdown doesn't fail
|
|
122
|
+
DBOS.setConfig({ name: 'drizzle-pool-test' });
|
|
123
|
+
await DBOS.launch();
|
|
124
|
+
});
|
|
125
|
+
});
|