@mastra/dsql 0.0.0 → 1.0.0-alpha.3
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/CHANGELOG.md +27 -0
- package/LICENSE.md +30 -0
- package/dist/index.cjs +4476 -0
- package/dist/index.cjs.map +1 -0
- package/dist/index.d.ts +4 -0
- package/dist/index.d.ts.map +1 -0
- package/dist/index.js +4467 -0
- package/dist/index.js.map +1 -0
- package/dist/shared/batch.d.ts +55 -0
- package/dist/shared/batch.d.ts.map +1 -0
- package/dist/shared/config.d.ts +129 -0
- package/dist/shared/config.d.ts.map +1 -0
- package/dist/shared/retry.d.ts +207 -0
- package/dist/shared/retry.d.ts.map +1 -0
- package/dist/storage/client.d.ts +91 -0
- package/dist/storage/client.d.ts.map +1 -0
- package/dist/storage/db/index.d.ts +179 -0
- package/dist/storage/db/index.d.ts.map +1 -0
- package/dist/storage/domains/agents/index.d.ts +47 -0
- package/dist/storage/domains/agents/index.d.ts.map +1 -0
- package/dist/storage/domains/memory/index.d.ts +80 -0
- package/dist/storage/domains/memory/index.d.ts.map +1 -0
- package/dist/storage/domains/observability/index.d.ts +37 -0
- package/dist/storage/domains/observability/index.d.ts.map +1 -0
- package/dist/storage/domains/scores/index.d.ts +53 -0
- package/dist/storage/domains/scores/index.d.ts.map +1 -0
- package/dist/storage/domains/utils.d.ts +14 -0
- package/dist/storage/domains/utils.d.ts.map +1 -0
- package/dist/storage/domains/workflows/index.d.ts +61 -0
- package/dist/storage/domains/workflows/index.d.ts.map +1 -0
- package/dist/storage/index.d.ts +46 -0
- package/dist/storage/index.d.ts.map +1 -0
- package/dist/storage/performance-indexes/performance-test.d.ts +47 -0
- package/dist/storage/performance-indexes/performance-test.d.ts.map +1 -0
- package/dist/storage/test-utils.d.ts +14 -0
- package/dist/storage/test-utils.d.ts.map +1 -0
- package/package.json +17 -17
package/dist/index.js
ADDED
|
@@ -0,0 +1,4467 @@
|
|
|
1
|
+
import { AuroraDSQLClient } from '@aws/aurora-dsql-node-postgres-connector';
|
|
2
|
+
import { MastraError, ErrorCategory, ErrorDomain } from '@mastra/core/error';
|
|
3
|
+
import { AgentsStorage, TABLE_AGENTS, TABLE_AGENT_VERSIONS, TABLE_SCHEMAS, createStorageErrorId, normalizePerPage, calculatePagination, MemoryStorage, TABLE_THREADS, TABLE_MESSAGES, TABLE_RESOURCES, ObservabilityStorage, TABLE_SPANS, listTracesArgsSchema, toTraceSpans, ScoresStorage, TABLE_SCORERS, WorkflowsStorage, TABLE_WORKFLOW_SNAPSHOT, MastraStorage, TraceStatus, getSqlType, getDefaultValue, transformScoreRow as transformScoreRow$1 } from '@mastra/core/storage';
|
|
4
|
+
import { Pool } from 'pg';
|
|
5
|
+
import { MastraBase } from '@mastra/core/base';
|
|
6
|
+
import { parseSqlIdentifier } from '@mastra/core/utils';
|
|
7
|
+
import { MessageList } from '@mastra/core/agent';
|
|
8
|
+
import { saveScorePayloadSchema } from '@mastra/core/evals';
|
|
9
|
+
|
|
10
|
+
// src/storage/index.ts
|
|
11
|
+
|
|
12
|
+
// src/shared/config.ts
|
|
13
|
+
var DSQL_POOL_DEFAULTS = {
|
|
14
|
+
/** Maximum connections in the pool */
|
|
15
|
+
max: 10,
|
|
16
|
+
/** Minimum connections in the pool */
|
|
17
|
+
min: 0,
|
|
18
|
+
/** Close idle connections after 10 minutes */
|
|
19
|
+
idleTimeoutMillis: 6e5,
|
|
20
|
+
/** Force connection rotation before DSQL's 60-minute limit */
|
|
21
|
+
maxLifetimeSeconds: 3300,
|
|
22
|
+
/** Connection acquisition timeout */
|
|
23
|
+
connectionTimeoutMillis: 5e3,
|
|
24
|
+
/** Allow process to exit when idle */
|
|
25
|
+
allowExitOnIdle: true
|
|
26
|
+
};
|
|
27
|
+
var isPoolConfig = (cfg) => {
|
|
28
|
+
return "pool" in cfg;
|
|
29
|
+
};
|
|
30
|
+
var isHostConfig = (cfg) => {
|
|
31
|
+
return "host" in cfg;
|
|
32
|
+
};
|
|
33
|
+
var validateConfig = (config) => {
|
|
34
|
+
if (!config.id || typeof config.id !== "string" || config.id.trim() === "") {
|
|
35
|
+
throw new Error("DSQLStore: id must be provided and cannot be empty.");
|
|
36
|
+
}
|
|
37
|
+
if (isPoolConfig(config)) {
|
|
38
|
+
if (!config.pool) {
|
|
39
|
+
throw new Error("DSQLStore: pool must be provided when using pool config.");
|
|
40
|
+
}
|
|
41
|
+
return;
|
|
42
|
+
}
|
|
43
|
+
if (!isHostConfig(config)) {
|
|
44
|
+
throw new Error("DSQLStore: host must be provided and cannot be empty.");
|
|
45
|
+
}
|
|
46
|
+
if (!config.host || config.host.trim() === "") {
|
|
47
|
+
throw new Error("DSQLStore: host must be provided and cannot be empty.");
|
|
48
|
+
}
|
|
49
|
+
if (config.maxLifetimeSeconds !== void 0 && config.maxLifetimeSeconds >= 3600) {
|
|
50
|
+
throw new Error(
|
|
51
|
+
"DSQLStore: maxLifetimeSeconds must be less than 3600 (60 minutes) due to Aurora DSQL connection duration limit."
|
|
52
|
+
);
|
|
53
|
+
}
|
|
54
|
+
};
|
|
55
|
+
var extractRegionFromHost = (host) => {
|
|
56
|
+
const match = host.match(/\.dsql\.([a-z0-9-]+)\.on\.aws$/);
|
|
57
|
+
return match?.[1];
|
|
58
|
+
};
|
|
59
|
+
var getEffectiveRegion = (config) => {
|
|
60
|
+
if (config.region) {
|
|
61
|
+
return config.region;
|
|
62
|
+
}
|
|
63
|
+
const extractedRegion = extractRegionFromHost(config.host);
|
|
64
|
+
if (extractedRegion) {
|
|
65
|
+
return extractedRegion;
|
|
66
|
+
}
|
|
67
|
+
throw new Error(
|
|
68
|
+
"DSQLStore: region could not be determined. Provide region in config or use a standard DSQL endpoint."
|
|
69
|
+
);
|
|
70
|
+
};
|
|
71
|
+
|
|
72
|
+
// src/storage/client.ts
|
|
73
|
+
function truncateQuery(query, maxLength = 100) {
|
|
74
|
+
const normalized = query.replace(/\s+/g, " ").trim();
|
|
75
|
+
if (normalized.length <= maxLength) {
|
|
76
|
+
return normalized;
|
|
77
|
+
}
|
|
78
|
+
return normalized.slice(0, maxLength) + "...";
|
|
79
|
+
}
|
|
80
|
+
var PoolAdapter = class {
|
|
81
|
+
constructor($pool) {
|
|
82
|
+
this.$pool = $pool;
|
|
83
|
+
}
|
|
84
|
+
connect() {
|
|
85
|
+
return this.$pool.connect();
|
|
86
|
+
}
|
|
87
|
+
async none(query, values) {
|
|
88
|
+
await this.$pool.query(query, values);
|
|
89
|
+
return null;
|
|
90
|
+
}
|
|
91
|
+
async one(query, values) {
|
|
92
|
+
const result = await this.$pool.query(query, values);
|
|
93
|
+
if (result.rows.length === 0) {
|
|
94
|
+
throw new Error(`No data returned from query: ${truncateQuery(query)}`);
|
|
95
|
+
}
|
|
96
|
+
if (result.rows.length > 1) {
|
|
97
|
+
throw new Error(`Multiple rows returned when one was expected: ${truncateQuery(query)}`);
|
|
98
|
+
}
|
|
99
|
+
return result.rows[0];
|
|
100
|
+
}
|
|
101
|
+
async oneOrNone(query, values) {
|
|
102
|
+
const result = await this.$pool.query(query, values);
|
|
103
|
+
if (result.rows.length === 0) {
|
|
104
|
+
return null;
|
|
105
|
+
}
|
|
106
|
+
if (result.rows.length > 1) {
|
|
107
|
+
throw new Error(`Multiple rows returned when one or none was expected: ${truncateQuery(query)}`);
|
|
108
|
+
}
|
|
109
|
+
return result.rows[0];
|
|
110
|
+
}
|
|
111
|
+
async any(query, values) {
|
|
112
|
+
const result = await this.$pool.query(query, values);
|
|
113
|
+
return result.rows;
|
|
114
|
+
}
|
|
115
|
+
async manyOrNone(query, values) {
|
|
116
|
+
return this.any(query, values);
|
|
117
|
+
}
|
|
118
|
+
async many(query, values) {
|
|
119
|
+
const result = await this.$pool.query(query, values);
|
|
120
|
+
if (result.rows.length === 0) {
|
|
121
|
+
throw new Error(`No data returned from query: ${truncateQuery(query)}`);
|
|
122
|
+
}
|
|
123
|
+
return result.rows;
|
|
124
|
+
}
|
|
125
|
+
async query(query, values) {
|
|
126
|
+
return this.$pool.query(query, values);
|
|
127
|
+
}
|
|
128
|
+
async tx(callback) {
|
|
129
|
+
const client = await this.$pool.connect();
|
|
130
|
+
try {
|
|
131
|
+
await client.query("BEGIN");
|
|
132
|
+
const txClient = new TransactionClient(client);
|
|
133
|
+
const result = await callback(txClient);
|
|
134
|
+
await client.query("COMMIT");
|
|
135
|
+
return result;
|
|
136
|
+
} catch (error) {
|
|
137
|
+
try {
|
|
138
|
+
await client.query("ROLLBACK");
|
|
139
|
+
} catch (rollbackError) {
|
|
140
|
+
console.error("Transaction rollback failed:", rollbackError);
|
|
141
|
+
}
|
|
142
|
+
throw error;
|
|
143
|
+
} finally {
|
|
144
|
+
client.release();
|
|
145
|
+
}
|
|
146
|
+
}
|
|
147
|
+
};
|
|
148
|
+
var TransactionClient = class {
|
|
149
|
+
constructor(client) {
|
|
150
|
+
this.client = client;
|
|
151
|
+
}
|
|
152
|
+
async none(query, values) {
|
|
153
|
+
await this.client.query(query, values);
|
|
154
|
+
return null;
|
|
155
|
+
}
|
|
156
|
+
async one(query, values) {
|
|
157
|
+
const result = await this.client.query(query, values);
|
|
158
|
+
if (result.rows.length === 0) {
|
|
159
|
+
throw new Error(`No data returned from query: ${truncateQuery(query)}`);
|
|
160
|
+
}
|
|
161
|
+
if (result.rows.length > 1) {
|
|
162
|
+
throw new Error(`Multiple rows returned when one was expected: ${truncateQuery(query)}`);
|
|
163
|
+
}
|
|
164
|
+
return result.rows[0];
|
|
165
|
+
}
|
|
166
|
+
async oneOrNone(query, values) {
|
|
167
|
+
const result = await this.client.query(query, values);
|
|
168
|
+
if (result.rows.length === 0) {
|
|
169
|
+
return null;
|
|
170
|
+
}
|
|
171
|
+
if (result.rows.length > 1) {
|
|
172
|
+
throw new Error(`Multiple rows returned when one or none was expected: ${truncateQuery(query)}`);
|
|
173
|
+
}
|
|
174
|
+
return result.rows[0];
|
|
175
|
+
}
|
|
176
|
+
async any(query, values) {
|
|
177
|
+
const result = await this.client.query(query, values);
|
|
178
|
+
return result.rows;
|
|
179
|
+
}
|
|
180
|
+
async manyOrNone(query, values) {
|
|
181
|
+
return this.any(query, values);
|
|
182
|
+
}
|
|
183
|
+
async many(query, values) {
|
|
184
|
+
const result = await this.client.query(query, values);
|
|
185
|
+
if (result.rows.length === 0) {
|
|
186
|
+
throw new Error(`No data returned from query: ${truncateQuery(query)}`);
|
|
187
|
+
}
|
|
188
|
+
return result.rows;
|
|
189
|
+
}
|
|
190
|
+
async query(query, values) {
|
|
191
|
+
return this.client.query(query, values);
|
|
192
|
+
}
|
|
193
|
+
async batch(promises) {
|
|
194
|
+
return Promise.all(promises);
|
|
195
|
+
}
|
|
196
|
+
};
|
|
197
|
+
|
|
198
|
+
// src/shared/retry.ts
|
|
199
|
+
var RETRIABLE_ERROR_CODES = {
|
|
200
|
+
/** Serialization failure - OCC conflict in Aurora DSQL (default retriable) */
|
|
201
|
+
SERIALIZATION_FAILURE: "40001"};
|
|
202
|
+
var DEFAULT_RETRIABLE_SQLSTATES = /* @__PURE__ */ new Set([RETRIABLE_ERROR_CODES.SERIALIZATION_FAILURE]);
|
|
203
|
+
var DEFAULT_RETRY_OPTIONS = {
|
|
204
|
+
maxAttempts: 5,
|
|
205
|
+
initialDelayMs: 100,
|
|
206
|
+
maxDelayMs: 2e3,
|
|
207
|
+
backoffMultiplier: 2,
|
|
208
|
+
jitter: true
|
|
209
|
+
};
|
|
210
|
+
function isPostgresError(error) {
|
|
211
|
+
return error instanceof Error && "code" in error && typeof error.code === "string";
|
|
212
|
+
}
|
|
213
|
+
function getErrorCode(error) {
|
|
214
|
+
if (isPostgresError(error)) {
|
|
215
|
+
return error.code;
|
|
216
|
+
}
|
|
217
|
+
return void 0;
|
|
218
|
+
}
|
|
219
|
+
var SQLSTATE_PATTERN = /^[0-9A-Z]{5}$/;
|
|
220
|
+
function getPostgresSqlStateCode(error) {
|
|
221
|
+
const raw = getErrorCode(error);
|
|
222
|
+
if (!raw) return void 0;
|
|
223
|
+
const code = raw.toUpperCase();
|
|
224
|
+
if (SQLSTATE_PATTERN.test(code)) {
|
|
225
|
+
return code;
|
|
226
|
+
}
|
|
227
|
+
return void 0;
|
|
228
|
+
}
|
|
229
|
+
function isRetriableError(error) {
|
|
230
|
+
const sqlstate = getPostgresSqlStateCode(error);
|
|
231
|
+
if (!sqlstate) {
|
|
232
|
+
return false;
|
|
233
|
+
}
|
|
234
|
+
return DEFAULT_RETRIABLE_SQLSTATES.has(sqlstate);
|
|
235
|
+
}
|
|
236
|
+
function calculateRetryDelay(attempt, options = {}) {
|
|
237
|
+
const {
|
|
238
|
+
initialDelayMs = DEFAULT_RETRY_OPTIONS.initialDelayMs,
|
|
239
|
+
maxDelayMs = DEFAULT_RETRY_OPTIONS.maxDelayMs,
|
|
240
|
+
backoffMultiplier = DEFAULT_RETRY_OPTIONS.backoffMultiplier,
|
|
241
|
+
jitter = DEFAULT_RETRY_OPTIONS.jitter
|
|
242
|
+
} = options;
|
|
243
|
+
const baseDelay = initialDelayMs * Math.pow(backoffMultiplier, attempt - 1);
|
|
244
|
+
const cappedDelay = Math.min(baseDelay, maxDelayMs);
|
|
245
|
+
if (jitter) {
|
|
246
|
+
return Math.floor(Math.random() * (cappedDelay + 1));
|
|
247
|
+
}
|
|
248
|
+
return cappedDelay;
|
|
249
|
+
}
|
|
250
|
+
function sleep(ms) {
|
|
251
|
+
return new Promise((resolve) => setTimeout(resolve, ms));
|
|
252
|
+
}
|
|
253
|
+
function validateRetryOptions(options) {
|
|
254
|
+
const { maxAttempts, initialDelayMs, maxDelayMs, backoffMultiplier } = options;
|
|
255
|
+
if (maxAttempts < 1) {
|
|
256
|
+
throw new Error(`Invalid retry option: maxAttempts must be >= 1, got ${maxAttempts}`);
|
|
257
|
+
}
|
|
258
|
+
if (initialDelayMs < 0) {
|
|
259
|
+
throw new Error(`Invalid retry option: initialDelayMs must be >= 0, got ${initialDelayMs}`);
|
|
260
|
+
}
|
|
261
|
+
if (maxDelayMs <= 0) {
|
|
262
|
+
throw new Error(`Invalid retry option: maxDelayMs must be > 0, got ${maxDelayMs}`);
|
|
263
|
+
}
|
|
264
|
+
if (backoffMultiplier < 1) {
|
|
265
|
+
throw new Error(`Invalid retry option: backoffMultiplier must be >= 1, got ${backoffMultiplier}`);
|
|
266
|
+
}
|
|
267
|
+
if (maxDelayMs < initialDelayMs) {
|
|
268
|
+
throw new Error(`Invalid retry option: maxDelayMs (${maxDelayMs}) must be >= initialDelayMs (${initialDelayMs})`);
|
|
269
|
+
}
|
|
270
|
+
}
|
|
271
|
+
async function withRetry(fn, options = {}) {
|
|
272
|
+
const {
|
|
273
|
+
maxAttempts = DEFAULT_RETRY_OPTIONS.maxAttempts,
|
|
274
|
+
initialDelayMs = DEFAULT_RETRY_OPTIONS.initialDelayMs,
|
|
275
|
+
maxDelayMs = DEFAULT_RETRY_OPTIONS.maxDelayMs,
|
|
276
|
+
backoffMultiplier = DEFAULT_RETRY_OPTIONS.backoffMultiplier,
|
|
277
|
+
jitter = DEFAULT_RETRY_OPTIONS.jitter,
|
|
278
|
+
onRetry,
|
|
279
|
+
isRetriable = isRetriableError
|
|
280
|
+
} = options;
|
|
281
|
+
validateRetryOptions({ maxAttempts, initialDelayMs, maxDelayMs, backoffMultiplier });
|
|
282
|
+
const startTime = Date.now();
|
|
283
|
+
let lastError;
|
|
284
|
+
for (let attempt = 1; attempt <= maxAttempts; attempt++) {
|
|
285
|
+
try {
|
|
286
|
+
const result = await fn();
|
|
287
|
+
return {
|
|
288
|
+
result,
|
|
289
|
+
attempts: attempt,
|
|
290
|
+
totalTimeMs: Date.now() - startTime
|
|
291
|
+
};
|
|
292
|
+
} catch (error) {
|
|
293
|
+
lastError = error instanceof Error ? error : new Error(String(error));
|
|
294
|
+
if (attempt === maxAttempts || !isRetriable(error)) {
|
|
295
|
+
throw lastError;
|
|
296
|
+
}
|
|
297
|
+
const delay = calculateRetryDelay(attempt, {
|
|
298
|
+
initialDelayMs,
|
|
299
|
+
maxDelayMs,
|
|
300
|
+
backoffMultiplier,
|
|
301
|
+
jitter
|
|
302
|
+
});
|
|
303
|
+
if (onRetry) {
|
|
304
|
+
onRetry(lastError, attempt, delay);
|
|
305
|
+
}
|
|
306
|
+
await sleep(delay);
|
|
307
|
+
}
|
|
308
|
+
}
|
|
309
|
+
throw lastError;
|
|
310
|
+
}
|
|
311
|
+
|
|
312
|
+
// src/shared/batch.ts
|
|
313
|
+
var DEFAULT_MAX_ROWS_PER_BATCH = 3e3;
|
|
314
|
+
function splitIntoBatches(records, options = {}) {
|
|
315
|
+
const maxRows = options.maxRows ?? DEFAULT_MAX_ROWS_PER_BATCH;
|
|
316
|
+
if (records.length === 0) {
|
|
317
|
+
return {
|
|
318
|
+
batches: [],
|
|
319
|
+
totalRecords: 0,
|
|
320
|
+
batchCount: 0
|
|
321
|
+
};
|
|
322
|
+
}
|
|
323
|
+
if (maxRows <= 0) {
|
|
324
|
+
throw new Error(`maxRows must be a positive number, got: ${maxRows}`);
|
|
325
|
+
}
|
|
326
|
+
const batches = [];
|
|
327
|
+
for (let i = 0; i < records.length; i += maxRows) {
|
|
328
|
+
batches.push(records.slice(i, i + maxRows));
|
|
329
|
+
}
|
|
330
|
+
return {
|
|
331
|
+
batches,
|
|
332
|
+
totalRecords: records.length,
|
|
333
|
+
batchCount: batches.length
|
|
334
|
+
};
|
|
335
|
+
}
|
|
336
|
+
|
|
337
|
+
// src/storage/db/index.ts
|
|
338
|
+
function resolveDsqlConfig(config) {
|
|
339
|
+
if ("client" in config) {
|
|
340
|
+
return {
|
|
341
|
+
client: config.client,
|
|
342
|
+
schemaName: config.schemaName,
|
|
343
|
+
skipDefaultIndexes: config.skipDefaultIndexes,
|
|
344
|
+
indexes: config.indexes
|
|
345
|
+
};
|
|
346
|
+
}
|
|
347
|
+
if ("pool" in config) {
|
|
348
|
+
return {
|
|
349
|
+
client: new PoolAdapter(config.pool),
|
|
350
|
+
schemaName: config.schemaName,
|
|
351
|
+
skipDefaultIndexes: config.skipDefaultIndexes,
|
|
352
|
+
indexes: config.indexes
|
|
353
|
+
};
|
|
354
|
+
}
|
|
355
|
+
const pool = new Pool({
|
|
356
|
+
host: config.host,
|
|
357
|
+
database: config.database,
|
|
358
|
+
user: config.user,
|
|
359
|
+
Client: AuroraDSQLClient
|
|
360
|
+
});
|
|
361
|
+
return {
|
|
362
|
+
client: new PoolAdapter(pool),
|
|
363
|
+
schemaName: config.schemaName,
|
|
364
|
+
skipDefaultIndexes: config.skipDefaultIndexes,
|
|
365
|
+
indexes: config.indexes
|
|
366
|
+
};
|
|
367
|
+
}
|
|
368
|
+
function getSchemaName(schema) {
|
|
369
|
+
return schema ? `"${parseSqlIdentifier(schema, "schema name")}"` : '"public"';
|
|
370
|
+
}
|
|
371
|
+
function getTableName({ indexName, schemaName }) {
|
|
372
|
+
const parsedIndexName = parseSqlIdentifier(indexName, "index name");
|
|
373
|
+
const quotedIndexName = `"${parsedIndexName}"`;
|
|
374
|
+
const quotedSchemaName = schemaName;
|
|
375
|
+
return quotedSchemaName ? `${quotedSchemaName}.${quotedIndexName}` : quotedIndexName;
|
|
376
|
+
}
|
|
377
|
+
var schemaSetupRegistry = /* @__PURE__ */ new Map();
|
|
378
|
+
var DsqlDB = class extends MastraBase {
|
|
379
|
+
client;
|
|
380
|
+
schemaName;
|
|
381
|
+
skipDefaultIndexes;
|
|
382
|
+
constructor(config) {
|
|
383
|
+
super({
|
|
384
|
+
component: "STORAGE",
|
|
385
|
+
name: "DSQL_DB_LAYER"
|
|
386
|
+
});
|
|
387
|
+
this.client = config.client;
|
|
388
|
+
this.schemaName = config.schemaName;
|
|
389
|
+
this.skipDefaultIndexes = config.skipDefaultIndexes;
|
|
390
|
+
}
|
|
391
|
+
async hasColumn(table, column) {
|
|
392
|
+
const schema = this.schemaName || "public";
|
|
393
|
+
const result = await this.client.oneOrNone(
|
|
394
|
+
`SELECT 1 FROM information_schema.columns WHERE table_schema = $1 AND table_name = $2 AND (column_name = $3 OR column_name = $4)`,
|
|
395
|
+
[schema, table, column, column.toLowerCase()]
|
|
396
|
+
);
|
|
397
|
+
return !!result;
|
|
398
|
+
}
|
|
399
|
+
/**
|
|
400
|
+
* Prepares values for insertion, handling TEXT columns storing JSON by stringifying them.
|
|
401
|
+
* Aurora DSQL does not support JSONB natively, so we store JSON as TEXT.
|
|
402
|
+
*/
|
|
403
|
+
prepareValuesForInsert(record, tableName) {
|
|
404
|
+
return Object.entries(record).map(([key, value]) => {
|
|
405
|
+
const schema = TABLE_SCHEMAS[tableName];
|
|
406
|
+
const columnSchema = schema?.[key];
|
|
407
|
+
if (columnSchema?.type === "jsonb" && value !== null && value !== void 0 && typeof value === "object") {
|
|
408
|
+
return JSON.stringify(value);
|
|
409
|
+
}
|
|
410
|
+
return value;
|
|
411
|
+
});
|
|
412
|
+
}
|
|
413
|
+
/**
|
|
414
|
+
* Adds timestamp Z columns to a record if timestamp columns exist
|
|
415
|
+
*/
|
|
416
|
+
addTimestampZColumns(record) {
|
|
417
|
+
if (record.createdAt) {
|
|
418
|
+
record.createdAtZ = record.createdAt;
|
|
419
|
+
}
|
|
420
|
+
if (record.created_at) {
|
|
421
|
+
record.created_atZ = record.created_at;
|
|
422
|
+
}
|
|
423
|
+
if (record.updatedAt) {
|
|
424
|
+
record.updatedAtZ = record.updatedAt;
|
|
425
|
+
}
|
|
426
|
+
}
|
|
427
|
+
/**
|
|
428
|
+
* Prepares a value for database operations, handling Date objects and JSON serialization.
|
|
429
|
+
* This is schema-aware and stringifies objects for TEXT columns storing JSON.
|
|
430
|
+
* Aurora DSQL: We use TEXT instead of JSONB and cast to ::jsonb when filtering.
|
|
431
|
+
*/
|
|
432
|
+
prepareValue(value, columnName, tableName) {
|
|
433
|
+
if (value === null || value === void 0) {
|
|
434
|
+
return value;
|
|
435
|
+
}
|
|
436
|
+
if (value instanceof Date) {
|
|
437
|
+
return value.toISOString();
|
|
438
|
+
}
|
|
439
|
+
const schema = TABLE_SCHEMAS[tableName];
|
|
440
|
+
const columnSchema = schema?.[columnName];
|
|
441
|
+
if (columnSchema?.type === "jsonb") {
|
|
442
|
+
if (typeof value === "object") {
|
|
443
|
+
return JSON.stringify(value);
|
|
444
|
+
}
|
|
445
|
+
}
|
|
446
|
+
if (typeof value === "object") {
|
|
447
|
+
return JSON.stringify(value);
|
|
448
|
+
}
|
|
449
|
+
return value;
|
|
450
|
+
}
|
|
451
|
+
async setupSchema() {
|
|
452
|
+
if (!this.schemaName) {
|
|
453
|
+
return;
|
|
454
|
+
}
|
|
455
|
+
let registryEntry = schemaSetupRegistry.get(this.schemaName);
|
|
456
|
+
if (registryEntry?.complete) {
|
|
457
|
+
return;
|
|
458
|
+
}
|
|
459
|
+
const quotedSchemaName = getSchemaName(this.schemaName);
|
|
460
|
+
if (!registryEntry?.promise) {
|
|
461
|
+
const schemaNameCapture = this.schemaName;
|
|
462
|
+
const setupPromise = (async () => {
|
|
463
|
+
try {
|
|
464
|
+
const schemaExists = await this.client.oneOrNone(
|
|
465
|
+
`
|
|
466
|
+
SELECT EXISTS (
|
|
467
|
+
SELECT 1 FROM information_schema.schemata
|
|
468
|
+
WHERE schema_name = $1
|
|
469
|
+
)
|
|
470
|
+
`,
|
|
471
|
+
[schemaNameCapture]
|
|
472
|
+
);
|
|
473
|
+
if (!schemaExists?.exists) {
|
|
474
|
+
try {
|
|
475
|
+
await this.client.none(`CREATE SCHEMA IF NOT EXISTS ${quotedSchemaName}`);
|
|
476
|
+
this.logger.info(`Schema "${schemaNameCapture}" created successfully`);
|
|
477
|
+
} catch (error) {
|
|
478
|
+
this.logger.error(`Failed to create schema "${schemaNameCapture}"`, { error });
|
|
479
|
+
throw new Error(
|
|
480
|
+
`Unable to create schema "${schemaNameCapture}". This requires CREATE privilege on the database. Either create the schema manually or grant CREATE privilege to the user.`
|
|
481
|
+
);
|
|
482
|
+
}
|
|
483
|
+
}
|
|
484
|
+
const entry = schemaSetupRegistry.get(schemaNameCapture);
|
|
485
|
+
if (entry) {
|
|
486
|
+
entry.complete = true;
|
|
487
|
+
}
|
|
488
|
+
this.logger.debug(`Schema "${quotedSchemaName}" is ready for use`);
|
|
489
|
+
} catch (error) {
|
|
490
|
+
schemaSetupRegistry.delete(schemaNameCapture);
|
|
491
|
+
throw error;
|
|
492
|
+
}
|
|
493
|
+
})();
|
|
494
|
+
schemaSetupRegistry.set(this.schemaName, { promise: setupPromise, complete: false });
|
|
495
|
+
registryEntry = schemaSetupRegistry.get(this.schemaName);
|
|
496
|
+
}
|
|
497
|
+
await registryEntry.promise;
|
|
498
|
+
}
|
|
499
|
+
/**
|
|
500
|
+
* Override getSqlType to map JSONB to TEXT for Aurora DSQL compatibility.
|
|
501
|
+
* Aurora DSQL does not fully support native JSONB, so we store JSON as TEXT
|
|
502
|
+
* and cast to ::jsonb only when filtering/querying.
|
|
503
|
+
*/
|
|
504
|
+
getSqlType(type) {
|
|
505
|
+
switch (type) {
|
|
506
|
+
case "jsonb":
|
|
507
|
+
return "TEXT";
|
|
508
|
+
case "uuid":
|
|
509
|
+
return "UUID";
|
|
510
|
+
case "boolean":
|
|
511
|
+
return "BOOLEAN";
|
|
512
|
+
default:
|
|
513
|
+
return getSqlType(type);
|
|
514
|
+
}
|
|
515
|
+
}
|
|
516
|
+
getDefaultValue(type) {
|
|
517
|
+
switch (type) {
|
|
518
|
+
case "timestamp":
|
|
519
|
+
return "DEFAULT NOW()";
|
|
520
|
+
case "jsonb":
|
|
521
|
+
return "DEFAULT '{}'";
|
|
522
|
+
default:
|
|
523
|
+
return getDefaultValue(type);
|
|
524
|
+
}
|
|
525
|
+
}
|
|
526
|
+
async insert({ tableName, record }) {
|
|
527
|
+
this.addTimestampZColumns(record);
|
|
528
|
+
const schemaName = getSchemaName(this.schemaName);
|
|
529
|
+
const columns = Object.keys(record).map((col) => parseSqlIdentifier(col, "column name"));
|
|
530
|
+
const values = this.prepareValuesForInsert(record, tableName);
|
|
531
|
+
const placeholders = values.map((_, i) => `$${i + 1}`).join(", ");
|
|
532
|
+
await withRetry(
|
|
533
|
+
async () => {
|
|
534
|
+
await this.client.none(
|
|
535
|
+
`INSERT INTO ${getTableName({ indexName: tableName, schemaName })} (${columns.map((c) => `"${c}"`).join(", ")}) VALUES (${placeholders})`,
|
|
536
|
+
values
|
|
537
|
+
);
|
|
538
|
+
},
|
|
539
|
+
{
|
|
540
|
+
onRetry: (error, attempt, delay) => {
|
|
541
|
+
this.logger?.warn?.(`insert retry ${attempt} for table ${tableName} after ${delay}ms: ${error.message}`);
|
|
542
|
+
}
|
|
543
|
+
}
|
|
544
|
+
).catch((error) => {
|
|
545
|
+
throw new MastraError(
|
|
546
|
+
{
|
|
547
|
+
id: createStorageErrorId("DSQL", "INSERT", "FAILED"),
|
|
548
|
+
domain: ErrorDomain.STORAGE,
|
|
549
|
+
category: ErrorCategory.THIRD_PARTY,
|
|
550
|
+
details: {
|
|
551
|
+
tableName
|
|
552
|
+
}
|
|
553
|
+
},
|
|
554
|
+
error
|
|
555
|
+
);
|
|
556
|
+
});
|
|
557
|
+
}
|
|
558
|
+
async clearTable({ tableName }) {
|
|
559
|
+
try {
|
|
560
|
+
await withRetry(
|
|
561
|
+
async () => {
|
|
562
|
+
const schemaName = getSchemaName(this.schemaName);
|
|
563
|
+
const tableNameWithSchema = getTableName({ indexName: tableName, schemaName });
|
|
564
|
+
await this.client.none(`DELETE FROM ${tableNameWithSchema}`);
|
|
565
|
+
},
|
|
566
|
+
{
|
|
567
|
+
onRetry: (error, attempt, delay) => {
|
|
568
|
+
this.logger?.warn?.(`clearTable retry ${attempt} for ${tableName} after ${delay}ms: ${error.message}`);
|
|
569
|
+
}
|
|
570
|
+
}
|
|
571
|
+
);
|
|
572
|
+
} catch (error) {
|
|
573
|
+
throw new MastraError(
|
|
574
|
+
{
|
|
575
|
+
id: createStorageErrorId("DSQL", "CLEAR_TABLE", "FAILED"),
|
|
576
|
+
domain: ErrorDomain.STORAGE,
|
|
577
|
+
category: ErrorCategory.THIRD_PARTY,
|
|
578
|
+
details: {
|
|
579
|
+
tableName
|
|
580
|
+
}
|
|
581
|
+
},
|
|
582
|
+
error
|
|
583
|
+
);
|
|
584
|
+
}
|
|
585
|
+
}
|
|
586
|
+
async createTable({
|
|
587
|
+
tableName,
|
|
588
|
+
schema
|
|
589
|
+
}) {
|
|
590
|
+
await withRetry(
|
|
591
|
+
async () => {
|
|
592
|
+
const timeZColumnNames = Object.entries(schema).filter(([_, def]) => def.type === "timestamp").map(([name]) => name);
|
|
593
|
+
const timeZColumns = Object.entries(schema).filter(([_, def]) => def.type === "timestamp").map(([name]) => {
|
|
594
|
+
const parsedName = parseSqlIdentifier(name, "column name");
|
|
595
|
+
return `"${parsedName}Z" TIMESTAMPTZ DEFAULT NOW()`;
|
|
596
|
+
});
|
|
597
|
+
const columns = Object.entries(schema).map(([name, def]) => {
|
|
598
|
+
const parsedName = parseSqlIdentifier(name, "column name");
|
|
599
|
+
const constraints = [];
|
|
600
|
+
if (def.primaryKey) constraints.push("PRIMARY KEY");
|
|
601
|
+
if (!def.nullable) constraints.push("NOT NULL");
|
|
602
|
+
const sqlType = this.getSqlType(def.type);
|
|
603
|
+
return `"${parsedName}" ${sqlType} ${constraints.join(" ")}`;
|
|
604
|
+
});
|
|
605
|
+
if (this.schemaName) {
|
|
606
|
+
await this.setupSchema();
|
|
607
|
+
}
|
|
608
|
+
const finalColumns = [...columns, ...timeZColumns].join(",\n");
|
|
609
|
+
const constraintPrefix = this.schemaName ? `${this.schemaName}_` : "";
|
|
610
|
+
const schemaName = getSchemaName(this.schemaName);
|
|
611
|
+
const createTableSql = `
|
|
612
|
+
CREATE TABLE IF NOT EXISTS ${getTableName({ indexName: tableName, schemaName })} (
|
|
613
|
+
${finalColumns}
|
|
614
|
+
);
|
|
615
|
+
`;
|
|
616
|
+
await this.client.none(createTableSql);
|
|
617
|
+
if (tableName === TABLE_WORKFLOW_SNAPSHOT) {
|
|
618
|
+
const indexName = `${constraintPrefix}mastra_workflow_snapshot_workflow_name_run_id_key`;
|
|
619
|
+
const fullTableName = getTableName({ indexName: tableName, schemaName });
|
|
620
|
+
try {
|
|
621
|
+
const indexExists = await this.client.oneOrNone(`SELECT 1 FROM pg_indexes WHERE indexname = $1`, [
|
|
622
|
+
indexName
|
|
623
|
+
]);
|
|
624
|
+
if (!indexExists) {
|
|
625
|
+
const result = await this.client.oneOrNone(
|
|
626
|
+
`CREATE UNIQUE INDEX ASYNC "${indexName}" ON ${fullTableName} ("workflow_name", "run_id")`
|
|
627
|
+
);
|
|
628
|
+
if (result?.job_uuid) {
|
|
629
|
+
await this.waitForDSQLJob(result.job_uuid);
|
|
630
|
+
}
|
|
631
|
+
this.logger?.debug?.(`Created unique index ${indexName} on ${fullTableName}`);
|
|
632
|
+
}
|
|
633
|
+
} catch (error) {
|
|
634
|
+
this.logger?.warn?.(`Failed to create unique index ${indexName}:`, error);
|
|
635
|
+
}
|
|
636
|
+
}
|
|
637
|
+
await this.alterTable({
|
|
638
|
+
tableName,
|
|
639
|
+
schema,
|
|
640
|
+
ifNotExists: timeZColumnNames
|
|
641
|
+
});
|
|
642
|
+
if (tableName === TABLE_SPANS) {
|
|
643
|
+
await this.setupTimestampTriggers(tableName);
|
|
644
|
+
await this.migrateSpansTable();
|
|
645
|
+
}
|
|
646
|
+
},
|
|
647
|
+
{
|
|
648
|
+
onRetry: (error, attempt, delay) => {
|
|
649
|
+
this.logger?.warn?.(`createTable retry ${attempt} for ${tableName} after ${delay}ms: ${error.message}`);
|
|
650
|
+
}
|
|
651
|
+
}
|
|
652
|
+
).catch((error) => {
|
|
653
|
+
throw new MastraError(
|
|
654
|
+
{
|
|
655
|
+
id: createStorageErrorId("DSQL", "CREATE_TABLE", "FAILED"),
|
|
656
|
+
domain: ErrorDomain.STORAGE,
|
|
657
|
+
category: ErrorCategory.THIRD_PARTY,
|
|
658
|
+
details: {
|
|
659
|
+
tableName
|
|
660
|
+
}
|
|
661
|
+
},
|
|
662
|
+
error
|
|
663
|
+
);
|
|
664
|
+
});
|
|
665
|
+
}
|
|
666
|
+
/**
|
|
667
|
+
* Set up timestamp triggers for a table to automatically manage createdAt/updatedAt
|
|
668
|
+
* Note: Aurora DSQL doesn't support triggers, PL/pgSQL, or CREATE FUNCTION.
|
|
669
|
+
* Timestamps are managed at the application level in insert/update operations.
|
|
670
|
+
* This method is kept as a no-op for API compatibility.
|
|
671
|
+
*/
|
|
672
|
+
async setupTimestampTriggers(_tableName) {
|
|
673
|
+
return;
|
|
674
|
+
}
|
|
675
|
+
/**
|
|
676
|
+
* Migrates the spans table schema from OLD_SPAN_SCHEMA to current SPAN_SCHEMA.
|
|
677
|
+
* This adds new columns that don't exist in old schema.
|
|
678
|
+
*/
|
|
679
|
+
async migrateSpansTable() {
|
|
680
|
+
const fullTableName = getTableName({ indexName: TABLE_SPANS, schemaName: getSchemaName(this.schemaName) });
|
|
681
|
+
const schema = TABLE_SCHEMAS[TABLE_SPANS];
|
|
682
|
+
try {
|
|
683
|
+
for (const [columnName, columnDef] of Object.entries(schema)) {
|
|
684
|
+
const columnExists = await this.hasColumn(TABLE_SPANS, columnName);
|
|
685
|
+
if (!columnExists) {
|
|
686
|
+
const parsedColumnName = parseSqlIdentifier(columnName, "column name");
|
|
687
|
+
const sqlType = this.getSqlType(columnDef.type);
|
|
688
|
+
const nullable = columnDef.nullable ? "" : "NOT NULL";
|
|
689
|
+
const defaultValue = !columnDef.nullable ? this.getDefaultValue(columnDef.type) : "";
|
|
690
|
+
const alterSql = `ALTER TABLE ${fullTableName} ADD COLUMN IF NOT EXISTS "${parsedColumnName}" ${sqlType} ${nullable} ${defaultValue}`.trim();
|
|
691
|
+
await this.client.none(alterSql);
|
|
692
|
+
this.logger?.debug?.(`Added column '${columnName}' to ${fullTableName}`);
|
|
693
|
+
if (sqlType === "TIMESTAMP") {
|
|
694
|
+
const timestampZSql = `ALTER TABLE ${fullTableName} ADD COLUMN IF NOT EXISTS "${parsedColumnName}Z" TIMESTAMPTZ DEFAULT NOW()`.trim();
|
|
695
|
+
await this.client.none(timestampZSql);
|
|
696
|
+
this.logger?.debug?.(`Added timezone column '${columnName}Z' to ${fullTableName}`);
|
|
697
|
+
}
|
|
698
|
+
}
|
|
699
|
+
}
|
|
700
|
+
for (const [columnName, columnDef] of Object.entries(schema)) {
|
|
701
|
+
if (columnDef.type === "timestamp") {
|
|
702
|
+
const tzColumnName = `${columnName}Z`;
|
|
703
|
+
const tzColumnExists = await this.hasColumn(TABLE_SPANS, tzColumnName);
|
|
704
|
+
if (!tzColumnExists) {
|
|
705
|
+
const parsedTzColumnName = parseSqlIdentifier(tzColumnName, "column name");
|
|
706
|
+
const timestampZSql = `ALTER TABLE ${fullTableName} ADD COLUMN IF NOT EXISTS "${parsedTzColumnName}" TIMESTAMPTZ DEFAULT NOW()`.trim();
|
|
707
|
+
await this.client.none(timestampZSql);
|
|
708
|
+
this.logger?.debug?.(`Added timezone column '${tzColumnName}' to ${fullTableName}`);
|
|
709
|
+
}
|
|
710
|
+
}
|
|
711
|
+
}
|
|
712
|
+
this.logger?.info?.(`Migration completed for ${fullTableName}`);
|
|
713
|
+
} catch (error) {
|
|
714
|
+
this.logger?.warn?.(`Failed to migrate spans table ${fullTableName}:`, error);
|
|
715
|
+
}
|
|
716
|
+
}
|
|
717
|
+
/**
|
|
718
|
+
* Alters table schema to add columns if they don't exist
|
|
719
|
+
* @param tableName Name of the table
|
|
720
|
+
* @param schema Schema of the table
|
|
721
|
+
* @param ifNotExists Array of column names to add if they don't exist
|
|
722
|
+
*/
|
|
723
|
+
async alterTable({
|
|
724
|
+
tableName,
|
|
725
|
+
schema,
|
|
726
|
+
ifNotExists
|
|
727
|
+
}) {
|
|
728
|
+
const fullTableName = getTableName({ indexName: tableName, schemaName: getSchemaName(this.schemaName) });
|
|
729
|
+
try {
|
|
730
|
+
for (const columnName of ifNotExists) {
|
|
731
|
+
if (schema[columnName]) {
|
|
732
|
+
const columnDef = schema[columnName];
|
|
733
|
+
const sqlType = this.getSqlType(columnDef.type);
|
|
734
|
+
const parsedColumnName = parseSqlIdentifier(columnName, "column name");
|
|
735
|
+
const alterSql = `ALTER TABLE ${fullTableName} ADD COLUMN IF NOT EXISTS "${parsedColumnName}" ${sqlType}`;
|
|
736
|
+
await this.client.none(alterSql);
|
|
737
|
+
if (sqlType === "TIMESTAMP") {
|
|
738
|
+
const alterSqlZ = `ALTER TABLE ${fullTableName} ADD COLUMN IF NOT EXISTS "${parsedColumnName}Z" TIMESTAMPTZ`;
|
|
739
|
+
await this.client.none(alterSqlZ);
|
|
740
|
+
}
|
|
741
|
+
this.logger?.debug?.(`Ensured column ${parsedColumnName} exists in table ${fullTableName}`);
|
|
742
|
+
}
|
|
743
|
+
}
|
|
744
|
+
} catch (error) {
|
|
745
|
+
throw new MastraError(
|
|
746
|
+
{
|
|
747
|
+
id: createStorageErrorId("DSQL", "ALTER_TABLE", "FAILED"),
|
|
748
|
+
domain: ErrorDomain.STORAGE,
|
|
749
|
+
category: ErrorCategory.THIRD_PARTY,
|
|
750
|
+
details: {
|
|
751
|
+
tableName
|
|
752
|
+
}
|
|
753
|
+
},
|
|
754
|
+
error
|
|
755
|
+
);
|
|
756
|
+
}
|
|
757
|
+
}
|
|
758
|
+
async load({ tableName, keys }) {
|
|
759
|
+
try {
|
|
760
|
+
const keyEntries = Object.entries(keys).map(([key, value]) => [parseSqlIdentifier(key, "column name"), value]);
|
|
761
|
+
const conditions = keyEntries.map(([key], index) => `"${key}" = $${index + 1}`).join(" AND ");
|
|
762
|
+
const values = keyEntries.map(([_, value]) => value);
|
|
763
|
+
const result = await this.client.oneOrNone(
|
|
764
|
+
`SELECT * FROM ${getTableName({ indexName: tableName, schemaName: getSchemaName(this.schemaName) })} WHERE ${conditions} ORDER BY "createdAt" DESC LIMIT 1`,
|
|
765
|
+
values
|
|
766
|
+
);
|
|
767
|
+
if (!result) {
|
|
768
|
+
return null;
|
|
769
|
+
}
|
|
770
|
+
if (tableName === TABLE_WORKFLOW_SNAPSHOT) {
|
|
771
|
+
const snapshot = result;
|
|
772
|
+
if (typeof snapshot.snapshot === "string") {
|
|
773
|
+
snapshot.snapshot = JSON.parse(snapshot.snapshot);
|
|
774
|
+
}
|
|
775
|
+
return snapshot;
|
|
776
|
+
}
|
|
777
|
+
return result;
|
|
778
|
+
} catch (error) {
|
|
779
|
+
throw new MastraError(
|
|
780
|
+
{
|
|
781
|
+
id: createStorageErrorId("DSQL", "LOAD", "FAILED"),
|
|
782
|
+
domain: ErrorDomain.STORAGE,
|
|
783
|
+
category: ErrorCategory.THIRD_PARTY,
|
|
784
|
+
details: {
|
|
785
|
+
tableName
|
|
786
|
+
}
|
|
787
|
+
},
|
|
788
|
+
error
|
|
789
|
+
);
|
|
790
|
+
}
|
|
791
|
+
}
|
|
792
|
+
async batchInsert({ tableName, records }) {
|
|
793
|
+
if (records.length === 0) {
|
|
794
|
+
return;
|
|
795
|
+
}
|
|
796
|
+
try {
|
|
797
|
+
const { batches } = splitIntoBatches(records, { maxRows: DEFAULT_MAX_ROWS_PER_BATCH });
|
|
798
|
+
for (const batch of batches) {
|
|
799
|
+
await withRetry(
|
|
800
|
+
async () => {
|
|
801
|
+
await this.client.tx(async (t) => {
|
|
802
|
+
for (const record of batch) {
|
|
803
|
+
this.addTimestampZColumns(record);
|
|
804
|
+
const schemaName = getSchemaName(this.schemaName);
|
|
805
|
+
const columns = Object.keys(record).map((col) => parseSqlIdentifier(col, "column name"));
|
|
806
|
+
const values = this.prepareValuesForInsert(record, tableName);
|
|
807
|
+
const placeholders = values.map((_, i) => `$${i + 1}`).join(", ");
|
|
808
|
+
await t.none(
|
|
809
|
+
`INSERT INTO ${getTableName({ indexName: tableName, schemaName })} (${columns.map((c) => `"${c}"`).join(", ")}) VALUES (${placeholders})`,
|
|
810
|
+
values
|
|
811
|
+
);
|
|
812
|
+
}
|
|
813
|
+
});
|
|
814
|
+
},
|
|
815
|
+
{
|
|
816
|
+
onRetry: (error, attempt, delay) => {
|
|
817
|
+
this.logger?.warn?.(
|
|
818
|
+
`Batch insert retry ${attempt} for table ${tableName} after ${delay}ms: ${error.message}`
|
|
819
|
+
);
|
|
820
|
+
}
|
|
821
|
+
}
|
|
822
|
+
);
|
|
823
|
+
}
|
|
824
|
+
} catch (error) {
|
|
825
|
+
throw new MastraError(
|
|
826
|
+
{
|
|
827
|
+
id: createStorageErrorId("DSQL", "BATCH_INSERT", "FAILED"),
|
|
828
|
+
domain: ErrorDomain.STORAGE,
|
|
829
|
+
category: ErrorCategory.THIRD_PARTY,
|
|
830
|
+
details: {
|
|
831
|
+
tableName,
|
|
832
|
+
numberOfRecords: records.length
|
|
833
|
+
}
|
|
834
|
+
},
|
|
835
|
+
error
|
|
836
|
+
);
|
|
837
|
+
}
|
|
838
|
+
}
|
|
839
|
+
async dropTable({ tableName }) {
|
|
840
|
+
try {
|
|
841
|
+
const schemaName = getSchemaName(this.schemaName);
|
|
842
|
+
const tableNameWithSchema = getTableName({ indexName: tableName, schemaName });
|
|
843
|
+
await this.client.none(`DROP TABLE IF EXISTS ${tableNameWithSchema}`);
|
|
844
|
+
} catch (error) {
|
|
845
|
+
throw new MastraError(
|
|
846
|
+
{
|
|
847
|
+
id: createStorageErrorId("DSQL", "DROP_TABLE", "FAILED"),
|
|
848
|
+
domain: ErrorDomain.STORAGE,
|
|
849
|
+
category: ErrorCategory.THIRD_PARTY,
|
|
850
|
+
details: {
|
|
851
|
+
tableName
|
|
852
|
+
}
|
|
853
|
+
},
|
|
854
|
+
error
|
|
855
|
+
);
|
|
856
|
+
}
|
|
857
|
+
}
|
|
858
|
+
/**
|
|
859
|
+
* Wait for an asynchronous DSQL job to complete.
|
|
860
|
+
* Aurora DSQL requires CREATE INDEX ASYNC and sys.wait_for_job() to wait for completion.
|
|
861
|
+
*/
|
|
862
|
+
async waitForDSQLJob(jobUuid, timeoutMs = 6e4) {
|
|
863
|
+
const pollIntervalMs = 1e3;
|
|
864
|
+
const startTime = Date.now();
|
|
865
|
+
while (Date.now() - startTime < timeoutMs) {
|
|
866
|
+
const result = await this.client.oneOrNone(`SELECT sys.wait_for_job($1, 1) as status`, [
|
|
867
|
+
jobUuid
|
|
868
|
+
]);
|
|
869
|
+
if (result?.status === "COMPLETED") {
|
|
870
|
+
return;
|
|
871
|
+
}
|
|
872
|
+
if (result?.status === "FAILED") {
|
|
873
|
+
throw new Error(`DSQL async job ${jobUuid} failed`);
|
|
874
|
+
}
|
|
875
|
+
await new Promise((resolve) => setTimeout(resolve, pollIntervalMs));
|
|
876
|
+
}
|
|
877
|
+
throw new Error(`DSQL async job ${jobUuid} timed out after ${timeoutMs}ms`);
|
|
878
|
+
}
|
|
879
|
+
async createIndex(options) {
|
|
880
|
+
try {
|
|
881
|
+
const {
|
|
882
|
+
name,
|
|
883
|
+
table,
|
|
884
|
+
columns,
|
|
885
|
+
unique = false,
|
|
886
|
+
// Note: 'concurrent' option is ignored in DSQL - always uses ASYNC
|
|
887
|
+
where,
|
|
888
|
+
method = "btree",
|
|
889
|
+
opclass,
|
|
890
|
+
storage,
|
|
891
|
+
tablespace
|
|
892
|
+
} = options;
|
|
893
|
+
const schemaName = this.schemaName || "public";
|
|
894
|
+
const fullTableName = getTableName({
|
|
895
|
+
indexName: table,
|
|
896
|
+
schemaName: getSchemaName(this.schemaName)
|
|
897
|
+
});
|
|
898
|
+
const indexExists = await this.client.oneOrNone(
|
|
899
|
+
`SELECT 1 as exists FROM pg_indexes
|
|
900
|
+
WHERE indexname = $1
|
|
901
|
+
AND schemaname = $2`,
|
|
902
|
+
[name, schemaName]
|
|
903
|
+
);
|
|
904
|
+
if (indexExists) {
|
|
905
|
+
return;
|
|
906
|
+
}
|
|
907
|
+
const uniqueStr = unique ? "UNIQUE " : "";
|
|
908
|
+
const methodStr = method !== "btree" ? `USING ${method} ` : "";
|
|
909
|
+
const columnsStr = columns.map((col) => {
|
|
910
|
+
const colName = col.replace(/\s+(DESC|ASC)$/i, "").trim();
|
|
911
|
+
const quotedCol = `"${parseSqlIdentifier(colName, "column name")}"`;
|
|
912
|
+
return opclass ? `${quotedCol} ${opclass}` : quotedCol;
|
|
913
|
+
}).join(", ");
|
|
914
|
+
const whereStr = where ? ` WHERE ${where}` : "";
|
|
915
|
+
const tablespaceStr = tablespace ? ` TABLESPACE ${tablespace}` : "";
|
|
916
|
+
let withStr = "";
|
|
917
|
+
if (storage && Object.keys(storage).length > 0) {
|
|
918
|
+
const storageParams = Object.entries(storage).map(([key, value]) => `${key} = ${value}`).join(", ");
|
|
919
|
+
withStr = ` WITH (${storageParams})`;
|
|
920
|
+
}
|
|
921
|
+
const quotedIndexName = `"${parseSqlIdentifier(name, "index name")}"`;
|
|
922
|
+
const sql = `CREATE ${uniqueStr}INDEX ASYNC ${quotedIndexName} ON ${fullTableName} ${methodStr}(${columnsStr})${withStr}${tablespaceStr}${whereStr}`;
|
|
923
|
+
const result = await this.client.oneOrNone(sql);
|
|
924
|
+
if (result?.job_uuid) {
|
|
925
|
+
await this.waitForDSQLJob(result.job_uuid);
|
|
926
|
+
}
|
|
927
|
+
} catch (error) {
|
|
928
|
+
throw new MastraError(
|
|
929
|
+
{
|
|
930
|
+
id: createStorageErrorId("DSQL", "INDEX_CREATE", "FAILED"),
|
|
931
|
+
domain: ErrorDomain.STORAGE,
|
|
932
|
+
category: ErrorCategory.THIRD_PARTY,
|
|
933
|
+
details: {
|
|
934
|
+
indexName: options.name,
|
|
935
|
+
tableName: options.table
|
|
936
|
+
}
|
|
937
|
+
},
|
|
938
|
+
error
|
|
939
|
+
);
|
|
940
|
+
}
|
|
941
|
+
}
|
|
942
|
+
async dropIndex(indexName) {
|
|
943
|
+
const schemaName = this.schemaName || "public";
|
|
944
|
+
await withRetry(
|
|
945
|
+
async () => {
|
|
946
|
+
const indexExists = await this.client.oneOrNone(
|
|
947
|
+
`SELECT 1 FROM pg_indexes
|
|
948
|
+
WHERE indexname = $1
|
|
949
|
+
AND schemaname = $2`,
|
|
950
|
+
[indexName, schemaName]
|
|
951
|
+
);
|
|
952
|
+
if (!indexExists) {
|
|
953
|
+
return;
|
|
954
|
+
}
|
|
955
|
+
const quotedIndexName = `"${parseSqlIdentifier(indexName, "index name")}"`;
|
|
956
|
+
const sql = `DROP INDEX IF EXISTS ${getSchemaName(this.schemaName)}.${quotedIndexName}`;
|
|
957
|
+
await this.client.none(sql);
|
|
958
|
+
},
|
|
959
|
+
{
|
|
960
|
+
onRetry: (error, attempt, delay) => {
|
|
961
|
+
this.logger?.warn?.(`dropIndex retry ${attempt} for ${indexName} after ${delay}ms: ${error.message}`);
|
|
962
|
+
}
|
|
963
|
+
}
|
|
964
|
+
).catch((error) => {
|
|
965
|
+
throw new MastraError(
|
|
966
|
+
{
|
|
967
|
+
id: createStorageErrorId("DSQL", "INDEX_DROP", "FAILED"),
|
|
968
|
+
domain: ErrorDomain.STORAGE,
|
|
969
|
+
category: ErrorCategory.THIRD_PARTY,
|
|
970
|
+
details: {
|
|
971
|
+
indexName
|
|
972
|
+
}
|
|
973
|
+
},
|
|
974
|
+
error
|
|
975
|
+
);
|
|
976
|
+
});
|
|
977
|
+
}
|
|
978
|
+
async listIndexes(tableName) {
|
|
979
|
+
try {
|
|
980
|
+
const schemaName = this.schemaName || "public";
|
|
981
|
+
let query;
|
|
982
|
+
let params;
|
|
983
|
+
if (tableName) {
|
|
984
|
+
query = `
|
|
985
|
+
SELECT
|
|
986
|
+
i.indexname as name,
|
|
987
|
+
i.tablename as table,
|
|
988
|
+
i.indexdef as definition,
|
|
989
|
+
ix.indisunique as is_unique,
|
|
990
|
+
pg_size_pretty(pg_relation_size(c.oid)) as size,
|
|
991
|
+
array_agg(a.attname ORDER BY array_position(ix.indkey, a.attnum)) as columns
|
|
992
|
+
FROM pg_indexes i
|
|
993
|
+
JOIN pg_class c ON c.relname = i.indexname AND c.relnamespace = (SELECT oid FROM pg_namespace WHERE nspname = i.schemaname)
|
|
994
|
+
JOIN pg_index ix ON ix.indexrelid = c.oid
|
|
995
|
+
JOIN pg_attribute a ON a.attrelid = ix.indrelid AND a.attnum = ANY(ix.indkey)
|
|
996
|
+
WHERE i.schemaname = $1
|
|
997
|
+
AND i.tablename = $2
|
|
998
|
+
GROUP BY i.indexname, i.tablename, i.indexdef, ix.indisunique, c.oid
|
|
999
|
+
`;
|
|
1000
|
+
params = [schemaName, tableName];
|
|
1001
|
+
} else {
|
|
1002
|
+
query = `
|
|
1003
|
+
SELECT
|
|
1004
|
+
i.indexname as name,
|
|
1005
|
+
i.tablename as table,
|
|
1006
|
+
i.indexdef as definition,
|
|
1007
|
+
ix.indisunique as is_unique,
|
|
1008
|
+
pg_size_pretty(pg_relation_size(c.oid)) as size,
|
|
1009
|
+
array_agg(a.attname ORDER BY array_position(ix.indkey, a.attnum)) as columns
|
|
1010
|
+
FROM pg_indexes i
|
|
1011
|
+
JOIN pg_class c ON c.relname = i.indexname AND c.relnamespace = (SELECT oid FROM pg_namespace WHERE nspname = i.schemaname)
|
|
1012
|
+
JOIN pg_index ix ON ix.indexrelid = c.oid
|
|
1013
|
+
JOIN pg_attribute a ON a.attrelid = ix.indrelid AND a.attnum = ANY(ix.indkey)
|
|
1014
|
+
WHERE i.schemaname = $1
|
|
1015
|
+
GROUP BY i.indexname, i.tablename, i.indexdef, ix.indisunique, c.oid
|
|
1016
|
+
`;
|
|
1017
|
+
params = [schemaName];
|
|
1018
|
+
}
|
|
1019
|
+
const results = await this.client.manyOrNone(query, params);
|
|
1020
|
+
return results.map((row) => {
|
|
1021
|
+
let columns = [];
|
|
1022
|
+
if (typeof row.columns === "string" && row.columns.startsWith("{") && row.columns.endsWith("}")) {
|
|
1023
|
+
const arrayContent = row.columns.slice(1, -1);
|
|
1024
|
+
columns = arrayContent ? arrayContent.split(",") : [];
|
|
1025
|
+
} else if (Array.isArray(row.columns)) {
|
|
1026
|
+
columns = row.columns;
|
|
1027
|
+
}
|
|
1028
|
+
return {
|
|
1029
|
+
name: row.name,
|
|
1030
|
+
table: row.table,
|
|
1031
|
+
columns,
|
|
1032
|
+
unique: row.is_unique || false,
|
|
1033
|
+
size: row.size || "0",
|
|
1034
|
+
definition: row.definition || ""
|
|
1035
|
+
};
|
|
1036
|
+
});
|
|
1037
|
+
} catch (error) {
|
|
1038
|
+
throw new MastraError(
|
|
1039
|
+
{
|
|
1040
|
+
id: createStorageErrorId("DSQL", "INDEX_LIST", "FAILED"),
|
|
1041
|
+
domain: ErrorDomain.STORAGE,
|
|
1042
|
+
category: ErrorCategory.THIRD_PARTY,
|
|
1043
|
+
details: tableName ? {
|
|
1044
|
+
tableName
|
|
1045
|
+
} : {}
|
|
1046
|
+
},
|
|
1047
|
+
error
|
|
1048
|
+
);
|
|
1049
|
+
}
|
|
1050
|
+
}
|
|
1051
|
+
async describeIndex(indexName) {
|
|
1052
|
+
try {
|
|
1053
|
+
const schemaName = this.schemaName || "public";
|
|
1054
|
+
const query = `
|
|
1055
|
+
SELECT
|
|
1056
|
+
i.indexname as name,
|
|
1057
|
+
i.tablename as table,
|
|
1058
|
+
i.indexdef as definition,
|
|
1059
|
+
ix.indisunique as is_unique,
|
|
1060
|
+
pg_size_pretty(pg_relation_size(c.oid)) as size,
|
|
1061
|
+
array_agg(a.attname ORDER BY array_position(ix.indkey, a.attnum)) as columns,
|
|
1062
|
+
am.amname as method,
|
|
1063
|
+
s.idx_scan as scans,
|
|
1064
|
+
s.idx_tup_read as tuples_read,
|
|
1065
|
+
s.idx_tup_fetch as tuples_fetched
|
|
1066
|
+
FROM pg_indexes i
|
|
1067
|
+
JOIN pg_class c ON c.relname = i.indexname AND c.relnamespace = (SELECT oid FROM pg_namespace WHERE nspname = i.schemaname)
|
|
1068
|
+
JOIN pg_index ix ON ix.indexrelid = c.oid
|
|
1069
|
+
JOIN pg_attribute a ON a.attrelid = ix.indrelid AND a.attnum = ANY(ix.indkey)
|
|
1070
|
+
JOIN pg_am am ON c.relam = am.oid
|
|
1071
|
+
LEFT JOIN pg_stat_user_indexes s ON s.indexrelname = i.indexname AND s.schemaname = i.schemaname
|
|
1072
|
+
WHERE i.schemaname = $1
|
|
1073
|
+
AND i.indexname = $2
|
|
1074
|
+
GROUP BY i.indexname, i.tablename, i.indexdef, ix.indisunique, c.oid, am.amname, s.idx_scan, s.idx_tup_read, s.idx_tup_fetch
|
|
1075
|
+
`;
|
|
1076
|
+
const result = await this.client.oneOrNone(query, [schemaName, indexName]);
|
|
1077
|
+
if (!result) {
|
|
1078
|
+
throw new Error(`Index "${indexName}" not found in schema "${schemaName}"`);
|
|
1079
|
+
}
|
|
1080
|
+
let columns = [];
|
|
1081
|
+
if (typeof result.columns === "string" && result.columns.startsWith("{") && result.columns.endsWith("}")) {
|
|
1082
|
+
const arrayContent = result.columns.slice(1, -1);
|
|
1083
|
+
columns = arrayContent ? arrayContent.split(",") : [];
|
|
1084
|
+
} else if (Array.isArray(result.columns)) {
|
|
1085
|
+
columns = result.columns;
|
|
1086
|
+
}
|
|
1087
|
+
const normalizedMethod = result.method === "btree_index" ? "btree" : result.method || "btree";
|
|
1088
|
+
return {
|
|
1089
|
+
name: result.name,
|
|
1090
|
+
table: result.table,
|
|
1091
|
+
columns,
|
|
1092
|
+
unique: result.is_unique || false,
|
|
1093
|
+
size: result.size || "0",
|
|
1094
|
+
definition: result.definition || "",
|
|
1095
|
+
method: normalizedMethod,
|
|
1096
|
+
scans: parseInt(result.scans) || 0,
|
|
1097
|
+
tuples_read: parseInt(result.tuples_read) || 0,
|
|
1098
|
+
tuples_fetched: parseInt(result.tuples_fetched) || 0
|
|
1099
|
+
};
|
|
1100
|
+
} catch (error) {
|
|
1101
|
+
throw new MastraError(
|
|
1102
|
+
{
|
|
1103
|
+
id: createStorageErrorId("DSQL", "INDEX_DESCRIBE", "FAILED"),
|
|
1104
|
+
domain: ErrorDomain.STORAGE,
|
|
1105
|
+
category: ErrorCategory.THIRD_PARTY,
|
|
1106
|
+
details: {
|
|
1107
|
+
indexName
|
|
1108
|
+
}
|
|
1109
|
+
},
|
|
1110
|
+
error
|
|
1111
|
+
);
|
|
1112
|
+
}
|
|
1113
|
+
}
|
|
1114
|
+
async update({
|
|
1115
|
+
tableName,
|
|
1116
|
+
keys,
|
|
1117
|
+
data
|
|
1118
|
+
}) {
|
|
1119
|
+
const setColumns = [];
|
|
1120
|
+
const setValues = [];
|
|
1121
|
+
let paramIndex = 1;
|
|
1122
|
+
const now = (/* @__PURE__ */ new Date()).toISOString();
|
|
1123
|
+
const dataWithTimestamp = {
|
|
1124
|
+
...data,
|
|
1125
|
+
updatedAt: now,
|
|
1126
|
+
updatedAtZ: now
|
|
1127
|
+
};
|
|
1128
|
+
Object.entries(dataWithTimestamp).forEach(([key, value]) => {
|
|
1129
|
+
const parsedKey = parseSqlIdentifier(key, "column name");
|
|
1130
|
+
setColumns.push(`"${parsedKey}" = $${paramIndex++}`);
|
|
1131
|
+
setValues.push(this.prepareValue(value, key, tableName));
|
|
1132
|
+
});
|
|
1133
|
+
const whereConditions = [];
|
|
1134
|
+
const whereValues = [];
|
|
1135
|
+
Object.entries(keys).forEach(([key, value]) => {
|
|
1136
|
+
const parsedKey = parseSqlIdentifier(key, "column name");
|
|
1137
|
+
whereConditions.push(`"${parsedKey}" = $${paramIndex++}`);
|
|
1138
|
+
whereValues.push(this.prepareValue(value, key, tableName));
|
|
1139
|
+
});
|
|
1140
|
+
const tableName_ = getTableName({
|
|
1141
|
+
indexName: tableName,
|
|
1142
|
+
schemaName: getSchemaName(this.schemaName)
|
|
1143
|
+
});
|
|
1144
|
+
const sql = `UPDATE ${tableName_} SET ${setColumns.join(", ")} WHERE ${whereConditions.join(" AND ")}`;
|
|
1145
|
+
const values = [...setValues, ...whereValues];
|
|
1146
|
+
await withRetry(
|
|
1147
|
+
async () => {
|
|
1148
|
+
await this.client.none(sql, values);
|
|
1149
|
+
},
|
|
1150
|
+
{
|
|
1151
|
+
onRetry: (error, attempt, delay) => {
|
|
1152
|
+
this.logger?.warn?.(`update retry ${attempt} for table ${tableName} after ${delay}ms: ${error.message}`);
|
|
1153
|
+
}
|
|
1154
|
+
}
|
|
1155
|
+
).catch((error) => {
|
|
1156
|
+
throw new MastraError(
|
|
1157
|
+
{
|
|
1158
|
+
id: createStorageErrorId("DSQL", "UPDATE", "FAILED"),
|
|
1159
|
+
domain: ErrorDomain.STORAGE,
|
|
1160
|
+
category: ErrorCategory.THIRD_PARTY,
|
|
1161
|
+
details: {
|
|
1162
|
+
tableName
|
|
1163
|
+
}
|
|
1164
|
+
},
|
|
1165
|
+
error
|
|
1166
|
+
);
|
|
1167
|
+
});
|
|
1168
|
+
}
|
|
1169
|
+
async batchUpdate({
|
|
1170
|
+
tableName,
|
|
1171
|
+
updates
|
|
1172
|
+
}) {
|
|
1173
|
+
if (updates.length === 0) {
|
|
1174
|
+
return;
|
|
1175
|
+
}
|
|
1176
|
+
try {
|
|
1177
|
+
const { batches } = splitIntoBatches(updates, { maxRows: DEFAULT_MAX_ROWS_PER_BATCH });
|
|
1178
|
+
for (const batch of batches) {
|
|
1179
|
+
await withRetry(
|
|
1180
|
+
async () => {
|
|
1181
|
+
await this.client.tx(async (t) => {
|
|
1182
|
+
for (const { keys, data } of batch) {
|
|
1183
|
+
const setClauses = [];
|
|
1184
|
+
const whereConditions = [];
|
|
1185
|
+
const values = [];
|
|
1186
|
+
let paramIndex = 1;
|
|
1187
|
+
const now = (/* @__PURE__ */ new Date()).toISOString();
|
|
1188
|
+
const dataWithTimestamp = {
|
|
1189
|
+
...data,
|
|
1190
|
+
updatedAt: now,
|
|
1191
|
+
updatedAtZ: now
|
|
1192
|
+
};
|
|
1193
|
+
Object.entries(dataWithTimestamp).forEach(([key, value]) => {
|
|
1194
|
+
const parsedKey = parseSqlIdentifier(key, "column name");
|
|
1195
|
+
const preparedValue = this.prepareValue(value, key, tableName);
|
|
1196
|
+
setClauses.push(`"${parsedKey}" = $${paramIndex++}`);
|
|
1197
|
+
values.push(preparedValue);
|
|
1198
|
+
});
|
|
1199
|
+
Object.entries(keys).forEach(([key, value]) => {
|
|
1200
|
+
const parsedKey = parseSqlIdentifier(key, "column name");
|
|
1201
|
+
whereConditions.push(`"${parsedKey}" = $${paramIndex++}`);
|
|
1202
|
+
values.push(value);
|
|
1203
|
+
});
|
|
1204
|
+
const tableName_ = getTableName({
|
|
1205
|
+
indexName: tableName,
|
|
1206
|
+
schemaName: getSchemaName(this.schemaName)
|
|
1207
|
+
});
|
|
1208
|
+
const sql = `UPDATE ${tableName_} SET ${setClauses.join(", ")} WHERE ${whereConditions.join(" AND ")}`;
|
|
1209
|
+
await t.none(sql, values);
|
|
1210
|
+
}
|
|
1211
|
+
});
|
|
1212
|
+
},
|
|
1213
|
+
{
|
|
1214
|
+
onRetry: (error, attempt, delay) => {
|
|
1215
|
+
this.logger?.warn?.(
|
|
1216
|
+
`Batch update retry ${attempt} for table ${tableName} after ${delay}ms: ${error.message}`
|
|
1217
|
+
);
|
|
1218
|
+
}
|
|
1219
|
+
}
|
|
1220
|
+
);
|
|
1221
|
+
}
|
|
1222
|
+
} catch (error) {
|
|
1223
|
+
throw new MastraError(
|
|
1224
|
+
{
|
|
1225
|
+
id: createStorageErrorId("DSQL", "BATCH_UPDATE", "FAILED"),
|
|
1226
|
+
domain: ErrorDomain.STORAGE,
|
|
1227
|
+
category: ErrorCategory.THIRD_PARTY,
|
|
1228
|
+
details: {
|
|
1229
|
+
tableName,
|
|
1230
|
+
numberOfRecords: updates.length
|
|
1231
|
+
}
|
|
1232
|
+
},
|
|
1233
|
+
error
|
|
1234
|
+
);
|
|
1235
|
+
}
|
|
1236
|
+
}
|
|
1237
|
+
async batchDelete({ tableName, keys }) {
|
|
1238
|
+
if (keys.length === 0) {
|
|
1239
|
+
return;
|
|
1240
|
+
}
|
|
1241
|
+
try {
|
|
1242
|
+
const tableName_ = getTableName({
|
|
1243
|
+
indexName: tableName,
|
|
1244
|
+
schemaName: getSchemaName(this.schemaName)
|
|
1245
|
+
});
|
|
1246
|
+
const { batches } = splitIntoBatches(keys, { maxRows: DEFAULT_MAX_ROWS_PER_BATCH });
|
|
1247
|
+
for (const batch of batches) {
|
|
1248
|
+
await withRetry(
|
|
1249
|
+
async () => {
|
|
1250
|
+
await this.client.tx(async (t) => {
|
|
1251
|
+
for (const keySet of batch) {
|
|
1252
|
+
const conditions = [];
|
|
1253
|
+
const values = [];
|
|
1254
|
+
let paramIndex = 1;
|
|
1255
|
+
Object.entries(keySet).forEach(([key, value]) => {
|
|
1256
|
+
const parsedKey = parseSqlIdentifier(key, "column name");
|
|
1257
|
+
conditions.push(`"${parsedKey}" = $${paramIndex++}`);
|
|
1258
|
+
values.push(value);
|
|
1259
|
+
});
|
|
1260
|
+
const sql = `DELETE FROM ${tableName_} WHERE ${conditions.join(" AND ")}`;
|
|
1261
|
+
await t.none(sql, values);
|
|
1262
|
+
}
|
|
1263
|
+
});
|
|
1264
|
+
},
|
|
1265
|
+
{
|
|
1266
|
+
onRetry: (error, attempt, delay) => {
|
|
1267
|
+
this.logger?.warn?.(
|
|
1268
|
+
`Batch delete retry ${attempt} for table ${tableName} after ${delay}ms: ${error.message}`
|
|
1269
|
+
);
|
|
1270
|
+
}
|
|
1271
|
+
}
|
|
1272
|
+
);
|
|
1273
|
+
}
|
|
1274
|
+
} catch (error) {
|
|
1275
|
+
throw new MastraError(
|
|
1276
|
+
{
|
|
1277
|
+
id: createStorageErrorId("DSQL", "BATCH_DELETE", "FAILED"),
|
|
1278
|
+
domain: ErrorDomain.STORAGE,
|
|
1279
|
+
category: ErrorCategory.THIRD_PARTY,
|
|
1280
|
+
details: {
|
|
1281
|
+
tableName,
|
|
1282
|
+
numberOfRecords: keys.length
|
|
1283
|
+
}
|
|
1284
|
+
},
|
|
1285
|
+
error
|
|
1286
|
+
);
|
|
1287
|
+
}
|
|
1288
|
+
}
|
|
1289
|
+
/**
|
|
1290
|
+
* Delete all data from a table (alias for clearTable for consistency with other stores)
|
|
1291
|
+
*/
|
|
1292
|
+
async deleteData({ tableName }) {
|
|
1293
|
+
return this.clearTable({ tableName });
|
|
1294
|
+
}
|
|
1295
|
+
};
|
|
1296
|
+
function getSchemaName2(schema) {
|
|
1297
|
+
return schema ? `"${parseSqlIdentifier(schema, "schema name")}"` : void 0;
|
|
1298
|
+
}
|
|
1299
|
+
function getTableName2({ indexName, schemaName }) {
|
|
1300
|
+
const parsedIndexName = parseSqlIdentifier(indexName, "index name");
|
|
1301
|
+
const quotedIndexName = `"${parsedIndexName}"`;
|
|
1302
|
+
const quotedSchemaName = schemaName;
|
|
1303
|
+
return quotedSchemaName ? `${quotedSchemaName}.${quotedIndexName}` : quotedIndexName;
|
|
1304
|
+
}
|
|
1305
|
+
function transformFromSqlRow({
|
|
1306
|
+
tableName,
|
|
1307
|
+
sqlRow
|
|
1308
|
+
}) {
|
|
1309
|
+
const schema = TABLE_SCHEMAS[tableName];
|
|
1310
|
+
const result = {};
|
|
1311
|
+
Object.entries(sqlRow).forEach(([key, value]) => {
|
|
1312
|
+
const columnSchema = schema?.[key];
|
|
1313
|
+
if (columnSchema?.type === "jsonb" && typeof value === "string") {
|
|
1314
|
+
try {
|
|
1315
|
+
result[key] = JSON.parse(value);
|
|
1316
|
+
} catch {
|
|
1317
|
+
result[key] = value;
|
|
1318
|
+
}
|
|
1319
|
+
} else if (columnSchema?.type === "timestamp" && value && typeof value === "string") {
|
|
1320
|
+
result[key] = new Date(value);
|
|
1321
|
+
} else if (columnSchema?.type === "timestamp" && value instanceof Date) {
|
|
1322
|
+
result[key] = value;
|
|
1323
|
+
} else if (columnSchema?.type === "boolean") {
|
|
1324
|
+
result[key] = Boolean(value);
|
|
1325
|
+
} else {
|
|
1326
|
+
result[key] = value;
|
|
1327
|
+
}
|
|
1328
|
+
});
|
|
1329
|
+
return result;
|
|
1330
|
+
}
|
|
1331
|
+
|
|
1332
|
+
// src/storage/domains/agents/index.ts
|
|
1333
|
+
var AgentsDSQL = class _AgentsDSQL extends AgentsStorage {
|
|
1334
|
+
#db;
|
|
1335
|
+
#schema;
|
|
1336
|
+
#skipDefaultIndexes;
|
|
1337
|
+
#indexes;
|
|
1338
|
+
/** Tables managed by this domain */
|
|
1339
|
+
static MANAGED_TABLES = [TABLE_AGENTS, TABLE_AGENT_VERSIONS];
|
|
1340
|
+
constructor(config) {
|
|
1341
|
+
super();
|
|
1342
|
+
const { client, schemaName, skipDefaultIndexes, indexes } = resolveDsqlConfig(config);
|
|
1343
|
+
this.#db = new DsqlDB({ client, schemaName, skipDefaultIndexes });
|
|
1344
|
+
this.#schema = schemaName || "public";
|
|
1345
|
+
this.#skipDefaultIndexes = skipDefaultIndexes;
|
|
1346
|
+
this.#indexes = indexes?.filter((idx) => _AgentsDSQL.MANAGED_TABLES.includes(idx.table));
|
|
1347
|
+
}
|
|
1348
|
+
/**
|
|
1349
|
+
* Returns default index definitions for the agents domain tables.
|
|
1350
|
+
* Currently no default indexes are defined for agents.
|
|
1351
|
+
*/
|
|
1352
|
+
getDefaultIndexDefinitions() {
|
|
1353
|
+
return [];
|
|
1354
|
+
}
|
|
1355
|
+
/**
|
|
1356
|
+
* Creates default indexes for optimal query performance.
|
|
1357
|
+
* Currently no default indexes are defined for agents.
|
|
1358
|
+
*/
|
|
1359
|
+
async createDefaultIndexes() {
|
|
1360
|
+
if (this.#skipDefaultIndexes) {
|
|
1361
|
+
return;
|
|
1362
|
+
}
|
|
1363
|
+
}
|
|
1364
|
+
async init() {
|
|
1365
|
+
await this.#db.createTable({ tableName: TABLE_AGENTS, schema: TABLE_SCHEMAS[TABLE_AGENTS] });
|
|
1366
|
+
await this.#db.createTable({ tableName: TABLE_AGENT_VERSIONS, schema: TABLE_SCHEMAS[TABLE_AGENT_VERSIONS] });
|
|
1367
|
+
await this.createDefaultIndexes();
|
|
1368
|
+
await this.createCustomIndexes();
|
|
1369
|
+
}
|
|
1370
|
+
/**
|
|
1371
|
+
* Creates custom user-defined indexes for this domain's tables.
|
|
1372
|
+
*/
|
|
1373
|
+
async createCustomIndexes() {
|
|
1374
|
+
if (!this.#indexes || this.#indexes.length === 0) {
|
|
1375
|
+
return;
|
|
1376
|
+
}
|
|
1377
|
+
for (const indexDef of this.#indexes) {
|
|
1378
|
+
try {
|
|
1379
|
+
await this.#db.createIndex(indexDef);
|
|
1380
|
+
} catch (error) {
|
|
1381
|
+
this.logger?.warn?.(`Failed to create custom index ${indexDef.name}:`, error);
|
|
1382
|
+
}
|
|
1383
|
+
}
|
|
1384
|
+
}
|
|
1385
|
+
async dangerouslyClearAll() {
|
|
1386
|
+
await this.#db.clearTable({ tableName: TABLE_AGENT_VERSIONS });
|
|
1387
|
+
await this.#db.clearTable({ tableName: TABLE_AGENTS });
|
|
1388
|
+
}
|
|
1389
|
+
parseJson(value, fieldName) {
|
|
1390
|
+
if (!value) return void 0;
|
|
1391
|
+
if (typeof value !== "string") return value;
|
|
1392
|
+
try {
|
|
1393
|
+
return JSON.parse(value);
|
|
1394
|
+
} catch (error) {
|
|
1395
|
+
const details = {
|
|
1396
|
+
value: value.length > 100 ? value.substring(0, 100) + "..." : value
|
|
1397
|
+
};
|
|
1398
|
+
if (fieldName) {
|
|
1399
|
+
details.field = fieldName;
|
|
1400
|
+
}
|
|
1401
|
+
throw new MastraError(
|
|
1402
|
+
{
|
|
1403
|
+
id: createStorageErrorId("DSQL", "PARSE_JSON", "INVALID_JSON"),
|
|
1404
|
+
domain: ErrorDomain.STORAGE,
|
|
1405
|
+
category: ErrorCategory.SYSTEM,
|
|
1406
|
+
text: `Failed to parse JSON${fieldName ? ` for field "${fieldName}"` : ""}: ${error instanceof Error ? error.message : "Unknown error"}`,
|
|
1407
|
+
details
|
|
1408
|
+
},
|
|
1409
|
+
error
|
|
1410
|
+
);
|
|
1411
|
+
}
|
|
1412
|
+
}
|
|
1413
|
+
parseRow(row) {
|
|
1414
|
+
return {
|
|
1415
|
+
id: row.id,
|
|
1416
|
+
status: row.status,
|
|
1417
|
+
activeVersionId: row.activeVersionId,
|
|
1418
|
+
authorId: row.authorId,
|
|
1419
|
+
metadata: this.parseJson(row.metadata, "metadata"),
|
|
1420
|
+
createdAt: row.createdAtZ || row.createdAt,
|
|
1421
|
+
updatedAt: row.updatedAtZ || row.updatedAt
|
|
1422
|
+
};
|
|
1423
|
+
}
|
|
1424
|
+
async getById(id) {
|
|
1425
|
+
try {
|
|
1426
|
+
const tableName = getTableName2({ indexName: TABLE_AGENTS, schemaName: getSchemaName2(this.#schema) });
|
|
1427
|
+
const result = await this.#db.client.oneOrNone(`SELECT * FROM ${tableName} WHERE id = $1`, [id]);
|
|
1428
|
+
if (!result) {
|
|
1429
|
+
return null;
|
|
1430
|
+
}
|
|
1431
|
+
return this.parseRow(result);
|
|
1432
|
+
} catch (error) {
|
|
1433
|
+
if (error instanceof MastraError) throw error;
|
|
1434
|
+
throw new MastraError(
|
|
1435
|
+
{
|
|
1436
|
+
id: createStorageErrorId("DSQL", "GET_AGENT_BY_ID", "FAILED"),
|
|
1437
|
+
domain: ErrorDomain.STORAGE,
|
|
1438
|
+
category: ErrorCategory.THIRD_PARTY,
|
|
1439
|
+
details: { agentId: id }
|
|
1440
|
+
},
|
|
1441
|
+
error
|
|
1442
|
+
);
|
|
1443
|
+
}
|
|
1444
|
+
}
|
|
1445
|
+
async create(input) {
|
|
1446
|
+
const { agent } = input;
|
|
1447
|
+
const tableName = getTableName2({ indexName: TABLE_AGENTS, schemaName: getSchemaName2(this.#schema) });
|
|
1448
|
+
const now = /* @__PURE__ */ new Date();
|
|
1449
|
+
const nowIso = now.toISOString();
|
|
1450
|
+
try {
|
|
1451
|
+
await withRetry(
|
|
1452
|
+
() => this.#db.client.none(
|
|
1453
|
+
`INSERT INTO ${tableName} (
|
|
1454
|
+
id, status, "authorId", metadata,
|
|
1455
|
+
"activeVersionId",
|
|
1456
|
+
"createdAt", "createdAtZ", "updatedAt", "updatedAtZ"
|
|
1457
|
+
) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9)`,
|
|
1458
|
+
[
|
|
1459
|
+
agent.id,
|
|
1460
|
+
"draft",
|
|
1461
|
+
agent.authorId ?? null,
|
|
1462
|
+
agent.metadata ? JSON.stringify(agent.metadata) : null,
|
|
1463
|
+
null,
|
|
1464
|
+
// activeVersionId starts as null
|
|
1465
|
+
nowIso,
|
|
1466
|
+
nowIso,
|
|
1467
|
+
nowIso,
|
|
1468
|
+
nowIso
|
|
1469
|
+
]
|
|
1470
|
+
)
|
|
1471
|
+
);
|
|
1472
|
+
const { id: _id, authorId: _authorId, metadata: _metadata, ...snapshotConfig } = agent;
|
|
1473
|
+
const versionId = crypto.randomUUID();
|
|
1474
|
+
await this.createVersion({
|
|
1475
|
+
id: versionId,
|
|
1476
|
+
agentId: agent.id,
|
|
1477
|
+
versionNumber: 1,
|
|
1478
|
+
...snapshotConfig,
|
|
1479
|
+
changedFields: Object.keys(snapshotConfig),
|
|
1480
|
+
changeMessage: "Initial version"
|
|
1481
|
+
});
|
|
1482
|
+
return {
|
|
1483
|
+
id: agent.id,
|
|
1484
|
+
status: "draft",
|
|
1485
|
+
activeVersionId: void 0,
|
|
1486
|
+
authorId: agent.authorId,
|
|
1487
|
+
metadata: agent.metadata,
|
|
1488
|
+
createdAt: now,
|
|
1489
|
+
updatedAt: now
|
|
1490
|
+
};
|
|
1491
|
+
} catch (error) {
|
|
1492
|
+
if (error instanceof MastraError) throw error;
|
|
1493
|
+
try {
|
|
1494
|
+
await this.#db.client.none(
|
|
1495
|
+
`DELETE FROM ${tableName} WHERE id = $1 AND status = 'draft' AND "activeVersionId" IS NULL`,
|
|
1496
|
+
[agent.id]
|
|
1497
|
+
);
|
|
1498
|
+
} catch {
|
|
1499
|
+
}
|
|
1500
|
+
throw new MastraError(
|
|
1501
|
+
{
|
|
1502
|
+
id: createStorageErrorId("DSQL", "CREATE_AGENT", "FAILED"),
|
|
1503
|
+
domain: ErrorDomain.STORAGE,
|
|
1504
|
+
category: ErrorCategory.THIRD_PARTY,
|
|
1505
|
+
details: { agentId: agent.id }
|
|
1506
|
+
},
|
|
1507
|
+
error
|
|
1508
|
+
);
|
|
1509
|
+
}
|
|
1510
|
+
}
|
|
1511
|
+
async update(input) {
|
|
1512
|
+
const { id, ...updates } = input;
|
|
1513
|
+
try {
|
|
1514
|
+
const tableName = getTableName2({ indexName: TABLE_AGENTS, schemaName: getSchemaName2(this.#schema) });
|
|
1515
|
+
const existingAgent = await this.getById(id);
|
|
1516
|
+
if (!existingAgent) {
|
|
1517
|
+
throw new MastraError({
|
|
1518
|
+
id: createStorageErrorId("DSQL", "UPDATE_AGENT", "NOT_FOUND"),
|
|
1519
|
+
domain: ErrorDomain.STORAGE,
|
|
1520
|
+
category: ErrorCategory.USER,
|
|
1521
|
+
text: `Agent ${id} not found`,
|
|
1522
|
+
details: { agentId: id }
|
|
1523
|
+
});
|
|
1524
|
+
}
|
|
1525
|
+
const { authorId, activeVersionId, metadata, status } = updates;
|
|
1526
|
+
const setClauses = [];
|
|
1527
|
+
const values = [];
|
|
1528
|
+
let paramIndex = 1;
|
|
1529
|
+
if (authorId !== void 0) {
|
|
1530
|
+
setClauses.push(`"authorId" = $${paramIndex++}`);
|
|
1531
|
+
values.push(authorId);
|
|
1532
|
+
}
|
|
1533
|
+
if (activeVersionId !== void 0) {
|
|
1534
|
+
setClauses.push(`"activeVersionId" = $${paramIndex++}`);
|
|
1535
|
+
values.push(activeVersionId);
|
|
1536
|
+
}
|
|
1537
|
+
if (status !== void 0) {
|
|
1538
|
+
setClauses.push(`status = $${paramIndex++}`);
|
|
1539
|
+
values.push(status);
|
|
1540
|
+
}
|
|
1541
|
+
if (metadata !== void 0) {
|
|
1542
|
+
setClauses.push(`metadata = $${paramIndex++}`);
|
|
1543
|
+
values.push(JSON.stringify(metadata));
|
|
1544
|
+
}
|
|
1545
|
+
const now = (/* @__PURE__ */ new Date()).toISOString();
|
|
1546
|
+
setClauses.push(`"updatedAt" = $${paramIndex++}`);
|
|
1547
|
+
values.push(now);
|
|
1548
|
+
setClauses.push(`"updatedAtZ" = $${paramIndex++}`);
|
|
1549
|
+
values.push(now);
|
|
1550
|
+
values.push(id);
|
|
1551
|
+
await withRetry(
|
|
1552
|
+
() => this.#db.client.none(`UPDATE ${tableName} SET ${setClauses.join(", ")} WHERE id = ${paramIndex}`, values)
|
|
1553
|
+
);
|
|
1554
|
+
const updatedAgent = await this.getById(id);
|
|
1555
|
+
if (!updatedAgent) {
|
|
1556
|
+
throw new MastraError({
|
|
1557
|
+
id: createStorageErrorId("DSQL", "UPDATE_AGENT", "NOT_FOUND_AFTER_UPDATE"),
|
|
1558
|
+
domain: ErrorDomain.STORAGE,
|
|
1559
|
+
category: ErrorCategory.SYSTEM,
|
|
1560
|
+
text: `Agent ${id} not found after update`,
|
|
1561
|
+
details: { agentId: id }
|
|
1562
|
+
});
|
|
1563
|
+
}
|
|
1564
|
+
return updatedAgent;
|
|
1565
|
+
} catch (error) {
|
|
1566
|
+
if (error instanceof MastraError) throw error;
|
|
1567
|
+
throw new MastraError(
|
|
1568
|
+
{
|
|
1569
|
+
id: createStorageErrorId("DSQL", "UPDATE_AGENT", "FAILED"),
|
|
1570
|
+
domain: ErrorDomain.STORAGE,
|
|
1571
|
+
category: ErrorCategory.THIRD_PARTY,
|
|
1572
|
+
details: { agentId: id }
|
|
1573
|
+
},
|
|
1574
|
+
error
|
|
1575
|
+
);
|
|
1576
|
+
}
|
|
1577
|
+
}
|
|
1578
|
+
async delete(id) {
|
|
1579
|
+
try {
|
|
1580
|
+
const tableName = getTableName2({ indexName: TABLE_AGENTS, schemaName: getSchemaName2(this.#schema) });
|
|
1581
|
+
await this.deleteVersionsByParentId(id);
|
|
1582
|
+
await this.#db.client.none(`DELETE FROM ${tableName} WHERE id = $1`, [id]);
|
|
1583
|
+
} catch (error) {
|
|
1584
|
+
if (error instanceof MastraError) throw error;
|
|
1585
|
+
throw new MastraError(
|
|
1586
|
+
{
|
|
1587
|
+
id: createStorageErrorId("DSQL", "DELETE_AGENT", "FAILED"),
|
|
1588
|
+
domain: ErrorDomain.STORAGE,
|
|
1589
|
+
category: ErrorCategory.THIRD_PARTY,
|
|
1590
|
+
details: { agentId: id }
|
|
1591
|
+
},
|
|
1592
|
+
error
|
|
1593
|
+
);
|
|
1594
|
+
}
|
|
1595
|
+
}
|
|
1596
|
+
async list(args) {
|
|
1597
|
+
const { page = 0, perPage: perPageInput, orderBy, authorId, metadata, status } = args || {};
|
|
1598
|
+
const { field, direction } = this.parseOrderBy(orderBy);
|
|
1599
|
+
if (page < 0) {
|
|
1600
|
+
throw new MastraError(
|
|
1601
|
+
{
|
|
1602
|
+
id: createStorageErrorId("DSQL", "LIST_AGENTS", "INVALID_PAGE"),
|
|
1603
|
+
domain: ErrorDomain.STORAGE,
|
|
1604
|
+
category: ErrorCategory.USER,
|
|
1605
|
+
details: { page }
|
|
1606
|
+
},
|
|
1607
|
+
new Error("page must be >= 0")
|
|
1608
|
+
);
|
|
1609
|
+
}
|
|
1610
|
+
const perPage = normalizePerPage(perPageInput, 100);
|
|
1611
|
+
const { offset, perPage: perPageForResponse } = calculatePagination(page, perPageInput, perPage);
|
|
1612
|
+
try {
|
|
1613
|
+
const tableName = getTableName2({ indexName: TABLE_AGENTS, schemaName: getSchemaName2(this.#schema) });
|
|
1614
|
+
const conditions = [];
|
|
1615
|
+
const queryParams = [];
|
|
1616
|
+
let paramIdx = 1;
|
|
1617
|
+
if (status) {
|
|
1618
|
+
conditions.push(`status = $${paramIdx++}`);
|
|
1619
|
+
queryParams.push(status);
|
|
1620
|
+
}
|
|
1621
|
+
if (authorId !== void 0) {
|
|
1622
|
+
conditions.push(`"authorId" = $${paramIdx++}`);
|
|
1623
|
+
queryParams.push(authorId);
|
|
1624
|
+
}
|
|
1625
|
+
if (metadata && Object.keys(metadata).length > 0) {
|
|
1626
|
+
conditions.push(`metadata::text = $${paramIdx++}`);
|
|
1627
|
+
queryParams.push(JSON.stringify(metadata));
|
|
1628
|
+
}
|
|
1629
|
+
const whereClause = conditions.length > 0 ? `WHERE ${conditions.join(" AND ")}` : "";
|
|
1630
|
+
const countResult = await this.#db.client.one(
|
|
1631
|
+
`SELECT COUNT(*) as count FROM ${tableName} ${whereClause}`,
|
|
1632
|
+
queryParams
|
|
1633
|
+
);
|
|
1634
|
+
const total = parseInt(countResult.count, 10);
|
|
1635
|
+
if (total === 0) {
|
|
1636
|
+
return {
|
|
1637
|
+
agents: [],
|
|
1638
|
+
total: 0,
|
|
1639
|
+
page,
|
|
1640
|
+
perPage: perPageForResponse,
|
|
1641
|
+
hasMore: false
|
|
1642
|
+
};
|
|
1643
|
+
}
|
|
1644
|
+
const limitValue = perPageInput === false ? total : perPage;
|
|
1645
|
+
const dataResult = await this.#db.client.manyOrNone(
|
|
1646
|
+
`SELECT * FROM ${tableName} ${whereClause} ORDER BY "${field}" ${direction} LIMIT $${paramIdx++} OFFSET $${paramIdx++}`,
|
|
1647
|
+
[...queryParams, limitValue, offset]
|
|
1648
|
+
);
|
|
1649
|
+
const agents = (dataResult || []).map((row) => this.parseRow(row));
|
|
1650
|
+
return {
|
|
1651
|
+
agents,
|
|
1652
|
+
total,
|
|
1653
|
+
page,
|
|
1654
|
+
perPage: perPageForResponse,
|
|
1655
|
+
hasMore: perPageInput === false ? false : offset + perPage < total
|
|
1656
|
+
};
|
|
1657
|
+
} catch (error) {
|
|
1658
|
+
if (error instanceof MastraError) throw error;
|
|
1659
|
+
throw new MastraError(
|
|
1660
|
+
{
|
|
1661
|
+
id: createStorageErrorId("DSQL", "LIST_AGENTS", "FAILED"),
|
|
1662
|
+
domain: ErrorDomain.STORAGE,
|
|
1663
|
+
category: ErrorCategory.THIRD_PARTY
|
|
1664
|
+
},
|
|
1665
|
+
error
|
|
1666
|
+
);
|
|
1667
|
+
}
|
|
1668
|
+
}
|
|
1669
|
+
// ==========================================================================
|
|
1670
|
+
// Agent Version Methods
|
|
1671
|
+
// ==========================================================================
|
|
1672
|
+
async createVersion(input) {
|
|
1673
|
+
try {
|
|
1674
|
+
const tableName = getTableName2({ indexName: TABLE_AGENT_VERSIONS, schemaName: getSchemaName2(this.#schema) });
|
|
1675
|
+
const now = /* @__PURE__ */ new Date();
|
|
1676
|
+
const nowIso = now.toISOString();
|
|
1677
|
+
await withRetry(
|
|
1678
|
+
() => this.#db.client.none(
|
|
1679
|
+
`INSERT INTO ${tableName} (
|
|
1680
|
+
id, "agentId", "versionNumber",
|
|
1681
|
+
name, description, instructions, model, tools,
|
|
1682
|
+
"defaultOptions", workflows, agents, "integrationTools",
|
|
1683
|
+
"inputProcessors", "outputProcessors", memory, scorers,
|
|
1684
|
+
"mcpClients", "requestContextSchema", workspace, skills, "skillsFormat",
|
|
1685
|
+
"changedFields", "changeMessage",
|
|
1686
|
+
"createdAt", "createdAtZ"
|
|
1687
|
+
) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17, $18, $19, $20, $21, $22, $23, $24, $25)`,
|
|
1688
|
+
[
|
|
1689
|
+
input.id,
|
|
1690
|
+
input.agentId,
|
|
1691
|
+
input.versionNumber,
|
|
1692
|
+
input.name,
|
|
1693
|
+
input.description ?? null,
|
|
1694
|
+
this.serializeInstructions(input.instructions),
|
|
1695
|
+
JSON.stringify(input.model),
|
|
1696
|
+
input.tools ? JSON.stringify(input.tools) : null,
|
|
1697
|
+
input.defaultOptions ? JSON.stringify(input.defaultOptions) : null,
|
|
1698
|
+
input.workflows ? JSON.stringify(input.workflows) : null,
|
|
1699
|
+
input.agents ? JSON.stringify(input.agents) : null,
|
|
1700
|
+
input.integrationTools ? JSON.stringify(input.integrationTools) : null,
|
|
1701
|
+
input.inputProcessors ? JSON.stringify(input.inputProcessors) : null,
|
|
1702
|
+
input.outputProcessors ? JSON.stringify(input.outputProcessors) : null,
|
|
1703
|
+
input.memory ? JSON.stringify(input.memory) : null,
|
|
1704
|
+
input.scorers ? JSON.stringify(input.scorers) : null,
|
|
1705
|
+
input.mcpClients ? JSON.stringify(input.mcpClients) : null,
|
|
1706
|
+
input.requestContextSchema ? JSON.stringify(input.requestContextSchema) : null,
|
|
1707
|
+
input.workspace ? JSON.stringify(input.workspace) : null,
|
|
1708
|
+
input.skills ? JSON.stringify(input.skills) : null,
|
|
1709
|
+
input.skillsFormat ?? null,
|
|
1710
|
+
input.changedFields ? JSON.stringify(input.changedFields) : null,
|
|
1711
|
+
input.changeMessage ?? null,
|
|
1712
|
+
nowIso,
|
|
1713
|
+
nowIso
|
|
1714
|
+
]
|
|
1715
|
+
)
|
|
1716
|
+
);
|
|
1717
|
+
return {
|
|
1718
|
+
...input,
|
|
1719
|
+
createdAt: now
|
|
1720
|
+
};
|
|
1721
|
+
} catch (error) {
|
|
1722
|
+
if (error instanceof MastraError) throw error;
|
|
1723
|
+
throw new MastraError(
|
|
1724
|
+
{
|
|
1725
|
+
id: createStorageErrorId("DSQL", "CREATE_VERSION", "FAILED"),
|
|
1726
|
+
domain: ErrorDomain.STORAGE,
|
|
1727
|
+
category: ErrorCategory.THIRD_PARTY,
|
|
1728
|
+
details: { versionId: input.id, agentId: input.agentId }
|
|
1729
|
+
},
|
|
1730
|
+
error
|
|
1731
|
+
);
|
|
1732
|
+
}
|
|
1733
|
+
}
|
|
1734
|
+
async getVersion(id) {
|
|
1735
|
+
try {
|
|
1736
|
+
const tableName = getTableName2({ indexName: TABLE_AGENT_VERSIONS, schemaName: getSchemaName2(this.#schema) });
|
|
1737
|
+
const result = await this.#db.client.oneOrNone(`SELECT * FROM ${tableName} WHERE id = $1`, [id]);
|
|
1738
|
+
if (!result) {
|
|
1739
|
+
return null;
|
|
1740
|
+
}
|
|
1741
|
+
return this.parseVersionRow(result);
|
|
1742
|
+
} catch (error) {
|
|
1743
|
+
if (error instanceof MastraError) throw error;
|
|
1744
|
+
throw new MastraError(
|
|
1745
|
+
{
|
|
1746
|
+
id: createStorageErrorId("DSQL", "GET_VERSION", "FAILED"),
|
|
1747
|
+
domain: ErrorDomain.STORAGE,
|
|
1748
|
+
category: ErrorCategory.THIRD_PARTY,
|
|
1749
|
+
details: { versionId: id }
|
|
1750
|
+
},
|
|
1751
|
+
error
|
|
1752
|
+
);
|
|
1753
|
+
}
|
|
1754
|
+
}
|
|
1755
|
+
async getVersionByNumber(agentId, versionNumber) {
|
|
1756
|
+
try {
|
|
1757
|
+
const tableName = getTableName2({ indexName: TABLE_AGENT_VERSIONS, schemaName: getSchemaName2(this.#schema) });
|
|
1758
|
+
const result = await this.#db.client.oneOrNone(
|
|
1759
|
+
`SELECT * FROM ${tableName} WHERE "agentId" = $1 AND "versionNumber" = $2`,
|
|
1760
|
+
[agentId, versionNumber]
|
|
1761
|
+
);
|
|
1762
|
+
if (!result) {
|
|
1763
|
+
return null;
|
|
1764
|
+
}
|
|
1765
|
+
return this.parseVersionRow(result);
|
|
1766
|
+
} catch (error) {
|
|
1767
|
+
if (error instanceof MastraError) throw error;
|
|
1768
|
+
throw new MastraError(
|
|
1769
|
+
{
|
|
1770
|
+
id: createStorageErrorId("DSQL", "GET_VERSION_BY_NUMBER", "FAILED"),
|
|
1771
|
+
domain: ErrorDomain.STORAGE,
|
|
1772
|
+
category: ErrorCategory.THIRD_PARTY,
|
|
1773
|
+
details: { agentId, versionNumber }
|
|
1774
|
+
},
|
|
1775
|
+
error
|
|
1776
|
+
);
|
|
1777
|
+
}
|
|
1778
|
+
}
|
|
1779
|
+
async getLatestVersion(agentId) {
|
|
1780
|
+
try {
|
|
1781
|
+
const tableName = getTableName2({ indexName: TABLE_AGENT_VERSIONS, schemaName: getSchemaName2(this.#schema) });
|
|
1782
|
+
const result = await this.#db.client.oneOrNone(
|
|
1783
|
+
`SELECT * FROM ${tableName} WHERE "agentId" = $1 ORDER BY "versionNumber" DESC LIMIT 1`,
|
|
1784
|
+
[agentId]
|
|
1785
|
+
);
|
|
1786
|
+
if (!result) {
|
|
1787
|
+
return null;
|
|
1788
|
+
}
|
|
1789
|
+
return this.parseVersionRow(result);
|
|
1790
|
+
} catch (error) {
|
|
1791
|
+
if (error instanceof MastraError) throw error;
|
|
1792
|
+
throw new MastraError(
|
|
1793
|
+
{
|
|
1794
|
+
id: createStorageErrorId("DSQL", "GET_LATEST_VERSION", "FAILED"),
|
|
1795
|
+
domain: ErrorDomain.STORAGE,
|
|
1796
|
+
category: ErrorCategory.THIRD_PARTY,
|
|
1797
|
+
details: { agentId }
|
|
1798
|
+
},
|
|
1799
|
+
error
|
|
1800
|
+
);
|
|
1801
|
+
}
|
|
1802
|
+
}
|
|
1803
|
+
async listVersions(input) {
|
|
1804
|
+
const { agentId, page = 0, perPage: perPageInput, orderBy } = input;
|
|
1805
|
+
const perPage = normalizePerPage(perPageInput, 100);
|
|
1806
|
+
const { offset, perPage: perPageForResponse } = calculatePagination(page, perPageInput, perPage);
|
|
1807
|
+
const sortField = orderBy?.field || "versionNumber";
|
|
1808
|
+
const sortDirection = orderBy?.direction || "DESC";
|
|
1809
|
+
try {
|
|
1810
|
+
const tableName = getTableName2({ indexName: TABLE_AGENT_VERSIONS, schemaName: getSchemaName2(this.#schema) });
|
|
1811
|
+
const countResult = await this.#db.client.one(`SELECT COUNT(*) as count FROM ${tableName} WHERE "agentId" = $1`, [
|
|
1812
|
+
agentId
|
|
1813
|
+
]);
|
|
1814
|
+
const total = parseInt(countResult.count, 10);
|
|
1815
|
+
if (total === 0) {
|
|
1816
|
+
return {
|
|
1817
|
+
versions: [],
|
|
1818
|
+
total: 0,
|
|
1819
|
+
page,
|
|
1820
|
+
perPage: perPageForResponse,
|
|
1821
|
+
hasMore: false
|
|
1822
|
+
};
|
|
1823
|
+
}
|
|
1824
|
+
const limitValue = perPageInput === false ? total : perPage;
|
|
1825
|
+
const rows = await this.#db.client.manyOrNone(
|
|
1826
|
+
`SELECT * FROM ${tableName} WHERE "agentId" = $1 ORDER BY "${sortField}" ${sortDirection} LIMIT $2 OFFSET $3`,
|
|
1827
|
+
[agentId, limitValue, offset]
|
|
1828
|
+
);
|
|
1829
|
+
const versions = (rows || []).map((row) => this.parseVersionRow(row));
|
|
1830
|
+
return {
|
|
1831
|
+
versions,
|
|
1832
|
+
total,
|
|
1833
|
+
page,
|
|
1834
|
+
perPage: perPageForResponse,
|
|
1835
|
+
hasMore: perPageInput === false ? false : offset + perPage < total
|
|
1836
|
+
};
|
|
1837
|
+
} catch (error) {
|
|
1838
|
+
if (error instanceof MastraError) throw error;
|
|
1839
|
+
throw new MastraError(
|
|
1840
|
+
{
|
|
1841
|
+
id: createStorageErrorId("DSQL", "LIST_VERSIONS", "FAILED"),
|
|
1842
|
+
domain: ErrorDomain.STORAGE,
|
|
1843
|
+
category: ErrorCategory.THIRD_PARTY,
|
|
1844
|
+
details: { agentId }
|
|
1845
|
+
},
|
|
1846
|
+
error
|
|
1847
|
+
);
|
|
1848
|
+
}
|
|
1849
|
+
}
|
|
1850
|
+
async deleteVersion(id) {
|
|
1851
|
+
try {
|
|
1852
|
+
const tableName = getTableName2({ indexName: TABLE_AGENT_VERSIONS, schemaName: getSchemaName2(this.#schema) });
|
|
1853
|
+
await this.#db.client.none(`DELETE FROM ${tableName} WHERE id = $1`, [id]);
|
|
1854
|
+
} catch (error) {
|
|
1855
|
+
if (error instanceof MastraError) throw error;
|
|
1856
|
+
throw new MastraError(
|
|
1857
|
+
{
|
|
1858
|
+
id: createStorageErrorId("DSQL", "DELETE_VERSION", "FAILED"),
|
|
1859
|
+
domain: ErrorDomain.STORAGE,
|
|
1860
|
+
category: ErrorCategory.THIRD_PARTY,
|
|
1861
|
+
details: { versionId: id }
|
|
1862
|
+
},
|
|
1863
|
+
error
|
|
1864
|
+
);
|
|
1865
|
+
}
|
|
1866
|
+
}
|
|
1867
|
+
async deleteVersionsByParentId(agentId) {
|
|
1868
|
+
try {
|
|
1869
|
+
const tableName = getTableName2({ indexName: TABLE_AGENT_VERSIONS, schemaName: getSchemaName2(this.#schema) });
|
|
1870
|
+
await this.#db.client.none(`DELETE FROM ${tableName} WHERE "agentId" = $1`, [agentId]);
|
|
1871
|
+
} catch (error) {
|
|
1872
|
+
if (error instanceof MastraError) throw error;
|
|
1873
|
+
throw new MastraError(
|
|
1874
|
+
{
|
|
1875
|
+
id: createStorageErrorId("DSQL", "DELETE_VERSIONS_BY_PARENT", "FAILED"),
|
|
1876
|
+
domain: ErrorDomain.STORAGE,
|
|
1877
|
+
category: ErrorCategory.THIRD_PARTY,
|
|
1878
|
+
details: { agentId }
|
|
1879
|
+
},
|
|
1880
|
+
error
|
|
1881
|
+
);
|
|
1882
|
+
}
|
|
1883
|
+
}
|
|
1884
|
+
async countVersions(agentId) {
|
|
1885
|
+
try {
|
|
1886
|
+
const tableName = getTableName2({ indexName: TABLE_AGENT_VERSIONS, schemaName: getSchemaName2(this.#schema) });
|
|
1887
|
+
const result = await this.#db.client.one(`SELECT COUNT(*) as count FROM ${tableName} WHERE "agentId" = $1`, [
|
|
1888
|
+
agentId
|
|
1889
|
+
]);
|
|
1890
|
+
return parseInt(result.count, 10);
|
|
1891
|
+
} catch (error) {
|
|
1892
|
+
if (error instanceof MastraError) throw error;
|
|
1893
|
+
throw new MastraError(
|
|
1894
|
+
{
|
|
1895
|
+
id: createStorageErrorId("DSQL", "COUNT_VERSIONS", "FAILED"),
|
|
1896
|
+
domain: ErrorDomain.STORAGE,
|
|
1897
|
+
category: ErrorCategory.THIRD_PARTY,
|
|
1898
|
+
details: { agentId }
|
|
1899
|
+
},
|
|
1900
|
+
error
|
|
1901
|
+
);
|
|
1902
|
+
}
|
|
1903
|
+
}
|
|
1904
|
+
// ==========================================================================
|
|
1905
|
+
// Private Helpers
|
|
1906
|
+
// ==========================================================================
|
|
1907
|
+
serializeInstructions(instructions) {
|
|
1908
|
+
if (instructions == null) return void 0;
|
|
1909
|
+
return Array.isArray(instructions) ? JSON.stringify(instructions) : instructions;
|
|
1910
|
+
}
|
|
1911
|
+
deserializeInstructions(raw) {
|
|
1912
|
+
if (!raw) return "";
|
|
1913
|
+
try {
|
|
1914
|
+
const parsed = JSON.parse(raw);
|
|
1915
|
+
if (Array.isArray(parsed)) return parsed;
|
|
1916
|
+
} catch {
|
|
1917
|
+
}
|
|
1918
|
+
return raw;
|
|
1919
|
+
}
|
|
1920
|
+
parseVersionRow(row) {
|
|
1921
|
+
return {
|
|
1922
|
+
id: row.id,
|
|
1923
|
+
agentId: row.agentId,
|
|
1924
|
+
versionNumber: row.versionNumber,
|
|
1925
|
+
name: row.name,
|
|
1926
|
+
description: row.description,
|
|
1927
|
+
instructions: this.deserializeInstructions(row.instructions),
|
|
1928
|
+
model: this.parseJson(row.model, "model"),
|
|
1929
|
+
tools: this.parseJson(row.tools, "tools"),
|
|
1930
|
+
defaultOptions: this.parseJson(row.defaultOptions, "defaultOptions"),
|
|
1931
|
+
workflows: this.parseJson(row.workflows, "workflows"),
|
|
1932
|
+
agents: this.parseJson(row.agents, "agents"),
|
|
1933
|
+
integrationTools: this.parseJson(row.integrationTools, "integrationTools"),
|
|
1934
|
+
inputProcessors: this.parseJson(row.inputProcessors, "inputProcessors"),
|
|
1935
|
+
outputProcessors: this.parseJson(row.outputProcessors, "outputProcessors"),
|
|
1936
|
+
memory: this.parseJson(row.memory, "memory"),
|
|
1937
|
+
scorers: this.parseJson(row.scorers, "scorers"),
|
|
1938
|
+
mcpClients: this.parseJson(row.mcpClients, "mcpClients"),
|
|
1939
|
+
requestContextSchema: this.parseJson(row.requestContextSchema, "requestContextSchema"),
|
|
1940
|
+
workspace: this.parseJson(row.workspace, "workspace"),
|
|
1941
|
+
skills: this.parseJson(row.skills, "skills"),
|
|
1942
|
+
skillsFormat: row.skillsFormat,
|
|
1943
|
+
changedFields: this.parseJson(row.changedFields, "changedFields"),
|
|
1944
|
+
changeMessage: row.changeMessage,
|
|
1945
|
+
createdAt: row.createdAtZ || row.createdAt
|
|
1946
|
+
};
|
|
1947
|
+
}
|
|
1948
|
+
};
|
|
1949
|
+
function inPlaceholders(count, startIndex = 1) {
|
|
1950
|
+
return Array.from({ length: count }, (_, i) => `$${i + startIndex}`).join(", ");
|
|
1951
|
+
}
|
|
1952
|
+
var MemoryDSQL = class _MemoryDSQL extends MemoryStorage {
|
|
1953
|
+
#db;
|
|
1954
|
+
#schema;
|
|
1955
|
+
#skipDefaultIndexes;
|
|
1956
|
+
#indexes;
|
|
1957
|
+
/** Tables managed by this domain */
|
|
1958
|
+
static MANAGED_TABLES = [TABLE_THREADS, TABLE_MESSAGES, TABLE_RESOURCES];
|
|
1959
|
+
constructor(config) {
|
|
1960
|
+
super();
|
|
1961
|
+
const { client, schemaName, skipDefaultIndexes, indexes } = resolveDsqlConfig(config);
|
|
1962
|
+
this.#db = new DsqlDB({ client, schemaName, skipDefaultIndexes });
|
|
1963
|
+
this.#schema = schemaName || "public";
|
|
1964
|
+
this.#skipDefaultIndexes = skipDefaultIndexes;
|
|
1965
|
+
this.#indexes = indexes?.filter((idx) => _MemoryDSQL.MANAGED_TABLES.includes(idx.table));
|
|
1966
|
+
}
|
|
1967
|
+
async init() {
|
|
1968
|
+
await this.#db.createTable({ tableName: TABLE_THREADS, schema: TABLE_SCHEMAS[TABLE_THREADS] });
|
|
1969
|
+
await this.#db.createTable({ tableName: TABLE_MESSAGES, schema: TABLE_SCHEMAS[TABLE_MESSAGES] });
|
|
1970
|
+
await this.#db.createTable({ tableName: TABLE_RESOURCES, schema: TABLE_SCHEMAS[TABLE_RESOURCES] });
|
|
1971
|
+
await this.#db.alterTable({
|
|
1972
|
+
tableName: TABLE_MESSAGES,
|
|
1973
|
+
schema: TABLE_SCHEMAS[TABLE_MESSAGES],
|
|
1974
|
+
ifNotExists: ["resourceId"]
|
|
1975
|
+
});
|
|
1976
|
+
await this.createDefaultIndexes();
|
|
1977
|
+
await this.createCustomIndexes();
|
|
1978
|
+
}
|
|
1979
|
+
/**
|
|
1980
|
+
* Returns default index definitions for the memory domain tables.
|
|
1981
|
+
* Note: Aurora DSQL does not support ASC/DESC in index columns.
|
|
1982
|
+
*/
|
|
1983
|
+
getDefaultIndexDefinitions() {
|
|
1984
|
+
const schemaPrefix = this.#schema !== "public" ? `${this.#schema}_` : "";
|
|
1985
|
+
return [
|
|
1986
|
+
{
|
|
1987
|
+
name: `${schemaPrefix}mastra_threads_resourceid_createdat_idx`,
|
|
1988
|
+
table: TABLE_THREADS,
|
|
1989
|
+
columns: ["resourceId", "createdAt"]
|
|
1990
|
+
},
|
|
1991
|
+
{
|
|
1992
|
+
name: `${schemaPrefix}mastra_messages_thread_id_createdat_idx`,
|
|
1993
|
+
table: TABLE_MESSAGES,
|
|
1994
|
+
columns: ["thread_id", "createdAt"]
|
|
1995
|
+
}
|
|
1996
|
+
];
|
|
1997
|
+
}
|
|
1998
|
+
/**
|
|
1999
|
+
* Creates default indexes for optimal query performance.
|
|
2000
|
+
*/
|
|
2001
|
+
async createDefaultIndexes() {
|
|
2002
|
+
if (this.#skipDefaultIndexes) {
|
|
2003
|
+
return;
|
|
2004
|
+
}
|
|
2005
|
+
for (const indexDef of this.getDefaultIndexDefinitions()) {
|
|
2006
|
+
try {
|
|
2007
|
+
await this.#db.createIndex(indexDef);
|
|
2008
|
+
} catch (error) {
|
|
2009
|
+
this.logger?.warn?.(`Failed to create index ${indexDef.name}:`, error);
|
|
2010
|
+
}
|
|
2011
|
+
}
|
|
2012
|
+
}
|
|
2013
|
+
/**
|
|
2014
|
+
* Creates custom user-defined indexes for this domain's tables.
|
|
2015
|
+
*/
|
|
2016
|
+
async createCustomIndexes() {
|
|
2017
|
+
if (!this.#indexes || this.#indexes.length === 0) {
|
|
2018
|
+
return;
|
|
2019
|
+
}
|
|
2020
|
+
for (const indexDef of this.#indexes) {
|
|
2021
|
+
try {
|
|
2022
|
+
await this.#db.createIndex(indexDef);
|
|
2023
|
+
} catch (error) {
|
|
2024
|
+
this.logger?.warn?.(`Failed to create custom index ${indexDef.name}:`, error);
|
|
2025
|
+
}
|
|
2026
|
+
}
|
|
2027
|
+
}
|
|
2028
|
+
async dangerouslyClearAll() {
|
|
2029
|
+
await this.#db.clearTable({ tableName: TABLE_MESSAGES });
|
|
2030
|
+
await this.#db.clearTable({ tableName: TABLE_THREADS });
|
|
2031
|
+
await this.#db.clearTable({ tableName: TABLE_RESOURCES });
|
|
2032
|
+
}
|
|
2033
|
+
/**
|
|
2034
|
+
* Normalizes message row from database by applying createdAtZ fallback
|
|
2035
|
+
*/
|
|
2036
|
+
normalizeMessageRow(row) {
|
|
2037
|
+
return {
|
|
2038
|
+
id: row.id,
|
|
2039
|
+
content: row.content,
|
|
2040
|
+
role: row.role,
|
|
2041
|
+
type: row.type,
|
|
2042
|
+
createdAt: row.createdAtZ || row.createdAt,
|
|
2043
|
+
threadId: row.threadId,
|
|
2044
|
+
resourceId: row.resourceId
|
|
2045
|
+
};
|
|
2046
|
+
}
|
|
2047
|
+
async getThreadById({ threadId }) {
|
|
2048
|
+
try {
|
|
2049
|
+
const tableName = getTableName2({ indexName: TABLE_THREADS, schemaName: getSchemaName2(this.#schema) });
|
|
2050
|
+
const thread = await this.#db.client.oneOrNone(
|
|
2051
|
+
`SELECT * FROM ${tableName} WHERE id = $1`,
|
|
2052
|
+
[threadId]
|
|
2053
|
+
);
|
|
2054
|
+
if (!thread) {
|
|
2055
|
+
return null;
|
|
2056
|
+
}
|
|
2057
|
+
return {
|
|
2058
|
+
id: thread.id,
|
|
2059
|
+
resourceId: thread.resourceId,
|
|
2060
|
+
title: thread.title,
|
|
2061
|
+
metadata: typeof thread.metadata === "string" ? JSON.parse(thread.metadata) : thread.metadata,
|
|
2062
|
+
createdAt: thread.createdAtZ || thread.createdAt,
|
|
2063
|
+
updatedAt: thread.updatedAtZ || thread.updatedAt
|
|
2064
|
+
};
|
|
2065
|
+
} catch (error) {
|
|
2066
|
+
throw new MastraError(
|
|
2067
|
+
{
|
|
2068
|
+
id: createStorageErrorId("DSQL", "GET_THREAD_BY_ID", "FAILED"),
|
|
2069
|
+
domain: ErrorDomain.STORAGE,
|
|
2070
|
+
category: ErrorCategory.THIRD_PARTY,
|
|
2071
|
+
details: {
|
|
2072
|
+
threadId
|
|
2073
|
+
}
|
|
2074
|
+
},
|
|
2075
|
+
error
|
|
2076
|
+
);
|
|
2077
|
+
}
|
|
2078
|
+
}
|
|
2079
|
+
async listThreads(args) {
|
|
2080
|
+
const { page = 0, perPage: perPageInput, orderBy, filter } = args;
|
|
2081
|
+
try {
|
|
2082
|
+
this.validatePaginationInput(page, perPageInput ?? 100);
|
|
2083
|
+
} catch (error) {
|
|
2084
|
+
throw new MastraError({
|
|
2085
|
+
id: createStorageErrorId("DSQL", "LIST_THREADS", "INVALID_PAGE"),
|
|
2086
|
+
domain: ErrorDomain.STORAGE,
|
|
2087
|
+
category: ErrorCategory.USER,
|
|
2088
|
+
text: error instanceof Error ? error.message : "Invalid pagination parameters",
|
|
2089
|
+
details: { page, ...perPageInput !== void 0 && { perPage: perPageInput } }
|
|
2090
|
+
});
|
|
2091
|
+
}
|
|
2092
|
+
const perPage = normalizePerPage(perPageInput, 100);
|
|
2093
|
+
try {
|
|
2094
|
+
this.validateMetadataKeys(filter?.metadata);
|
|
2095
|
+
} catch (error) {
|
|
2096
|
+
throw new MastraError({
|
|
2097
|
+
id: createStorageErrorId("DSQL", "LIST_THREADS", "INVALID_METADATA_KEY"),
|
|
2098
|
+
domain: ErrorDomain.STORAGE,
|
|
2099
|
+
category: ErrorCategory.USER,
|
|
2100
|
+
text: error instanceof Error ? error.message : "Invalid metadata key",
|
|
2101
|
+
details: { metadataKeys: filter?.metadata ? Object.keys(filter.metadata).join(", ") : "" }
|
|
2102
|
+
});
|
|
2103
|
+
}
|
|
2104
|
+
const { field, direction } = this.parseOrderBy(orderBy);
|
|
2105
|
+
const { offset, perPage: perPageForResponse } = calculatePagination(page, perPageInput, perPage);
|
|
2106
|
+
try {
|
|
2107
|
+
const tableName = getTableName2({ indexName: TABLE_THREADS, schemaName: getSchemaName2(this.#schema) });
|
|
2108
|
+
const whereClauses = [];
|
|
2109
|
+
const queryParams = [];
|
|
2110
|
+
let paramIndex = 1;
|
|
2111
|
+
if (filter?.resourceId) {
|
|
2112
|
+
whereClauses.push(`"resourceId" = ${paramIndex}`);
|
|
2113
|
+
queryParams.push(filter.resourceId);
|
|
2114
|
+
paramIndex++;
|
|
2115
|
+
}
|
|
2116
|
+
if (filter?.metadata && Object.keys(filter.metadata).length > 0) {
|
|
2117
|
+
for (const [key, value] of Object.entries(filter.metadata)) {
|
|
2118
|
+
whereClauses.push(`metadata::jsonb @> ${paramIndex}::jsonb`);
|
|
2119
|
+
queryParams.push(JSON.stringify({ [key]: value }));
|
|
2120
|
+
paramIndex++;
|
|
2121
|
+
}
|
|
2122
|
+
}
|
|
2123
|
+
const whereClause = whereClauses.length > 0 ? `WHERE ${whereClauses.join(" AND ")}` : "";
|
|
2124
|
+
const baseQuery = `FROM ${tableName} ${whereClause}`;
|
|
2125
|
+
const countQuery = `SELECT COUNT(*) ${baseQuery}`;
|
|
2126
|
+
const countResult = await this.#db.client.one(countQuery, queryParams);
|
|
2127
|
+
const total = parseInt(countResult.count, 10);
|
|
2128
|
+
if (total === 0) {
|
|
2129
|
+
return {
|
|
2130
|
+
threads: [],
|
|
2131
|
+
total: 0,
|
|
2132
|
+
page,
|
|
2133
|
+
perPage: perPageForResponse,
|
|
2134
|
+
hasMore: false
|
|
2135
|
+
};
|
|
2136
|
+
}
|
|
2137
|
+
const limitValue = perPageInput === false ? total : perPage;
|
|
2138
|
+
const dataQuery = `SELECT id, "resourceId", title, metadata, "createdAt", "createdAtZ", "updatedAt", "updatedAtZ" ${baseQuery} ORDER BY "${field}" ${direction} LIMIT ${paramIndex} OFFSET ${paramIndex + 1}`;
|
|
2139
|
+
const rows = await this.#db.client.manyOrNone(
|
|
2140
|
+
dataQuery,
|
|
2141
|
+
[...queryParams, limitValue, offset]
|
|
2142
|
+
);
|
|
2143
|
+
const threads = (rows || []).map((thread) => ({
|
|
2144
|
+
id: thread.id,
|
|
2145
|
+
resourceId: thread.resourceId,
|
|
2146
|
+
title: thread.title,
|
|
2147
|
+
metadata: typeof thread.metadata === "string" ? JSON.parse(thread.metadata) : thread.metadata,
|
|
2148
|
+
createdAt: thread.createdAtZ || thread.createdAt,
|
|
2149
|
+
updatedAt: thread.updatedAtZ || thread.updatedAt
|
|
2150
|
+
}));
|
|
2151
|
+
return {
|
|
2152
|
+
threads,
|
|
2153
|
+
total,
|
|
2154
|
+
page,
|
|
2155
|
+
perPage: perPageForResponse,
|
|
2156
|
+
hasMore: perPageInput === false ? false : offset + perPage < total
|
|
2157
|
+
};
|
|
2158
|
+
} catch (error) {
|
|
2159
|
+
const mastraError = new MastraError(
|
|
2160
|
+
{
|
|
2161
|
+
id: createStorageErrorId("DSQL", "LIST_THREADS", "FAILED"),
|
|
2162
|
+
domain: ErrorDomain.STORAGE,
|
|
2163
|
+
category: ErrorCategory.THIRD_PARTY,
|
|
2164
|
+
details: {
|
|
2165
|
+
...filter?.resourceId && { resourceId: filter.resourceId },
|
|
2166
|
+
hasMetadataFilter: !!filter?.metadata,
|
|
2167
|
+
page
|
|
2168
|
+
}
|
|
2169
|
+
},
|
|
2170
|
+
error
|
|
2171
|
+
);
|
|
2172
|
+
this.logger?.error?.(mastraError.toString());
|
|
2173
|
+
this.logger?.trackException(mastraError);
|
|
2174
|
+
return {
|
|
2175
|
+
threads: [],
|
|
2176
|
+
total: 0,
|
|
2177
|
+
page,
|
|
2178
|
+
perPage: perPageForResponse,
|
|
2179
|
+
hasMore: false
|
|
2180
|
+
};
|
|
2181
|
+
}
|
|
2182
|
+
}
|
|
2183
|
+
async saveThread({ thread }) {
|
|
2184
|
+
const tableName = getTableName2({ indexName: TABLE_THREADS, schemaName: getSchemaName2(this.#schema) });
|
|
2185
|
+
await withRetry(
|
|
2186
|
+
async () => {
|
|
2187
|
+
await this.#db.client.none(
|
|
2188
|
+
`INSERT INTO ${tableName} (
|
|
2189
|
+
id,
|
|
2190
|
+
"resourceId",
|
|
2191
|
+
title,
|
|
2192
|
+
metadata,
|
|
2193
|
+
"createdAt",
|
|
2194
|
+
"createdAtZ",
|
|
2195
|
+
"updatedAt",
|
|
2196
|
+
"updatedAtZ"
|
|
2197
|
+
) VALUES ($1, $2, $3, $4, $5, $6, $7, $8)
|
|
2198
|
+
ON CONFLICT (id) DO UPDATE SET
|
|
2199
|
+
"resourceId" = EXCLUDED."resourceId",
|
|
2200
|
+
title = EXCLUDED.title,
|
|
2201
|
+
metadata = EXCLUDED.metadata,
|
|
2202
|
+
"createdAt" = EXCLUDED."createdAt",
|
|
2203
|
+
"createdAtZ" = EXCLUDED."createdAtZ",
|
|
2204
|
+
"updatedAt" = EXCLUDED."updatedAt",
|
|
2205
|
+
"updatedAtZ" = EXCLUDED."updatedAtZ"`,
|
|
2206
|
+
[
|
|
2207
|
+
thread.id,
|
|
2208
|
+
thread.resourceId,
|
|
2209
|
+
thread.title,
|
|
2210
|
+
thread.metadata ? JSON.stringify(thread.metadata) : null,
|
|
2211
|
+
thread.createdAt,
|
|
2212
|
+
thread.createdAt,
|
|
2213
|
+
thread.updatedAt,
|
|
2214
|
+
thread.updatedAt
|
|
2215
|
+
]
|
|
2216
|
+
);
|
|
2217
|
+
},
|
|
2218
|
+
{
|
|
2219
|
+
onRetry: (error, attempt, delay) => {
|
|
2220
|
+
this.logger?.warn?.(`saveThread retry ${attempt} for ${thread.id} after ${delay}ms: ${error.message}`);
|
|
2221
|
+
}
|
|
2222
|
+
}
|
|
2223
|
+
).catch((error) => {
|
|
2224
|
+
throw new MastraError(
|
|
2225
|
+
{
|
|
2226
|
+
id: createStorageErrorId("DSQL", "SAVE_THREAD", "FAILED"),
|
|
2227
|
+
domain: ErrorDomain.STORAGE,
|
|
2228
|
+
category: ErrorCategory.THIRD_PARTY,
|
|
2229
|
+
details: {
|
|
2230
|
+
threadId: thread.id
|
|
2231
|
+
}
|
|
2232
|
+
},
|
|
2233
|
+
error
|
|
2234
|
+
);
|
|
2235
|
+
});
|
|
2236
|
+
return thread;
|
|
2237
|
+
}
|
|
2238
|
+
async updateThread({
|
|
2239
|
+
id,
|
|
2240
|
+
title,
|
|
2241
|
+
metadata
|
|
2242
|
+
}) {
|
|
2243
|
+
const threadTableName = getTableName2({ indexName: TABLE_THREADS, schemaName: getSchemaName2(this.#schema) });
|
|
2244
|
+
const { result } = await withRetry(
|
|
2245
|
+
async () => {
|
|
2246
|
+
const existingThread = await this.getThreadById({ threadId: id });
|
|
2247
|
+
if (!existingThread) {
|
|
2248
|
+
throw new MastraError({
|
|
2249
|
+
id: createStorageErrorId("DSQL", "UPDATE_THREAD", "NOT_FOUND"),
|
|
2250
|
+
domain: ErrorDomain.STORAGE,
|
|
2251
|
+
category: ErrorCategory.USER,
|
|
2252
|
+
text: `Thread ${id} not found`,
|
|
2253
|
+
details: {
|
|
2254
|
+
threadId: id,
|
|
2255
|
+
title
|
|
2256
|
+
}
|
|
2257
|
+
});
|
|
2258
|
+
}
|
|
2259
|
+
const mergedMetadata = {
|
|
2260
|
+
...existingThread.metadata,
|
|
2261
|
+
...metadata
|
|
2262
|
+
};
|
|
2263
|
+
const now = (/* @__PURE__ */ new Date()).toISOString();
|
|
2264
|
+
const thread = await this.#db.client.one(
|
|
2265
|
+
`UPDATE ${threadTableName}
|
|
2266
|
+
SET
|
|
2267
|
+
title = $1,
|
|
2268
|
+
metadata = $2,
|
|
2269
|
+
"updatedAt" = $3::timestamp,
|
|
2270
|
+
"updatedAtZ" = $4::timestamptz
|
|
2271
|
+
WHERE id = $5
|
|
2272
|
+
RETURNING *
|
|
2273
|
+
`,
|
|
2274
|
+
[title, JSON.stringify(mergedMetadata), now, now, id]
|
|
2275
|
+
);
|
|
2276
|
+
return {
|
|
2277
|
+
id: thread.id,
|
|
2278
|
+
resourceId: thread.resourceId,
|
|
2279
|
+
title: thread.title,
|
|
2280
|
+
metadata: typeof thread.metadata === "string" ? JSON.parse(thread.metadata) : thread.metadata,
|
|
2281
|
+
createdAt: thread.createdAtZ || thread.createdAt,
|
|
2282
|
+
updatedAt: thread.updatedAtZ || thread.updatedAt
|
|
2283
|
+
};
|
|
2284
|
+
},
|
|
2285
|
+
{
|
|
2286
|
+
onRetry: (error, attempt, delay) => {
|
|
2287
|
+
this.logger?.warn?.(`updateThread retry ${attempt} for ${id} after ${delay}ms: ${error.message}`);
|
|
2288
|
+
}
|
|
2289
|
+
}
|
|
2290
|
+
).catch((error) => {
|
|
2291
|
+
if (error instanceof MastraError) {
|
|
2292
|
+
throw error;
|
|
2293
|
+
}
|
|
2294
|
+
throw new MastraError(
|
|
2295
|
+
{
|
|
2296
|
+
id: createStorageErrorId("DSQL", "UPDATE_THREAD", "FAILED"),
|
|
2297
|
+
domain: ErrorDomain.STORAGE,
|
|
2298
|
+
category: ErrorCategory.THIRD_PARTY,
|
|
2299
|
+
details: {
|
|
2300
|
+
threadId: id,
|
|
2301
|
+
title
|
|
2302
|
+
}
|
|
2303
|
+
},
|
|
2304
|
+
error
|
|
2305
|
+
);
|
|
2306
|
+
});
|
|
2307
|
+
return result;
|
|
2308
|
+
}
|
|
2309
|
+
async deleteThread({ threadId }) {
|
|
2310
|
+
const tableName = getTableName2({ indexName: TABLE_MESSAGES, schemaName: getSchemaName2(this.#schema) });
|
|
2311
|
+
const threadTableName = getTableName2({ indexName: TABLE_THREADS, schemaName: getSchemaName2(this.#schema) });
|
|
2312
|
+
await withRetry(
|
|
2313
|
+
async () => {
|
|
2314
|
+
await this.#db.client.tx(async (t) => {
|
|
2315
|
+
await t.none(`DELETE FROM ${tableName} WHERE thread_id = $1`, [threadId]);
|
|
2316
|
+
await t.none(`DELETE FROM ${threadTableName} WHERE id = $1`, [threadId]);
|
|
2317
|
+
});
|
|
2318
|
+
},
|
|
2319
|
+
{
|
|
2320
|
+
onRetry: (error, attempt, delay) => {
|
|
2321
|
+
this.logger?.warn?.(`deleteThread retry ${attempt} for ${threadId} after ${delay}ms: ${error.message}`);
|
|
2322
|
+
}
|
|
2323
|
+
}
|
|
2324
|
+
).catch((error) => {
|
|
2325
|
+
throw new MastraError(
|
|
2326
|
+
{
|
|
2327
|
+
id: createStorageErrorId("DSQL", "DELETE_THREAD", "FAILED"),
|
|
2328
|
+
domain: ErrorDomain.STORAGE,
|
|
2329
|
+
category: ErrorCategory.THIRD_PARTY,
|
|
2330
|
+
details: {
|
|
2331
|
+
threadId
|
|
2332
|
+
}
|
|
2333
|
+
},
|
|
2334
|
+
error
|
|
2335
|
+
);
|
|
2336
|
+
});
|
|
2337
|
+
}
|
|
2338
|
+
async _getIncludedMessages({ include }) {
|
|
2339
|
+
if (!include || include.length === 0) return null;
|
|
2340
|
+
const unionQueries = [];
|
|
2341
|
+
const params = [];
|
|
2342
|
+
let paramIdx = 1;
|
|
2343
|
+
const tableName = getTableName2({ indexName: TABLE_MESSAGES, schemaName: getSchemaName2(this.#schema) });
|
|
2344
|
+
for (const inc of include) {
|
|
2345
|
+
const { id, withPreviousMessages = 0, withNextMessages = 0 } = inc;
|
|
2346
|
+
unionQueries.push(
|
|
2347
|
+
`
|
|
2348
|
+
SELECT * FROM (
|
|
2349
|
+
WITH target_thread AS (
|
|
2350
|
+
SELECT thread_id FROM ${tableName} WHERE id = $${paramIdx}
|
|
2351
|
+
),
|
|
2352
|
+
ordered_messages AS (
|
|
2353
|
+
SELECT
|
|
2354
|
+
*,
|
|
2355
|
+
ROW_NUMBER() OVER (ORDER BY "createdAt" ASC) as row_num
|
|
2356
|
+
FROM ${tableName}
|
|
2357
|
+
WHERE thread_id = (SELECT thread_id FROM target_thread)
|
|
2358
|
+
)
|
|
2359
|
+
SELECT
|
|
2360
|
+
m.id,
|
|
2361
|
+
m.content,
|
|
2362
|
+
m.role,
|
|
2363
|
+
m.type,
|
|
2364
|
+
m."createdAt",
|
|
2365
|
+
m."createdAtZ",
|
|
2366
|
+
m.thread_id AS "threadId",
|
|
2367
|
+
m."resourceId"
|
|
2368
|
+
FROM ordered_messages m
|
|
2369
|
+
WHERE m.id = $${paramIdx}
|
|
2370
|
+
OR EXISTS (
|
|
2371
|
+
SELECT 1 FROM ordered_messages target
|
|
2372
|
+
WHERE target.id = $${paramIdx}
|
|
2373
|
+
AND (
|
|
2374
|
+
(m.row_num < target.row_num AND m.row_num >= target.row_num - $${paramIdx + 1})
|
|
2375
|
+
OR
|
|
2376
|
+
(m.row_num > target.row_num AND m.row_num <= target.row_num + $${paramIdx + 2})
|
|
2377
|
+
)
|
|
2378
|
+
)
|
|
2379
|
+
) AS query_${paramIdx}
|
|
2380
|
+
`
|
|
2381
|
+
);
|
|
2382
|
+
params.push(id, withPreviousMessages, withNextMessages);
|
|
2383
|
+
paramIdx += 3;
|
|
2384
|
+
}
|
|
2385
|
+
const finalQuery = unionQueries.join(" UNION ALL ") + ' ORDER BY "createdAt" ASC';
|
|
2386
|
+
const includedRows = await this.#db.client.manyOrNone(finalQuery, params);
|
|
2387
|
+
const seen = /* @__PURE__ */ new Set();
|
|
2388
|
+
const dedupedRows = includedRows.filter((row) => {
|
|
2389
|
+
if (seen.has(row.id)) return false;
|
|
2390
|
+
seen.add(row.id);
|
|
2391
|
+
return true;
|
|
2392
|
+
});
|
|
2393
|
+
return dedupedRows;
|
|
2394
|
+
}
|
|
2395
|
+
parseRow(row) {
|
|
2396
|
+
const normalized = this.normalizeMessageRow(row);
|
|
2397
|
+
let content = normalized.content;
|
|
2398
|
+
try {
|
|
2399
|
+
content = JSON.parse(normalized.content);
|
|
2400
|
+
} catch {
|
|
2401
|
+
}
|
|
2402
|
+
return {
|
|
2403
|
+
id: normalized.id,
|
|
2404
|
+
content,
|
|
2405
|
+
role: normalized.role,
|
|
2406
|
+
createdAt: new Date(normalized.createdAt),
|
|
2407
|
+
threadId: normalized.threadId,
|
|
2408
|
+
resourceId: normalized.resourceId,
|
|
2409
|
+
...normalized.type && normalized.type !== "v2" ? { type: normalized.type } : {}
|
|
2410
|
+
};
|
|
2411
|
+
}
|
|
2412
|
+
async listMessagesById({ messageIds }) {
|
|
2413
|
+
if (messageIds.length === 0) return { messages: [] };
|
|
2414
|
+
const selectStatement = `SELECT id, content, role, type, "createdAt", "createdAtZ", thread_id AS "threadId", "resourceId"`;
|
|
2415
|
+
try {
|
|
2416
|
+
const tableName = getTableName2({ indexName: TABLE_MESSAGES, schemaName: getSchemaName2(this.#schema) });
|
|
2417
|
+
const query = `
|
|
2418
|
+
${selectStatement} FROM ${tableName}
|
|
2419
|
+
WHERE id IN (${inPlaceholders(messageIds.length)})
|
|
2420
|
+
ORDER BY "createdAt" DESC
|
|
2421
|
+
`;
|
|
2422
|
+
const resultRows = await this.#db.client.manyOrNone(query, messageIds);
|
|
2423
|
+
const list = new MessageList().add(
|
|
2424
|
+
resultRows.map((row) => this.parseRow(row)),
|
|
2425
|
+
"memory"
|
|
2426
|
+
);
|
|
2427
|
+
return { messages: list.get.all.db() };
|
|
2428
|
+
} catch (error) {
|
|
2429
|
+
const mastraError = new MastraError(
|
|
2430
|
+
{
|
|
2431
|
+
id: createStorageErrorId("DSQL", "LIST_MESSAGES_BY_ID", "FAILED"),
|
|
2432
|
+
domain: ErrorDomain.STORAGE,
|
|
2433
|
+
category: ErrorCategory.THIRD_PARTY,
|
|
2434
|
+
details: {
|
|
2435
|
+
messageIds: JSON.stringify(messageIds)
|
|
2436
|
+
}
|
|
2437
|
+
},
|
|
2438
|
+
error
|
|
2439
|
+
);
|
|
2440
|
+
this.logger?.error?.(mastraError.toString());
|
|
2441
|
+
this.logger?.trackException(mastraError);
|
|
2442
|
+
return { messages: [] };
|
|
2443
|
+
}
|
|
2444
|
+
}
|
|
2445
|
+
async listMessages(args) {
|
|
2446
|
+
const { threadId, resourceId, include, filter, perPage: perPageInput, page = 0, orderBy } = args;
|
|
2447
|
+
const threadIds = (Array.isArray(threadId) ? threadId : [threadId]).filter(
|
|
2448
|
+
(id) => typeof id === "string"
|
|
2449
|
+
);
|
|
2450
|
+
if (threadIds.length === 0 || threadIds.some((id) => !id.trim())) {
|
|
2451
|
+
throw new MastraError(
|
|
2452
|
+
{
|
|
2453
|
+
id: createStorageErrorId("DSQL", "LIST_MESSAGES", "INVALID_THREAD_ID"),
|
|
2454
|
+
domain: ErrorDomain.STORAGE,
|
|
2455
|
+
category: ErrorCategory.USER,
|
|
2456
|
+
details: { threadId: Array.isArray(threadId) ? String(threadId) : String(threadId) }
|
|
2457
|
+
},
|
|
2458
|
+
new Error("threadId must be a non-empty string or array of non-empty strings")
|
|
2459
|
+
);
|
|
2460
|
+
}
|
|
2461
|
+
if (page < 0) {
|
|
2462
|
+
throw new MastraError({
|
|
2463
|
+
id: createStorageErrorId("DSQL", "LIST_MESSAGES", "INVALID_PAGE"),
|
|
2464
|
+
domain: ErrorDomain.STORAGE,
|
|
2465
|
+
category: ErrorCategory.USER,
|
|
2466
|
+
text: "Page number must be non-negative",
|
|
2467
|
+
details: {
|
|
2468
|
+
threadId: Array.isArray(threadId) ? threadId.join(",") : threadId,
|
|
2469
|
+
page
|
|
2470
|
+
}
|
|
2471
|
+
});
|
|
2472
|
+
}
|
|
2473
|
+
const perPage = normalizePerPage(perPageInput, 40);
|
|
2474
|
+
const { offset, perPage: perPageForResponse } = calculatePagination(page, perPageInput, perPage);
|
|
2475
|
+
try {
|
|
2476
|
+
const { field, direction } = this.parseOrderBy(orderBy, "ASC");
|
|
2477
|
+
const orderByStatement = `ORDER BY "${field}" ${direction}`;
|
|
2478
|
+
const selectStatement = `SELECT id, content, role, type, "createdAt", "createdAtZ", thread_id AS "threadId", "resourceId"`;
|
|
2479
|
+
const tableName = getTableName2({ indexName: TABLE_MESSAGES, schemaName: getSchemaName2(this.#schema) });
|
|
2480
|
+
const conditions = [`thread_id IN (${inPlaceholders(threadIds.length)})`];
|
|
2481
|
+
const queryParams = [...threadIds];
|
|
2482
|
+
let paramIndex = threadIds.length + 1;
|
|
2483
|
+
if (resourceId) {
|
|
2484
|
+
conditions.push(`"resourceId" = $${paramIndex++}`);
|
|
2485
|
+
queryParams.push(resourceId);
|
|
2486
|
+
}
|
|
2487
|
+
if (filter?.dateRange?.start) {
|
|
2488
|
+
conditions.push(`"createdAtZ" >= $${paramIndex++}::timestamptz`);
|
|
2489
|
+
queryParams.push(filter.dateRange.start);
|
|
2490
|
+
}
|
|
2491
|
+
if (filter?.dateRange?.end) {
|
|
2492
|
+
conditions.push(`"createdAtZ" <= $${paramIndex++}::timestamptz`);
|
|
2493
|
+
queryParams.push(filter.dateRange.end);
|
|
2494
|
+
}
|
|
2495
|
+
const whereClause = conditions.length > 0 ? `WHERE ${conditions.join(" AND ")}` : "";
|
|
2496
|
+
const countQuery = `SELECT COUNT(*) FROM ${tableName} ${whereClause}`;
|
|
2497
|
+
const countResult = await this.#db.client.one(countQuery, queryParams);
|
|
2498
|
+
const total = parseInt(countResult.count, 10);
|
|
2499
|
+
const limitValue = perPageInput === false ? total : perPage;
|
|
2500
|
+
const dataQuery = `${selectStatement} FROM ${tableName} ${whereClause} ${orderByStatement} LIMIT $${paramIndex++} OFFSET $${paramIndex++}`;
|
|
2501
|
+
const rows = await this.#db.client.manyOrNone(dataQuery, [...queryParams, limitValue, offset]);
|
|
2502
|
+
const messages = [...rows || []];
|
|
2503
|
+
if (total === 0 && messages.length === 0 && (!include || include.length === 0)) {
|
|
2504
|
+
return {
|
|
2505
|
+
messages: [],
|
|
2506
|
+
total: 0,
|
|
2507
|
+
page,
|
|
2508
|
+
perPage: perPageForResponse,
|
|
2509
|
+
hasMore: false
|
|
2510
|
+
};
|
|
2511
|
+
}
|
|
2512
|
+
const messageIds = new Set(messages.map((m) => m.id));
|
|
2513
|
+
if (include && include.length > 0) {
|
|
2514
|
+
const includeMessages = await this._getIncludedMessages({ include });
|
|
2515
|
+
if (includeMessages) {
|
|
2516
|
+
for (const includeMsg of includeMessages) {
|
|
2517
|
+
if (!messageIds.has(includeMsg.id)) {
|
|
2518
|
+
messages.push(includeMsg);
|
|
2519
|
+
messageIds.add(includeMsg.id);
|
|
2520
|
+
}
|
|
2521
|
+
}
|
|
2522
|
+
}
|
|
2523
|
+
}
|
|
2524
|
+
const messagesWithParsedContent = messages.map((row) => this.parseRow(row));
|
|
2525
|
+
const list = new MessageList().add(messagesWithParsedContent, "memory");
|
|
2526
|
+
let finalMessages = list.get.all.db();
|
|
2527
|
+
finalMessages = finalMessages.sort((a, b) => {
|
|
2528
|
+
const aValue = field === "createdAt" ? new Date(a.createdAt).getTime() : a[field];
|
|
2529
|
+
const bValue = field === "createdAt" ? new Date(b.createdAt).getTime() : b[field];
|
|
2530
|
+
if (aValue == null && bValue == null) return a.id.localeCompare(b.id);
|
|
2531
|
+
if (aValue == null) return 1;
|
|
2532
|
+
if (bValue == null) return -1;
|
|
2533
|
+
if (aValue === bValue) {
|
|
2534
|
+
return a.id.localeCompare(b.id);
|
|
2535
|
+
}
|
|
2536
|
+
if (typeof aValue === "number" && typeof bValue === "number") {
|
|
2537
|
+
return direction === "ASC" ? aValue - bValue : bValue - aValue;
|
|
2538
|
+
}
|
|
2539
|
+
return direction === "ASC" ? String(aValue).localeCompare(String(bValue)) : String(bValue).localeCompare(String(aValue));
|
|
2540
|
+
});
|
|
2541
|
+
const threadIdSet = new Set(threadIds);
|
|
2542
|
+
const returnedThreadMessageIds = new Set(
|
|
2543
|
+
finalMessages.filter((m) => m.threadId && threadIdSet.has(m.threadId)).map((m) => m.id)
|
|
2544
|
+
);
|
|
2545
|
+
const allThreadMessagesReturned = returnedThreadMessageIds.size >= total;
|
|
2546
|
+
const hasMore = perPageInput !== false && !allThreadMessagesReturned && offset + perPage < total;
|
|
2547
|
+
return {
|
|
2548
|
+
messages: finalMessages,
|
|
2549
|
+
total,
|
|
2550
|
+
page,
|
|
2551
|
+
perPage: perPageForResponse,
|
|
2552
|
+
hasMore
|
|
2553
|
+
};
|
|
2554
|
+
} catch (error) {
|
|
2555
|
+
const mastraError = new MastraError(
|
|
2556
|
+
{
|
|
2557
|
+
id: createStorageErrorId("DSQL", "LIST_MESSAGES", "FAILED"),
|
|
2558
|
+
domain: ErrorDomain.STORAGE,
|
|
2559
|
+
category: ErrorCategory.THIRD_PARTY,
|
|
2560
|
+
details: {
|
|
2561
|
+
threadId: Array.isArray(threadId) ? threadId.join(",") : threadId,
|
|
2562
|
+
resourceId: resourceId ?? ""
|
|
2563
|
+
}
|
|
2564
|
+
},
|
|
2565
|
+
error
|
|
2566
|
+
);
|
|
2567
|
+
this.logger?.error?.(mastraError.toString());
|
|
2568
|
+
this.logger?.trackException(mastraError);
|
|
2569
|
+
return {
|
|
2570
|
+
messages: [],
|
|
2571
|
+
total: 0,
|
|
2572
|
+
page,
|
|
2573
|
+
perPage: perPageForResponse,
|
|
2574
|
+
hasMore: false
|
|
2575
|
+
};
|
|
2576
|
+
}
|
|
2577
|
+
}
|
|
2578
|
+
async saveMessages({ messages }) {
|
|
2579
|
+
if (messages.length === 0) return { messages: [] };
|
|
2580
|
+
const threadId = messages[0]?.threadId;
|
|
2581
|
+
if (!threadId) {
|
|
2582
|
+
throw new MastraError({
|
|
2583
|
+
id: createStorageErrorId("DSQL", "SAVE_MESSAGES", "FAILED"),
|
|
2584
|
+
domain: ErrorDomain.STORAGE,
|
|
2585
|
+
category: ErrorCategory.USER,
|
|
2586
|
+
text: `Thread ID is required`
|
|
2587
|
+
});
|
|
2588
|
+
}
|
|
2589
|
+
const thread = await this.getThreadById({ threadId });
|
|
2590
|
+
if (!thread) {
|
|
2591
|
+
throw new MastraError({
|
|
2592
|
+
id: createStorageErrorId("DSQL", "SAVE_MESSAGES", "FAILED"),
|
|
2593
|
+
domain: ErrorDomain.STORAGE,
|
|
2594
|
+
category: ErrorCategory.USER,
|
|
2595
|
+
text: `Thread ${threadId} not found`,
|
|
2596
|
+
details: {
|
|
2597
|
+
threadId
|
|
2598
|
+
}
|
|
2599
|
+
});
|
|
2600
|
+
}
|
|
2601
|
+
const tableName = getTableName2({ indexName: TABLE_MESSAGES, schemaName: getSchemaName2(this.#schema) });
|
|
2602
|
+
const threadTableName = getTableName2({ indexName: TABLE_THREADS, schemaName: getSchemaName2(this.#schema) });
|
|
2603
|
+
await withRetry(
|
|
2604
|
+
async () => {
|
|
2605
|
+
await this.#db.client.tx(async (t) => {
|
|
2606
|
+
const messageInserts = messages.map((message) => {
|
|
2607
|
+
if (!message.threadId) {
|
|
2608
|
+
throw new Error(
|
|
2609
|
+
`Expected to find a threadId for message, but couldn't find one. An unexpected error has occurred.`
|
|
2610
|
+
);
|
|
2611
|
+
}
|
|
2612
|
+
if (!message.resourceId) {
|
|
2613
|
+
throw new Error(
|
|
2614
|
+
`Expected to find a resourceId for message, but couldn't find one. An unexpected error has occurred.`
|
|
2615
|
+
);
|
|
2616
|
+
}
|
|
2617
|
+
const createdAtIso = message.createdAt ? new Date(message.createdAt).toISOString() : (/* @__PURE__ */ new Date()).toISOString();
|
|
2618
|
+
return t.none(
|
|
2619
|
+
`INSERT INTO ${tableName} (id, thread_id, content, "createdAt", "createdAtZ", role, type, "resourceId")
|
|
2620
|
+
VALUES ($1, $2, $3, $4::timestamp, $5::timestamptz, $6, $7, $8)
|
|
2621
|
+
ON CONFLICT (id) DO UPDATE SET
|
|
2622
|
+
thread_id = EXCLUDED.thread_id,
|
|
2623
|
+
content = EXCLUDED.content,
|
|
2624
|
+
role = EXCLUDED.role,
|
|
2625
|
+
type = EXCLUDED.type,
|
|
2626
|
+
"resourceId" = EXCLUDED."resourceId"`,
|
|
2627
|
+
[
|
|
2628
|
+
message.id,
|
|
2629
|
+
message.threadId,
|
|
2630
|
+
typeof message.content === "string" ? message.content : JSON.stringify(message.content),
|
|
2631
|
+
createdAtIso,
|
|
2632
|
+
createdAtIso,
|
|
2633
|
+
message.role,
|
|
2634
|
+
message.type || "v2",
|
|
2635
|
+
message.resourceId
|
|
2636
|
+
]
|
|
2637
|
+
);
|
|
2638
|
+
});
|
|
2639
|
+
const nowIso = (/* @__PURE__ */ new Date()).toISOString();
|
|
2640
|
+
const threadUpdate = t.none(
|
|
2641
|
+
`UPDATE ${threadTableName}
|
|
2642
|
+
SET
|
|
2643
|
+
"updatedAt" = $1::timestamp,
|
|
2644
|
+
"updatedAtZ" = $2::timestamptz
|
|
2645
|
+
WHERE id = $3
|
|
2646
|
+
`,
|
|
2647
|
+
[nowIso, nowIso, threadId]
|
|
2648
|
+
);
|
|
2649
|
+
await Promise.all([...messageInserts, threadUpdate]);
|
|
2650
|
+
});
|
|
2651
|
+
},
|
|
2652
|
+
{
|
|
2653
|
+
onRetry: (error, attempt, delay) => {
|
|
2654
|
+
this.logger?.warn?.(
|
|
2655
|
+
`saveMessages retry ${attempt} for thread ${threadId} after ${delay}ms: ${error.message}`
|
|
2656
|
+
);
|
|
2657
|
+
}
|
|
2658
|
+
}
|
|
2659
|
+
).catch((error) => {
|
|
2660
|
+
throw new MastraError(
|
|
2661
|
+
{
|
|
2662
|
+
id: createStorageErrorId("DSQL", "SAVE_MESSAGES", "FAILED"),
|
|
2663
|
+
domain: ErrorDomain.STORAGE,
|
|
2664
|
+
category: ErrorCategory.THIRD_PARTY,
|
|
2665
|
+
details: {
|
|
2666
|
+
threadId
|
|
2667
|
+
}
|
|
2668
|
+
},
|
|
2669
|
+
error
|
|
2670
|
+
);
|
|
2671
|
+
});
|
|
2672
|
+
const messagesWithParsedContent = messages.map((message) => {
|
|
2673
|
+
if (typeof message.content === "string") {
|
|
2674
|
+
try {
|
|
2675
|
+
return { ...message, content: JSON.parse(message.content) };
|
|
2676
|
+
} catch {
|
|
2677
|
+
return message;
|
|
2678
|
+
}
|
|
2679
|
+
}
|
|
2680
|
+
return message;
|
|
2681
|
+
});
|
|
2682
|
+
const list = new MessageList().add(messagesWithParsedContent, "memory");
|
|
2683
|
+
return { messages: list.get.all.db() };
|
|
2684
|
+
}
|
|
2685
|
+
async updateMessages({
|
|
2686
|
+
messages
|
|
2687
|
+
}) {
|
|
2688
|
+
if (messages.length === 0) {
|
|
2689
|
+
return [];
|
|
2690
|
+
}
|
|
2691
|
+
const messageIds = messages.map((m) => m.id);
|
|
2692
|
+
const selectQuery = `SELECT id, content, role, type, "createdAt", "createdAtZ", thread_id AS "threadId", "resourceId" FROM ${getTableName2({ indexName: TABLE_MESSAGES, schemaName: getSchemaName2(this.#schema) })} WHERE id IN (${inPlaceholders(messageIds.length)})`;
|
|
2693
|
+
const existingMessagesDb = await this.#db.client.manyOrNone(selectQuery, messageIds);
|
|
2694
|
+
if (existingMessagesDb.length === 0) {
|
|
2695
|
+
return [];
|
|
2696
|
+
}
|
|
2697
|
+
const existingMessages = existingMessagesDb.map((msg) => {
|
|
2698
|
+
if (typeof msg.content === "string") {
|
|
2699
|
+
try {
|
|
2700
|
+
msg.content = JSON.parse(msg.content);
|
|
2701
|
+
} catch {
|
|
2702
|
+
}
|
|
2703
|
+
}
|
|
2704
|
+
return msg;
|
|
2705
|
+
});
|
|
2706
|
+
const threadIdsToUpdate = /* @__PURE__ */ new Set();
|
|
2707
|
+
await withRetry(
|
|
2708
|
+
async () => {
|
|
2709
|
+
await this.#db.client.tx(async (t) => {
|
|
2710
|
+
const queries = [];
|
|
2711
|
+
const columnMapping = {
|
|
2712
|
+
threadId: "thread_id"
|
|
2713
|
+
};
|
|
2714
|
+
for (const existingMessage of existingMessages) {
|
|
2715
|
+
const updatePayload = messages.find((m) => m.id === existingMessage.id);
|
|
2716
|
+
if (!updatePayload) continue;
|
|
2717
|
+
const { id, ...fieldsToUpdate } = updatePayload;
|
|
2718
|
+
if (Object.keys(fieldsToUpdate).length === 0) continue;
|
|
2719
|
+
threadIdsToUpdate.add(existingMessage.threadId);
|
|
2720
|
+
if (updatePayload.threadId && updatePayload.threadId !== existingMessage.threadId) {
|
|
2721
|
+
threadIdsToUpdate.add(updatePayload.threadId);
|
|
2722
|
+
}
|
|
2723
|
+
const setClauses = [];
|
|
2724
|
+
const values = [];
|
|
2725
|
+
let paramIndex = 1;
|
|
2726
|
+
const updatableFields = { ...fieldsToUpdate };
|
|
2727
|
+
if (updatableFields.content) {
|
|
2728
|
+
const newContent = {
|
|
2729
|
+
...existingMessage.content,
|
|
2730
|
+
...updatableFields.content,
|
|
2731
|
+
...existingMessage.content?.metadata && updatableFields.content.metadata ? {
|
|
2732
|
+
metadata: {
|
|
2733
|
+
...existingMessage.content.metadata,
|
|
2734
|
+
...updatableFields.content.metadata
|
|
2735
|
+
}
|
|
2736
|
+
} : {}
|
|
2737
|
+
};
|
|
2738
|
+
setClauses.push(`content = $${paramIndex++}`);
|
|
2739
|
+
values.push(newContent);
|
|
2740
|
+
delete updatableFields.content;
|
|
2741
|
+
}
|
|
2742
|
+
for (const key in updatableFields) {
|
|
2743
|
+
if (Object.prototype.hasOwnProperty.call(updatableFields, key)) {
|
|
2744
|
+
const dbColumn = columnMapping[key] || key;
|
|
2745
|
+
setClauses.push(`"${dbColumn}" = $${paramIndex++}`);
|
|
2746
|
+
values.push(updatableFields[key]);
|
|
2747
|
+
}
|
|
2748
|
+
}
|
|
2749
|
+
if (setClauses.length > 0) {
|
|
2750
|
+
values.push(id);
|
|
2751
|
+
const sql = `UPDATE ${getTableName2({ indexName: TABLE_MESSAGES, schemaName: getSchemaName2(this.#schema) })} SET ${setClauses.join(", ")} WHERE id = $${paramIndex}`;
|
|
2752
|
+
queries.push(t.none(sql, values));
|
|
2753
|
+
}
|
|
2754
|
+
}
|
|
2755
|
+
if (threadIdsToUpdate.size > 0) {
|
|
2756
|
+
const threadIds = Array.from(threadIdsToUpdate);
|
|
2757
|
+
queries.push(
|
|
2758
|
+
t.none(
|
|
2759
|
+
`UPDATE ${getTableName2({ indexName: TABLE_THREADS, schemaName: getSchemaName2(this.#schema) })} SET "updatedAt" = NOW(), "updatedAtZ" = NOW() WHERE id IN (${inPlaceholders(threadIds.length)})`,
|
|
2760
|
+
threadIds
|
|
2761
|
+
)
|
|
2762
|
+
);
|
|
2763
|
+
}
|
|
2764
|
+
if (queries.length > 0) {
|
|
2765
|
+
await t.batch(queries);
|
|
2766
|
+
}
|
|
2767
|
+
});
|
|
2768
|
+
},
|
|
2769
|
+
{
|
|
2770
|
+
onRetry: (error, attempt, delay) => {
|
|
2771
|
+
this.logger?.warn?.(
|
|
2772
|
+
`updateMessages retry ${attempt} for ${messageIds.length} messages after ${delay}ms: ${error.message}`
|
|
2773
|
+
);
|
|
2774
|
+
}
|
|
2775
|
+
}
|
|
2776
|
+
).catch((error) => {
|
|
2777
|
+
throw new MastraError(
|
|
2778
|
+
{
|
|
2779
|
+
id: createStorageErrorId("DSQL", "UPDATE_MESSAGES", "FAILED"),
|
|
2780
|
+
domain: ErrorDomain.STORAGE,
|
|
2781
|
+
category: ErrorCategory.THIRD_PARTY,
|
|
2782
|
+
details: {
|
|
2783
|
+
messageIdsLength: messageIds.length
|
|
2784
|
+
}
|
|
2785
|
+
},
|
|
2786
|
+
error
|
|
2787
|
+
);
|
|
2788
|
+
});
|
|
2789
|
+
const updatedMessages = await this.#db.client.manyOrNone(selectQuery, messageIds);
|
|
2790
|
+
return (updatedMessages || []).map((row) => {
|
|
2791
|
+
const message = this.normalizeMessageRow(row);
|
|
2792
|
+
if (typeof message.content === "string") {
|
|
2793
|
+
try {
|
|
2794
|
+
return { ...message, content: JSON.parse(message.content) };
|
|
2795
|
+
} catch {
|
|
2796
|
+
}
|
|
2797
|
+
}
|
|
2798
|
+
return message;
|
|
2799
|
+
});
|
|
2800
|
+
}
|
|
2801
|
+
async deleteMessages(messageIds) {
|
|
2802
|
+
if (!messageIds || messageIds.length === 0) {
|
|
2803
|
+
return;
|
|
2804
|
+
}
|
|
2805
|
+
const messageTableName = getTableName2({ indexName: TABLE_MESSAGES, schemaName: getSchemaName2(this.#schema) });
|
|
2806
|
+
const threadTableName = getTableName2({ indexName: TABLE_THREADS, schemaName: getSchemaName2(this.#schema) });
|
|
2807
|
+
await withRetry(
|
|
2808
|
+
async () => {
|
|
2809
|
+
await this.#db.client.tx(async (t) => {
|
|
2810
|
+
const placeholders = messageIds.map((_, idx) => `$${idx + 1}`).join(",");
|
|
2811
|
+
const messages = await t.manyOrNone(
|
|
2812
|
+
`SELECT DISTINCT thread_id FROM ${messageTableName} WHERE id IN (${placeholders})`,
|
|
2813
|
+
messageIds
|
|
2814
|
+
);
|
|
2815
|
+
const threadIds = messages?.map((msg) => msg.thread_id).filter(Boolean) || [];
|
|
2816
|
+
await t.none(`DELETE FROM ${messageTableName} WHERE id IN (${placeholders})`, messageIds);
|
|
2817
|
+
if (threadIds.length > 0) {
|
|
2818
|
+
const updatePromises = threadIds.map(
|
|
2819
|
+
(threadId) => t.none(`UPDATE ${threadTableName} SET "updatedAt" = NOW(), "updatedAtZ" = NOW() WHERE id = $1`, [
|
|
2820
|
+
threadId
|
|
2821
|
+
])
|
|
2822
|
+
);
|
|
2823
|
+
await Promise.all(updatePromises);
|
|
2824
|
+
}
|
|
2825
|
+
});
|
|
2826
|
+
},
|
|
2827
|
+
{
|
|
2828
|
+
onRetry: (error, attempt, delay) => {
|
|
2829
|
+
this.logger?.warn?.(
|
|
2830
|
+
`deleteMessages retry ${attempt} for ${messageIds.length} messages after ${delay}ms: ${error.message}`
|
|
2831
|
+
);
|
|
2832
|
+
}
|
|
2833
|
+
}
|
|
2834
|
+
).catch((error) => {
|
|
2835
|
+
throw new MastraError(
|
|
2836
|
+
{
|
|
2837
|
+
id: createStorageErrorId("DSQL", "DELETE_MESSAGES", "FAILED"),
|
|
2838
|
+
domain: ErrorDomain.STORAGE,
|
|
2839
|
+
category: ErrorCategory.THIRD_PARTY,
|
|
2840
|
+
details: { messageIds: messageIds.join(", ") }
|
|
2841
|
+
},
|
|
2842
|
+
error
|
|
2843
|
+
);
|
|
2844
|
+
});
|
|
2845
|
+
}
|
|
2846
|
+
async getResourceById({ resourceId }) {
|
|
2847
|
+
try {
|
|
2848
|
+
const tableName = getTableName2({ indexName: TABLE_RESOURCES, schemaName: getSchemaName2(this.#schema) });
|
|
2849
|
+
const result = await this.#db.client.oneOrNone(
|
|
2850
|
+
`SELECT * FROM ${tableName} WHERE id = $1`,
|
|
2851
|
+
[resourceId]
|
|
2852
|
+
);
|
|
2853
|
+
if (!result) {
|
|
2854
|
+
return null;
|
|
2855
|
+
}
|
|
2856
|
+
return {
|
|
2857
|
+
id: result.id,
|
|
2858
|
+
createdAt: result.createdAtZ || result.createdAt,
|
|
2859
|
+
updatedAt: result.updatedAtZ || result.updatedAt,
|
|
2860
|
+
workingMemory: result.workingMemory,
|
|
2861
|
+
metadata: typeof result.metadata === "string" ? JSON.parse(result.metadata) : result.metadata
|
|
2862
|
+
};
|
|
2863
|
+
} catch (error) {
|
|
2864
|
+
throw new MastraError(
|
|
2865
|
+
{
|
|
2866
|
+
id: createStorageErrorId("DSQL", "GET_RESOURCE_BY_ID", "FAILED"),
|
|
2867
|
+
domain: ErrorDomain.STORAGE,
|
|
2868
|
+
category: ErrorCategory.THIRD_PARTY,
|
|
2869
|
+
details: { resourceId }
|
|
2870
|
+
},
|
|
2871
|
+
error
|
|
2872
|
+
);
|
|
2873
|
+
}
|
|
2874
|
+
}
|
|
2875
|
+
async saveResource({ resource }) {
|
|
2876
|
+
await this.#db.insert({
|
|
2877
|
+
tableName: TABLE_RESOURCES,
|
|
2878
|
+
record: {
|
|
2879
|
+
...resource,
|
|
2880
|
+
metadata: JSON.stringify(resource.metadata)
|
|
2881
|
+
}
|
|
2882
|
+
});
|
|
2883
|
+
return resource;
|
|
2884
|
+
}
|
|
2885
|
+
async updateResource({
|
|
2886
|
+
resourceId,
|
|
2887
|
+
workingMemory,
|
|
2888
|
+
metadata
|
|
2889
|
+
}) {
|
|
2890
|
+
const tableName = getTableName2({ indexName: TABLE_RESOURCES, schemaName: getSchemaName2(this.#schema) });
|
|
2891
|
+
const { result } = await withRetry(
|
|
2892
|
+
async () => {
|
|
2893
|
+
const existingResource = await this.getResourceById({ resourceId });
|
|
2894
|
+
if (!existingResource) {
|
|
2895
|
+
const newResource = {
|
|
2896
|
+
id: resourceId,
|
|
2897
|
+
workingMemory,
|
|
2898
|
+
metadata: metadata || {},
|
|
2899
|
+
createdAt: /* @__PURE__ */ new Date(),
|
|
2900
|
+
updatedAt: /* @__PURE__ */ new Date()
|
|
2901
|
+
};
|
|
2902
|
+
return this.saveResource({ resource: newResource });
|
|
2903
|
+
}
|
|
2904
|
+
const updatedResource = {
|
|
2905
|
+
...existingResource,
|
|
2906
|
+
workingMemory: workingMemory !== void 0 ? workingMemory : existingResource.workingMemory,
|
|
2907
|
+
metadata: {
|
|
2908
|
+
...existingResource.metadata,
|
|
2909
|
+
...metadata
|
|
2910
|
+
},
|
|
2911
|
+
updatedAt: /* @__PURE__ */ new Date()
|
|
2912
|
+
};
|
|
2913
|
+
const updates = [];
|
|
2914
|
+
const values = [];
|
|
2915
|
+
let paramIndex = 1;
|
|
2916
|
+
if (workingMemory !== void 0) {
|
|
2917
|
+
updates.push(`"workingMemory" = $${paramIndex}`);
|
|
2918
|
+
values.push(workingMemory);
|
|
2919
|
+
paramIndex++;
|
|
2920
|
+
}
|
|
2921
|
+
if (metadata) {
|
|
2922
|
+
updates.push(`metadata = $${paramIndex}`);
|
|
2923
|
+
values.push(JSON.stringify(updatedResource.metadata));
|
|
2924
|
+
paramIndex++;
|
|
2925
|
+
}
|
|
2926
|
+
updates.push(`"updatedAt" = $${paramIndex}`);
|
|
2927
|
+
values.push(updatedResource.updatedAt.toISOString());
|
|
2928
|
+
paramIndex++;
|
|
2929
|
+
updates.push(`"updatedAtZ" = $${paramIndex}`);
|
|
2930
|
+
values.push(updatedResource.updatedAt.toISOString());
|
|
2931
|
+
paramIndex++;
|
|
2932
|
+
values.push(resourceId);
|
|
2933
|
+
await this.#db.client.none(`UPDATE ${tableName} SET ${updates.join(", ")} WHERE id = $${paramIndex}`, values);
|
|
2934
|
+
return updatedResource;
|
|
2935
|
+
},
|
|
2936
|
+
{
|
|
2937
|
+
onRetry: (error, attempt, delay) => {
|
|
2938
|
+
this.logger?.warn?.(`updateResource retry ${attempt} for ${resourceId} after ${delay}ms: ${error.message}`);
|
|
2939
|
+
}
|
|
2940
|
+
}
|
|
2941
|
+
).catch((error) => {
|
|
2942
|
+
if (error instanceof MastraError) {
|
|
2943
|
+
throw error;
|
|
2944
|
+
}
|
|
2945
|
+
throw new MastraError(
|
|
2946
|
+
{
|
|
2947
|
+
id: createStorageErrorId("DSQL", "UPDATE_RESOURCE", "FAILED"),
|
|
2948
|
+
domain: ErrorDomain.STORAGE,
|
|
2949
|
+
category: ErrorCategory.THIRD_PARTY,
|
|
2950
|
+
details: {
|
|
2951
|
+
resourceId
|
|
2952
|
+
}
|
|
2953
|
+
},
|
|
2954
|
+
error
|
|
2955
|
+
);
|
|
2956
|
+
});
|
|
2957
|
+
return result;
|
|
2958
|
+
}
|
|
2959
|
+
};
|
|
2960
|
+
var ObservabilityDSQL = class _ObservabilityDSQL extends ObservabilityStorage {
|
|
2961
|
+
#db;
|
|
2962
|
+
#schema;
|
|
2963
|
+
#skipDefaultIndexes;
|
|
2964
|
+
#indexes;
|
|
2965
|
+
/** Tables managed by this domain */
|
|
2966
|
+
static MANAGED_TABLES = [TABLE_SPANS];
|
|
2967
|
+
constructor(config) {
|
|
2968
|
+
super();
|
|
2969
|
+
const { client, schemaName, skipDefaultIndexes, indexes } = resolveDsqlConfig(config);
|
|
2970
|
+
this.#db = new DsqlDB({ client, schemaName, skipDefaultIndexes });
|
|
2971
|
+
this.#schema = schemaName || "public";
|
|
2972
|
+
this.#skipDefaultIndexes = skipDefaultIndexes;
|
|
2973
|
+
this.#indexes = indexes?.filter((idx) => _ObservabilityDSQL.MANAGED_TABLES.includes(idx.table));
|
|
2974
|
+
}
|
|
2975
|
+
async init() {
|
|
2976
|
+
await this.#db.createTable({ tableName: TABLE_SPANS, schema: TABLE_SCHEMAS[TABLE_SPANS] });
|
|
2977
|
+
await this.createDefaultIndexes();
|
|
2978
|
+
await this.createCustomIndexes();
|
|
2979
|
+
}
|
|
2980
|
+
/**
|
|
2981
|
+
* Returns default index definitions for the observability domain tables.
|
|
2982
|
+
*/
|
|
2983
|
+
getDefaultIndexDefinitions() {
|
|
2984
|
+
const schemaPrefix = this.#schema !== "public" ? `${this.#schema}_` : "";
|
|
2985
|
+
return [
|
|
2986
|
+
{
|
|
2987
|
+
name: `${schemaPrefix}mastra_ai_spans_traceid_startedat_idx`,
|
|
2988
|
+
table: TABLE_SPANS,
|
|
2989
|
+
columns: ["traceId", "startedAt"]
|
|
2990
|
+
},
|
|
2991
|
+
{
|
|
2992
|
+
name: `${schemaPrefix}mastra_ai_spans_parentspanid_startedat_idx`,
|
|
2993
|
+
table: TABLE_SPANS,
|
|
2994
|
+
columns: ["parentSpanId", "startedAt"]
|
|
2995
|
+
},
|
|
2996
|
+
{
|
|
2997
|
+
name: `${schemaPrefix}mastra_ai_spans_name_idx`,
|
|
2998
|
+
table: TABLE_SPANS,
|
|
2999
|
+
columns: ["name"]
|
|
3000
|
+
},
|
|
3001
|
+
{
|
|
3002
|
+
name: `${schemaPrefix}mastra_ai_spans_spantype_startedat_idx`,
|
|
3003
|
+
table: TABLE_SPANS,
|
|
3004
|
+
columns: ["spanType", "startedAt"]
|
|
3005
|
+
},
|
|
3006
|
+
// Entity identification indexes - common filtering patterns
|
|
3007
|
+
{
|
|
3008
|
+
name: `${schemaPrefix}mastra_ai_spans_entitytype_entityid_idx`,
|
|
3009
|
+
table: TABLE_SPANS,
|
|
3010
|
+
columns: ["entityType", "entityId"]
|
|
3011
|
+
},
|
|
3012
|
+
{
|
|
3013
|
+
name: `${schemaPrefix}mastra_ai_spans_entitytype_entityname_idx`,
|
|
3014
|
+
table: TABLE_SPANS,
|
|
3015
|
+
columns: ["entityType", "entityName"]
|
|
3016
|
+
},
|
|
3017
|
+
// Multi-tenant filtering - organizationId + userId
|
|
3018
|
+
{
|
|
3019
|
+
name: `${schemaPrefix}mastra_ai_spans_orgid_userid_idx`,
|
|
3020
|
+
table: TABLE_SPANS,
|
|
3021
|
+
columns: ["organizationId", "userId"]
|
|
3022
|
+
}
|
|
3023
|
+
];
|
|
3024
|
+
}
|
|
3025
|
+
/**
|
|
3026
|
+
* Creates default indexes for optimal query performance.
|
|
3027
|
+
*/
|
|
3028
|
+
async createDefaultIndexes() {
|
|
3029
|
+
if (this.#skipDefaultIndexes) {
|
|
3030
|
+
return;
|
|
3031
|
+
}
|
|
3032
|
+
for (const indexDef of this.getDefaultIndexDefinitions()) {
|
|
3033
|
+
try {
|
|
3034
|
+
await this.#db.createIndex(indexDef);
|
|
3035
|
+
} catch (error) {
|
|
3036
|
+
this.logger?.warn?.(`Failed to create index ${indexDef.name}:`, error);
|
|
3037
|
+
}
|
|
3038
|
+
}
|
|
3039
|
+
}
|
|
3040
|
+
/**
|
|
3041
|
+
* Creates custom user-defined indexes for this domain's tables.
|
|
3042
|
+
*/
|
|
3043
|
+
async createCustomIndexes() {
|
|
3044
|
+
if (!this.#indexes || this.#indexes.length === 0) {
|
|
3045
|
+
return;
|
|
3046
|
+
}
|
|
3047
|
+
for (const indexDef of this.#indexes) {
|
|
3048
|
+
try {
|
|
3049
|
+
await this.#db.createIndex(indexDef);
|
|
3050
|
+
} catch (error) {
|
|
3051
|
+
this.logger?.warn?.(`Failed to create custom index ${indexDef.name}:`, error);
|
|
3052
|
+
}
|
|
3053
|
+
}
|
|
3054
|
+
}
|
|
3055
|
+
async dangerouslyClearAll() {
|
|
3056
|
+
await this.#db.clearTable({ tableName: TABLE_SPANS });
|
|
3057
|
+
}
|
|
3058
|
+
get tracingStrategy() {
|
|
3059
|
+
return {
|
|
3060
|
+
preferred: "batch-with-updates",
|
|
3061
|
+
supported: ["batch-with-updates", "insert-only"]
|
|
3062
|
+
};
|
|
3063
|
+
}
|
|
3064
|
+
async createSpan(args) {
|
|
3065
|
+
const { span } = args;
|
|
3066
|
+
try {
|
|
3067
|
+
const startedAt = span.startedAt instanceof Date ? span.startedAt.toISOString() : span.startedAt;
|
|
3068
|
+
const endedAt = span.endedAt instanceof Date ? span.endedAt.toISOString() : span.endedAt;
|
|
3069
|
+
const now = (/* @__PURE__ */ new Date()).toISOString();
|
|
3070
|
+
const record = {
|
|
3071
|
+
...span,
|
|
3072
|
+
startedAt,
|
|
3073
|
+
endedAt,
|
|
3074
|
+
startedAtZ: startedAt,
|
|
3075
|
+
endedAtZ: endedAt,
|
|
3076
|
+
// Aurora DSQL doesn't support triggers, so we set timestamps explicitly
|
|
3077
|
+
createdAt: now,
|
|
3078
|
+
updatedAt: now
|
|
3079
|
+
};
|
|
3080
|
+
await this.#db.insert({ tableName: TABLE_SPANS, record });
|
|
3081
|
+
} catch (error) {
|
|
3082
|
+
throw new MastraError(
|
|
3083
|
+
{
|
|
3084
|
+
id: createStorageErrorId("DSQL", "CREATE_SPAN", "FAILED"),
|
|
3085
|
+
domain: ErrorDomain.STORAGE,
|
|
3086
|
+
category: ErrorCategory.USER,
|
|
3087
|
+
details: {
|
|
3088
|
+
spanId: span.spanId,
|
|
3089
|
+
traceId: span.traceId,
|
|
3090
|
+
spanType: span.spanType,
|
|
3091
|
+
spanName: span.name
|
|
3092
|
+
}
|
|
3093
|
+
},
|
|
3094
|
+
error
|
|
3095
|
+
);
|
|
3096
|
+
}
|
|
3097
|
+
}
|
|
3098
|
+
async getSpan(args) {
|
|
3099
|
+
const { traceId, spanId } = args;
|
|
3100
|
+
try {
|
|
3101
|
+
const tableName = getTableName2({
|
|
3102
|
+
indexName: TABLE_SPANS,
|
|
3103
|
+
schemaName: getSchemaName2(this.#schema)
|
|
3104
|
+
});
|
|
3105
|
+
const row = await this.#db.client.oneOrNone(
|
|
3106
|
+
`SELECT
|
|
3107
|
+
"traceId", "spanId", "parentSpanId", "name",
|
|
3108
|
+
"entityType", "entityId", "entityName",
|
|
3109
|
+
"userId", "organizationId", "resourceId",
|
|
3110
|
+
"runId", "sessionId", "threadId", "requestId",
|
|
3111
|
+
"environment", "source", "serviceName", "scope",
|
|
3112
|
+
"spanType", "attributes", "metadata", "tags", "links",
|
|
3113
|
+
"input", "output", "error", "isEvent",
|
|
3114
|
+
"startedAtZ" as "startedAt", "endedAtZ" as "endedAt",
|
|
3115
|
+
"createdAtZ" as "createdAt", "updatedAtZ" as "updatedAt"
|
|
3116
|
+
FROM ${tableName}
|
|
3117
|
+
WHERE "traceId" = $1 AND "spanId" = $2`,
|
|
3118
|
+
[traceId, spanId]
|
|
3119
|
+
);
|
|
3120
|
+
if (!row) {
|
|
3121
|
+
return null;
|
|
3122
|
+
}
|
|
3123
|
+
return {
|
|
3124
|
+
span: transformFromSqlRow({
|
|
3125
|
+
tableName: TABLE_SPANS,
|
|
3126
|
+
sqlRow: row
|
|
3127
|
+
})
|
|
3128
|
+
};
|
|
3129
|
+
} catch (error) {
|
|
3130
|
+
throw new MastraError(
|
|
3131
|
+
{
|
|
3132
|
+
id: createStorageErrorId("DSQL", "GET_SPAN", "FAILED"),
|
|
3133
|
+
domain: ErrorDomain.STORAGE,
|
|
3134
|
+
category: ErrorCategory.USER,
|
|
3135
|
+
details: { traceId, spanId }
|
|
3136
|
+
},
|
|
3137
|
+
error
|
|
3138
|
+
);
|
|
3139
|
+
}
|
|
3140
|
+
}
|
|
3141
|
+
async getRootSpan(args) {
|
|
3142
|
+
const { traceId } = args;
|
|
3143
|
+
try {
|
|
3144
|
+
const tableName = getTableName2({
|
|
3145
|
+
indexName: TABLE_SPANS,
|
|
3146
|
+
schemaName: getSchemaName2(this.#schema)
|
|
3147
|
+
});
|
|
3148
|
+
const row = await this.#db.client.oneOrNone(
|
|
3149
|
+
`SELECT
|
|
3150
|
+
"traceId", "spanId", "parentSpanId", "name",
|
|
3151
|
+
"entityType", "entityId", "entityName",
|
|
3152
|
+
"userId", "organizationId", "resourceId",
|
|
3153
|
+
"runId", "sessionId", "threadId", "requestId",
|
|
3154
|
+
"environment", "source", "serviceName", "scope",
|
|
3155
|
+
"spanType", "attributes", "metadata", "tags", "links",
|
|
3156
|
+
"input", "output", "error", "isEvent",
|
|
3157
|
+
"startedAtZ" as "startedAt", "endedAtZ" as "endedAt",
|
|
3158
|
+
"createdAtZ" as "createdAt", "updatedAtZ" as "updatedAt"
|
|
3159
|
+
FROM ${tableName}
|
|
3160
|
+
WHERE "traceId" = $1 AND "parentSpanId" IS NULL`,
|
|
3161
|
+
[traceId]
|
|
3162
|
+
);
|
|
3163
|
+
if (!row) {
|
|
3164
|
+
return null;
|
|
3165
|
+
}
|
|
3166
|
+
return {
|
|
3167
|
+
span: transformFromSqlRow({
|
|
3168
|
+
tableName: TABLE_SPANS,
|
|
3169
|
+
sqlRow: row
|
|
3170
|
+
})
|
|
3171
|
+
};
|
|
3172
|
+
} catch (error) {
|
|
3173
|
+
throw new MastraError(
|
|
3174
|
+
{
|
|
3175
|
+
id: createStorageErrorId("DSQL", "GET_ROOT_SPAN", "FAILED"),
|
|
3176
|
+
domain: ErrorDomain.STORAGE,
|
|
3177
|
+
category: ErrorCategory.USER,
|
|
3178
|
+
details: { traceId }
|
|
3179
|
+
},
|
|
3180
|
+
error
|
|
3181
|
+
);
|
|
3182
|
+
}
|
|
3183
|
+
}
|
|
3184
|
+
async getTrace(args) {
|
|
3185
|
+
const { traceId } = args;
|
|
3186
|
+
try {
|
|
3187
|
+
const tableName = getTableName2({
|
|
3188
|
+
indexName: TABLE_SPANS,
|
|
3189
|
+
schemaName: getSchemaName2(this.#schema)
|
|
3190
|
+
});
|
|
3191
|
+
const spans = await this.#db.client.manyOrNone(
|
|
3192
|
+
`SELECT
|
|
3193
|
+
"traceId", "spanId", "parentSpanId", "name",
|
|
3194
|
+
"entityType", "entityId", "entityName",
|
|
3195
|
+
"userId", "organizationId", "resourceId",
|
|
3196
|
+
"runId", "sessionId", "threadId", "requestId",
|
|
3197
|
+
"environment", "source", "serviceName", "scope",
|
|
3198
|
+
"spanType", "attributes", "metadata", "tags", "links",
|
|
3199
|
+
"input", "output", "error", "isEvent",
|
|
3200
|
+
"startedAtZ" as "startedAt", "endedAtZ" as "endedAt",
|
|
3201
|
+
"createdAtZ" as "createdAt", "updatedAtZ" as "updatedAt"
|
|
3202
|
+
FROM ${tableName}
|
|
3203
|
+
WHERE "traceId" = $1
|
|
3204
|
+
ORDER BY "startedAtZ" ASC`,
|
|
3205
|
+
[traceId]
|
|
3206
|
+
);
|
|
3207
|
+
if (!spans || spans.length === 0) {
|
|
3208
|
+
return null;
|
|
3209
|
+
}
|
|
3210
|
+
return {
|
|
3211
|
+
traceId,
|
|
3212
|
+
spans: spans.map(
|
|
3213
|
+
(span) => transformFromSqlRow({
|
|
3214
|
+
tableName: TABLE_SPANS,
|
|
3215
|
+
sqlRow: span
|
|
3216
|
+
})
|
|
3217
|
+
)
|
|
3218
|
+
};
|
|
3219
|
+
} catch (error) {
|
|
3220
|
+
throw new MastraError(
|
|
3221
|
+
{
|
|
3222
|
+
id: createStorageErrorId("DSQL", "GET_TRACE", "FAILED"),
|
|
3223
|
+
domain: ErrorDomain.STORAGE,
|
|
3224
|
+
category: ErrorCategory.USER,
|
|
3225
|
+
details: {
|
|
3226
|
+
traceId
|
|
3227
|
+
}
|
|
3228
|
+
},
|
|
3229
|
+
error
|
|
3230
|
+
);
|
|
3231
|
+
}
|
|
3232
|
+
}
|
|
3233
|
+
async updateSpan(args) {
|
|
3234
|
+
const { traceId, spanId, updates } = args;
|
|
3235
|
+
try {
|
|
3236
|
+
const data = { ...updates };
|
|
3237
|
+
if (data.endedAt instanceof Date) {
|
|
3238
|
+
const endedAt = data.endedAt.toISOString();
|
|
3239
|
+
data.endedAt = endedAt;
|
|
3240
|
+
data.endedAtZ = endedAt;
|
|
3241
|
+
}
|
|
3242
|
+
if (data.startedAt instanceof Date) {
|
|
3243
|
+
const startedAt = data.startedAt.toISOString();
|
|
3244
|
+
data.startedAt = startedAt;
|
|
3245
|
+
data.startedAtZ = startedAt;
|
|
3246
|
+
}
|
|
3247
|
+
await this.#db.update({
|
|
3248
|
+
tableName: TABLE_SPANS,
|
|
3249
|
+
keys: { spanId, traceId },
|
|
3250
|
+
data
|
|
3251
|
+
});
|
|
3252
|
+
} catch (error) {
|
|
3253
|
+
throw new MastraError(
|
|
3254
|
+
{
|
|
3255
|
+
id: createStorageErrorId("DSQL", "UPDATE_SPAN", "FAILED"),
|
|
3256
|
+
domain: ErrorDomain.STORAGE,
|
|
3257
|
+
category: ErrorCategory.USER,
|
|
3258
|
+
details: {
|
|
3259
|
+
spanId,
|
|
3260
|
+
traceId
|
|
3261
|
+
}
|
|
3262
|
+
},
|
|
3263
|
+
error
|
|
3264
|
+
);
|
|
3265
|
+
}
|
|
3266
|
+
}
|
|
3267
|
+
async listTraces(args) {
|
|
3268
|
+
const { filters, pagination, orderBy } = listTracesArgsSchema.parse(args);
|
|
3269
|
+
const page = pagination?.page ?? 0;
|
|
3270
|
+
const perPage = pagination?.perPage ?? 100;
|
|
3271
|
+
const tableName = getTableName2({
|
|
3272
|
+
indexName: TABLE_SPANS,
|
|
3273
|
+
schemaName: getSchemaName2(this.#schema)
|
|
3274
|
+
});
|
|
3275
|
+
try {
|
|
3276
|
+
const conditions = ['r."parentSpanId" IS NULL'];
|
|
3277
|
+
const params = [];
|
|
3278
|
+
let paramIndex = 1;
|
|
3279
|
+
if (filters) {
|
|
3280
|
+
if (filters.startedAt?.start) {
|
|
3281
|
+
conditions.push(`r."startedAtZ" >= $${paramIndex++}`);
|
|
3282
|
+
params.push(filters.startedAt.start.toISOString());
|
|
3283
|
+
}
|
|
3284
|
+
if (filters.startedAt?.end) {
|
|
3285
|
+
conditions.push(`r."startedAtZ" <= $${paramIndex++}`);
|
|
3286
|
+
params.push(filters.startedAt.end.toISOString());
|
|
3287
|
+
}
|
|
3288
|
+
if (filters.endedAt?.start) {
|
|
3289
|
+
conditions.push(`r."endedAtZ" >= $${paramIndex++}`);
|
|
3290
|
+
params.push(filters.endedAt.start.toISOString());
|
|
3291
|
+
}
|
|
3292
|
+
if (filters.endedAt?.end) {
|
|
3293
|
+
conditions.push(`r."endedAtZ" <= $${paramIndex++}`);
|
|
3294
|
+
params.push(filters.endedAt.end.toISOString());
|
|
3295
|
+
}
|
|
3296
|
+
if (filters.spanType !== void 0) {
|
|
3297
|
+
conditions.push(`r."spanType" = $${paramIndex++}`);
|
|
3298
|
+
params.push(filters.spanType);
|
|
3299
|
+
}
|
|
3300
|
+
if (filters.entityType !== void 0) {
|
|
3301
|
+
conditions.push(`r."entityType" = $${paramIndex++}`);
|
|
3302
|
+
params.push(filters.entityType);
|
|
3303
|
+
}
|
|
3304
|
+
if (filters.entityId !== void 0) {
|
|
3305
|
+
conditions.push(`r."entityId" = $${paramIndex++}`);
|
|
3306
|
+
params.push(filters.entityId);
|
|
3307
|
+
}
|
|
3308
|
+
if (filters.entityName !== void 0) {
|
|
3309
|
+
conditions.push(`r."entityName" = $${paramIndex++}`);
|
|
3310
|
+
params.push(filters.entityName);
|
|
3311
|
+
}
|
|
3312
|
+
if (filters.userId !== void 0) {
|
|
3313
|
+
conditions.push(`r."userId" = $${paramIndex++}`);
|
|
3314
|
+
params.push(filters.userId);
|
|
3315
|
+
}
|
|
3316
|
+
if (filters.organizationId !== void 0) {
|
|
3317
|
+
conditions.push(`r."organizationId" = $${paramIndex++}`);
|
|
3318
|
+
params.push(filters.organizationId);
|
|
3319
|
+
}
|
|
3320
|
+
if (filters.resourceId !== void 0) {
|
|
3321
|
+
conditions.push(`r."resourceId" = $${paramIndex++}`);
|
|
3322
|
+
params.push(filters.resourceId);
|
|
3323
|
+
}
|
|
3324
|
+
if (filters.runId !== void 0) {
|
|
3325
|
+
conditions.push(`r."runId" = $${paramIndex++}`);
|
|
3326
|
+
params.push(filters.runId);
|
|
3327
|
+
}
|
|
3328
|
+
if (filters.sessionId !== void 0) {
|
|
3329
|
+
conditions.push(`r."sessionId" = $${paramIndex++}`);
|
|
3330
|
+
params.push(filters.sessionId);
|
|
3331
|
+
}
|
|
3332
|
+
if (filters.threadId !== void 0) {
|
|
3333
|
+
conditions.push(`r."threadId" = $${paramIndex++}`);
|
|
3334
|
+
params.push(filters.threadId);
|
|
3335
|
+
}
|
|
3336
|
+
if (filters.requestId !== void 0) {
|
|
3337
|
+
conditions.push(`r."requestId" = $${paramIndex++}`);
|
|
3338
|
+
params.push(filters.requestId);
|
|
3339
|
+
}
|
|
3340
|
+
if (filters.environment !== void 0) {
|
|
3341
|
+
conditions.push(`r."environment" = $${paramIndex++}`);
|
|
3342
|
+
params.push(filters.environment);
|
|
3343
|
+
}
|
|
3344
|
+
if (filters.source !== void 0) {
|
|
3345
|
+
conditions.push(`r."source" = $${paramIndex++}`);
|
|
3346
|
+
params.push(filters.source);
|
|
3347
|
+
}
|
|
3348
|
+
if (filters.serviceName !== void 0) {
|
|
3349
|
+
conditions.push(`r."serviceName" = $${paramIndex++}`);
|
|
3350
|
+
params.push(filters.serviceName);
|
|
3351
|
+
}
|
|
3352
|
+
if (filters.scope != null) {
|
|
3353
|
+
conditions.push(`r."scope"::text = $${paramIndex++}`);
|
|
3354
|
+
params.push(JSON.stringify(filters.scope));
|
|
3355
|
+
}
|
|
3356
|
+
if (filters.metadata != null) {
|
|
3357
|
+
conditions.push(`r."metadata"::text = $${paramIndex++}`);
|
|
3358
|
+
params.push(JSON.stringify(filters.metadata));
|
|
3359
|
+
}
|
|
3360
|
+
if (filters.tags != null && filters.tags.length > 0) {
|
|
3361
|
+
conditions.push(`r."tags"::text = $${paramIndex++}`);
|
|
3362
|
+
params.push(JSON.stringify(filters.tags));
|
|
3363
|
+
}
|
|
3364
|
+
if (filters.status !== void 0) {
|
|
3365
|
+
switch (filters.status) {
|
|
3366
|
+
case TraceStatus.ERROR:
|
|
3367
|
+
conditions.push(`r."error" IS NOT NULL`);
|
|
3368
|
+
break;
|
|
3369
|
+
case TraceStatus.RUNNING:
|
|
3370
|
+
conditions.push(`r."endedAtZ" IS NULL AND r."error" IS NULL`);
|
|
3371
|
+
break;
|
|
3372
|
+
case TraceStatus.SUCCESS:
|
|
3373
|
+
conditions.push(`r."endedAtZ" IS NOT NULL AND r."error" IS NULL`);
|
|
3374
|
+
break;
|
|
3375
|
+
}
|
|
3376
|
+
}
|
|
3377
|
+
if (filters.hasChildError !== void 0) {
|
|
3378
|
+
if (filters.hasChildError) {
|
|
3379
|
+
conditions.push(`EXISTS (
|
|
3380
|
+
SELECT 1 FROM ${tableName} c
|
|
3381
|
+
WHERE c."traceId" = r."traceId" AND c."error" IS NOT NULL
|
|
3382
|
+
)`);
|
|
3383
|
+
} else {
|
|
3384
|
+
conditions.push(`NOT EXISTS (
|
|
3385
|
+
SELECT 1 FROM ${tableName} c
|
|
3386
|
+
WHERE c."traceId" = r."traceId" AND c."error" IS NOT NULL
|
|
3387
|
+
)`);
|
|
3388
|
+
}
|
|
3389
|
+
}
|
|
3390
|
+
}
|
|
3391
|
+
const whereClause = conditions.length > 0 ? `WHERE ${conditions.join(" AND ")}` : "";
|
|
3392
|
+
const sortField = `${orderBy?.field ?? "startedAt"}Z`;
|
|
3393
|
+
const sortDirection = orderBy?.direction ?? "DESC";
|
|
3394
|
+
const orderClause = `ORDER BY r."${sortField}" ${sortDirection}`;
|
|
3395
|
+
const countResult = await this.#db.client.oneOrNone(
|
|
3396
|
+
`SELECT COUNT(*) FROM ${tableName} r ${whereClause}`,
|
|
3397
|
+
params
|
|
3398
|
+
);
|
|
3399
|
+
const count = Number(countResult?.count ?? 0);
|
|
3400
|
+
if (count === 0) {
|
|
3401
|
+
return {
|
|
3402
|
+
pagination: {
|
|
3403
|
+
total: 0,
|
|
3404
|
+
page,
|
|
3405
|
+
perPage,
|
|
3406
|
+
hasMore: false
|
|
3407
|
+
},
|
|
3408
|
+
spans: []
|
|
3409
|
+
};
|
|
3410
|
+
}
|
|
3411
|
+
const spans = await this.#db.client.manyOrNone(
|
|
3412
|
+
`SELECT
|
|
3413
|
+
r."traceId", r."spanId", r."parentSpanId", r."name",
|
|
3414
|
+
r."entityType", r."entityId", r."entityName",
|
|
3415
|
+
r."userId", r."organizationId", r."resourceId",
|
|
3416
|
+
r."runId", r."sessionId", r."threadId", r."requestId",
|
|
3417
|
+
r."environment", r."source", r."serviceName", r."scope",
|
|
3418
|
+
r."spanType", r."attributes", r."metadata", r."tags", r."links",
|
|
3419
|
+
r."input", r."output", r."error", r."isEvent",
|
|
3420
|
+
r."startedAtZ" as "startedAt", r."endedAtZ" as "endedAt",
|
|
3421
|
+
r."createdAtZ" as "createdAt", r."updatedAtZ" as "updatedAt"
|
|
3422
|
+
FROM ${tableName} r
|
|
3423
|
+
${whereClause}
|
|
3424
|
+
${orderClause}
|
|
3425
|
+
LIMIT $${paramIndex} OFFSET $${paramIndex + 1}`,
|
|
3426
|
+
[...params, perPage, page * perPage]
|
|
3427
|
+
);
|
|
3428
|
+
return {
|
|
3429
|
+
pagination: {
|
|
3430
|
+
total: count,
|
|
3431
|
+
page,
|
|
3432
|
+
perPage,
|
|
3433
|
+
hasMore: (page + 1) * perPage < count
|
|
3434
|
+
},
|
|
3435
|
+
spans: toTraceSpans(
|
|
3436
|
+
spans.map(
|
|
3437
|
+
(span) => transformFromSqlRow({
|
|
3438
|
+
tableName: TABLE_SPANS,
|
|
3439
|
+
sqlRow: span
|
|
3440
|
+
})
|
|
3441
|
+
)
|
|
3442
|
+
)
|
|
3443
|
+
};
|
|
3444
|
+
} catch (error) {
|
|
3445
|
+
throw new MastraError(
|
|
3446
|
+
{
|
|
3447
|
+
id: createStorageErrorId("DSQL", "LIST_TRACES", "FAILED"),
|
|
3448
|
+
domain: ErrorDomain.STORAGE,
|
|
3449
|
+
category: ErrorCategory.USER
|
|
3450
|
+
},
|
|
3451
|
+
error
|
|
3452
|
+
);
|
|
3453
|
+
}
|
|
3454
|
+
}
|
|
3455
|
+
async batchCreateSpans(args) {
|
|
3456
|
+
try {
|
|
3457
|
+
const now = (/* @__PURE__ */ new Date()).toISOString();
|
|
3458
|
+
const records = args.records.map((record) => {
|
|
3459
|
+
const startedAt = record.startedAt instanceof Date ? record.startedAt.toISOString() : record.startedAt;
|
|
3460
|
+
const endedAt = record.endedAt instanceof Date ? record.endedAt.toISOString() : record.endedAt;
|
|
3461
|
+
return {
|
|
3462
|
+
...record,
|
|
3463
|
+
startedAt,
|
|
3464
|
+
endedAt,
|
|
3465
|
+
startedAtZ: startedAt,
|
|
3466
|
+
endedAtZ: endedAt,
|
|
3467
|
+
// Aurora DSQL doesn't support triggers, so we set timestamps explicitly
|
|
3468
|
+
createdAt: now,
|
|
3469
|
+
updatedAt: now
|
|
3470
|
+
};
|
|
3471
|
+
});
|
|
3472
|
+
await this.#db.batchInsert({
|
|
3473
|
+
tableName: TABLE_SPANS,
|
|
3474
|
+
records
|
|
3475
|
+
});
|
|
3476
|
+
} catch (error) {
|
|
3477
|
+
throw new MastraError(
|
|
3478
|
+
{
|
|
3479
|
+
id: createStorageErrorId("DSQL", "BATCH_CREATE_SPANS", "FAILED"),
|
|
3480
|
+
domain: ErrorDomain.STORAGE,
|
|
3481
|
+
category: ErrorCategory.USER
|
|
3482
|
+
},
|
|
3483
|
+
error
|
|
3484
|
+
);
|
|
3485
|
+
}
|
|
3486
|
+
}
|
|
3487
|
+
async batchUpdateSpans(args) {
|
|
3488
|
+
try {
|
|
3489
|
+
await this.#db.batchUpdate({
|
|
3490
|
+
tableName: TABLE_SPANS,
|
|
3491
|
+
updates: args.records.map((record) => {
|
|
3492
|
+
const data = { ...record.updates };
|
|
3493
|
+
if (data.endedAt instanceof Date) {
|
|
3494
|
+
const endedAt = data.endedAt.toISOString();
|
|
3495
|
+
data.endedAt = endedAt;
|
|
3496
|
+
data.endedAtZ = endedAt;
|
|
3497
|
+
}
|
|
3498
|
+
if (data.startedAt instanceof Date) {
|
|
3499
|
+
const startedAt = data.startedAt.toISOString();
|
|
3500
|
+
data.startedAt = startedAt;
|
|
3501
|
+
data.startedAtZ = startedAt;
|
|
3502
|
+
}
|
|
3503
|
+
return {
|
|
3504
|
+
keys: { spanId: record.spanId, traceId: record.traceId },
|
|
3505
|
+
data
|
|
3506
|
+
};
|
|
3507
|
+
})
|
|
3508
|
+
});
|
|
3509
|
+
} catch (error) {
|
|
3510
|
+
throw new MastraError(
|
|
3511
|
+
{
|
|
3512
|
+
id: createStorageErrorId("DSQL", "BATCH_UPDATE_SPANS", "FAILED"),
|
|
3513
|
+
domain: ErrorDomain.STORAGE,
|
|
3514
|
+
category: ErrorCategory.USER
|
|
3515
|
+
},
|
|
3516
|
+
error
|
|
3517
|
+
);
|
|
3518
|
+
}
|
|
3519
|
+
}
|
|
3520
|
+
async batchDeleteTraces(args) {
|
|
3521
|
+
const { batches } = splitIntoBatches(args.traceIds, { maxRows: DEFAULT_MAX_ROWS_PER_BATCH });
|
|
3522
|
+
const tableName = getTableName2({
|
|
3523
|
+
indexName: TABLE_SPANS,
|
|
3524
|
+
schemaName: getSchemaName2(this.#schema)
|
|
3525
|
+
});
|
|
3526
|
+
for (const batchTraceIds of batches) {
|
|
3527
|
+
const placeholders = batchTraceIds.map((_, i) => `$${i + 1}`).join(", ");
|
|
3528
|
+
await withRetry(
|
|
3529
|
+
async () => {
|
|
3530
|
+
await this.#db.client.none(`DELETE FROM ${tableName} WHERE "traceId" IN (${placeholders})`, batchTraceIds);
|
|
3531
|
+
},
|
|
3532
|
+
{
|
|
3533
|
+
onRetry: (error, attempt, delay) => {
|
|
3534
|
+
this.logger?.warn?.(
|
|
3535
|
+
`batchDeleteTraces retry ${attempt} for ${batchTraceIds.length} traces after ${delay}ms: ${error.message}`
|
|
3536
|
+
);
|
|
3537
|
+
}
|
|
3538
|
+
}
|
|
3539
|
+
).catch((error) => {
|
|
3540
|
+
throw new MastraError(
|
|
3541
|
+
{
|
|
3542
|
+
id: createStorageErrorId("DSQL", "BATCH_DELETE_TRACES", "FAILED"),
|
|
3543
|
+
domain: ErrorDomain.STORAGE,
|
|
3544
|
+
category: ErrorCategory.USER
|
|
3545
|
+
},
|
|
3546
|
+
error
|
|
3547
|
+
);
|
|
3548
|
+
});
|
|
3549
|
+
}
|
|
3550
|
+
}
|
|
3551
|
+
};
|
|
3552
|
+
function transformScoreRow(row) {
|
|
3553
|
+
return transformScoreRow$1(row, {
|
|
3554
|
+
preferredTimestampFields: {
|
|
3555
|
+
createdAt: "createdAtZ",
|
|
3556
|
+
updatedAt: "updatedAtZ"
|
|
3557
|
+
}
|
|
3558
|
+
});
|
|
3559
|
+
}
|
|
3560
|
+
var ScoresDSQL = class _ScoresDSQL extends ScoresStorage {
|
|
3561
|
+
#db;
|
|
3562
|
+
#schema;
|
|
3563
|
+
#skipDefaultIndexes;
|
|
3564
|
+
#indexes;
|
|
3565
|
+
/** Tables managed by this domain */
|
|
3566
|
+
static MANAGED_TABLES = [TABLE_SCORERS];
|
|
3567
|
+
constructor(config) {
|
|
3568
|
+
super();
|
|
3569
|
+
const { client, schemaName, skipDefaultIndexes, indexes } = resolveDsqlConfig(config);
|
|
3570
|
+
this.#db = new DsqlDB({ client, schemaName, skipDefaultIndexes });
|
|
3571
|
+
this.#schema = schemaName || "public";
|
|
3572
|
+
this.#skipDefaultIndexes = skipDefaultIndexes;
|
|
3573
|
+
this.#indexes = indexes?.filter((idx) => _ScoresDSQL.MANAGED_TABLES.includes(idx.table));
|
|
3574
|
+
}
|
|
3575
|
+
async init() {
|
|
3576
|
+
await this.#db.createTable({ tableName: TABLE_SCORERS, schema: TABLE_SCHEMAS[TABLE_SCORERS] });
|
|
3577
|
+
await this.createDefaultIndexes();
|
|
3578
|
+
await this.createCustomIndexes();
|
|
3579
|
+
}
|
|
3580
|
+
/**
|
|
3581
|
+
* Returns default index definitions for the scores domain tables.
|
|
3582
|
+
* Note: Aurora DSQL does not support ASC/DESC in index columns.
|
|
3583
|
+
*/
|
|
3584
|
+
getDefaultIndexDefinitions() {
|
|
3585
|
+
const schemaPrefix = this.#schema !== "public" ? `${this.#schema}_` : "";
|
|
3586
|
+
return [
|
|
3587
|
+
{
|
|
3588
|
+
name: `${schemaPrefix}mastra_scores_trace_id_span_id_created_at_idx`,
|
|
3589
|
+
table: TABLE_SCORERS,
|
|
3590
|
+
columns: ["traceId", "spanId", "createdAt"]
|
|
3591
|
+
}
|
|
3592
|
+
];
|
|
3593
|
+
}
|
|
3594
|
+
/**
|
|
3595
|
+
* Creates default indexes for optimal query performance.
|
|
3596
|
+
*/
|
|
3597
|
+
async createDefaultIndexes() {
|
|
3598
|
+
if (this.#skipDefaultIndexes) {
|
|
3599
|
+
return;
|
|
3600
|
+
}
|
|
3601
|
+
for (const indexDef of this.getDefaultIndexDefinitions()) {
|
|
3602
|
+
try {
|
|
3603
|
+
await this.#db.createIndex(indexDef);
|
|
3604
|
+
} catch (error) {
|
|
3605
|
+
this.logger?.warn?.(`Failed to create index ${indexDef.name}:`, error);
|
|
3606
|
+
}
|
|
3607
|
+
}
|
|
3608
|
+
}
|
|
3609
|
+
/**
|
|
3610
|
+
* Creates custom user-defined indexes for this domain's tables.
|
|
3611
|
+
*/
|
|
3612
|
+
async createCustomIndexes() {
|
|
3613
|
+
if (!this.#indexes || this.#indexes.length === 0) {
|
|
3614
|
+
return;
|
|
3615
|
+
}
|
|
3616
|
+
for (const indexDef of this.#indexes) {
|
|
3617
|
+
try {
|
|
3618
|
+
await this.#db.createIndex(indexDef);
|
|
3619
|
+
} catch (error) {
|
|
3620
|
+
this.logger?.warn?.(`Failed to create custom index ${indexDef.name}:`, error);
|
|
3621
|
+
}
|
|
3622
|
+
}
|
|
3623
|
+
}
|
|
3624
|
+
async dangerouslyClearAll() {
|
|
3625
|
+
await this.#db.clearTable({ tableName: TABLE_SCORERS });
|
|
3626
|
+
}
|
|
3627
|
+
async getScoreById({ id }) {
|
|
3628
|
+
try {
|
|
3629
|
+
const result = await this.#db.client.oneOrNone(
|
|
3630
|
+
`SELECT * FROM ${getTableName2({ indexName: TABLE_SCORERS, schemaName: getSchemaName2(this.#schema) })} WHERE id = $1`,
|
|
3631
|
+
[id]
|
|
3632
|
+
);
|
|
3633
|
+
return result ? transformScoreRow(result) : null;
|
|
3634
|
+
} catch (error) {
|
|
3635
|
+
throw new MastraError(
|
|
3636
|
+
{
|
|
3637
|
+
id: createStorageErrorId("DSQL", "GET_SCORE_BY_ID", "FAILED"),
|
|
3638
|
+
domain: ErrorDomain.STORAGE,
|
|
3639
|
+
category: ErrorCategory.THIRD_PARTY
|
|
3640
|
+
},
|
|
3641
|
+
error
|
|
3642
|
+
);
|
|
3643
|
+
}
|
|
3644
|
+
}
|
|
3645
|
+
async listScoresByScorerId({
|
|
3646
|
+
scorerId,
|
|
3647
|
+
pagination,
|
|
3648
|
+
entityId,
|
|
3649
|
+
entityType,
|
|
3650
|
+
source
|
|
3651
|
+
}) {
|
|
3652
|
+
try {
|
|
3653
|
+
const conditions = [`"scorerId" = $1`];
|
|
3654
|
+
const queryParams = [scorerId];
|
|
3655
|
+
let paramIndex = 2;
|
|
3656
|
+
if (entityId) {
|
|
3657
|
+
conditions.push(`"entityId" = $${paramIndex++}`);
|
|
3658
|
+
queryParams.push(entityId);
|
|
3659
|
+
}
|
|
3660
|
+
if (entityType) {
|
|
3661
|
+
conditions.push(`"entityType" = $${paramIndex++}`);
|
|
3662
|
+
queryParams.push(entityType);
|
|
3663
|
+
}
|
|
3664
|
+
if (source) {
|
|
3665
|
+
conditions.push(`"source" = $${paramIndex++}`);
|
|
3666
|
+
queryParams.push(source);
|
|
3667
|
+
}
|
|
3668
|
+
const whereClause = conditions.join(" AND ");
|
|
3669
|
+
const total = await this.#db.client.oneOrNone(
|
|
3670
|
+
`SELECT COUNT(*) FROM ${getTableName2({ indexName: TABLE_SCORERS, schemaName: getSchemaName2(this.#schema) })} WHERE ${whereClause}`,
|
|
3671
|
+
queryParams
|
|
3672
|
+
);
|
|
3673
|
+
const { page, perPage: perPageInput } = pagination;
|
|
3674
|
+
const perPage = normalizePerPage(perPageInput, 100);
|
|
3675
|
+
const { offset: start, perPage: perPageForResponse } = calculatePagination(page, perPageInput, perPage);
|
|
3676
|
+
if (total?.count === "0" || !total?.count) {
|
|
3677
|
+
return {
|
|
3678
|
+
pagination: {
|
|
3679
|
+
total: 0,
|
|
3680
|
+
page,
|
|
3681
|
+
perPage: perPageForResponse,
|
|
3682
|
+
hasMore: false
|
|
3683
|
+
},
|
|
3684
|
+
scores: []
|
|
3685
|
+
};
|
|
3686
|
+
}
|
|
3687
|
+
const limitValue = perPageInput === false ? Number(total?.count) : perPage;
|
|
3688
|
+
const end = perPageInput === false ? Number(total?.count) : start + perPage;
|
|
3689
|
+
const result = await this.#db.client.manyOrNone(
|
|
3690
|
+
`SELECT * FROM ${getTableName2({ indexName: TABLE_SCORERS, schemaName: getSchemaName2(this.#schema) })} WHERE ${whereClause} ORDER BY "createdAt" DESC LIMIT $${paramIndex++} OFFSET $${paramIndex++}`,
|
|
3691
|
+
[...queryParams, limitValue, start]
|
|
3692
|
+
);
|
|
3693
|
+
return {
|
|
3694
|
+
pagination: {
|
|
3695
|
+
total: Number(total?.count) || 0,
|
|
3696
|
+
page,
|
|
3697
|
+
perPage: perPageForResponse,
|
|
3698
|
+
hasMore: end < Number(total?.count)
|
|
3699
|
+
},
|
|
3700
|
+
scores: result.map(transformScoreRow)
|
|
3701
|
+
};
|
|
3702
|
+
} catch (error) {
|
|
3703
|
+
throw new MastraError(
|
|
3704
|
+
{
|
|
3705
|
+
id: createStorageErrorId("DSQL", "LIST_SCORES_BY_SCORER_ID", "FAILED"),
|
|
3706
|
+
domain: ErrorDomain.STORAGE,
|
|
3707
|
+
category: ErrorCategory.THIRD_PARTY
|
|
3708
|
+
},
|
|
3709
|
+
error
|
|
3710
|
+
);
|
|
3711
|
+
}
|
|
3712
|
+
}
|
|
3713
|
+
async saveScore(score) {
|
|
3714
|
+
let parsedScore;
|
|
3715
|
+
try {
|
|
3716
|
+
parsedScore = saveScorePayloadSchema.parse(score);
|
|
3717
|
+
} catch (error) {
|
|
3718
|
+
throw new MastraError(
|
|
3719
|
+
{
|
|
3720
|
+
id: createStorageErrorId("DSQL", "SAVE_SCORE", "VALIDATION_FAILED"),
|
|
3721
|
+
domain: ErrorDomain.STORAGE,
|
|
3722
|
+
category: ErrorCategory.USER,
|
|
3723
|
+
details: {
|
|
3724
|
+
scorer: typeof score.scorer?.id === "string" ? score.scorer.id : String(score.scorer?.id ?? "unknown"),
|
|
3725
|
+
entityId: score.entityId ?? "unknown",
|
|
3726
|
+
entityType: score.entityType ?? "unknown",
|
|
3727
|
+
traceId: score.traceId ?? "",
|
|
3728
|
+
spanId: score.spanId ?? ""
|
|
3729
|
+
}
|
|
3730
|
+
},
|
|
3731
|
+
error
|
|
3732
|
+
);
|
|
3733
|
+
}
|
|
3734
|
+
try {
|
|
3735
|
+
const id = crypto.randomUUID();
|
|
3736
|
+
const now = /* @__PURE__ */ new Date();
|
|
3737
|
+
const {
|
|
3738
|
+
scorer,
|
|
3739
|
+
preprocessStepResult,
|
|
3740
|
+
analyzeStepResult,
|
|
3741
|
+
metadata,
|
|
3742
|
+
input,
|
|
3743
|
+
output,
|
|
3744
|
+
additionalContext,
|
|
3745
|
+
requestContext,
|
|
3746
|
+
entity,
|
|
3747
|
+
...rest
|
|
3748
|
+
} = parsedScore;
|
|
3749
|
+
await this.#db.insert({
|
|
3750
|
+
tableName: TABLE_SCORERS,
|
|
3751
|
+
record: {
|
|
3752
|
+
id,
|
|
3753
|
+
...rest,
|
|
3754
|
+
input: JSON.stringify(input) || "",
|
|
3755
|
+
output: JSON.stringify(output) || "",
|
|
3756
|
+
scorer: scorer ? JSON.stringify(scorer) : null,
|
|
3757
|
+
preprocessStepResult: preprocessStepResult ? JSON.stringify(preprocessStepResult) : null,
|
|
3758
|
+
analyzeStepResult: analyzeStepResult ? JSON.stringify(analyzeStepResult) : null,
|
|
3759
|
+
metadata: metadata ? JSON.stringify(metadata) : null,
|
|
3760
|
+
additionalContext: additionalContext ? JSON.stringify(additionalContext) : null,
|
|
3761
|
+
requestContext: requestContext ? JSON.stringify(requestContext) : null,
|
|
3762
|
+
entity: entity ? JSON.stringify(entity) : null,
|
|
3763
|
+
createdAt: now.toISOString(),
|
|
3764
|
+
updatedAt: now.toISOString()
|
|
3765
|
+
}
|
|
3766
|
+
});
|
|
3767
|
+
return { score: { ...parsedScore, id, createdAt: now, updatedAt: now } };
|
|
3768
|
+
} catch (error) {
|
|
3769
|
+
throw new MastraError(
|
|
3770
|
+
{
|
|
3771
|
+
id: createStorageErrorId("DSQL", "SAVE_SCORE", "FAILED"),
|
|
3772
|
+
domain: ErrorDomain.STORAGE,
|
|
3773
|
+
category: ErrorCategory.THIRD_PARTY
|
|
3774
|
+
},
|
|
3775
|
+
error
|
|
3776
|
+
);
|
|
3777
|
+
}
|
|
3778
|
+
}
|
|
3779
|
+
async listScoresByRunId({
|
|
3780
|
+
runId,
|
|
3781
|
+
pagination
|
|
3782
|
+
}) {
|
|
3783
|
+
try {
|
|
3784
|
+
const total = await this.#db.client.oneOrNone(
|
|
3785
|
+
`SELECT COUNT(*) FROM ${getTableName2({ indexName: TABLE_SCORERS, schemaName: getSchemaName2(this.#schema) })} WHERE "runId" = $1`,
|
|
3786
|
+
[runId]
|
|
3787
|
+
);
|
|
3788
|
+
const { page, perPage: perPageInput } = pagination;
|
|
3789
|
+
const perPage = normalizePerPage(perPageInput, 100);
|
|
3790
|
+
const { offset: start, perPage: perPageForResponse } = calculatePagination(page, perPageInput, perPage);
|
|
3791
|
+
if (total?.count === "0" || !total?.count) {
|
|
3792
|
+
return {
|
|
3793
|
+
pagination: {
|
|
3794
|
+
total: 0,
|
|
3795
|
+
page,
|
|
3796
|
+
perPage: perPageForResponse,
|
|
3797
|
+
hasMore: false
|
|
3798
|
+
},
|
|
3799
|
+
scores: []
|
|
3800
|
+
};
|
|
3801
|
+
}
|
|
3802
|
+
const limitValue = perPageInput === false ? Number(total?.count) : perPage;
|
|
3803
|
+
const end = perPageInput === false ? Number(total?.count) : start + perPage;
|
|
3804
|
+
const result = await this.#db.client.manyOrNone(
|
|
3805
|
+
`SELECT * FROM ${getTableName2({ indexName: TABLE_SCORERS, schemaName: getSchemaName2(this.#schema) })} WHERE "runId" = $1 ORDER BY "createdAt" DESC LIMIT $2 OFFSET $3`,
|
|
3806
|
+
[runId, limitValue, start]
|
|
3807
|
+
);
|
|
3808
|
+
return {
|
|
3809
|
+
pagination: {
|
|
3810
|
+
total: Number(total?.count) || 0,
|
|
3811
|
+
page,
|
|
3812
|
+
perPage: perPageForResponse,
|
|
3813
|
+
hasMore: end < Number(total?.count)
|
|
3814
|
+
},
|
|
3815
|
+
scores: result.map(transformScoreRow)
|
|
3816
|
+
};
|
|
3817
|
+
} catch (error) {
|
|
3818
|
+
throw new MastraError(
|
|
3819
|
+
{
|
|
3820
|
+
id: createStorageErrorId("DSQL", "LIST_SCORES_BY_RUN_ID", "FAILED"),
|
|
3821
|
+
domain: ErrorDomain.STORAGE,
|
|
3822
|
+
category: ErrorCategory.THIRD_PARTY
|
|
3823
|
+
},
|
|
3824
|
+
error
|
|
3825
|
+
);
|
|
3826
|
+
}
|
|
3827
|
+
}
|
|
3828
|
+
async listScoresByEntityId({
|
|
3829
|
+
entityId,
|
|
3830
|
+
entityType,
|
|
3831
|
+
pagination
|
|
3832
|
+
}) {
|
|
3833
|
+
try {
|
|
3834
|
+
const total = await this.#db.client.oneOrNone(
|
|
3835
|
+
`SELECT COUNT(*) FROM ${getTableName2({ indexName: TABLE_SCORERS, schemaName: getSchemaName2(this.#schema) })} WHERE "entityId" = $1 AND "entityType" = $2`,
|
|
3836
|
+
[entityId, entityType]
|
|
3837
|
+
);
|
|
3838
|
+
const { page, perPage: perPageInput } = pagination;
|
|
3839
|
+
const perPage = normalizePerPage(perPageInput, 100);
|
|
3840
|
+
const { offset: start, perPage: perPageForResponse } = calculatePagination(page, perPageInput, perPage);
|
|
3841
|
+
if (total?.count === "0" || !total?.count) {
|
|
3842
|
+
return {
|
|
3843
|
+
pagination: {
|
|
3844
|
+
total: 0,
|
|
3845
|
+
page,
|
|
3846
|
+
perPage: perPageForResponse,
|
|
3847
|
+
hasMore: false
|
|
3848
|
+
},
|
|
3849
|
+
scores: []
|
|
3850
|
+
};
|
|
3851
|
+
}
|
|
3852
|
+
const limitValue = perPageInput === false ? Number(total?.count) : perPage;
|
|
3853
|
+
const end = perPageInput === false ? Number(total?.count) : start + perPage;
|
|
3854
|
+
const result = await this.#db.client.manyOrNone(
|
|
3855
|
+
`SELECT * FROM ${getTableName2({ indexName: TABLE_SCORERS, schemaName: getSchemaName2(this.#schema) })} WHERE "entityId" = $1 AND "entityType" = $2 ORDER BY "createdAt" DESC LIMIT $3 OFFSET $4`,
|
|
3856
|
+
[entityId, entityType, limitValue, start]
|
|
3857
|
+
);
|
|
3858
|
+
return {
|
|
3859
|
+
pagination: {
|
|
3860
|
+
total: Number(total?.count) || 0,
|
|
3861
|
+
page,
|
|
3862
|
+
perPage: perPageForResponse,
|
|
3863
|
+
hasMore: end < Number(total?.count)
|
|
3864
|
+
},
|
|
3865
|
+
scores: result.map(transformScoreRow)
|
|
3866
|
+
};
|
|
3867
|
+
} catch (error) {
|
|
3868
|
+
throw new MastraError(
|
|
3869
|
+
{
|
|
3870
|
+
id: createStorageErrorId("DSQL", "LIST_SCORES_BY_ENTITY_ID", "FAILED"),
|
|
3871
|
+
domain: ErrorDomain.STORAGE,
|
|
3872
|
+
category: ErrorCategory.THIRD_PARTY
|
|
3873
|
+
},
|
|
3874
|
+
error
|
|
3875
|
+
);
|
|
3876
|
+
}
|
|
3877
|
+
}
|
|
3878
|
+
async listScoresBySpan({
|
|
3879
|
+
traceId,
|
|
3880
|
+
spanId,
|
|
3881
|
+
pagination
|
|
3882
|
+
}) {
|
|
3883
|
+
try {
|
|
3884
|
+
const tableName = getTableName2({ indexName: TABLE_SCORERS, schemaName: getSchemaName2(this.#schema) });
|
|
3885
|
+
const countSQLResult = await this.#db.client.oneOrNone(
|
|
3886
|
+
`SELECT COUNT(*) as count FROM ${tableName} WHERE "traceId" = $1 AND "spanId" = $2`,
|
|
3887
|
+
[traceId, spanId]
|
|
3888
|
+
);
|
|
3889
|
+
const total = Number(countSQLResult?.count ?? 0);
|
|
3890
|
+
const { page, perPage: perPageInput } = pagination;
|
|
3891
|
+
const perPage = normalizePerPage(perPageInput, 100);
|
|
3892
|
+
const { offset: start, perPage: perPageForResponse } = calculatePagination(page, perPageInput, perPage);
|
|
3893
|
+
const limitValue = perPageInput === false ? total : perPage;
|
|
3894
|
+
const end = perPageInput === false ? total : start + perPage;
|
|
3895
|
+
const result = await this.#db.client.manyOrNone(
|
|
3896
|
+
`SELECT * FROM ${tableName} WHERE "traceId" = $1 AND "spanId" = $2 ORDER BY "createdAt" DESC LIMIT $3 OFFSET $4`,
|
|
3897
|
+
[traceId, spanId, limitValue, start]
|
|
3898
|
+
);
|
|
3899
|
+
const hasMore = end < total;
|
|
3900
|
+
const scores = result.map((row) => transformScoreRow(row)) ?? [];
|
|
3901
|
+
return {
|
|
3902
|
+
scores,
|
|
3903
|
+
pagination: {
|
|
3904
|
+
total,
|
|
3905
|
+
page,
|
|
3906
|
+
perPage: perPageForResponse,
|
|
3907
|
+
hasMore
|
|
3908
|
+
}
|
|
3909
|
+
};
|
|
3910
|
+
} catch (error) {
|
|
3911
|
+
throw new MastraError(
|
|
3912
|
+
{
|
|
3913
|
+
id: createStorageErrorId("DSQL", "LIST_SCORES_BY_SPAN", "FAILED"),
|
|
3914
|
+
domain: ErrorDomain.STORAGE,
|
|
3915
|
+
category: ErrorCategory.THIRD_PARTY
|
|
3916
|
+
},
|
|
3917
|
+
error
|
|
3918
|
+
);
|
|
3919
|
+
}
|
|
3920
|
+
}
|
|
3921
|
+
};
|
|
3922
|
+
function parseWorkflowRun(row) {
|
|
3923
|
+
let parsedSnapshot = row.snapshot;
|
|
3924
|
+
if (typeof parsedSnapshot === "string") {
|
|
3925
|
+
try {
|
|
3926
|
+
parsedSnapshot = JSON.parse(row.snapshot);
|
|
3927
|
+
} catch (e) {
|
|
3928
|
+
console.warn(`Failed to parse snapshot for workflow ${row.workflow_name}: ${e}`);
|
|
3929
|
+
}
|
|
3930
|
+
}
|
|
3931
|
+
return {
|
|
3932
|
+
workflowName: row.workflow_name,
|
|
3933
|
+
runId: row.run_id,
|
|
3934
|
+
snapshot: parsedSnapshot,
|
|
3935
|
+
resourceId: row.resourceId,
|
|
3936
|
+
createdAt: new Date(row.createdAtZ || row.createdAt),
|
|
3937
|
+
updatedAt: new Date(row.updatedAtZ || row.updatedAt)
|
|
3938
|
+
};
|
|
3939
|
+
}
|
|
3940
|
+
var WorkflowsDSQL = class _WorkflowsDSQL extends WorkflowsStorage {
|
|
3941
|
+
#db;
|
|
3942
|
+
#schema;
|
|
3943
|
+
#skipDefaultIndexes;
|
|
3944
|
+
#indexes;
|
|
3945
|
+
/** Tables managed by this domain */
|
|
3946
|
+
static MANAGED_TABLES = [TABLE_WORKFLOW_SNAPSHOT];
|
|
3947
|
+
constructor(config) {
|
|
3948
|
+
super();
|
|
3949
|
+
const { client, schemaName, skipDefaultIndexes, indexes } = resolveDsqlConfig(config);
|
|
3950
|
+
this.#db = new DsqlDB({ client, schemaName });
|
|
3951
|
+
this.#schema = schemaName || "public";
|
|
3952
|
+
this.#skipDefaultIndexes = skipDefaultIndexes;
|
|
3953
|
+
this.#indexes = indexes?.filter((idx) => _WorkflowsDSQL.MANAGED_TABLES.includes(idx.table));
|
|
3954
|
+
}
|
|
3955
|
+
supportsConcurrentUpdates() {
|
|
3956
|
+
return true;
|
|
3957
|
+
}
|
|
3958
|
+
/**
|
|
3959
|
+
* Returns default index definitions for the workflows domain tables.
|
|
3960
|
+
* Currently no default indexes are defined for workflows.
|
|
3961
|
+
*/
|
|
3962
|
+
getDefaultIndexDefinitions() {
|
|
3963
|
+
return [];
|
|
3964
|
+
}
|
|
3965
|
+
/**
|
|
3966
|
+
* Creates default indexes for optimal query performance.
|
|
3967
|
+
* Currently no default indexes are defined for workflows.
|
|
3968
|
+
*/
|
|
3969
|
+
async createDefaultIndexes() {
|
|
3970
|
+
if (this.#skipDefaultIndexes) {
|
|
3971
|
+
return;
|
|
3972
|
+
}
|
|
3973
|
+
}
|
|
3974
|
+
async init() {
|
|
3975
|
+
await this.#db.createTable({ tableName: TABLE_WORKFLOW_SNAPSHOT, schema: TABLE_SCHEMAS[TABLE_WORKFLOW_SNAPSHOT] });
|
|
3976
|
+
await this.#db.alterTable({
|
|
3977
|
+
tableName: TABLE_WORKFLOW_SNAPSHOT,
|
|
3978
|
+
schema: TABLE_SCHEMAS[TABLE_WORKFLOW_SNAPSHOT],
|
|
3979
|
+
ifNotExists: ["resourceId"]
|
|
3980
|
+
});
|
|
3981
|
+
await this.createDefaultIndexes();
|
|
3982
|
+
await this.createCustomIndexes();
|
|
3983
|
+
}
|
|
3984
|
+
/**
|
|
3985
|
+
* Creates custom user-defined indexes for this domain's tables.
|
|
3986
|
+
*/
|
|
3987
|
+
async createCustomIndexes() {
|
|
3988
|
+
if (!this.#indexes || this.#indexes.length === 0) {
|
|
3989
|
+
return;
|
|
3990
|
+
}
|
|
3991
|
+
for (const indexDef of this.#indexes) {
|
|
3992
|
+
try {
|
|
3993
|
+
await this.#db.createIndex(indexDef);
|
|
3994
|
+
} catch (error) {
|
|
3995
|
+
this.logger?.warn?.(`Failed to create custom index ${indexDef.name}:`, error);
|
|
3996
|
+
}
|
|
3997
|
+
}
|
|
3998
|
+
}
|
|
3999
|
+
async dangerouslyClearAll() {
|
|
4000
|
+
await this.#db.clearTable({ tableName: TABLE_WORKFLOW_SNAPSHOT });
|
|
4001
|
+
}
|
|
4002
|
+
async updateWorkflowResults({
|
|
4003
|
+
workflowName,
|
|
4004
|
+
runId,
|
|
4005
|
+
stepId,
|
|
4006
|
+
result,
|
|
4007
|
+
requestContext
|
|
4008
|
+
}) {
|
|
4009
|
+
try {
|
|
4010
|
+
const { result: context } = await withRetry(
|
|
4011
|
+
async () => {
|
|
4012
|
+
return this.#db.client.tx(async (t) => {
|
|
4013
|
+
const tableName = getTableName2({
|
|
4014
|
+
indexName: TABLE_WORKFLOW_SNAPSHOT,
|
|
4015
|
+
schemaName: getSchemaName2(this.#schema)
|
|
4016
|
+
});
|
|
4017
|
+
const existingSnapshotResult = await t.oneOrNone(
|
|
4018
|
+
`SELECT snapshot FROM ${tableName} WHERE workflow_name = $1 AND run_id = $2`,
|
|
4019
|
+
[workflowName, runId]
|
|
4020
|
+
);
|
|
4021
|
+
let snapshot;
|
|
4022
|
+
if (!existingSnapshotResult) {
|
|
4023
|
+
snapshot = {
|
|
4024
|
+
context: {},
|
|
4025
|
+
activePaths: [],
|
|
4026
|
+
timestamp: Date.now(),
|
|
4027
|
+
suspendedPaths: {},
|
|
4028
|
+
activeStepsPath: {},
|
|
4029
|
+
resumeLabels: {},
|
|
4030
|
+
serializedStepGraph: [],
|
|
4031
|
+
status: "pending",
|
|
4032
|
+
value: {},
|
|
4033
|
+
waitingPaths: {},
|
|
4034
|
+
runId,
|
|
4035
|
+
requestContext: {}
|
|
4036
|
+
};
|
|
4037
|
+
} else {
|
|
4038
|
+
const existingSnapshot = existingSnapshotResult.snapshot;
|
|
4039
|
+
snapshot = typeof existingSnapshot === "string" ? JSON.parse(existingSnapshot) : existingSnapshot;
|
|
4040
|
+
}
|
|
4041
|
+
snapshot.context[stepId] = result;
|
|
4042
|
+
snapshot.requestContext = { ...snapshot.requestContext, ...requestContext };
|
|
4043
|
+
const now = /* @__PURE__ */ new Date();
|
|
4044
|
+
await t.none(
|
|
4045
|
+
`INSERT INTO ${tableName} (workflow_name, run_id, snapshot, "createdAt", "updatedAt")
|
|
4046
|
+
VALUES ($1, $2, $3, $4, $5)
|
|
4047
|
+
ON CONFLICT (workflow_name, run_id) DO UPDATE
|
|
4048
|
+
SET snapshot = $3, "updatedAt" = $5`,
|
|
4049
|
+
[workflowName, runId, JSON.stringify(snapshot), now, now]
|
|
4050
|
+
);
|
|
4051
|
+
return snapshot.context;
|
|
4052
|
+
});
|
|
4053
|
+
},
|
|
4054
|
+
{
|
|
4055
|
+
onRetry: (error, attempt, delay) => {
|
|
4056
|
+
this.logger?.warn?.(
|
|
4057
|
+
`updateWorkflowResults retry ${attempt} for workflow ${workflowName}/${runId} after ${delay}ms: ${error.message}`
|
|
4058
|
+
);
|
|
4059
|
+
}
|
|
4060
|
+
}
|
|
4061
|
+
);
|
|
4062
|
+
return context;
|
|
4063
|
+
} catch (error) {
|
|
4064
|
+
throw new MastraError(
|
|
4065
|
+
{
|
|
4066
|
+
id: createStorageErrorId("DSQL", "UPDATE_WORKFLOW_RESULTS", "FAILED"),
|
|
4067
|
+
domain: ErrorDomain.STORAGE,
|
|
4068
|
+
category: ErrorCategory.THIRD_PARTY,
|
|
4069
|
+
details: {
|
|
4070
|
+
workflowName,
|
|
4071
|
+
runId,
|
|
4072
|
+
stepId
|
|
4073
|
+
}
|
|
4074
|
+
},
|
|
4075
|
+
error
|
|
4076
|
+
);
|
|
4077
|
+
}
|
|
4078
|
+
}
|
|
4079
|
+
async updateWorkflowState({
|
|
4080
|
+
workflowName,
|
|
4081
|
+
runId,
|
|
4082
|
+
opts
|
|
4083
|
+
}) {
|
|
4084
|
+
try {
|
|
4085
|
+
const { result } = await withRetry(
|
|
4086
|
+
async () => {
|
|
4087
|
+
return this.#db.client.tx(async (t) => {
|
|
4088
|
+
const tableName = getTableName2({
|
|
4089
|
+
indexName: TABLE_WORKFLOW_SNAPSHOT,
|
|
4090
|
+
schemaName: getSchemaName2(this.#schema)
|
|
4091
|
+
});
|
|
4092
|
+
const existingSnapshotResult = await t.oneOrNone(
|
|
4093
|
+
`SELECT snapshot FROM ${tableName} WHERE workflow_name = $1 AND run_id = $2`,
|
|
4094
|
+
[workflowName, runId]
|
|
4095
|
+
);
|
|
4096
|
+
if (!existingSnapshotResult) {
|
|
4097
|
+
return void 0;
|
|
4098
|
+
}
|
|
4099
|
+
const existingSnapshot = existingSnapshotResult.snapshot;
|
|
4100
|
+
const snapshot = typeof existingSnapshot === "string" ? JSON.parse(existingSnapshot) : existingSnapshot;
|
|
4101
|
+
if (!snapshot || !snapshot?.context) {
|
|
4102
|
+
throw new Error(`Snapshot not found for runId ${runId}`);
|
|
4103
|
+
}
|
|
4104
|
+
const updatedSnapshot = { ...snapshot, ...opts };
|
|
4105
|
+
await t.none(
|
|
4106
|
+
`UPDATE ${tableName} SET snapshot = $1, "updatedAt" = $2 WHERE workflow_name = $3 AND run_id = $4`,
|
|
4107
|
+
[JSON.stringify(updatedSnapshot), /* @__PURE__ */ new Date(), workflowName, runId]
|
|
4108
|
+
);
|
|
4109
|
+
return updatedSnapshot;
|
|
4110
|
+
});
|
|
4111
|
+
},
|
|
4112
|
+
{
|
|
4113
|
+
onRetry: (error, attempt, delay) => {
|
|
4114
|
+
this.logger?.warn?.(
|
|
4115
|
+
`updateWorkflowState retry ${attempt} for workflow ${workflowName}/${runId} after ${delay}ms: ${error.message}`
|
|
4116
|
+
);
|
|
4117
|
+
}
|
|
4118
|
+
}
|
|
4119
|
+
);
|
|
4120
|
+
return result;
|
|
4121
|
+
} catch (error) {
|
|
4122
|
+
throw new MastraError(
|
|
4123
|
+
{
|
|
4124
|
+
id: createStorageErrorId("DSQL", "UPDATE_WORKFLOW_STATE", "FAILED"),
|
|
4125
|
+
domain: ErrorDomain.STORAGE,
|
|
4126
|
+
category: ErrorCategory.THIRD_PARTY,
|
|
4127
|
+
details: {
|
|
4128
|
+
workflowName,
|
|
4129
|
+
runId
|
|
4130
|
+
}
|
|
4131
|
+
},
|
|
4132
|
+
error
|
|
4133
|
+
);
|
|
4134
|
+
}
|
|
4135
|
+
}
|
|
4136
|
+
async persistWorkflowSnapshot({
|
|
4137
|
+
workflowName,
|
|
4138
|
+
runId,
|
|
4139
|
+
resourceId,
|
|
4140
|
+
snapshot,
|
|
4141
|
+
createdAt,
|
|
4142
|
+
updatedAt
|
|
4143
|
+
}) {
|
|
4144
|
+
try {
|
|
4145
|
+
const now = /* @__PURE__ */ new Date();
|
|
4146
|
+
const createdAtValue = createdAt ? createdAt : now;
|
|
4147
|
+
const updatedAtValue = updatedAt ? updatedAt : now;
|
|
4148
|
+
await this.#db.client.none(
|
|
4149
|
+
`INSERT INTO ${getTableName2({ indexName: TABLE_WORKFLOW_SNAPSHOT, schemaName: getSchemaName2(this.#schema) })} (workflow_name, run_id, "resourceId", snapshot, "createdAt", "updatedAt")
|
|
4150
|
+
VALUES ($1, $2, $3, $4, $5, $6)
|
|
4151
|
+
ON CONFLICT (workflow_name, run_id) DO UPDATE
|
|
4152
|
+
SET "resourceId" = $3, snapshot = $4, "updatedAt" = $6`,
|
|
4153
|
+
[
|
|
4154
|
+
workflowName,
|
|
4155
|
+
runId,
|
|
4156
|
+
resourceId,
|
|
4157
|
+
JSON.stringify(snapshot),
|
|
4158
|
+
createdAtValue.toISOString(),
|
|
4159
|
+
updatedAtValue.toISOString()
|
|
4160
|
+
]
|
|
4161
|
+
);
|
|
4162
|
+
} catch (error) {
|
|
4163
|
+
throw new MastraError(
|
|
4164
|
+
{
|
|
4165
|
+
id: createStorageErrorId("DSQL", "PERSIST_WORKFLOW_SNAPSHOT", "FAILED"),
|
|
4166
|
+
domain: ErrorDomain.STORAGE,
|
|
4167
|
+
category: ErrorCategory.THIRD_PARTY
|
|
4168
|
+
},
|
|
4169
|
+
error
|
|
4170
|
+
);
|
|
4171
|
+
}
|
|
4172
|
+
}
|
|
4173
|
+
async loadWorkflowSnapshot({
|
|
4174
|
+
workflowName,
|
|
4175
|
+
runId
|
|
4176
|
+
}) {
|
|
4177
|
+
try {
|
|
4178
|
+
const result = await this.#db.load({
|
|
4179
|
+
tableName: TABLE_WORKFLOW_SNAPSHOT,
|
|
4180
|
+
keys: { workflow_name: workflowName, run_id: runId }
|
|
4181
|
+
});
|
|
4182
|
+
return result ? result.snapshot : null;
|
|
4183
|
+
} catch (error) {
|
|
4184
|
+
throw new MastraError(
|
|
4185
|
+
{
|
|
4186
|
+
id: createStorageErrorId("DSQL", "LOAD_WORKFLOW_SNAPSHOT", "FAILED"),
|
|
4187
|
+
domain: ErrorDomain.STORAGE,
|
|
4188
|
+
category: ErrorCategory.THIRD_PARTY
|
|
4189
|
+
},
|
|
4190
|
+
error
|
|
4191
|
+
);
|
|
4192
|
+
}
|
|
4193
|
+
}
|
|
4194
|
+
async getWorkflowRunById({
|
|
4195
|
+
runId,
|
|
4196
|
+
workflowName
|
|
4197
|
+
}) {
|
|
4198
|
+
try {
|
|
4199
|
+
const conditions = [];
|
|
4200
|
+
const values = [];
|
|
4201
|
+
let paramIndex = 1;
|
|
4202
|
+
if (runId) {
|
|
4203
|
+
conditions.push(`run_id = $${paramIndex}`);
|
|
4204
|
+
values.push(runId);
|
|
4205
|
+
paramIndex++;
|
|
4206
|
+
}
|
|
4207
|
+
if (workflowName) {
|
|
4208
|
+
conditions.push(`workflow_name = $${paramIndex}`);
|
|
4209
|
+
values.push(workflowName);
|
|
4210
|
+
paramIndex++;
|
|
4211
|
+
}
|
|
4212
|
+
const whereClause = conditions.length > 0 ? `WHERE ${conditions.join(" AND ")}` : "";
|
|
4213
|
+
const query = `
|
|
4214
|
+
SELECT * FROM ${getTableName2({ indexName: TABLE_WORKFLOW_SNAPSHOT, schemaName: getSchemaName2(this.#schema) })}
|
|
4215
|
+
${whereClause}
|
|
4216
|
+
ORDER BY "createdAt" DESC LIMIT 1
|
|
4217
|
+
`;
|
|
4218
|
+
const queryValues = values;
|
|
4219
|
+
const result = await this.#db.client.oneOrNone(query, queryValues);
|
|
4220
|
+
if (!result) {
|
|
4221
|
+
return null;
|
|
4222
|
+
}
|
|
4223
|
+
return parseWorkflowRun(result);
|
|
4224
|
+
} catch (error) {
|
|
4225
|
+
throw new MastraError(
|
|
4226
|
+
{
|
|
4227
|
+
id: createStorageErrorId("DSQL", "GET_WORKFLOW_RUN_BY_ID", "FAILED"),
|
|
4228
|
+
domain: ErrorDomain.STORAGE,
|
|
4229
|
+
category: ErrorCategory.THIRD_PARTY,
|
|
4230
|
+
details: {
|
|
4231
|
+
runId,
|
|
4232
|
+
workflowName: workflowName || ""
|
|
4233
|
+
}
|
|
4234
|
+
},
|
|
4235
|
+
error
|
|
4236
|
+
);
|
|
4237
|
+
}
|
|
4238
|
+
}
|
|
4239
|
+
async deleteWorkflowRunById({ runId, workflowName }) {
|
|
4240
|
+
try {
|
|
4241
|
+
await this.#db.client.none(
|
|
4242
|
+
`DELETE FROM ${getTableName2({ indexName: TABLE_WORKFLOW_SNAPSHOT, schemaName: getSchemaName2(this.#schema) })} WHERE run_id = $1 AND workflow_name = $2`,
|
|
4243
|
+
[runId, workflowName]
|
|
4244
|
+
);
|
|
4245
|
+
} catch (error) {
|
|
4246
|
+
throw new MastraError(
|
|
4247
|
+
{
|
|
4248
|
+
id: createStorageErrorId("DSQL", "DELETE_WORKFLOW_RUN_BY_ID", "FAILED"),
|
|
4249
|
+
domain: ErrorDomain.STORAGE,
|
|
4250
|
+
category: ErrorCategory.THIRD_PARTY,
|
|
4251
|
+
details: {
|
|
4252
|
+
runId,
|
|
4253
|
+
workflowName
|
|
4254
|
+
}
|
|
4255
|
+
},
|
|
4256
|
+
error
|
|
4257
|
+
);
|
|
4258
|
+
}
|
|
4259
|
+
}
|
|
4260
|
+
async listWorkflowRuns({
|
|
4261
|
+
workflowName,
|
|
4262
|
+
fromDate,
|
|
4263
|
+
toDate,
|
|
4264
|
+
perPage,
|
|
4265
|
+
page,
|
|
4266
|
+
resourceId,
|
|
4267
|
+
status
|
|
4268
|
+
} = {}) {
|
|
4269
|
+
try {
|
|
4270
|
+
const conditions = [];
|
|
4271
|
+
const values = [];
|
|
4272
|
+
let paramIndex = 1;
|
|
4273
|
+
if (workflowName) {
|
|
4274
|
+
conditions.push(`workflow_name = $${paramIndex}`);
|
|
4275
|
+
values.push(workflowName);
|
|
4276
|
+
paramIndex++;
|
|
4277
|
+
}
|
|
4278
|
+
if (status) {
|
|
4279
|
+
conditions.push(`snapshot::jsonb ->> 'status' = $${paramIndex}`);
|
|
4280
|
+
values.push(status);
|
|
4281
|
+
paramIndex++;
|
|
4282
|
+
}
|
|
4283
|
+
if (resourceId) {
|
|
4284
|
+
const hasResourceId = await this.#db.hasColumn(TABLE_WORKFLOW_SNAPSHOT, "resourceId");
|
|
4285
|
+
if (hasResourceId) {
|
|
4286
|
+
conditions.push(`"resourceId" = $${paramIndex}`);
|
|
4287
|
+
values.push(resourceId);
|
|
4288
|
+
paramIndex++;
|
|
4289
|
+
} else {
|
|
4290
|
+
console.warn(`[${TABLE_WORKFLOW_SNAPSHOT}] resourceId column not found. Skipping resourceId filter.`);
|
|
4291
|
+
}
|
|
4292
|
+
}
|
|
4293
|
+
if (fromDate) {
|
|
4294
|
+
conditions.push(`"createdAt" >= $${paramIndex}`);
|
|
4295
|
+
values.push(fromDate instanceof Date ? fromDate.toISOString() : fromDate);
|
|
4296
|
+
paramIndex++;
|
|
4297
|
+
}
|
|
4298
|
+
if (toDate) {
|
|
4299
|
+
conditions.push(`"createdAt" <= $${paramIndex}`);
|
|
4300
|
+
values.push(toDate instanceof Date ? toDate.toISOString() : toDate);
|
|
4301
|
+
paramIndex++;
|
|
4302
|
+
}
|
|
4303
|
+
const whereClause = conditions.length > 0 ? `WHERE ${conditions.join(" AND ")}` : "";
|
|
4304
|
+
let total = 0;
|
|
4305
|
+
const usePagination = typeof perPage === "number" && typeof page === "number";
|
|
4306
|
+
if (usePagination) {
|
|
4307
|
+
const countResult = await this.#db.client.one(
|
|
4308
|
+
`SELECT COUNT(*) as count FROM ${getTableName2({ indexName: TABLE_WORKFLOW_SNAPSHOT, schemaName: getSchemaName2(this.#schema) })} ${whereClause}`,
|
|
4309
|
+
values
|
|
4310
|
+
);
|
|
4311
|
+
total = Number(countResult.count);
|
|
4312
|
+
}
|
|
4313
|
+
const normalizedPerPage = usePagination ? normalizePerPage(perPage, Number.MAX_SAFE_INTEGER) : 0;
|
|
4314
|
+
const offset = usePagination ? page * normalizedPerPage : void 0;
|
|
4315
|
+
const query = `
|
|
4316
|
+
SELECT * FROM ${getTableName2({ indexName: TABLE_WORKFLOW_SNAPSHOT, schemaName: getSchemaName2(this.#schema) })}
|
|
4317
|
+
${whereClause}
|
|
4318
|
+
ORDER BY "createdAt" DESC
|
|
4319
|
+
${usePagination ? ` LIMIT $${paramIndex} OFFSET $${paramIndex + 1}` : ""}
|
|
4320
|
+
`;
|
|
4321
|
+
const queryValues = usePagination ? [...values, normalizedPerPage, offset] : values;
|
|
4322
|
+
const result = await this.#db.client.manyOrNone(query, queryValues);
|
|
4323
|
+
const runs = (result || []).map((row) => {
|
|
4324
|
+
return parseWorkflowRun(row);
|
|
4325
|
+
});
|
|
4326
|
+
return { runs, total: total || runs.length };
|
|
4327
|
+
} catch (error) {
|
|
4328
|
+
throw new MastraError(
|
|
4329
|
+
{
|
|
4330
|
+
id: createStorageErrorId("DSQL", "LIST_WORKFLOW_RUNS", "FAILED"),
|
|
4331
|
+
domain: ErrorDomain.STORAGE,
|
|
4332
|
+
category: ErrorCategory.THIRD_PARTY,
|
|
4333
|
+
details: {
|
|
4334
|
+
workflowName: workflowName || "all"
|
|
4335
|
+
}
|
|
4336
|
+
},
|
|
4337
|
+
error
|
|
4338
|
+
);
|
|
4339
|
+
}
|
|
4340
|
+
}
|
|
4341
|
+
};
|
|
4342
|
+
|
|
4343
|
+
// src/storage/index.ts
|
|
4344
|
+
var DSQLStore = class extends MastraStorage {
|
|
4345
|
+
#pool;
|
|
4346
|
+
#db;
|
|
4347
|
+
#ownsPool;
|
|
4348
|
+
schema;
|
|
4349
|
+
isInitialized = false;
|
|
4350
|
+
stores;
|
|
4351
|
+
constructor(config) {
|
|
4352
|
+
try {
|
|
4353
|
+
validateConfig(config);
|
|
4354
|
+
super({ id: config.id, name: "DSQLStore", disableInit: config.disableInit });
|
|
4355
|
+
this.schema = config.schemaName || "public";
|
|
4356
|
+
if (isPoolConfig(config)) {
|
|
4357
|
+
this.#pool = config.pool;
|
|
4358
|
+
this.#ownsPool = false;
|
|
4359
|
+
} else {
|
|
4360
|
+
this.#pool = this.createPool(config);
|
|
4361
|
+
this.#ownsPool = true;
|
|
4362
|
+
}
|
|
4363
|
+
this.#db = new PoolAdapter(this.#pool);
|
|
4364
|
+
const domainConfig = {
|
|
4365
|
+
client: this.#db,
|
|
4366
|
+
schemaName: this.schema,
|
|
4367
|
+
skipDefaultIndexes: config.skipDefaultIndexes,
|
|
4368
|
+
indexes: config.indexes
|
|
4369
|
+
};
|
|
4370
|
+
this.stores = {
|
|
4371
|
+
scores: new ScoresDSQL(domainConfig),
|
|
4372
|
+
workflows: new WorkflowsDSQL(domainConfig),
|
|
4373
|
+
memory: new MemoryDSQL(domainConfig),
|
|
4374
|
+
observability: new ObservabilityDSQL(domainConfig),
|
|
4375
|
+
agents: new AgentsDSQL(domainConfig)
|
|
4376
|
+
};
|
|
4377
|
+
} catch (e) {
|
|
4378
|
+
throw new MastraError(
|
|
4379
|
+
{
|
|
4380
|
+
id: createStorageErrorId("DSQL", "INITIALIZATION", "FAILED"),
|
|
4381
|
+
domain: ErrorDomain.STORAGE,
|
|
4382
|
+
category: ErrorCategory.USER
|
|
4383
|
+
},
|
|
4384
|
+
e
|
|
4385
|
+
);
|
|
4386
|
+
}
|
|
4387
|
+
}
|
|
4388
|
+
/**
|
|
4389
|
+
* Creates a connection pool with AuroraDSQLClient for IAM authentication.
|
|
4390
|
+
*/
|
|
4391
|
+
createPool(config) {
|
|
4392
|
+
if (!isHostConfig(config)) {
|
|
4393
|
+
throw new Error("DSQLStore: Invalid configuration for creating pool.");
|
|
4394
|
+
}
|
|
4395
|
+
const region = getEffectiveRegion(config);
|
|
4396
|
+
const poolConfig = {
|
|
4397
|
+
host: config.host,
|
|
4398
|
+
user: config.user ?? "admin",
|
|
4399
|
+
database: config.database ?? "postgres",
|
|
4400
|
+
// Use AuroraDSQLClient for automatic IAM token generation
|
|
4401
|
+
Client: AuroraDSQLClient,
|
|
4402
|
+
// Pass region for IAM token generation
|
|
4403
|
+
region,
|
|
4404
|
+
// Custom credentials provider (optional)
|
|
4405
|
+
customCredentialsProvider: config.customCredentialsProvider,
|
|
4406
|
+
// Pool settings optimized for Aurora DSQL
|
|
4407
|
+
max: config.max ?? DSQL_POOL_DEFAULTS.max,
|
|
4408
|
+
min: config.min ?? DSQL_POOL_DEFAULTS.min,
|
|
4409
|
+
idleTimeoutMillis: config.idleTimeoutMillis ?? DSQL_POOL_DEFAULTS.idleTimeoutMillis,
|
|
4410
|
+
maxLifetimeSeconds: config.maxLifetimeSeconds ?? DSQL_POOL_DEFAULTS.maxLifetimeSeconds,
|
|
4411
|
+
connectionTimeoutMillis: config.connectionTimeoutMillis ?? DSQL_POOL_DEFAULTS.connectionTimeoutMillis,
|
|
4412
|
+
allowExitOnIdle: config.allowExitOnIdle ?? DSQL_POOL_DEFAULTS.allowExitOnIdle
|
|
4413
|
+
};
|
|
4414
|
+
return new Pool(poolConfig);
|
|
4415
|
+
}
|
|
4416
|
+
async init() {
|
|
4417
|
+
if (this.isInitialized) {
|
|
4418
|
+
return;
|
|
4419
|
+
}
|
|
4420
|
+
try {
|
|
4421
|
+
this.isInitialized = true;
|
|
4422
|
+
await super.init();
|
|
4423
|
+
} catch (error) {
|
|
4424
|
+
this.isInitialized = false;
|
|
4425
|
+
throw new MastraError(
|
|
4426
|
+
{
|
|
4427
|
+
id: createStorageErrorId("DSQL", "INIT", "FAILED"),
|
|
4428
|
+
domain: ErrorDomain.STORAGE,
|
|
4429
|
+
category: ErrorCategory.THIRD_PARTY
|
|
4430
|
+
},
|
|
4431
|
+
error
|
|
4432
|
+
);
|
|
4433
|
+
}
|
|
4434
|
+
}
|
|
4435
|
+
/**
|
|
4436
|
+
* Database client for executing queries.
|
|
4437
|
+
*
|
|
4438
|
+
* @example
|
|
4439
|
+
* ```typescript
|
|
4440
|
+
* const rows = await store.db.any('SELECT * FROM users WHERE active = $1', [true]);
|
|
4441
|
+
* const user = await store.db.one('SELECT * FROM users WHERE id = $1', [userId]);
|
|
4442
|
+
* ```
|
|
4443
|
+
*/
|
|
4444
|
+
get db() {
|
|
4445
|
+
return this.#db;
|
|
4446
|
+
}
|
|
4447
|
+
/**
|
|
4448
|
+
* The underlying pg.Pool for direct database access or ORM integration.
|
|
4449
|
+
*/
|
|
4450
|
+
get pool() {
|
|
4451
|
+
return this.#pool;
|
|
4452
|
+
}
|
|
4453
|
+
/**
|
|
4454
|
+
* Closes the connection pool if it was created by this store.
|
|
4455
|
+
* If a pool was passed in via config, it will not be closed.
|
|
4456
|
+
*/
|
|
4457
|
+
async close() {
|
|
4458
|
+
if (this.#ownsPool) {
|
|
4459
|
+
await this.#pool.end();
|
|
4460
|
+
}
|
|
4461
|
+
this.isInitialized = false;
|
|
4462
|
+
}
|
|
4463
|
+
};
|
|
4464
|
+
|
|
4465
|
+
export { AgentsDSQL, DSQLStore, DSQL_POOL_DEFAULTS, MemoryDSQL, ObservabilityDSQL, PoolAdapter, ScoresDSQL, WorkflowsDSQL };
|
|
4466
|
+
//# sourceMappingURL=index.js.map
|
|
4467
|
+
//# sourceMappingURL=index.js.map
|