@workglow/postgres 0.2.28
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/job-queue/PostgresQueueStorage.d.ts +155 -0
- package/dist/job-queue/PostgresQueueStorage.d.ts.map +1 -0
- package/dist/job-queue/PostgresRateLimiterStorage.d.ts +57 -0
- package/dist/job-queue/PostgresRateLimiterStorage.d.ts.map +1 -0
- package/dist/job-queue/browser.d.ts +7 -0
- package/dist/job-queue/browser.d.ts.map +1 -0
- package/dist/job-queue/browser.js +732 -0
- package/dist/job-queue/browser.js.map +11 -0
- package/dist/job-queue/bun.d.ts +7 -0
- package/dist/job-queue/bun.d.ts.map +1 -0
- package/dist/job-queue/common.d.ts +8 -0
- package/dist/job-queue/common.d.ts.map +1 -0
- package/dist/job-queue/node.d.ts +7 -0
- package/dist/job-queue/node.d.ts.map +1 -0
- package/dist/job-queue/node.js +732 -0
- package/dist/job-queue/node.js.map +11 -0
- package/dist/storage/PostgresKvStorage.d.ts +27 -0
- package/dist/storage/PostgresKvStorage.d.ts.map +1 -0
- package/dist/storage/PostgresTabularStorage.d.ts +194 -0
- package/dist/storage/PostgresTabularStorage.d.ts.map +1 -0
- package/dist/storage/PostgresVectorStorage.d.ts +39 -0
- package/dist/storage/PostgresVectorStorage.d.ts.map +1 -0
- package/dist/storage/_postgres/browser.d.ts +32 -0
- package/dist/storage/_postgres/browser.d.ts.map +1 -0
- package/dist/storage/_postgres/node-bun.d.ts +26 -0
- package/dist/storage/_postgres/node-bun.d.ts.map +1 -0
- package/dist/storage/_postgres/pglite-pool.d.ts +21 -0
- package/dist/storage/_postgres/pglite-pool.d.ts.map +1 -0
- package/dist/storage/browser.d.ts +10 -0
- package/dist/storage/browser.d.ts.map +1 -0
- package/dist/storage/browser.js +951 -0
- package/dist/storage/browser.js.map +14 -0
- package/dist/storage/bun.d.ts +7 -0
- package/dist/storage/bun.d.ts.map +1 -0
- package/dist/storage/common.d.ts +10 -0
- package/dist/storage/common.d.ts.map +1 -0
- package/dist/storage/node.d.ts +7 -0
- package/dist/storage/node.d.ts.map +1 -0
- package/dist/storage/node.js +842 -0
- package/dist/storage/node.js.map +13 -0
- package/package.json +78 -0
|
@@ -0,0 +1,732 @@
|
|
|
1
|
+
// src/job-queue/PostgresQueueStorage.ts
|
|
2
|
+
import { createHash } from "node:crypto";
|
|
3
|
+
import { createServiceToken, getLogger, makeFingerprint, uuid4 } from "@workglow/util";
|
|
4
|
+
import { JobStatus } from "@workglow/job-queue";
|
|
5
|
+
var POSTGRES_QUEUE_STORAGE = createServiceToken("jobqueue.storage.postgres");
|
|
6
|
+
var SAFE_IDENTIFIER = /^[a-zA-Z][a-zA-Z0-9_]*$/;
|
|
7
|
+
|
|
8
|
+
class PostgresQueueStorage {
|
|
9
|
+
db;
|
|
10
|
+
queueName;
|
|
11
|
+
scope = "cluster";
|
|
12
|
+
prefixes;
|
|
13
|
+
prefixValues;
|
|
14
|
+
tableName;
|
|
15
|
+
constructor(db, queueName, options) {
|
|
16
|
+
this.db = db;
|
|
17
|
+
this.queueName = queueName;
|
|
18
|
+
this.prefixes = options?.prefixes ?? [];
|
|
19
|
+
this.prefixValues = options?.prefixValues ?? {};
|
|
20
|
+
for (const prefix of this.prefixes) {
|
|
21
|
+
if (!SAFE_IDENTIFIER.test(prefix.name)) {
|
|
22
|
+
throw new Error(`Prefix column name must start with a letter and contain only letters, digits, and underscores, got: ${prefix.name}`);
|
|
23
|
+
}
|
|
24
|
+
}
|
|
25
|
+
if (this.prefixes.length > 0) {
|
|
26
|
+
const prefixNames = this.prefixes.map((p) => p.name).join("_");
|
|
27
|
+
this.tableName = `job_queue_${prefixNames}`;
|
|
28
|
+
} else {
|
|
29
|
+
this.tableName = "job_queue";
|
|
30
|
+
}
|
|
31
|
+
}
|
|
32
|
+
getPrefixColumnType(type) {
|
|
33
|
+
return type === "uuid" ? "UUID" : "INTEGER";
|
|
34
|
+
}
|
|
35
|
+
buildPrefixColumnsSql() {
|
|
36
|
+
if (this.prefixes.length === 0)
|
|
37
|
+
return "";
|
|
38
|
+
return this.prefixes.map((p) => `${p.name} ${this.getPrefixColumnType(p.type)} NOT NULL`).join(`,
|
|
39
|
+
`) + `,
|
|
40
|
+
`;
|
|
41
|
+
}
|
|
42
|
+
getPrefixColumnNames() {
|
|
43
|
+
return this.prefixes.map((p) => p.name);
|
|
44
|
+
}
|
|
45
|
+
buildPrefixWhereClause(startParam) {
|
|
46
|
+
if (this.prefixes.length === 0) {
|
|
47
|
+
return { conditions: "", params: [] };
|
|
48
|
+
}
|
|
49
|
+
const conditions = this.prefixes.map((p, i) => `${p.name} = $${startParam + i}`).join(" AND ");
|
|
50
|
+
const params = this.prefixes.map((p) => this.prefixValues[p.name]);
|
|
51
|
+
return { conditions: " AND " + conditions, params };
|
|
52
|
+
}
|
|
53
|
+
getPrefixParamValues() {
|
|
54
|
+
return this.prefixes.map((p) => this.prefixValues[p.name]);
|
|
55
|
+
}
|
|
56
|
+
async setupDatabase() {
|
|
57
|
+
let sql;
|
|
58
|
+
try {
|
|
59
|
+
const enumValues = Object.values(JobStatus);
|
|
60
|
+
for (const v of enumValues) {
|
|
61
|
+
if (!SAFE_IDENTIFIER.test(v)) {
|
|
62
|
+
throw new Error(`Invalid JobStatus enum value: ${v}`);
|
|
63
|
+
}
|
|
64
|
+
}
|
|
65
|
+
sql = `CREATE TYPE job_status AS ENUM (${enumValues.map((v) => `'${v}'`).join(",")})`;
|
|
66
|
+
await this.db.query(sql);
|
|
67
|
+
} catch (e) {
|
|
68
|
+
if (e.code !== "42710")
|
|
69
|
+
throw e;
|
|
70
|
+
}
|
|
71
|
+
const prefixColumnsSql = this.buildPrefixColumnsSql();
|
|
72
|
+
const prefixColumnNames = this.getPrefixColumnNames();
|
|
73
|
+
const prefixIndexPrefix = prefixColumnNames.length > 0 ? prefixColumnNames.join(", ") + ", " : "";
|
|
74
|
+
sql = `
|
|
75
|
+
CREATE TABLE IF NOT EXISTS ${this.tableName} (
|
|
76
|
+
id SERIAL NOT NULL,
|
|
77
|
+
${prefixColumnsSql}fingerprint text NOT NULL,
|
|
78
|
+
queue text NOT NULL,
|
|
79
|
+
job_run_id text NOT NULL,
|
|
80
|
+
status job_status NOT NULL default 'PENDING',
|
|
81
|
+
input jsonb NOT NULL,
|
|
82
|
+
output jsonb,
|
|
83
|
+
run_attempts integer default 0,
|
|
84
|
+
max_retries integer default 20,
|
|
85
|
+
run_after timestamp with time zone DEFAULT now(),
|
|
86
|
+
last_ran_at timestamp with time zone,
|
|
87
|
+
created_at timestamp with time zone DEFAULT now(),
|
|
88
|
+
deadline_at timestamp with time zone,
|
|
89
|
+
completed_at timestamp with time zone,
|
|
90
|
+
error text,
|
|
91
|
+
error_code text,
|
|
92
|
+
progress real DEFAULT 0,
|
|
93
|
+
progress_message text DEFAULT '',
|
|
94
|
+
progress_details jsonb,
|
|
95
|
+
worker_id text
|
|
96
|
+
)`;
|
|
97
|
+
await this.db.query(sql);
|
|
98
|
+
const indexSuffix = prefixColumnNames.length > 0 ? "_" + prefixColumnNames.join("_") : "";
|
|
99
|
+
sql = `
|
|
100
|
+
CREATE INDEX IF NOT EXISTS job_fetcher${indexSuffix}_idx
|
|
101
|
+
ON ${this.tableName} (${prefixIndexPrefix}id, status, run_after)`;
|
|
102
|
+
await this.db.query(sql);
|
|
103
|
+
sql = `
|
|
104
|
+
CREATE INDEX IF NOT EXISTS job_queue_fetcher${indexSuffix}_idx
|
|
105
|
+
ON ${this.tableName} (${prefixIndexPrefix}queue, status, run_after)`;
|
|
106
|
+
await this.db.query(sql);
|
|
107
|
+
sql = `
|
|
108
|
+
CREATE INDEX IF NOT EXISTS jobs_fingerprint${indexSuffix}_unique_idx
|
|
109
|
+
ON ${this.tableName} (${prefixIndexPrefix}queue, fingerprint, status)`;
|
|
110
|
+
await this.db.query(sql);
|
|
111
|
+
const fnName = `${this.tableName}_notify`;
|
|
112
|
+
const trgName = `${this.tableName}_notify_trg`;
|
|
113
|
+
try {
|
|
114
|
+
await this.db.query(`
|
|
115
|
+
CREATE OR REPLACE FUNCTION ${fnName}() RETURNS trigger AS $fn$
|
|
116
|
+
DECLARE
|
|
117
|
+
channel TEXT := 'wglw_q_' || md5('${this.tableName}' || COALESCE(NEW.queue, OLD.queue));
|
|
118
|
+
payload TEXT;
|
|
119
|
+
BEGIN
|
|
120
|
+
payload := json_build_object(
|
|
121
|
+
'op', TG_OP,
|
|
122
|
+
'id', COALESCE(NEW.id, OLD.id),
|
|
123
|
+
'queue', COALESCE(NEW.queue, OLD.queue),
|
|
124
|
+
'status', COALESCE(NEW.status::text, OLD.status::text)
|
|
125
|
+
)::text;
|
|
126
|
+
PERFORM pg_notify(channel, payload);
|
|
127
|
+
RETURN NULL;
|
|
128
|
+
END;
|
|
129
|
+
$fn$ LANGUAGE plpgsql;
|
|
130
|
+
`);
|
|
131
|
+
await this.db.query(`DROP TRIGGER IF EXISTS ${trgName} ON ${this.tableName}`);
|
|
132
|
+
await this.db.query(`
|
|
133
|
+
CREATE TRIGGER ${trgName}
|
|
134
|
+
AFTER INSERT OR UPDATE ON ${this.tableName}
|
|
135
|
+
FOR EACH ROW EXECUTE FUNCTION ${fnName}();
|
|
136
|
+
`);
|
|
137
|
+
} catch {}
|
|
138
|
+
}
|
|
139
|
+
notifyChannelName() {
|
|
140
|
+
const tableAndQueue = `${this.tableName}${this.queueName}`;
|
|
141
|
+
const hash = createHash("md5").update(tableAndQueue).digest("hex");
|
|
142
|
+
return `wglw_q_${hash}`;
|
|
143
|
+
}
|
|
144
|
+
async add(job) {
|
|
145
|
+
const now = new Date().toISOString();
|
|
146
|
+
job.queue = this.queueName;
|
|
147
|
+
job.job_run_id = job.job_run_id ?? uuid4();
|
|
148
|
+
job.fingerprint = await makeFingerprint(job.input);
|
|
149
|
+
job.status = JobStatus.PENDING;
|
|
150
|
+
job.progress = 0;
|
|
151
|
+
job.progress_message = "";
|
|
152
|
+
job.progress_details = null;
|
|
153
|
+
job.created_at = now;
|
|
154
|
+
job.run_after = now;
|
|
155
|
+
const prefixColumnNames = this.getPrefixColumnNames();
|
|
156
|
+
const prefixColumnsInsert = prefixColumnNames.length > 0 ? prefixColumnNames.join(", ") + ", " : "";
|
|
157
|
+
const prefixParamValues = this.getPrefixParamValues();
|
|
158
|
+
const prefixParamPlaceholders = prefixColumnNames.length > 0 ? prefixColumnNames.map((_, i) => `$${i + 1}`).join(",") + "," : "";
|
|
159
|
+
const baseParamStart = prefixColumnNames.length + 1;
|
|
160
|
+
const sql = `
|
|
161
|
+
INSERT INTO ${this.tableName}(
|
|
162
|
+
${prefixColumnsInsert}queue,
|
|
163
|
+
fingerprint,
|
|
164
|
+
input,
|
|
165
|
+
run_after,
|
|
166
|
+
created_at,
|
|
167
|
+
deadline_at,
|
|
168
|
+
max_retries,
|
|
169
|
+
job_run_id,
|
|
170
|
+
progress,
|
|
171
|
+
progress_message,
|
|
172
|
+
progress_details
|
|
173
|
+
)
|
|
174
|
+
VALUES
|
|
175
|
+
(${prefixParamPlaceholders}$${baseParamStart},$${baseParamStart + 1},$${baseParamStart + 2},$${baseParamStart + 3},$${baseParamStart + 4},$${baseParamStart + 5},$${baseParamStart + 6},$${baseParamStart + 7},$${baseParamStart + 8},$${baseParamStart + 9},$${baseParamStart + 10})
|
|
176
|
+
RETURNING id`;
|
|
177
|
+
const params = [
|
|
178
|
+
...prefixParamValues,
|
|
179
|
+
job.queue,
|
|
180
|
+
job.fingerprint,
|
|
181
|
+
JSON.stringify(job.input),
|
|
182
|
+
job.run_after,
|
|
183
|
+
job.created_at,
|
|
184
|
+
job.deadline_at,
|
|
185
|
+
job.max_retries,
|
|
186
|
+
job.job_run_id,
|
|
187
|
+
job.progress,
|
|
188
|
+
job.progress_message,
|
|
189
|
+
job.progress_details ? JSON.stringify(job.progress_details) : null
|
|
190
|
+
];
|
|
191
|
+
const result = await this.db.query(sql, params);
|
|
192
|
+
if (!result)
|
|
193
|
+
throw new Error("Failed to add to queue");
|
|
194
|
+
job.id = result.rows[0].id;
|
|
195
|
+
return job.id;
|
|
196
|
+
}
|
|
197
|
+
async get(id) {
|
|
198
|
+
const { conditions: prefixConditions, params: prefixParams } = this.buildPrefixWhereClause(3);
|
|
199
|
+
const result = await this.db.query(`SELECT *
|
|
200
|
+
FROM ${this.tableName}
|
|
201
|
+
WHERE id = $1 AND queue = $2${prefixConditions}
|
|
202
|
+
FOR UPDATE SKIP LOCKED
|
|
203
|
+
LIMIT 1`, [id, this.queueName, ...prefixParams]);
|
|
204
|
+
if (!result || result.rows.length === 0)
|
|
205
|
+
return;
|
|
206
|
+
return result.rows[0];
|
|
207
|
+
}
|
|
208
|
+
async peek(status = JobStatus.PENDING, num = 100) {
|
|
209
|
+
num = Number(num) || 100;
|
|
210
|
+
const { conditions: prefixConditions, params: prefixParams } = this.buildPrefixWhereClause(4);
|
|
211
|
+
const result = await this.db.query(`
|
|
212
|
+
SELECT *
|
|
213
|
+
FROM ${this.tableName}
|
|
214
|
+
WHERE queue = $1
|
|
215
|
+
AND status = $2${prefixConditions}
|
|
216
|
+
ORDER BY run_after ASC
|
|
217
|
+
LIMIT $3
|
|
218
|
+
FOR UPDATE SKIP LOCKED`, [this.queueName, status, num, ...prefixParams]);
|
|
219
|
+
if (!result)
|
|
220
|
+
return [];
|
|
221
|
+
return result.rows;
|
|
222
|
+
}
|
|
223
|
+
async next(workerId) {
|
|
224
|
+
const { conditions: prefixConditions, params: prefixParams } = this.buildPrefixWhereClause(5);
|
|
225
|
+
const result = await this.db.query(`
|
|
226
|
+
UPDATE ${this.tableName}
|
|
227
|
+
SET status = $1, last_ran_at = NOW() AT TIME ZONE 'UTC', worker_id = $4
|
|
228
|
+
WHERE id = (
|
|
229
|
+
SELECT id
|
|
230
|
+
FROM ${this.tableName}
|
|
231
|
+
WHERE queue = $2
|
|
232
|
+
AND status = $3
|
|
233
|
+
${prefixConditions}
|
|
234
|
+
AND run_after <= NOW() AT TIME ZONE 'UTC'
|
|
235
|
+
ORDER BY run_after ASC
|
|
236
|
+
FOR UPDATE SKIP LOCKED
|
|
237
|
+
LIMIT 1
|
|
238
|
+
)
|
|
239
|
+
RETURNING *`, [JobStatus.PROCESSING, this.queueName, JobStatus.PENDING, workerId, ...prefixParams]);
|
|
240
|
+
return result?.rows?.[0] ?? undefined;
|
|
241
|
+
}
|
|
242
|
+
async size(status = JobStatus.PENDING) {
|
|
243
|
+
const { conditions: prefixConditions, params: prefixParams } = this.buildPrefixWhereClause(3);
|
|
244
|
+
const result = await this.db.query(`
|
|
245
|
+
SELECT COUNT(*) as count
|
|
246
|
+
FROM ${this.tableName}
|
|
247
|
+
WHERE queue = $1
|
|
248
|
+
AND status = $2${prefixConditions}`, [this.queueName, status, ...prefixParams]);
|
|
249
|
+
if (!result)
|
|
250
|
+
return 0;
|
|
251
|
+
return parseInt(result.rows[0].count, 10);
|
|
252
|
+
}
|
|
253
|
+
async complete(jobDetails) {
|
|
254
|
+
const prefixParams = this.getPrefixParamValues();
|
|
255
|
+
if (jobDetails.status === JobStatus.DISABLED) {
|
|
256
|
+
const { conditions: prefixConditions } = this.buildPrefixWhereClause(4);
|
|
257
|
+
await this.db.query(`UPDATE ${this.tableName}
|
|
258
|
+
SET
|
|
259
|
+
status = $1,
|
|
260
|
+
progress = 100,
|
|
261
|
+
progress_message = '',
|
|
262
|
+
progress_details = NULL,
|
|
263
|
+
completed_at = NOW() AT TIME ZONE 'UTC'
|
|
264
|
+
WHERE id = $2 AND queue = $3${prefixConditions}`, [jobDetails.status, jobDetails.id, this.queueName, ...prefixParams]);
|
|
265
|
+
} else if (jobDetails.status === JobStatus.PENDING) {
|
|
266
|
+
const { conditions: prefixConditions } = this.buildPrefixWhereClause(7);
|
|
267
|
+
await this.db.query(`UPDATE ${this.tableName}
|
|
268
|
+
SET
|
|
269
|
+
error = $1,
|
|
270
|
+
error_code = $2,
|
|
271
|
+
status = $3,
|
|
272
|
+
run_after = $4,
|
|
273
|
+
progress = 0,
|
|
274
|
+
progress_message = '',
|
|
275
|
+
progress_details = NULL,
|
|
276
|
+
run_attempts = run_attempts + 1,
|
|
277
|
+
last_ran_at = NOW() AT TIME ZONE 'UTC'
|
|
278
|
+
WHERE id = $5 AND queue = $6${prefixConditions}`, [
|
|
279
|
+
jobDetails.error,
|
|
280
|
+
jobDetails.error_code,
|
|
281
|
+
jobDetails.status,
|
|
282
|
+
jobDetails.run_after,
|
|
283
|
+
jobDetails.id,
|
|
284
|
+
this.queueName,
|
|
285
|
+
...prefixParams
|
|
286
|
+
]);
|
|
287
|
+
} else {
|
|
288
|
+
const { conditions: prefixConditions } = this.buildPrefixWhereClause(7);
|
|
289
|
+
await this.db.query(`
|
|
290
|
+
UPDATE ${this.tableName}
|
|
291
|
+
SET
|
|
292
|
+
output = $1,
|
|
293
|
+
error = $2,
|
|
294
|
+
error_code = $3,
|
|
295
|
+
status = $4,
|
|
296
|
+
progress = 100,
|
|
297
|
+
progress_message = '',
|
|
298
|
+
progress_details = NULL,
|
|
299
|
+
run_attempts = run_attempts + 1,
|
|
300
|
+
completed_at = NOW() AT TIME ZONE 'UTC',
|
|
301
|
+
last_ran_at = NOW() AT TIME ZONE 'UTC'
|
|
302
|
+
WHERE id = $5 AND queue = $6${prefixConditions}`, [
|
|
303
|
+
jobDetails.output ? JSON.stringify(jobDetails.output) : null,
|
|
304
|
+
jobDetails.error ?? null,
|
|
305
|
+
jobDetails.error_code ?? null,
|
|
306
|
+
jobDetails.status,
|
|
307
|
+
jobDetails.id,
|
|
308
|
+
this.queueName,
|
|
309
|
+
...prefixParams
|
|
310
|
+
]);
|
|
311
|
+
}
|
|
312
|
+
}
|
|
313
|
+
async deleteAll() {
|
|
314
|
+
const { conditions: prefixConditions, params: prefixParams } = this.buildPrefixWhereClause(2);
|
|
315
|
+
await this.db.query(`
|
|
316
|
+
DELETE FROM ${this.tableName}
|
|
317
|
+
WHERE queue = $1${prefixConditions}`, [this.queueName, ...prefixParams]);
|
|
318
|
+
}
|
|
319
|
+
async outputForInput(input) {
|
|
320
|
+
const fingerprint = await makeFingerprint(input);
|
|
321
|
+
const { conditions: prefixConditions, params: prefixParams } = this.buildPrefixWhereClause(3);
|
|
322
|
+
const result = await this.db.query(`
|
|
323
|
+
SELECT output
|
|
324
|
+
FROM ${this.tableName}
|
|
325
|
+
WHERE fingerprint = $1 AND queue = $2 AND status = 'COMPLETED'${prefixConditions}`, [fingerprint, this.queueName, ...prefixParams]);
|
|
326
|
+
if (!result || result.rows.length === 0)
|
|
327
|
+
return null;
|
|
328
|
+
return result.rows[0].output;
|
|
329
|
+
}
|
|
330
|
+
async abort(jobId) {
|
|
331
|
+
const { conditions: prefixConditions, params: prefixParams } = this.buildPrefixWhereClause(3);
|
|
332
|
+
await this.db.query(`
|
|
333
|
+
UPDATE ${this.tableName}
|
|
334
|
+
SET status = 'ABORTING'
|
|
335
|
+
WHERE id = $1 AND queue = $2${prefixConditions}`, [jobId, this.queueName, ...prefixParams]);
|
|
336
|
+
}
|
|
337
|
+
async release(jobId) {
|
|
338
|
+
const { conditions: prefixConditions, params: prefixParams } = this.buildPrefixWhereClause(3);
|
|
339
|
+
await this.db.query(`
|
|
340
|
+
UPDATE ${this.tableName}
|
|
341
|
+
SET status = 'PENDING',
|
|
342
|
+
worker_id = NULL,
|
|
343
|
+
progress = 0,
|
|
344
|
+
progress_message = '',
|
|
345
|
+
progress_details = NULL
|
|
346
|
+
WHERE id = $1 AND queue = $2${prefixConditions}`, [jobId, this.queueName, ...prefixParams]);
|
|
347
|
+
}
|
|
348
|
+
async getByRunId(job_run_id) {
|
|
349
|
+
const { conditions: prefixConditions, params: prefixParams } = this.buildPrefixWhereClause(3);
|
|
350
|
+
const result = await this.db.query(`
|
|
351
|
+
SELECT * FROM ${this.tableName} WHERE job_run_id = $1 AND queue = $2${prefixConditions}`, [job_run_id, this.queueName, ...prefixParams]);
|
|
352
|
+
if (!result)
|
|
353
|
+
return [];
|
|
354
|
+
return result.rows;
|
|
355
|
+
}
|
|
356
|
+
async saveProgress(jobId, progress, message, details) {
|
|
357
|
+
const { conditions: prefixConditions, params: prefixParams } = this.buildPrefixWhereClause(6);
|
|
358
|
+
await this.db.query(`
|
|
359
|
+
UPDATE ${this.tableName}
|
|
360
|
+
SET progress = $1,
|
|
361
|
+
progress_message = $2,
|
|
362
|
+
progress_details = $3
|
|
363
|
+
WHERE id = $4 AND queue = $5${prefixConditions}`, [
|
|
364
|
+
progress,
|
|
365
|
+
message,
|
|
366
|
+
details ? JSON.stringify(details) : null,
|
|
367
|
+
jobId,
|
|
368
|
+
this.queueName,
|
|
369
|
+
...prefixParams
|
|
370
|
+
]);
|
|
371
|
+
}
|
|
372
|
+
async delete(jobId) {
|
|
373
|
+
const { conditions: prefixConditions, params: prefixParams } = this.buildPrefixWhereClause(3);
|
|
374
|
+
await this.db.query(`DELETE FROM ${this.tableName} WHERE id = $1 AND queue = $2${prefixConditions}`, [jobId, this.queueName, ...prefixParams]);
|
|
375
|
+
}
|
|
376
|
+
async deleteJobsByStatusAndAge(status, olderThanMs) {
|
|
377
|
+
const cutoffDate = new Date(Date.now() - olderThanMs).toISOString();
|
|
378
|
+
const { conditions: prefixConditions, params: prefixParams } = this.buildPrefixWhereClause(4);
|
|
379
|
+
await this.db.query(`DELETE FROM ${this.tableName}
|
|
380
|
+
WHERE queue = $1
|
|
381
|
+
AND status = $2
|
|
382
|
+
AND completed_at IS NOT NULL
|
|
383
|
+
AND completed_at <= $3${prefixConditions}`, [this.queueName, status, cutoffDate, ...prefixParams]);
|
|
384
|
+
}
|
|
385
|
+
subscribeToChanges(callback, options) {
|
|
386
|
+
const poolMaybe = this.db;
|
|
387
|
+
if (typeof poolMaybe.connect !== "function") {
|
|
388
|
+
throw new Error("PostgresQueueStorage.subscribeToChanges requires a pg.Pool (got a single-connection wrapper)");
|
|
389
|
+
}
|
|
390
|
+
const pool = poolMaybe;
|
|
391
|
+
const channel = this.notifyChannelName();
|
|
392
|
+
const effectivePrefixFilter = options?.prefixFilter === undefined ? this.prefixes.length > 0 ? this.prefixValues : null : Object.keys(options.prefixFilter).length === 0 ? null : options.prefixFilter;
|
|
393
|
+
let unsubscribed = false;
|
|
394
|
+
let activeClient = null;
|
|
395
|
+
let reconnectTimer = null;
|
|
396
|
+
let backoffMs = 250;
|
|
397
|
+
const dispatch = (change) => {
|
|
398
|
+
try {
|
|
399
|
+
callback(change);
|
|
400
|
+
} catch (err) {
|
|
401
|
+
getLogger().debug("PostgresQueueStorage subscribe callback threw", {
|
|
402
|
+
channel,
|
|
403
|
+
changeType: change.type,
|
|
404
|
+
error: err
|
|
405
|
+
});
|
|
406
|
+
}
|
|
407
|
+
};
|
|
408
|
+
const matchesPrefix = (change, filter) => {
|
|
409
|
+
const row = change.new ?? change.old;
|
|
410
|
+
if (!row)
|
|
411
|
+
return false;
|
|
412
|
+
for (const [k, v] of Object.entries(filter)) {
|
|
413
|
+
if (row[k] !== v)
|
|
414
|
+
return false;
|
|
415
|
+
}
|
|
416
|
+
return true;
|
|
417
|
+
};
|
|
418
|
+
const hydrate = async (id) => {
|
|
419
|
+
try {
|
|
420
|
+
return await this.get(id);
|
|
421
|
+
} catch {
|
|
422
|
+
return;
|
|
423
|
+
}
|
|
424
|
+
};
|
|
425
|
+
const handleNotification = (msg) => {
|
|
426
|
+
if (msg.channel !== channel || !msg.payload)
|
|
427
|
+
return;
|
|
428
|
+
let parsed;
|
|
429
|
+
try {
|
|
430
|
+
parsed = JSON.parse(msg.payload);
|
|
431
|
+
} catch {
|
|
432
|
+
return;
|
|
433
|
+
}
|
|
434
|
+
(async () => {
|
|
435
|
+
const op = parsed.op;
|
|
436
|
+
const fallback = {
|
|
437
|
+
id: parsed.id,
|
|
438
|
+
queue: parsed.queue,
|
|
439
|
+
status: parsed.status
|
|
440
|
+
};
|
|
441
|
+
const fullRow = op === "DELETE" ? undefined : await hydrate(parsed.id);
|
|
442
|
+
const change = {
|
|
443
|
+
type: op === "INSERT" ? "INSERT" : op === "DELETE" ? "DELETE" : "UPDATE",
|
|
444
|
+
new: op === "DELETE" ? undefined : fullRow ?? fallback,
|
|
445
|
+
old: op === "DELETE" ? fallback : undefined
|
|
446
|
+
};
|
|
447
|
+
if (effectivePrefixFilter && !matchesPrefix(change, effectivePrefixFilter))
|
|
448
|
+
return;
|
|
449
|
+
dispatch(change);
|
|
450
|
+
})();
|
|
451
|
+
};
|
|
452
|
+
const connect = async () => {
|
|
453
|
+
if (unsubscribed)
|
|
454
|
+
return;
|
|
455
|
+
try {
|
|
456
|
+
const client = await pool.connect();
|
|
457
|
+
if (unsubscribed) {
|
|
458
|
+
client.release();
|
|
459
|
+
return;
|
|
460
|
+
}
|
|
461
|
+
activeClient = client;
|
|
462
|
+
client.on("notification", handleNotification);
|
|
463
|
+
client.on("error", () => scheduleReconnect());
|
|
464
|
+
await client.query(`LISTEN ${channel}`);
|
|
465
|
+
backoffMs = 250;
|
|
466
|
+
dispatch({ type: "RESYNC" });
|
|
467
|
+
} catch {
|
|
468
|
+
scheduleReconnect();
|
|
469
|
+
}
|
|
470
|
+
};
|
|
471
|
+
const scheduleReconnect = () => {
|
|
472
|
+
if (unsubscribed || reconnectTimer)
|
|
473
|
+
return;
|
|
474
|
+
const c = activeClient;
|
|
475
|
+
activeClient = null;
|
|
476
|
+
try {
|
|
477
|
+
c?.removeAllListeners?.("notification");
|
|
478
|
+
c?.removeAllListeners?.("error");
|
|
479
|
+
c?.release();
|
|
480
|
+
} catch {}
|
|
481
|
+
const delay = Math.min(backoffMs, 30000);
|
|
482
|
+
backoffMs = Math.min(backoffMs * 2, 30000);
|
|
483
|
+
reconnectTimer = setTimeout(() => {
|
|
484
|
+
reconnectTimer = null;
|
|
485
|
+
connect();
|
|
486
|
+
}, delay);
|
|
487
|
+
};
|
|
488
|
+
connect();
|
|
489
|
+
return () => {
|
|
490
|
+
unsubscribed = true;
|
|
491
|
+
if (reconnectTimer) {
|
|
492
|
+
clearTimeout(reconnectTimer);
|
|
493
|
+
reconnectTimer = null;
|
|
494
|
+
}
|
|
495
|
+
const c = activeClient;
|
|
496
|
+
activeClient = null;
|
|
497
|
+
if (c) {
|
|
498
|
+
c.query(`UNLISTEN ${channel}`).catch(() => {}).finally(() => {
|
|
499
|
+
try {
|
|
500
|
+
c.removeAllListeners?.("notification");
|
|
501
|
+
c.removeAllListeners?.("error");
|
|
502
|
+
c.release();
|
|
503
|
+
} catch {}
|
|
504
|
+
});
|
|
505
|
+
}
|
|
506
|
+
};
|
|
507
|
+
}
|
|
508
|
+
}
|
|
509
|
+
// src/job-queue/PostgresRateLimiterStorage.ts
|
|
510
|
+
import { createServiceToken as createServiceToken2 } from "@workglow/util";
|
|
511
|
+
var POSTGRES_RATE_LIMITER_STORAGE = createServiceToken2("ratelimiter.storage.postgres");
|
|
512
|
+
|
|
513
|
+
class PostgresRateLimiterStorage {
|
|
514
|
+
db;
|
|
515
|
+
scope = "cluster";
|
|
516
|
+
prefixes;
|
|
517
|
+
prefixValues;
|
|
518
|
+
executionTableName;
|
|
519
|
+
nextAvailableTableName;
|
|
520
|
+
constructor(db, options) {
|
|
521
|
+
this.db = db;
|
|
522
|
+
this.prefixes = options?.prefixes ?? [];
|
|
523
|
+
this.prefixValues = options?.prefixValues ?? {};
|
|
524
|
+
if (this.prefixes.length > 0) {
|
|
525
|
+
const prefixNames = this.prefixes.map((p) => p.name).join("_");
|
|
526
|
+
this.executionTableName = `rate_limit_executions_${prefixNames}`;
|
|
527
|
+
this.nextAvailableTableName = `rate_limit_next_available_${prefixNames}`;
|
|
528
|
+
} else {
|
|
529
|
+
this.executionTableName = "rate_limit_executions";
|
|
530
|
+
this.nextAvailableTableName = "rate_limit_next_available";
|
|
531
|
+
}
|
|
532
|
+
}
|
|
533
|
+
getPrefixColumnType(type) {
|
|
534
|
+
return type === "uuid" ? "UUID" : "INTEGER";
|
|
535
|
+
}
|
|
536
|
+
buildPrefixColumnsSql() {
|
|
537
|
+
if (this.prefixes.length === 0)
|
|
538
|
+
return "";
|
|
539
|
+
return this.prefixes.map((p) => `${p.name} ${this.getPrefixColumnType(p.type)} NOT NULL`).join(`,
|
|
540
|
+
`) + `,
|
|
541
|
+
`;
|
|
542
|
+
}
|
|
543
|
+
getPrefixColumnNames() {
|
|
544
|
+
return this.prefixes.map((p) => p.name);
|
|
545
|
+
}
|
|
546
|
+
buildPrefixWhereClause(startParam) {
|
|
547
|
+
if (this.prefixes.length === 0) {
|
|
548
|
+
return { conditions: "", params: [] };
|
|
549
|
+
}
|
|
550
|
+
const conditions = this.prefixes.map((p, i) => `${p.name} = $${startParam + i}`).join(" AND ");
|
|
551
|
+
const params = this.prefixes.map((p) => this.prefixValues[p.name]);
|
|
552
|
+
return { conditions: " AND " + conditions, params };
|
|
553
|
+
}
|
|
554
|
+
getPrefixParamValues() {
|
|
555
|
+
return this.prefixes.map((p) => this.prefixValues[p.name]);
|
|
556
|
+
}
|
|
557
|
+
async setupDatabase() {
|
|
558
|
+
const prefixColumnsSql = this.buildPrefixColumnsSql();
|
|
559
|
+
const prefixColumnNames = this.getPrefixColumnNames();
|
|
560
|
+
const prefixIndexPrefix = prefixColumnNames.length > 0 ? prefixColumnNames.join(", ") + ", " : "";
|
|
561
|
+
const indexSuffix = prefixColumnNames.length > 0 ? "_" + prefixColumnNames.join("_") : "";
|
|
562
|
+
await this.db.query(`
|
|
563
|
+
CREATE TABLE IF NOT EXISTS ${this.executionTableName} (
|
|
564
|
+
id SERIAL PRIMARY KEY,
|
|
565
|
+
${prefixColumnsSql}queue_name TEXT NOT NULL,
|
|
566
|
+
executed_at TIMESTAMP WITH TIME ZONE DEFAULT NOW()
|
|
567
|
+
)
|
|
568
|
+
`);
|
|
569
|
+
await this.db.query(`
|
|
570
|
+
CREATE INDEX IF NOT EXISTS rate_limit_exec_queue${indexSuffix}_idx
|
|
571
|
+
ON ${this.executionTableName} (${prefixIndexPrefix}queue_name, executed_at)
|
|
572
|
+
`);
|
|
573
|
+
const primaryKeyColumns = prefixColumnNames.length > 0 ? `${prefixColumnNames.join(", ")}, queue_name` : "queue_name";
|
|
574
|
+
await this.db.query(`
|
|
575
|
+
CREATE TABLE IF NOT EXISTS ${this.nextAvailableTableName} (
|
|
576
|
+
${prefixColumnsSql}queue_name TEXT NOT NULL,
|
|
577
|
+
next_available_at TIMESTAMP WITH TIME ZONE,
|
|
578
|
+
PRIMARY KEY (${primaryKeyColumns})
|
|
579
|
+
)
|
|
580
|
+
`);
|
|
581
|
+
}
|
|
582
|
+
async tryReserveExecution(queueName, maxExecutions, windowMs) {
|
|
583
|
+
const prefixColumnNames = this.getPrefixColumnNames();
|
|
584
|
+
const prefixParamValues = this.getPrefixParamValues();
|
|
585
|
+
const prefixCount = prefixColumnNames.length;
|
|
586
|
+
const queueParam = `$${prefixCount + 1}`;
|
|
587
|
+
const windowStartParam = `$${prefixCount + 2}`;
|
|
588
|
+
const lockKeyParts = [`'${this.executionTableName}'`];
|
|
589
|
+
for (let i = 0;i < prefixCount; i++) {
|
|
590
|
+
lockKeyParts.push(`$${i + 1}::text`);
|
|
591
|
+
}
|
|
592
|
+
lockKeyParts.push(`${queueParam}::text`);
|
|
593
|
+
const lockKeyExpr = `hashtextextended(${lockKeyParts.join(" || '|' || ")}, 0)`;
|
|
594
|
+
const prefixWhere = prefixCount > 0 ? " AND " + prefixColumnNames.map((p, i) => `${p} = $${i + 1}`).join(" AND ") : "";
|
|
595
|
+
const prefixInsertCols = prefixCount > 0 ? prefixColumnNames.join(", ") + ", " : "";
|
|
596
|
+
const prefixInsertPlaceholders = prefixCount > 0 ? prefixColumnNames.map((_, i) => `$${i + 1}`).join(", ") + ", " : "";
|
|
597
|
+
const windowStart = new Date(Date.now() - windowMs).toISOString();
|
|
598
|
+
const supportsConnect = typeof this.db.connect === "function";
|
|
599
|
+
if (!supportsConnect) {
|
|
600
|
+
const dbAny = this.db;
|
|
601
|
+
const ctorName = dbAny.constructor?.name;
|
|
602
|
+
const looksLikePGlite = typeof dbAny.exec === "function" && dbAny.waitReady !== undefined;
|
|
603
|
+
const looksLikePGLitePool = ctorName === "PGLitePool";
|
|
604
|
+
if (!looksLikePGlite && !looksLikePGLitePool) {
|
|
605
|
+
throw new Error(`PostgresRateLimiterStorage.tryReserveExecution requires a pg.Pool with connect() or a known single-connection wrapper (PGLitePool, PGlite); got ${ctorName ?? typeof this.db}. A multi-connection pool without connect() would dispatch the advisory lock and the INSERT to different sessions, breaking atomicity.`);
|
|
606
|
+
}
|
|
607
|
+
}
|
|
608
|
+
const conn = supportsConnect ? await this.db.connect() : { query: this.db.query.bind(this.db), release: () => {} };
|
|
609
|
+
try {
|
|
610
|
+
await conn.query("BEGIN");
|
|
611
|
+
try {
|
|
612
|
+
await conn.query(`SELECT pg_advisory_xact_lock(${lockKeyExpr})`, [
|
|
613
|
+
...prefixParamValues,
|
|
614
|
+
queueName
|
|
615
|
+
]);
|
|
616
|
+
const countResult = await conn.query(`
|
|
617
|
+
SELECT COUNT(*)::int AS n
|
|
618
|
+
FROM ${this.executionTableName}
|
|
619
|
+
WHERE queue_name = ${queueParam} AND executed_at > ${windowStartParam}${prefixWhere}
|
|
620
|
+
`, [...prefixParamValues, queueName, windowStart]);
|
|
621
|
+
const n = countResult.rows[0]?.n ?? 0;
|
|
622
|
+
if (n >= maxExecutions) {
|
|
623
|
+
await conn.query("COMMIT");
|
|
624
|
+
return null;
|
|
625
|
+
}
|
|
626
|
+
const naResult = await conn.query(`
|
|
627
|
+
SELECT next_available_at
|
|
628
|
+
FROM ${this.nextAvailableTableName}
|
|
629
|
+
WHERE queue_name = ${queueParam}${prefixWhere}
|
|
630
|
+
`, [...prefixParamValues, queueName]);
|
|
631
|
+
const nextAvailableAt = naResult.rows[0]?.next_available_at ?? null;
|
|
632
|
+
if (nextAvailableAt && new Date(nextAvailableAt).getTime() > Date.now()) {
|
|
633
|
+
await conn.query("COMMIT");
|
|
634
|
+
return null;
|
|
635
|
+
}
|
|
636
|
+
const insertResult = await conn.query(`
|
|
637
|
+
INSERT INTO ${this.executionTableName} (${prefixInsertCols}queue_name)
|
|
638
|
+
VALUES (${prefixInsertPlaceholders}${queueParam})
|
|
639
|
+
RETURNING id
|
|
640
|
+
`, [...prefixParamValues, queueName]);
|
|
641
|
+
await conn.query("COMMIT");
|
|
642
|
+
return insertResult.rows[0]?.id ?? null;
|
|
643
|
+
} catch (err) {
|
|
644
|
+
try {
|
|
645
|
+
await conn.query("ROLLBACK");
|
|
646
|
+
} catch {}
|
|
647
|
+
throw err;
|
|
648
|
+
}
|
|
649
|
+
} finally {
|
|
650
|
+
conn.release();
|
|
651
|
+
}
|
|
652
|
+
}
|
|
653
|
+
async releaseExecution(queueName, token) {
|
|
654
|
+
if (token === null || token === undefined)
|
|
655
|
+
return;
|
|
656
|
+
const { conditions: prefixConditions, params: prefixParams } = this.buildPrefixWhereClause(3);
|
|
657
|
+
await this.db.query(`DELETE FROM ${this.executionTableName} WHERE id = $1 AND queue_name = $2${prefixConditions}`, [token, queueName, ...prefixParams]);
|
|
658
|
+
}
|
|
659
|
+
async recordExecution(queueName) {
|
|
660
|
+
const prefixColumnNames = this.getPrefixColumnNames();
|
|
661
|
+
const prefixColumnsInsert = prefixColumnNames.length > 0 ? prefixColumnNames.join(", ") + ", " : "";
|
|
662
|
+
const prefixParamValues = this.getPrefixParamValues();
|
|
663
|
+
const prefixParamPlaceholders = prefixColumnNames.length > 0 ? prefixColumnNames.map((_, i) => `$${i + 1}`).join(", ") + ", " : "";
|
|
664
|
+
const queueParamNum = prefixColumnNames.length + 1;
|
|
665
|
+
await this.db.query(`
|
|
666
|
+
INSERT INTO ${this.executionTableName} (${prefixColumnsInsert}queue_name)
|
|
667
|
+
VALUES (${prefixParamPlaceholders}$${queueParamNum})
|
|
668
|
+
`, [...prefixParamValues, queueName]);
|
|
669
|
+
}
|
|
670
|
+
async getExecutionCount(queueName, windowStartTime) {
|
|
671
|
+
const { conditions: prefixConditions, params: prefixParams } = this.buildPrefixWhereClause(3);
|
|
672
|
+
const result = await this.db.query(`
|
|
673
|
+
SELECT COUNT(*) AS count
|
|
674
|
+
FROM ${this.executionTableName}
|
|
675
|
+
WHERE queue_name = $1 AND executed_at > $2${prefixConditions}
|
|
676
|
+
`, [queueName, windowStartTime, ...prefixParams]);
|
|
677
|
+
return parseInt(result.rows[0]?.count ?? "0", 10);
|
|
678
|
+
}
|
|
679
|
+
async getOldestExecutionAtOffset(queueName, offset) {
|
|
680
|
+
const { conditions: prefixConditions, params: prefixParams } = this.buildPrefixWhereClause(3);
|
|
681
|
+
const result = await this.db.query(`
|
|
682
|
+
SELECT executed_at
|
|
683
|
+
FROM ${this.executionTableName}
|
|
684
|
+
WHERE queue_name = $1${prefixConditions}
|
|
685
|
+
ORDER BY executed_at ASC
|
|
686
|
+
LIMIT 1 OFFSET $2
|
|
687
|
+
`, [queueName, offset, ...prefixParams]);
|
|
688
|
+
const executedAt = result.rows[0]?.executed_at;
|
|
689
|
+
if (!executedAt)
|
|
690
|
+
return;
|
|
691
|
+
return new Date(executedAt).toISOString();
|
|
692
|
+
}
|
|
693
|
+
async getNextAvailableTime(queueName) {
|
|
694
|
+
const { conditions: prefixConditions, params: prefixParams } = this.buildPrefixWhereClause(2);
|
|
695
|
+
const result = await this.db.query(`
|
|
696
|
+
SELECT next_available_at
|
|
697
|
+
FROM ${this.nextAvailableTableName}
|
|
698
|
+
WHERE queue_name = $1${prefixConditions}
|
|
699
|
+
`, [queueName, ...prefixParams]);
|
|
700
|
+
const nextAvailableAt = result.rows[0]?.next_available_at;
|
|
701
|
+
if (!nextAvailableAt)
|
|
702
|
+
return;
|
|
703
|
+
return new Date(nextAvailableAt).toISOString();
|
|
704
|
+
}
|
|
705
|
+
async setNextAvailableTime(queueName, nextAvailableAt) {
|
|
706
|
+
const prefixColumnNames = this.getPrefixColumnNames();
|
|
707
|
+
const prefixColumnsInsert = prefixColumnNames.length > 0 ? prefixColumnNames.join(", ") + ", " : "";
|
|
708
|
+
const prefixParamValues = this.getPrefixParamValues();
|
|
709
|
+
const prefixParamPlaceholders = prefixColumnNames.length > 0 ? prefixColumnNames.map((_, i) => `$${i + 1}`).join(", ") + ", " : "";
|
|
710
|
+
const baseParamStart = prefixColumnNames.length + 1;
|
|
711
|
+
const conflictColumns = prefixColumnNames.length > 0 ? `${prefixColumnNames.join(", ")}, queue_name` : "queue_name";
|
|
712
|
+
await this.db.query(`
|
|
713
|
+
INSERT INTO ${this.nextAvailableTableName} (${prefixColumnsInsert}queue_name, next_available_at)
|
|
714
|
+
VALUES (${prefixParamPlaceholders}$${baseParamStart}, $${baseParamStart + 1})
|
|
715
|
+
ON CONFLICT (${conflictColumns})
|
|
716
|
+
DO UPDATE SET next_available_at = EXCLUDED.next_available_at
|
|
717
|
+
`, [...prefixParamValues, queueName, nextAvailableAt]);
|
|
718
|
+
}
|
|
719
|
+
async clear(queueName) {
|
|
720
|
+
const { conditions: prefixConditions, params: prefixParams } = this.buildPrefixWhereClause(2);
|
|
721
|
+
await this.db.query(`DELETE FROM ${this.executionTableName} WHERE queue_name = $1${prefixConditions}`, [queueName, ...prefixParams]);
|
|
722
|
+
await this.db.query(`DELETE FROM ${this.nextAvailableTableName} WHERE queue_name = $1${prefixConditions}`, [queueName, ...prefixParams]);
|
|
723
|
+
}
|
|
724
|
+
}
|
|
725
|
+
export {
|
|
726
|
+
PostgresRateLimiterStorage,
|
|
727
|
+
PostgresQueueStorage,
|
|
728
|
+
POSTGRES_RATE_LIMITER_STORAGE,
|
|
729
|
+
POSTGRES_QUEUE_STORAGE
|
|
730
|
+
};
|
|
731
|
+
|
|
732
|
+
//# debugId=561DECA09B89A19664756E2164756E21
|