@mastra/pg 0.1.0-alpha.2

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.js ADDED
@@ -0,0 +1,886 @@
1
+ import { MastraVector } from '@mastra/core/vector';
2
+ import pg from 'pg';
3
+ import { BaseFilterTranslator } from '@mastra/core/filter';
4
+ import '@mastra/core/memory';
5
+ import { MastraStorage } from '@mastra/core/storage';
6
+ import '@mastra/core/workflows';
7
+ import pgPromise from 'pg-promise';
8
+
9
+ // src/vector/index.ts
10
+ var PGFilterTranslator = class extends BaseFilterTranslator {
11
+ getSupportedOperators() {
12
+ return {
13
+ ...BaseFilterTranslator.DEFAULT_OPERATORS,
14
+ custom: ["$contains", "$size"]
15
+ };
16
+ }
17
+ translate(filter) {
18
+ if (this.isEmpty(filter)) {
19
+ return filter;
20
+ }
21
+ this.validateFilter(filter);
22
+ return this.translateNode(filter);
23
+ }
24
+ translateNode(node, currentPath = "") {
25
+ const withPath = (result2) => currentPath ? { [currentPath]: result2 } : result2;
26
+ if (this.isPrimitive(node)) {
27
+ return withPath({ $eq: this.normalizeComparisonValue(node) });
28
+ }
29
+ if (Array.isArray(node)) {
30
+ return withPath({ $in: this.normalizeArrayValues(node) });
31
+ }
32
+ if (node instanceof RegExp) {
33
+ return withPath(this.translateRegexPattern(node.source, node.flags));
34
+ }
35
+ const entries = Object.entries(node);
36
+ const result = {};
37
+ if ("$options" in node && !("$regex" in node)) {
38
+ throw new Error("$options is not valid without $regex");
39
+ }
40
+ if ("$regex" in node) {
41
+ const options = node.$options || "";
42
+ return withPath(this.translateRegexPattern(node.$regex, options));
43
+ }
44
+ for (const [key, value] of entries) {
45
+ if (key === "$options") continue;
46
+ const newPath = currentPath ? `${currentPath}.${key}` : key;
47
+ if (this.isLogicalOperator(key)) {
48
+ result[key] = Array.isArray(value) ? value.map((filter) => this.translateNode(filter)) : this.translateNode(value);
49
+ } else if (this.isOperator(key)) {
50
+ if (this.isArrayOperator(key) && !Array.isArray(value) && key !== "$elemMatch") {
51
+ result[key] = [value];
52
+ } else if (this.isBasicOperator(key) && Array.isArray(value)) {
53
+ result[key] = JSON.stringify(value);
54
+ } else {
55
+ result[key] = value;
56
+ }
57
+ } else if (typeof value === "object" && value !== null) {
58
+ const hasOperators = Object.keys(value).some((k) => this.isOperator(k));
59
+ if (hasOperators) {
60
+ result[newPath] = this.translateNode(value);
61
+ } else {
62
+ Object.assign(result, this.translateNode(value, newPath));
63
+ }
64
+ } else {
65
+ result[newPath] = this.translateNode(value);
66
+ }
67
+ }
68
+ return result;
69
+ }
70
+ translateRegexPattern(pattern, options = "") {
71
+ if (!options) return { $regex: pattern };
72
+ const flags = options.split("").filter((f) => "imsux".includes(f)).join("");
73
+ return { $regex: flags ? `(?${flags})${pattern}` : pattern };
74
+ }
75
+ };
76
+ var createBasicOperator = (symbol) => {
77
+ return (key, paramIndex) => ({
78
+ sql: `CASE
79
+ WHEN $${paramIndex}::text IS NULL THEN metadata#>>'{${handleKey(key)}}' IS ${symbol === "=" ? "" : "NOT"} NULL
80
+ ELSE metadata#>>'{${handleKey(key)}}' ${symbol} $${paramIndex}::text
81
+ END`,
82
+ needsValue: true
83
+ });
84
+ };
85
+ var createNumericOperator = (symbol) => {
86
+ return (key, paramIndex) => ({
87
+ sql: `(metadata#>>'{${handleKey(key)}}')::numeric ${symbol} $${paramIndex}`,
88
+ needsValue: true
89
+ });
90
+ };
91
+ function buildElemMatchConditions(value, paramIndex) {
92
+ if (typeof value !== "object" || Array.isArray(value)) {
93
+ throw new Error("$elemMatch requires an object with conditions");
94
+ }
95
+ const conditions = [];
96
+ const values = [];
97
+ Object.entries(value).forEach(([field, val]) => {
98
+ const nextParamIndex = paramIndex + values.length;
99
+ let paramOperator;
100
+ let paramKey;
101
+ let paramValue;
102
+ if (field.startsWith("$")) {
103
+ paramOperator = field;
104
+ paramKey = "";
105
+ paramValue = val;
106
+ } else if (typeof val === "object" && !Array.isArray(val)) {
107
+ const [op, opValue] = Object.entries(val || {})[0] || [];
108
+ paramOperator = op;
109
+ paramKey = field;
110
+ paramValue = opValue;
111
+ } else {
112
+ paramOperator = "$eq";
113
+ paramKey = field;
114
+ paramValue = val;
115
+ }
116
+ const operatorFn = FILTER_OPERATORS[paramOperator];
117
+ if (!operatorFn) {
118
+ throw new Error(`Invalid operator: ${paramOperator}`);
119
+ }
120
+ const result = operatorFn(paramKey, nextParamIndex, paramValue);
121
+ const sql = result.sql.replaceAll("metadata#>>", "elem#>>");
122
+ conditions.push(sql);
123
+ if (result.needsValue) {
124
+ values.push(paramValue);
125
+ }
126
+ });
127
+ return {
128
+ sql: conditions.join(" AND "),
129
+ values
130
+ };
131
+ }
132
+ var FILTER_OPERATORS = {
133
+ $eq: createBasicOperator("="),
134
+ $ne: createBasicOperator("!="),
135
+ $gt: createNumericOperator(">"),
136
+ $gte: createNumericOperator(">="),
137
+ $lt: createNumericOperator("<"),
138
+ $lte: createNumericOperator("<="),
139
+ // Array Operators
140
+ $in: (key, paramIndex) => ({
141
+ sql: `metadata#>>'{${handleKey(key)}}' = ANY($${paramIndex}::text[])`,
142
+ needsValue: true
143
+ }),
144
+ $nin: (key, paramIndex) => ({
145
+ sql: `metadata#>>'{${handleKey(key)}}' != ALL($${paramIndex}::text[])`,
146
+ needsValue: true
147
+ }),
148
+ $all: (key, paramIndex) => ({
149
+ sql: `CASE WHEN array_length($${paramIndex}::text[], 1) IS NULL THEN false
150
+ ELSE (metadata#>'{${handleKey(key)}}')::jsonb ?& $${paramIndex}::text[] END`,
151
+ needsValue: true
152
+ }),
153
+ $elemMatch: (key, paramIndex, value) => {
154
+ const { sql, values } = buildElemMatchConditions(value, paramIndex);
155
+ return {
156
+ sql: `(
157
+ CASE
158
+ WHEN jsonb_typeof(metadata->'${handleKey(key)}') = 'array' THEN
159
+ EXISTS (
160
+ SELECT 1
161
+ FROM jsonb_array_elements(metadata->'${handleKey(key)}') as elem
162
+ WHERE ${sql}
163
+ )
164
+ ELSE FALSE
165
+ END
166
+ )`,
167
+ needsValue: true,
168
+ transformValue: () => values
169
+ };
170
+ },
171
+ // Element Operators
172
+ $exists: (key) => ({
173
+ sql: `metadata ? '${key}'`,
174
+ needsValue: false
175
+ }),
176
+ // Logical Operators
177
+ $and: (key) => ({ sql: `(${key})`, needsValue: false }),
178
+ $or: (key) => ({ sql: `(${key})`, needsValue: false }),
179
+ $not: (key) => ({ sql: `NOT (${key})`, needsValue: false }),
180
+ $nor: (key) => ({ sql: `NOT (${key})`, needsValue: false }),
181
+ // Regex Operators
182
+ $regex: (key, paramIndex) => ({
183
+ sql: `metadata#>>'{${handleKey(key)}}' ~ $${paramIndex}`,
184
+ needsValue: true
185
+ }),
186
+ $contains: (key, paramIndex) => ({
187
+ sql: `metadata @> $${paramIndex}::jsonb`,
188
+ needsValue: true,
189
+ transformValue: (value) => {
190
+ const parts = key.split(".");
191
+ return JSON.stringify(parts.reduceRight((value2, key2) => ({ [key2]: value2 }), value));
192
+ }
193
+ }),
194
+ $size: (key, paramIndex) => ({
195
+ sql: `(
196
+ CASE
197
+ WHEN jsonb_typeof(metadata#>'{${handleKey(key)}}') = 'array' THEN
198
+ jsonb_array_length(metadata#>'{${handleKey(key)}}') = $${paramIndex}
199
+ ELSE FALSE
200
+ END
201
+ )`,
202
+ needsValue: true
203
+ })
204
+ };
205
+ var handleKey = (key) => {
206
+ return key.replace(/\./g, ",");
207
+ };
208
+ function buildFilterQuery(filter, minScore) {
209
+ const values = [minScore];
210
+ function buildCondition(key, value, parentPath) {
211
+ if (["$and", "$or", "$not", "$nor"].includes(key)) {
212
+ return handleLogicalOperator(key, value);
213
+ }
214
+ if (!value || typeof value !== "object") {
215
+ values.push(value);
216
+ return `metadata#>>'{${handleKey(key)}}' = $${values.length}`;
217
+ }
218
+ const [[operator, operatorValue] = []] = Object.entries(value);
219
+ if (operator === "$not") {
220
+ const entries = Object.entries(operatorValue);
221
+ const conditions2 = entries.map(([nestedOp, nestedValue]) => {
222
+ if (!FILTER_OPERATORS[nestedOp]) {
223
+ throw new Error(`Invalid operator in $not condition: ${nestedOp}`);
224
+ }
225
+ const operatorFn2 = FILTER_OPERATORS[nestedOp];
226
+ const operatorResult2 = operatorFn2(key, values.length + 1);
227
+ if (operatorResult2.needsValue) {
228
+ values.push(nestedValue);
229
+ }
230
+ return operatorResult2.sql;
231
+ }).join(" AND ");
232
+ return `NOT (${conditions2})`;
233
+ }
234
+ const operatorFn = FILTER_OPERATORS[operator];
235
+ const operatorResult = operatorFn(key, values.length + 1, operatorValue);
236
+ if (operatorResult.needsValue) {
237
+ const transformedValue = operatorResult.transformValue ? operatorResult.transformValue(operatorValue) : operatorValue;
238
+ if (Array.isArray(transformedValue) && operator === "$elemMatch") {
239
+ values.push(...transformedValue);
240
+ } else {
241
+ values.push(transformedValue);
242
+ }
243
+ }
244
+ return operatorResult.sql;
245
+ }
246
+ function handleLogicalOperator(key, value, parentPath) {
247
+ if (key === "$not") {
248
+ const entries = Object.entries(value);
249
+ const conditions3 = entries.map(([fieldKey, fieldValue]) => buildCondition(fieldKey, fieldValue)).join(" AND ");
250
+ return `NOT (${conditions3})`;
251
+ }
252
+ if (!value || value.length === 0) {
253
+ switch (key) {
254
+ case "$and":
255
+ case "$nor":
256
+ return "true";
257
+ // Empty $and/$nor match everything
258
+ case "$or":
259
+ return "false";
260
+ // Empty $or matches nothing
261
+ default:
262
+ return "true";
263
+ }
264
+ }
265
+ const joinOperator = key === "$or" || key === "$nor" ? "OR" : "AND";
266
+ const conditions2 = value.map((f) => {
267
+ const entries = Object.entries(f);
268
+ if (entries.length === 0) return "";
269
+ const [firstKey, firstValue] = entries[0] || [];
270
+ if (["$and", "$or", "$not", "$nor"].includes(firstKey)) {
271
+ return buildCondition(firstKey, firstValue);
272
+ }
273
+ return entries.map(([k, v]) => buildCondition(k, v)).join(` ${joinOperator} `);
274
+ });
275
+ const joined = conditions2.join(` ${joinOperator} `);
276
+ const operatorFn = FILTER_OPERATORS[key];
277
+ return operatorFn(joined, 0, value).sql;
278
+ }
279
+ if (!filter) {
280
+ return { sql: "", values };
281
+ }
282
+ const conditions = Object.entries(filter).map(([key, value]) => buildCondition(key, value)).filter(Boolean).join(" AND ");
283
+ return { sql: conditions ? `WHERE ${conditions}` : "", values };
284
+ }
285
+
286
+ // src/vector/index.ts
287
+ var PgVector = class extends MastraVector {
288
+ constructor(connectionString) {
289
+ super();
290
+ const basePool = new pg.Pool({
291
+ connectionString,
292
+ max: 20,
293
+ // Maximum number of clients in the pool
294
+ idleTimeoutMillis: 3e4,
295
+ // Close idle connections after 30 seconds
296
+ connectionTimeoutMillis: 2e3
297
+ // Fail fast if can't connect
298
+ });
299
+ const telemetry = this.__getTelemetry();
300
+ this.pool = telemetry?.traceClass(basePool, {
301
+ spanNamePrefix: "pg-vector",
302
+ attributes: {
303
+ "vector.type": "postgres"
304
+ }
305
+ }) ?? basePool;
306
+ }
307
+ transformFilter(filter) {
308
+ const pgFilter = new PGFilterTranslator();
309
+ const translatedFilter = pgFilter.translate(filter ?? {});
310
+ return translatedFilter;
311
+ }
312
+ async query(indexName, queryVector, topK = 10, filter, includeVector = false, minScore = 0) {
313
+ const client = await this.pool.connect();
314
+ try {
315
+ const vectorStr = `[${queryVector.join(",")}]`;
316
+ const translatedFilter = this.transformFilter(filter);
317
+ const { sql: filterQuery, values: filterValues } = buildFilterQuery(translatedFilter, minScore);
318
+ const query = `
319
+ WITH vector_scores AS (
320
+ SELECT
321
+ vector_id as id,
322
+ 1 - (embedding <=> '${vectorStr}'::vector) as score,
323
+ metadata
324
+ ${includeVector ? ", embedding" : ""}
325
+ FROM ${indexName}
326
+ ${filterQuery}
327
+ )
328
+ SELECT *
329
+ FROM vector_scores
330
+ WHERE score > $1
331
+ ORDER BY score DESC
332
+ LIMIT ${topK}`;
333
+ const result = await client.query(query, filterValues);
334
+ return result.rows.map(({ id, score, metadata, embedding }) => ({
335
+ id,
336
+ score,
337
+ metadata,
338
+ ...includeVector && embedding && { vector: JSON.parse(embedding) }
339
+ }));
340
+ } finally {
341
+ client.release();
342
+ }
343
+ }
344
+ async upsert(indexName, vectors, metadata, ids) {
345
+ const client = await this.pool.connect();
346
+ try {
347
+ await client.query("BEGIN");
348
+ const vectorIds = ids || vectors.map(() => crypto.randomUUID());
349
+ for (let i = 0; i < vectors.length; i++) {
350
+ const query = `
351
+ INSERT INTO ${indexName} (vector_id, embedding, metadata)
352
+ VALUES ($1, $2::vector, $3::jsonb)
353
+ ON CONFLICT (vector_id)
354
+ DO UPDATE SET
355
+ embedding = $2::vector,
356
+ metadata = $3::jsonb
357
+ RETURNING embedding::text
358
+ `;
359
+ await client.query(query, [vectorIds[i], `[${vectors[i]?.join(",")}]`, JSON.stringify(metadata?.[i] || {})]);
360
+ }
361
+ await client.query("COMMIT");
362
+ return vectorIds;
363
+ } catch (error) {
364
+ await client.query("ROLLBACK");
365
+ throw error;
366
+ } finally {
367
+ client.release();
368
+ }
369
+ }
370
+ async createIndex(indexName, dimension, metric = "cosine") {
371
+ const client = await this.pool.connect();
372
+ try {
373
+ if (!indexName.match(/^[a-zA-Z_][a-zA-Z0-9_]*$/)) {
374
+ throw new Error("Invalid index name format");
375
+ }
376
+ if (!Number.isInteger(dimension) || dimension <= 0) {
377
+ throw new Error("Dimension must be a positive integer");
378
+ }
379
+ const extensionCheck = await client.query(`
380
+ SELECT EXISTS (
381
+ SELECT 1 FROM pg_available_extensions WHERE name = 'vector'
382
+ );
383
+ `);
384
+ if (!extensionCheck.rows[0].exists) {
385
+ throw new Error("PostgreSQL vector extension is not available. Please install it first.");
386
+ }
387
+ await client.query("CREATE EXTENSION IF NOT EXISTS vector");
388
+ await client.query(`
389
+ CREATE TABLE IF NOT EXISTS ${indexName} (
390
+ id SERIAL PRIMARY KEY,
391
+ vector_id TEXT UNIQUE NOT NULL,
392
+ embedding vector(${dimension}),
393
+ metadata JSONB DEFAULT '{}'::jsonb
394
+ );
395
+ `);
396
+ const indexMethod = metric === "cosine" ? "vector_cosine_ops" : metric === "euclidean" ? "vector_l2_ops" : "vector_ip_ops";
397
+ await client.query(`
398
+ CREATE INDEX IF NOT EXISTS ${indexName}_vector_idx
399
+ ON public.${indexName}
400
+ USING ivfflat (embedding ${indexMethod})
401
+ WITH (lists = 100);
402
+ `);
403
+ } catch (error) {
404
+ console.error("Failed to create vector table:", error);
405
+ throw error;
406
+ } finally {
407
+ client.release();
408
+ }
409
+ }
410
+ async listIndexes() {
411
+ const client = await this.pool.connect();
412
+ try {
413
+ const vectorTablesQuery = `
414
+ SELECT DISTINCT table_name
415
+ FROM information_schema.columns
416
+ WHERE table_schema = 'public'
417
+ AND udt_name = 'vector';
418
+ `;
419
+ const vectorTables = await client.query(vectorTablesQuery);
420
+ return vectorTables.rows.map((row) => row.table_name);
421
+ } finally {
422
+ client.release();
423
+ }
424
+ }
425
+ async describeIndex(indexName) {
426
+ const client = await this.pool.connect();
427
+ try {
428
+ const dimensionQuery = `
429
+ SELECT atttypmod as dimension
430
+ FROM pg_attribute
431
+ WHERE attrelid = $1::regclass
432
+ AND attname = 'embedding';
433
+ `;
434
+ const countQuery = `
435
+ SELECT COUNT(*) as count
436
+ FROM ${indexName};
437
+ `;
438
+ const metricQuery = `
439
+ SELECT
440
+ am.amname as index_method,
441
+ opclass.opcname as operator_class
442
+ FROM pg_index i
443
+ JOIN pg_class c ON i.indexrelid = c.oid
444
+ JOIN pg_am am ON c.relam = am.oid
445
+ JOIN pg_opclass opclass ON i.indclass[0] = opclass.oid
446
+ WHERE c.relname = '${indexName}_vector_idx';
447
+ `;
448
+ const [dimResult, countResult, metricResult] = await Promise.all([
449
+ client.query(dimensionQuery, [indexName]),
450
+ client.query(countQuery),
451
+ client.query(metricQuery)
452
+ ]);
453
+ let metric = "cosine";
454
+ if (metricResult.rows.length > 0) {
455
+ const operatorClass = metricResult.rows[0].operator_class;
456
+ if (operatorClass.includes("l2")) {
457
+ metric = "euclidean";
458
+ } else if (operatorClass.includes("ip")) {
459
+ metric = "dotproduct";
460
+ } else if (operatorClass.includes("cosine")) {
461
+ metric = "cosine";
462
+ }
463
+ }
464
+ return {
465
+ dimension: dimResult.rows[0].dimension,
466
+ count: parseInt(countResult.rows[0].count),
467
+ metric
468
+ };
469
+ } catch (e) {
470
+ await client.query("ROLLBACK");
471
+ throw new Error(`Failed to describe vector table: ${e.message}`);
472
+ } finally {
473
+ client.release();
474
+ }
475
+ }
476
+ async deleteIndex(indexName) {
477
+ const client = await this.pool.connect();
478
+ try {
479
+ await client.query(`DROP TABLE IF EXISTS ${indexName} CASCADE`);
480
+ } catch (error) {
481
+ await client.query("ROLLBACK");
482
+ throw new Error(`Failed to delete vector table: ${error.message}`);
483
+ } finally {
484
+ client.release();
485
+ }
486
+ }
487
+ async truncateIndex(indexName) {
488
+ const client = await this.pool.connect();
489
+ try {
490
+ await client.query(`TRUNCATE ${indexName}`);
491
+ } catch (e) {
492
+ await client.query("ROLLBACK");
493
+ throw new Error(`Failed to truncate vector table: ${e.message}`);
494
+ } finally {
495
+ client.release();
496
+ }
497
+ }
498
+ async disconnect() {
499
+ await this.pool.end();
500
+ }
501
+ };
502
+ var PostgresStore = class extends MastraStorage {
503
+ constructor(config) {
504
+ super({ name: "PostgresStore" });
505
+ this.pgp = pgPromise();
506
+ this.db = this.pgp(
507
+ `connectionString` in config ? { connectionString: config.connectionString } : {
508
+ host: config.host,
509
+ port: config.port,
510
+ database: config.database,
511
+ user: config.user,
512
+ password: config.password
513
+ }
514
+ );
515
+ }
516
+ async createTable({
517
+ tableName,
518
+ schema
519
+ }) {
520
+ try {
521
+ const columns = Object.entries(schema).map(([name, def]) => {
522
+ const constraints = [];
523
+ if (def.primaryKey) constraints.push("PRIMARY KEY");
524
+ if (!def.nullable) constraints.push("NOT NULL");
525
+ return `"${name}" ${def.type.toUpperCase()} ${constraints.join(" ")}`;
526
+ }).join(",\n");
527
+ const sql = `
528
+ CREATE TABLE IF NOT EXISTS ${tableName} (
529
+ ${columns}
530
+ );
531
+ ${tableName === MastraStorage.TABLE_WORKFLOW_SNAPSHOT ? `
532
+ DO $$ BEGIN
533
+ IF NOT EXISTS (
534
+ SELECT 1 FROM pg_constraint WHERE conname = 'mastra_workflow_snapshot_workflow_name_run_id_key'
535
+ ) THEN
536
+ ALTER TABLE ${tableName}
537
+ ADD CONSTRAINT mastra_workflow_snapshot_workflow_name_run_id_key
538
+ UNIQUE (workflow_name, run_id);
539
+ END IF;
540
+ END $$;
541
+ ` : ""}
542
+ `;
543
+ await this.db.none(sql);
544
+ } catch (error) {
545
+ console.error(`Error creating table ${tableName}:`, error);
546
+ throw error;
547
+ }
548
+ }
549
+ async clearTable({ tableName }) {
550
+ try {
551
+ await this.db.none(`TRUNCATE TABLE ${tableName} CASCADE`);
552
+ } catch (error) {
553
+ console.error(`Error clearing table ${tableName}:`, error);
554
+ throw error;
555
+ }
556
+ }
557
+ async insert({ tableName, record }) {
558
+ try {
559
+ const columns = Object.keys(record);
560
+ const values = Object.values(record);
561
+ const placeholders = values.map((_, i) => `$${i + 1}`).join(", ");
562
+ await this.db.none(
563
+ `INSERT INTO ${tableName} (${columns.map((c) => `"${c}"`).join(", ")}) VALUES (${placeholders})`,
564
+ values
565
+ );
566
+ } catch (error) {
567
+ console.error(`Error inserting into ${tableName}:`, error);
568
+ throw error;
569
+ }
570
+ }
571
+ async load({ tableName, keys }) {
572
+ try {
573
+ const keyEntries = Object.entries(keys);
574
+ const conditions = keyEntries.map(([key], index) => `"${key}" = $${index + 1}`).join(" AND ");
575
+ const values = keyEntries.map(([_, value]) => value);
576
+ const result = await this.db.oneOrNone(`SELECT * FROM ${tableName} WHERE ${conditions}`, values);
577
+ if (!result) {
578
+ return null;
579
+ }
580
+ if (tableName === MastraStorage.TABLE_WORKFLOW_SNAPSHOT) {
581
+ const snapshot = result;
582
+ if (typeof snapshot.snapshot === "string") {
583
+ snapshot.snapshot = JSON.parse(snapshot.snapshot);
584
+ }
585
+ return snapshot;
586
+ }
587
+ return result;
588
+ } catch (error) {
589
+ console.error(`Error loading from ${tableName}:`, error);
590
+ throw error;
591
+ }
592
+ }
593
+ async getThreadById({ threadId }) {
594
+ try {
595
+ const thread = await this.db.oneOrNone(
596
+ `SELECT
597
+ id,
598
+ "resourceId",
599
+ title,
600
+ metadata,
601
+ "createdAt",
602
+ "updatedAt"
603
+ FROM "${MastraStorage.TABLE_THREADS}"
604
+ WHERE id = $1`,
605
+ [threadId]
606
+ );
607
+ if (!thread) {
608
+ return null;
609
+ }
610
+ return {
611
+ ...thread,
612
+ metadata: typeof thread.metadata === "string" ? JSON.parse(thread.metadata) : thread.metadata,
613
+ createdAt: thread.createdAt,
614
+ updatedAt: thread.updatedAt
615
+ };
616
+ } catch (error) {
617
+ console.error(`Error getting thread ${threadId}:`, error);
618
+ throw error;
619
+ }
620
+ }
621
+ async getThreadsByResourceId({ resourceId }) {
622
+ try {
623
+ const threads = await this.db.manyOrNone(
624
+ `SELECT
625
+ id,
626
+ "resourceId",
627
+ title,
628
+ metadata,
629
+ "createdAt",
630
+ "updatedAt"
631
+ FROM "${MastraStorage.TABLE_THREADS}"
632
+ WHERE "resourceId" = $1`,
633
+ [resourceId]
634
+ );
635
+ return threads.map((thread) => ({
636
+ ...thread,
637
+ metadata: typeof thread.metadata === "string" ? JSON.parse(thread.metadata) : thread.metadata,
638
+ createdAt: thread.createdAt,
639
+ updatedAt: thread.updatedAt
640
+ }));
641
+ } catch (error) {
642
+ console.error(`Error getting threads for resource ${resourceId}:`, error);
643
+ throw error;
644
+ }
645
+ }
646
+ async saveThread({ thread }) {
647
+ try {
648
+ await this.db.none(
649
+ `INSERT INTO "${MastraStorage.TABLE_THREADS}" (
650
+ id,
651
+ "resourceId",
652
+ title,
653
+ metadata,
654
+ "createdAt",
655
+ "updatedAt"
656
+ ) VALUES ($1, $2, $3, $4, $5, $6)
657
+ ON CONFLICT (id) DO UPDATE SET
658
+ "resourceId" = EXCLUDED."resourceId",
659
+ title = EXCLUDED.title,
660
+ metadata = EXCLUDED.metadata,
661
+ "createdAt" = EXCLUDED."createdAt",
662
+ "updatedAt" = EXCLUDED."updatedAt"`,
663
+ [
664
+ thread.id,
665
+ thread.resourceId,
666
+ thread.title,
667
+ thread.metadata ? JSON.stringify(thread.metadata) : null,
668
+ thread.createdAt,
669
+ thread.updatedAt
670
+ ]
671
+ );
672
+ return thread;
673
+ } catch (error) {
674
+ console.error("Error saving thread:", error);
675
+ throw error;
676
+ }
677
+ }
678
+ async updateThread({
679
+ id,
680
+ title,
681
+ metadata
682
+ }) {
683
+ try {
684
+ const existingThread = await this.getThreadById({ threadId: id });
685
+ if (!existingThread) {
686
+ throw new Error(`Thread ${id} not found`);
687
+ }
688
+ const mergedMetadata = {
689
+ ...existingThread.metadata,
690
+ ...metadata
691
+ };
692
+ const thread = await this.db.one(
693
+ `UPDATE "${MastraStorage.TABLE_THREADS}"
694
+ SET title = $1,
695
+ metadata = $2,
696
+ "updatedAt" = $3
697
+ WHERE id = $4
698
+ RETURNING *`,
699
+ [title, mergedMetadata, (/* @__PURE__ */ new Date()).toISOString(), id]
700
+ );
701
+ return {
702
+ ...thread,
703
+ metadata: typeof thread.metadata === "string" ? JSON.parse(thread.metadata) : thread.metadata,
704
+ createdAt: thread.createdAt,
705
+ updatedAt: thread.updatedAt
706
+ };
707
+ } catch (error) {
708
+ console.error("Error updating thread:", error);
709
+ throw error;
710
+ }
711
+ }
712
+ async deleteThread({ threadId }) {
713
+ try {
714
+ await this.db.tx(async (t) => {
715
+ await t.none(`DELETE FROM "${MastraStorage.TABLE_MESSAGES}" WHERE thread_id = $1`, [threadId]);
716
+ await t.none(`DELETE FROM "${MastraStorage.TABLE_THREADS}" WHERE id = $1`, [threadId]);
717
+ });
718
+ } catch (error) {
719
+ console.error("Error deleting thread:", error);
720
+ throw error;
721
+ }
722
+ }
723
+ async getMessages({ threadId, selectBy }) {
724
+ try {
725
+ const messages = [];
726
+ const limit = typeof selectBy?.last === `number` ? selectBy.last : 40;
727
+ const include = selectBy?.include || [];
728
+ if (include.length) {
729
+ const includeResult = await this.db.manyOrNone(
730
+ `
731
+ WITH ordered_messages AS (
732
+ SELECT
733
+ *,
734
+ ROW_NUMBER() OVER (ORDER BY "createdAt") as row_num
735
+ FROM "${MastraStorage.TABLE_MESSAGES}"
736
+ WHERE thread_id = $1
737
+ )
738
+ SELECT DISTINCT ON (m.id)
739
+ m.id,
740
+ m.content,
741
+ m.role,
742
+ m.type,
743
+ m."createdAt",
744
+ m.thread_id AS "threadId"
745
+ FROM ordered_messages m
746
+ WHERE m.id = ANY($2)
747
+ OR EXISTS (
748
+ SELECT 1 FROM ordered_messages target
749
+ WHERE target.id = ANY($2)
750
+ AND (
751
+ -- Get previous messages based on the max withPreviousMessages
752
+ (m.row_num >= target.row_num - $3 AND m.row_num < target.row_num)
753
+ OR
754
+ -- Get next messages based on the max withNextMessages
755
+ (m.row_num <= target.row_num + $4 AND m.row_num > target.row_num)
756
+ )
757
+ )
758
+ ORDER BY m.id, m."createdAt"
759
+ `,
760
+ [
761
+ threadId,
762
+ include.map((i) => i.id),
763
+ Math.max(...include.map((i) => i.withPreviousMessages || 0)),
764
+ Math.max(...include.map((i) => i.withNextMessages || 0))
765
+ ]
766
+ );
767
+ messages.push(...includeResult);
768
+ }
769
+ const result = await this.db.manyOrNone(
770
+ `
771
+ SELECT
772
+ id,
773
+ content,
774
+ role,
775
+ type,
776
+ "createdAt",
777
+ thread_id AS "threadId"
778
+ FROM "${MastraStorage.TABLE_MESSAGES}"
779
+ WHERE thread_id = $1
780
+ AND id != ALL($2)
781
+ ORDER BY "createdAt" DESC
782
+ LIMIT $3
783
+ `,
784
+ [threadId, messages.map((m) => m.id), limit]
785
+ );
786
+ messages.push(...result);
787
+ messages.sort((a, b) => new Date(a.createdAt).getTime() - new Date(b.createdAt).getTime());
788
+ messages.forEach((message) => {
789
+ if (typeof message.content === "string") {
790
+ try {
791
+ message.content = JSON.parse(message.content);
792
+ } catch (e) {
793
+ }
794
+ }
795
+ });
796
+ return messages;
797
+ } catch (error) {
798
+ console.error("Error getting messages:", error);
799
+ throw error;
800
+ }
801
+ }
802
+ async saveMessages({ messages }) {
803
+ if (messages.length === 0) return messages;
804
+ try {
805
+ const threadId = messages[0]?.threadId;
806
+ if (!threadId) {
807
+ throw new Error("Thread ID is required");
808
+ }
809
+ const thread = await this.getThreadById({ threadId });
810
+ if (!thread) {
811
+ throw new Error(`Thread ${threadId} not found`);
812
+ }
813
+ await this.db.tx(async (t) => {
814
+ for (const message of messages) {
815
+ await t.none(
816
+ `INSERT INTO "${MastraStorage.TABLE_MESSAGES}" (id, thread_id, content, "createdAt", role, type)
817
+ VALUES ($1, $2, $3, $4, $5, $6)`,
818
+ [
819
+ message.id,
820
+ threadId,
821
+ typeof message.content === "string" ? message.content : JSON.stringify(message.content),
822
+ message.createdAt || (/* @__PURE__ */ new Date()).toISOString(),
823
+ message.role,
824
+ message.type
825
+ ]
826
+ );
827
+ }
828
+ });
829
+ return messages;
830
+ } catch (error) {
831
+ console.error("Error saving messages:", error);
832
+ throw error;
833
+ }
834
+ }
835
+ async persistWorkflowSnapshot({
836
+ workflowName,
837
+ runId,
838
+ snapshot
839
+ }) {
840
+ try {
841
+ const now = (/* @__PURE__ */ new Date()).toISOString();
842
+ await this.db.none(
843
+ `INSERT INTO "${MastraStorage.TABLE_WORKFLOW_SNAPSHOT}" (
844
+ workflow_name,
845
+ run_id,
846
+ snapshot,
847
+ "createdAt",
848
+ "updatedAt"
849
+ ) VALUES ($1, $2, $3, $4, $5)
850
+ ON CONFLICT (workflow_name, run_id) DO UPDATE
851
+ SET snapshot = EXCLUDED.snapshot,
852
+ "updatedAt" = EXCLUDED."updatedAt"`,
853
+ [workflowName, runId, JSON.stringify(snapshot), now, now]
854
+ );
855
+ } catch (error) {
856
+ console.error("Error persisting workflow snapshot:", error);
857
+ throw error;
858
+ }
859
+ }
860
+ async loadWorkflowSnapshot({
861
+ workflowName,
862
+ runId
863
+ }) {
864
+ try {
865
+ const result = await this.load({
866
+ tableName: MastraStorage.TABLE_WORKFLOW_SNAPSHOT,
867
+ keys: {
868
+ workflow_name: workflowName,
869
+ run_id: runId
870
+ }
871
+ });
872
+ if (!result) {
873
+ return null;
874
+ }
875
+ return result.snapshot;
876
+ } catch (error) {
877
+ console.error("Error loading workflow snapshot:", error);
878
+ throw error;
879
+ }
880
+ }
881
+ async close() {
882
+ this.pgp.end();
883
+ }
884
+ };
885
+
886
+ export { PgVector, PostgresStore };