@mastra/clickhouse 1.0.0-beta.5 → 1.0.0-beta.7

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
package/dist/index.cjs CHANGED
@@ -4,6 +4,7 @@ var client = require('@clickhouse/client');
4
4
  var error = require('@mastra/core/error');
5
5
  var storage = require('@mastra/core/storage');
6
6
  var agent = require('@mastra/core/agent');
7
+ var base = require('@mastra/core/base');
7
8
  var evals = require('@mastra/core/evals');
8
9
 
9
10
  // src/storage/index.ts
@@ -28,6 +29,8 @@ var COLUMN_TYPES = {
28
29
  bigint: "Int64",
29
30
  boolean: "Bool"
30
31
  };
32
+ var JSON_FIELDS = ["content", "attributes", "metadata", "input", "output", "error", "scope", "links"];
33
+ var NULLABLE_STRING_FIELDS = ["parentSpanId", "error"];
31
34
  function transformRow(row) {
32
35
  if (!row) {
33
36
  return row;
@@ -38,8 +41,21 @@ function transformRow(row) {
38
41
  if (row.updatedAt) {
39
42
  row.updatedAt = new Date(row.updatedAt);
40
43
  }
41
- if (row.content && typeof row.content === "string") {
42
- row.content = storage.safelyParseJSON(row.content);
44
+ if (row.startedAt) {
45
+ row.startedAt = new Date(row.startedAt);
46
+ }
47
+ if (row.endedAt) {
48
+ row.endedAt = new Date(row.endedAt);
49
+ }
50
+ for (const field of JSON_FIELDS) {
51
+ if (row[field] && typeof row[field] === "string") {
52
+ row[field] = storage.safelyParseJSON(row[field]);
53
+ }
54
+ }
55
+ for (const field of NULLABLE_STRING_FIELDS) {
56
+ if (row[field] === "") {
57
+ row[field] = null;
58
+ }
43
59
  }
44
60
  return row;
45
61
  }
@@ -47,6 +63,326 @@ function transformRows(rows) {
47
63
  return rows.map((row) => transformRow(row));
48
64
  }
49
65
 
66
+ // src/storage/db/index.ts
67
+ function resolveClickhouseConfig(config) {
68
+ if ("client" in config) {
69
+ return { client: config.client, ttl: config.ttl };
70
+ }
71
+ const client$1 = client.createClient({
72
+ url: config.url,
73
+ username: config.username,
74
+ password: config.password,
75
+ clickhouse_settings: {
76
+ date_time_input_format: "best_effort",
77
+ date_time_output_format: "iso",
78
+ use_client_time_zone: 1,
79
+ output_format_json_quote_64bit_integers: 0
80
+ }
81
+ });
82
+ return { client: client$1, ttl: config.ttl };
83
+ }
84
+ var ClickhouseDB = class extends base.MastraBase {
85
+ ttl;
86
+ client;
87
+ constructor({ client, ttl }) {
88
+ super({
89
+ name: "CLICKHOUSE_DB"
90
+ });
91
+ this.ttl = ttl;
92
+ this.client = client;
93
+ }
94
+ async hasColumn(table, column) {
95
+ const result = await this.client.query({
96
+ query: `DESCRIBE TABLE ${table}`,
97
+ format: "JSONEachRow"
98
+ });
99
+ const columns = await result.json();
100
+ return columns.some((c) => c.name === column);
101
+ }
102
+ getSqlType(type) {
103
+ switch (type) {
104
+ case "text":
105
+ case "uuid":
106
+ case "jsonb":
107
+ return "String";
108
+ case "timestamp":
109
+ return "DateTime64(3)";
110
+ case "integer":
111
+ case "bigint":
112
+ return "Int64";
113
+ case "float":
114
+ return "Float64";
115
+ case "boolean":
116
+ return "Bool";
117
+ default:
118
+ return storage.getSqlType(type);
119
+ }
120
+ }
121
+ async createTable({
122
+ tableName,
123
+ schema
124
+ }) {
125
+ try {
126
+ const columns = Object.entries(schema).map(([name, def]) => {
127
+ let sqlType = this.getSqlType(def.type);
128
+ const isNullable = def.nullable === true;
129
+ if (isNullable) {
130
+ sqlType = `Nullable(${sqlType})`;
131
+ }
132
+ const constraints = [];
133
+ if (name === "metadata" && def.type === "text" && isNullable) {
134
+ constraints.push("DEFAULT '{}'");
135
+ }
136
+ const columnTtl = this.ttl?.[tableName]?.columns?.[name];
137
+ return `"${name}" ${sqlType} ${constraints.join(" ")} ${columnTtl ? `TTL toDateTime(${columnTtl.ttlKey ?? "createdAt"}) + INTERVAL ${columnTtl.interval} ${columnTtl.unit}` : ""}`;
138
+ }).join(",\n");
139
+ const rowTtl = this.ttl?.[tableName]?.row;
140
+ let sql;
141
+ if (tableName === storage.TABLE_WORKFLOW_SNAPSHOT) {
142
+ sql = `
143
+ CREATE TABLE IF NOT EXISTS ${tableName} (
144
+ ${["id String"].concat(columns)}
145
+ )
146
+ ENGINE = ${TABLE_ENGINES[tableName] ?? "MergeTree()"}
147
+ PRIMARY KEY (createdAt, run_id, workflow_name)
148
+ ORDER BY (createdAt, run_id, workflow_name)
149
+ ${rowTtl ? `TTL toDateTime(${rowTtl.ttlKey ?? "createdAt"}) + INTERVAL ${rowTtl.interval} ${rowTtl.unit}` : ""}
150
+ SETTINGS index_granularity = 8192
151
+ `;
152
+ } else if (tableName === storage.TABLE_SPANS) {
153
+ sql = `
154
+ CREATE TABLE IF NOT EXISTS ${tableName} (
155
+ ${columns}
156
+ )
157
+ ENGINE = ${TABLE_ENGINES[tableName] ?? "MergeTree()"}
158
+ PRIMARY KEY (createdAt, traceId, spanId)
159
+ ORDER BY (createdAt, traceId, spanId)
160
+ ${rowTtl ? `TTL toDateTime(${rowTtl.ttlKey ?? "createdAt"}) + INTERVAL ${rowTtl.interval} ${rowTtl.unit}` : ""}
161
+ SETTINGS index_granularity = 8192
162
+ `;
163
+ } else {
164
+ sql = `
165
+ CREATE TABLE IF NOT EXISTS ${tableName} (
166
+ ${columns}
167
+ )
168
+ ENGINE = ${TABLE_ENGINES[tableName] ?? "MergeTree()"}
169
+ PRIMARY KEY (createdAt, ${"id"})
170
+ ORDER BY (createdAt, ${"id"})
171
+ ${this.ttl?.[tableName]?.row ? `TTL toDateTime(createdAt) + INTERVAL ${this.ttl[tableName].row.interval} ${this.ttl[tableName].row.unit}` : ""}
172
+ SETTINGS index_granularity = 8192
173
+ `;
174
+ }
175
+ await this.client.query({
176
+ query: sql,
177
+ clickhouse_settings: {
178
+ // Allows to insert serialized JS Dates (such as '2023-12-06T10:54:48.000Z')
179
+ date_time_input_format: "best_effort",
180
+ date_time_output_format: "iso",
181
+ use_client_time_zone: 1,
182
+ output_format_json_quote_64bit_integers: 0
183
+ }
184
+ });
185
+ } catch (error$1) {
186
+ throw new error.MastraError(
187
+ {
188
+ id: storage.createStorageErrorId("CLICKHOUSE", "CREATE_TABLE", "FAILED"),
189
+ domain: error.ErrorDomain.STORAGE,
190
+ category: error.ErrorCategory.THIRD_PARTY,
191
+ details: { tableName }
192
+ },
193
+ error$1
194
+ );
195
+ }
196
+ }
197
+ async alterTable({
198
+ tableName,
199
+ schema,
200
+ ifNotExists
201
+ }) {
202
+ try {
203
+ const describeSql = `DESCRIBE TABLE ${tableName}`;
204
+ const result = await this.client.query({
205
+ query: describeSql
206
+ });
207
+ const rows = await result.json();
208
+ const existingColumnNames = new Set(rows.data.map((row) => row.name.toLowerCase()));
209
+ for (const columnName of ifNotExists) {
210
+ if (!existingColumnNames.has(columnName.toLowerCase()) && schema[columnName]) {
211
+ const columnDef = schema[columnName];
212
+ let sqlType = this.getSqlType(columnDef.type);
213
+ if (columnDef.nullable !== false) {
214
+ sqlType = `Nullable(${sqlType})`;
215
+ }
216
+ const defaultValue = columnDef.nullable === false ? storage.getDefaultValue(columnDef.type) : "";
217
+ const alterSql = `ALTER TABLE ${tableName} ADD COLUMN IF NOT EXISTS "${columnName}" ${sqlType} ${defaultValue}`.trim();
218
+ await this.client.query({
219
+ query: alterSql
220
+ });
221
+ this.logger?.debug?.(`Added column ${columnName} to table ${tableName}`);
222
+ }
223
+ }
224
+ } catch (error$1) {
225
+ throw new error.MastraError(
226
+ {
227
+ id: storage.createStorageErrorId("CLICKHOUSE", "ALTER_TABLE", "FAILED"),
228
+ domain: error.ErrorDomain.STORAGE,
229
+ category: error.ErrorCategory.THIRD_PARTY,
230
+ details: { tableName }
231
+ },
232
+ error$1
233
+ );
234
+ }
235
+ }
236
+ async clearTable({ tableName }) {
237
+ try {
238
+ await this.client.query({
239
+ query: `TRUNCATE TABLE ${tableName}`,
240
+ clickhouse_settings: {
241
+ // Allows to insert serialized JS Dates (such as '2023-12-06T10:54:48.000Z')
242
+ date_time_input_format: "best_effort",
243
+ date_time_output_format: "iso",
244
+ use_client_time_zone: 1,
245
+ output_format_json_quote_64bit_integers: 0
246
+ }
247
+ });
248
+ } catch (error$1) {
249
+ throw new error.MastraError(
250
+ {
251
+ id: storage.createStorageErrorId("CLICKHOUSE", "CLEAR_TABLE", "FAILED"),
252
+ domain: error.ErrorDomain.STORAGE,
253
+ category: error.ErrorCategory.THIRD_PARTY,
254
+ details: { tableName }
255
+ },
256
+ error$1
257
+ );
258
+ }
259
+ }
260
+ async dropTable({ tableName }) {
261
+ await this.client.query({
262
+ query: `DROP TABLE IF EXISTS ${tableName}`
263
+ });
264
+ }
265
+ async insert({ tableName, record }) {
266
+ const rawCreatedAt = record.createdAt || record.created_at || /* @__PURE__ */ new Date();
267
+ const rawUpdatedAt = record.updatedAt || /* @__PURE__ */ new Date();
268
+ const createdAt = rawCreatedAt instanceof Date ? rawCreatedAt.toISOString() : rawCreatedAt;
269
+ const updatedAt = rawUpdatedAt instanceof Date ? rawUpdatedAt.toISOString() : rawUpdatedAt;
270
+ try {
271
+ await this.client.insert({
272
+ table: tableName,
273
+ values: [
274
+ {
275
+ ...record,
276
+ createdAt,
277
+ updatedAt
278
+ }
279
+ ],
280
+ format: "JSONEachRow",
281
+ clickhouse_settings: {
282
+ // Allows to insert serialized JS Dates (such as '2023-12-06T10:54:48.000Z')
283
+ output_format_json_quote_64bit_integers: 0,
284
+ date_time_input_format: "best_effort",
285
+ use_client_time_zone: 1
286
+ }
287
+ });
288
+ } catch (error$1) {
289
+ throw new error.MastraError(
290
+ {
291
+ id: storage.createStorageErrorId("CLICKHOUSE", "INSERT", "FAILED"),
292
+ domain: error.ErrorDomain.STORAGE,
293
+ category: error.ErrorCategory.THIRD_PARTY,
294
+ details: { tableName }
295
+ },
296
+ error$1
297
+ );
298
+ }
299
+ }
300
+ async batchInsert({ tableName, records }) {
301
+ const recordsToBeInserted = records.map((record) => ({
302
+ ...Object.fromEntries(
303
+ Object.entries(record).map(([key, value]) => [
304
+ key,
305
+ storage.TABLE_SCHEMAS[tableName]?.[key]?.type === "timestamp" ? new Date(value).toISOString() : value
306
+ ])
307
+ )
308
+ }));
309
+ try {
310
+ await this.client.insert({
311
+ table: tableName,
312
+ values: recordsToBeInserted,
313
+ format: "JSONEachRow",
314
+ clickhouse_settings: {
315
+ // Allows to insert serialized JS Dates (such as '2023-12-06T10:54:48.000Z')
316
+ date_time_input_format: "best_effort",
317
+ use_client_time_zone: 1,
318
+ output_format_json_quote_64bit_integers: 0
319
+ }
320
+ });
321
+ } catch (error$1) {
322
+ throw new error.MastraError(
323
+ {
324
+ id: storage.createStorageErrorId("CLICKHOUSE", "BATCH_INSERT", "FAILED"),
325
+ domain: error.ErrorDomain.STORAGE,
326
+ category: error.ErrorCategory.THIRD_PARTY,
327
+ details: { tableName }
328
+ },
329
+ error$1
330
+ );
331
+ }
332
+ }
333
+ async load({ tableName, keys }) {
334
+ try {
335
+ const engine = TABLE_ENGINES[tableName] ?? "MergeTree()";
336
+ const keyEntries = Object.entries(keys);
337
+ const conditions = keyEntries.map(
338
+ ([key]) => `"${key}" = {var_${key}:${this.getSqlType(storage.TABLE_SCHEMAS[tableName]?.[key]?.type ?? "text")}}`
339
+ ).join(" AND ");
340
+ const values = keyEntries.reduce((acc, [key, value]) => {
341
+ return { ...acc, [`var_${key}`]: value };
342
+ }, {});
343
+ const hasUpdatedAt = storage.TABLE_SCHEMAS[tableName]?.updatedAt;
344
+ const selectClause = `SELECT *, toDateTime64(createdAt, 3) as createdAt${hasUpdatedAt ? ", toDateTime64(updatedAt, 3) as updatedAt" : ""}`;
345
+ const result = await this.client.query({
346
+ query: `${selectClause} FROM ${tableName} ${engine.startsWith("ReplacingMergeTree") ? "FINAL" : ""} WHERE ${conditions} ORDER BY createdAt DESC LIMIT 1`,
347
+ query_params: values,
348
+ clickhouse_settings: {
349
+ // Allows to insert serialized JS Dates (such as '2023-12-06T10:54:48.000Z')
350
+ date_time_input_format: "best_effort",
351
+ date_time_output_format: "iso",
352
+ use_client_time_zone: 1,
353
+ output_format_json_quote_64bit_integers: 0
354
+ }
355
+ });
356
+ if (!result) {
357
+ return null;
358
+ }
359
+ const rows = await result.json();
360
+ if (tableName === storage.TABLE_WORKFLOW_SNAPSHOT) {
361
+ const snapshot = rows.data[0];
362
+ if (!snapshot) {
363
+ return null;
364
+ }
365
+ if (typeof snapshot.snapshot === "string") {
366
+ snapshot.snapshot = JSON.parse(snapshot.snapshot);
367
+ }
368
+ return transformRow(snapshot);
369
+ }
370
+ const data = transformRow(rows.data[0]);
371
+ return data;
372
+ } catch (error$1) {
373
+ throw new error.MastraError(
374
+ {
375
+ id: storage.createStorageErrorId("CLICKHOUSE", "LOAD", "FAILED"),
376
+ domain: error.ErrorDomain.STORAGE,
377
+ category: error.ErrorCategory.THIRD_PARTY,
378
+ details: { tableName }
379
+ },
380
+ error$1
381
+ );
382
+ }
383
+ }
384
+ };
385
+
50
386
  // src/storage/domains/memory/index.ts
51
387
  function serializeMetadata(metadata) {
52
388
  if (!metadata || Object.keys(metadata).length === 0) {
@@ -68,11 +404,60 @@ function parseMetadata(metadata) {
68
404
  }
69
405
  var MemoryStorageClickhouse = class extends storage.MemoryStorage {
70
406
  client;
71
- operations;
72
- constructor({ client, operations }) {
407
+ #db;
408
+ constructor(config) {
73
409
  super();
410
+ const { client, ttl } = resolveClickhouseConfig(config);
74
411
  this.client = client;
75
- this.operations = operations;
412
+ this.#db = new ClickhouseDB({ client, ttl });
413
+ }
414
+ async init() {
415
+ await this.#db.createTable({ tableName: storage.TABLE_THREADS, schema: storage.TABLE_SCHEMAS[storage.TABLE_THREADS] });
416
+ await this.#db.createTable({ tableName: storage.TABLE_MESSAGES, schema: storage.TABLE_SCHEMAS[storage.TABLE_MESSAGES] });
417
+ await this.#db.createTable({ tableName: storage.TABLE_RESOURCES, schema: storage.TABLE_SCHEMAS[storage.TABLE_RESOURCES] });
418
+ await this.#db.alterTable({
419
+ tableName: storage.TABLE_MESSAGES,
420
+ schema: storage.TABLE_SCHEMAS[storage.TABLE_MESSAGES],
421
+ ifNotExists: ["resourceId"]
422
+ });
423
+ }
424
+ async dangerouslyClearAll() {
425
+ await this.#db.clearTable({ tableName: storage.TABLE_MESSAGES });
426
+ await this.#db.clearTable({ tableName: storage.TABLE_RESOURCES });
427
+ await this.#db.clearTable({ tableName: storage.TABLE_THREADS });
428
+ }
429
+ async deleteMessages(messageIds) {
430
+ if (!messageIds || messageIds.length === 0) return;
431
+ try {
432
+ const result = await this.client.query({
433
+ query: `SELECT DISTINCT thread_id FROM ${storage.TABLE_MESSAGES} WHERE id IN {messageIds:Array(String)}`,
434
+ query_params: { messageIds },
435
+ format: "JSONEachRow"
436
+ });
437
+ const rows = await result.json();
438
+ const threadIds = rows.map((r) => r.thread_id);
439
+ await this.client.command({
440
+ query: `DELETE FROM ${storage.TABLE_MESSAGES} WHERE id IN {messageIds:Array(String)}`,
441
+ query_params: { messageIds }
442
+ });
443
+ if (threadIds.length > 0) {
444
+ const now = (/* @__PURE__ */ new Date()).toISOString().replace("Z", "");
445
+ await this.client.command({
446
+ query: `ALTER TABLE ${storage.TABLE_THREADS} UPDATE updatedAt = {now:DateTime64(3)} WHERE id IN {threadIds:Array(String)}`,
447
+ query_params: { now, threadIds }
448
+ });
449
+ }
450
+ } catch (error$1) {
451
+ throw new error.MastraError(
452
+ {
453
+ id: storage.createStorageErrorId("CLICKHOUSE", "DELETE_MESSAGES", "FAILED"),
454
+ domain: error.ErrorDomain.STORAGE,
455
+ category: error.ErrorCategory.THIRD_PARTY,
456
+ details: { count: messageIds.length }
457
+ },
458
+ error$1
459
+ );
460
+ }
76
461
  }
77
462
  async listMessagesById({ messageIds }) {
78
463
  if (messageIds.length === 0) return { messages: [] };
@@ -1197,274 +1582,526 @@ var MemoryStorageClickhouse = class extends storage.MemoryStorage {
1197
1582
  }
1198
1583
  }
1199
1584
  };
1200
- var StoreOperationsClickhouse = class extends storage.StoreOperations {
1201
- ttl;
1585
+ var ObservabilityStorageClickhouse = class extends storage.ObservabilityStorage {
1202
1586
  client;
1203
- constructor({ client, ttl }) {
1587
+ #db;
1588
+ constructor(config) {
1204
1589
  super();
1205
- this.ttl = ttl;
1590
+ const { client, ttl } = resolveClickhouseConfig(config);
1206
1591
  this.client = client;
1592
+ this.#db = new ClickhouseDB({ client, ttl });
1207
1593
  }
1208
- async hasColumn(table, column) {
1209
- const result = await this.client.query({
1210
- query: `DESCRIBE TABLE ${table}`,
1211
- format: "JSONEachRow"
1212
- });
1213
- const columns = await result.json();
1214
- return columns.some((c) => c.name === column);
1594
+ async init() {
1595
+ await this.#db.createTable({ tableName: storage.TABLE_SPANS, schema: storage.SPAN_SCHEMA });
1215
1596
  }
1216
- getSqlType(type) {
1217
- switch (type) {
1218
- case "text":
1219
- return "String";
1220
- case "timestamp":
1221
- return "DateTime64(3)";
1222
- case "integer":
1223
- case "bigint":
1224
- return "Int64";
1225
- case "jsonb":
1226
- return "String";
1227
- default:
1228
- return super.getSqlType(type);
1597
+ async dangerouslyClearAll() {
1598
+ await this.#db.clearTable({ tableName: storage.TABLE_SPANS });
1599
+ }
1600
+ get tracingStrategy() {
1601
+ return {
1602
+ preferred: "insert-only",
1603
+ supported: ["insert-only"]
1604
+ };
1605
+ }
1606
+ async createSpan(args) {
1607
+ const { span } = args;
1608
+ try {
1609
+ const now = Date.now();
1610
+ const record = {
1611
+ ...span,
1612
+ // Convert Date objects to millisecond timestamps for DateTime64(3)
1613
+ startedAt: span.startedAt instanceof Date ? span.startedAt.getTime() : span.startedAt,
1614
+ endedAt: span.endedAt instanceof Date ? span.endedAt.getTime() : span.endedAt,
1615
+ createdAt: now,
1616
+ updatedAt: now
1617
+ };
1618
+ await this.#db.insert({ tableName: storage.TABLE_SPANS, record });
1619
+ } catch (error$1) {
1620
+ throw new error.MastraError(
1621
+ {
1622
+ id: storage.createStorageErrorId("CLICKHOUSE", "CREATE_SPAN", "FAILED"),
1623
+ domain: error.ErrorDomain.STORAGE,
1624
+ category: error.ErrorCategory.THIRD_PARTY,
1625
+ details: {
1626
+ spanId: span.spanId,
1627
+ traceId: span.traceId,
1628
+ spanType: span.spanType,
1629
+ spanName: span.name
1630
+ }
1631
+ },
1632
+ error$1
1633
+ );
1229
1634
  }
1230
1635
  }
1231
- async createTable({
1232
- tableName,
1233
- schema
1234
- }) {
1636
+ async getSpan(args) {
1637
+ const { traceId, spanId } = args;
1235
1638
  try {
1236
- const columns = Object.entries(schema).map(([name, def]) => {
1237
- const constraints = [];
1238
- if (!def.nullable) constraints.push("NOT NULL");
1239
- if (name === "metadata" && def.type === "text" && def.nullable) {
1240
- constraints.push("DEFAULT '{}'");
1241
- }
1242
- const columnTtl = this.ttl?.[tableName]?.columns?.[name];
1243
- return `"${name}" ${COLUMN_TYPES[def.type]} ${constraints.join(" ")} ${columnTtl ? `TTL toDateTime(${columnTtl.ttlKey ?? "createdAt"}) + INTERVAL ${columnTtl.interval} ${columnTtl.unit}` : ""}`;
1244
- }).join(",\n");
1245
- const rowTtl = this.ttl?.[tableName]?.row;
1246
- const sql = tableName === storage.TABLE_WORKFLOW_SNAPSHOT ? `
1247
- CREATE TABLE IF NOT EXISTS ${tableName} (
1248
- ${["id String"].concat(columns)}
1249
- )
1250
- ENGINE = ${TABLE_ENGINES[tableName] ?? "MergeTree()"}
1251
- PRIMARY KEY (createdAt, run_id, workflow_name)
1252
- ORDER BY (createdAt, run_id, workflow_name)
1253
- ${rowTtl ? `TTL toDateTime(${rowTtl.ttlKey ?? "createdAt"}) + INTERVAL ${rowTtl.interval} ${rowTtl.unit}` : ""}
1254
- SETTINGS index_granularity = 8192
1255
- ` : `
1256
- CREATE TABLE IF NOT EXISTS ${tableName} (
1257
- ${columns}
1258
- )
1259
- ENGINE = ${TABLE_ENGINES[tableName] ?? "MergeTree()"}
1260
- PRIMARY KEY (createdAt, ${"id"})
1261
- ORDER BY (createdAt, ${"id"})
1262
- ${this.ttl?.[tableName]?.row ? `TTL toDateTime(createdAt) + INTERVAL ${this.ttl[tableName].row.interval} ${this.ttl[tableName].row.unit}` : ""}
1263
- SETTINGS index_granularity = 8192
1264
- `;
1265
- await this.client.query({
1266
- query: sql,
1639
+ const engine = TABLE_ENGINES[storage.TABLE_SPANS] ?? "MergeTree()";
1640
+ const result = await this.client.query({
1641
+ query: `
1642
+ SELECT *
1643
+ FROM ${storage.TABLE_SPANS} ${engine.startsWith("ReplacingMergeTree") ? "FINAL" : ""}
1644
+ WHERE traceId = {traceId:String} AND spanId = {spanId:String}
1645
+ LIMIT 1
1646
+ `,
1647
+ query_params: { traceId, spanId },
1648
+ format: "JSONEachRow",
1267
1649
  clickhouse_settings: {
1268
- // Allows to insert serialized JS Dates (such as '2023-12-06T10:54:48.000Z')
1269
1650
  date_time_input_format: "best_effort",
1270
1651
  date_time_output_format: "iso",
1271
1652
  use_client_time_zone: 1,
1272
1653
  output_format_json_quote_64bit_integers: 0
1273
1654
  }
1274
1655
  });
1656
+ const rows = await result.json();
1657
+ if (!rows || rows.length === 0) {
1658
+ return null;
1659
+ }
1660
+ const spans = transformRows(rows);
1661
+ const span = spans[0];
1662
+ if (!span) {
1663
+ return null;
1664
+ }
1665
+ return { span };
1275
1666
  } catch (error$1) {
1276
1667
  throw new error.MastraError(
1277
1668
  {
1278
- id: storage.createStorageErrorId("CLICKHOUSE", "CREATE_TABLE", "FAILED"),
1669
+ id: storage.createStorageErrorId("CLICKHOUSE", "GET_SPAN", "FAILED"),
1279
1670
  domain: error.ErrorDomain.STORAGE,
1280
1671
  category: error.ErrorCategory.THIRD_PARTY,
1281
- details: { tableName }
1672
+ details: { traceId, spanId }
1282
1673
  },
1283
1674
  error$1
1284
1675
  );
1285
1676
  }
1286
1677
  }
1287
- async alterTable({
1288
- tableName,
1289
- schema,
1290
- ifNotExists
1291
- }) {
1678
+ async getRootSpan(args) {
1679
+ const { traceId } = args;
1292
1680
  try {
1293
- const describeSql = `DESCRIBE TABLE ${tableName}`;
1681
+ const engine = TABLE_ENGINES[storage.TABLE_SPANS] ?? "MergeTree()";
1294
1682
  const result = await this.client.query({
1295
- query: describeSql
1683
+ query: `
1684
+ SELECT *
1685
+ FROM ${storage.TABLE_SPANS} ${engine.startsWith("ReplacingMergeTree") ? "FINAL" : ""}
1686
+ WHERE traceId = {traceId:String} AND (parentSpanId IS NULL OR parentSpanId = '')
1687
+ LIMIT 1
1688
+ `,
1689
+ query_params: { traceId },
1690
+ format: "JSONEachRow",
1691
+ clickhouse_settings: {
1692
+ date_time_input_format: "best_effort",
1693
+ date_time_output_format: "iso",
1694
+ use_client_time_zone: 1,
1695
+ output_format_json_quote_64bit_integers: 0
1696
+ }
1296
1697
  });
1297
1698
  const rows = await result.json();
1298
- const existingColumnNames = new Set(rows.data.map((row) => row.name.toLowerCase()));
1299
- for (const columnName of ifNotExists) {
1300
- if (!existingColumnNames.has(columnName.toLowerCase()) && schema[columnName]) {
1301
- const columnDef = schema[columnName];
1302
- let sqlType = this.getSqlType(columnDef.type);
1303
- if (columnDef.nullable !== false) {
1304
- sqlType = `Nullable(${sqlType})`;
1305
- }
1306
- const defaultValue = columnDef.nullable === false ? this.getDefaultValue(columnDef.type) : "";
1307
- const alterSql = `ALTER TABLE ${tableName} ADD COLUMN IF NOT EXISTS "${columnName}" ${sqlType} ${defaultValue}`.trim();
1308
- await this.client.query({
1309
- query: alterSql
1310
- });
1311
- this.logger?.debug?.(`Added column ${columnName} to table ${tableName}`);
1699
+ if (!rows || rows.length === 0) {
1700
+ return null;
1701
+ }
1702
+ const spans = transformRows(rows);
1703
+ const span = spans[0];
1704
+ if (!span) {
1705
+ return null;
1706
+ }
1707
+ return { span };
1708
+ } catch (error$1) {
1709
+ throw new error.MastraError(
1710
+ {
1711
+ id: storage.createStorageErrorId("CLICKHOUSE", "GET_ROOT_SPAN", "FAILED"),
1712
+ domain: error.ErrorDomain.STORAGE,
1713
+ category: error.ErrorCategory.THIRD_PARTY,
1714
+ details: { traceId }
1715
+ },
1716
+ error$1
1717
+ );
1718
+ }
1719
+ }
1720
+ async getTrace(args) {
1721
+ const { traceId } = args;
1722
+ try {
1723
+ const engine = TABLE_ENGINES[storage.TABLE_SPANS] ?? "MergeTree()";
1724
+ const result = await this.client.query({
1725
+ query: `
1726
+ SELECT *
1727
+ FROM ${storage.TABLE_SPANS} ${engine.startsWith("ReplacingMergeTree") ? "FINAL" : ""}
1728
+ WHERE traceId = {traceId:String}
1729
+ ORDER BY startedAt DESC
1730
+ `,
1731
+ query_params: { traceId },
1732
+ format: "JSONEachRow",
1733
+ clickhouse_settings: {
1734
+ date_time_input_format: "best_effort",
1735
+ date_time_output_format: "iso",
1736
+ use_client_time_zone: 1,
1737
+ output_format_json_quote_64bit_integers: 0
1312
1738
  }
1739
+ });
1740
+ const rows = await result.json();
1741
+ if (!rows || rows.length === 0) {
1742
+ return null;
1313
1743
  }
1744
+ return {
1745
+ traceId,
1746
+ spans: transformRows(rows)
1747
+ };
1314
1748
  } catch (error$1) {
1315
1749
  throw new error.MastraError(
1316
1750
  {
1317
- id: storage.createStorageErrorId("CLICKHOUSE", "ALTER_TABLE", "FAILED"),
1751
+ id: storage.createStorageErrorId("CLICKHOUSE", "GET_TRACE", "FAILED"),
1318
1752
  domain: error.ErrorDomain.STORAGE,
1319
1753
  category: error.ErrorCategory.THIRD_PARTY,
1320
- details: { tableName }
1754
+ details: { traceId }
1321
1755
  },
1322
1756
  error$1
1323
1757
  );
1324
1758
  }
1325
1759
  }
1326
- async clearTable({ tableName }) {
1760
+ async updateSpan(args) {
1761
+ const { traceId, spanId, updates } = args;
1327
1762
  try {
1328
- await this.client.query({
1329
- query: `TRUNCATE TABLE ${tableName}`,
1763
+ const existing = await this.#db.load({
1764
+ tableName: storage.TABLE_SPANS,
1765
+ keys: { spanId, traceId }
1766
+ });
1767
+ if (!existing) {
1768
+ throw new error.MastraError({
1769
+ id: storage.createStorageErrorId("CLICKHOUSE", "UPDATE_SPAN", "NOT_FOUND"),
1770
+ domain: error.ErrorDomain.STORAGE,
1771
+ category: error.ErrorCategory.THIRD_PARTY,
1772
+ details: { spanId, traceId }
1773
+ });
1774
+ }
1775
+ const data = { ...updates };
1776
+ if (data.endedAt instanceof Date) {
1777
+ data.endedAt = data.endedAt.getTime();
1778
+ }
1779
+ if (data.startedAt instanceof Date) {
1780
+ data.startedAt = data.startedAt.getTime();
1781
+ }
1782
+ const updated = {
1783
+ ...existing,
1784
+ ...data,
1785
+ updatedAt: Date.now()
1786
+ };
1787
+ await this.client.insert({
1788
+ table: storage.TABLE_SPANS,
1789
+ values: [updated],
1790
+ format: "JSONEachRow",
1791
+ clickhouse_settings: {
1792
+ date_time_input_format: "best_effort",
1793
+ use_client_time_zone: 1,
1794
+ output_format_json_quote_64bit_integers: 0
1795
+ }
1796
+ });
1797
+ } catch (error$1) {
1798
+ if (error$1 instanceof error.MastraError) throw error$1;
1799
+ throw new error.MastraError(
1800
+ {
1801
+ id: storage.createStorageErrorId("CLICKHOUSE", "UPDATE_SPAN", "FAILED"),
1802
+ domain: error.ErrorDomain.STORAGE,
1803
+ category: error.ErrorCategory.THIRD_PARTY,
1804
+ details: { spanId, traceId }
1805
+ },
1806
+ error$1
1807
+ );
1808
+ }
1809
+ }
1810
+ async listTraces(args) {
1811
+ const { filters, pagination, orderBy } = storage.listTracesArgsSchema.parse(args);
1812
+ const { page, perPage } = pagination;
1813
+ try {
1814
+ const conditions = [`(parentSpanId IS NULL OR parentSpanId = '')`];
1815
+ const values = {};
1816
+ let paramIndex = 0;
1817
+ if (filters) {
1818
+ if (filters.startedAt?.start) {
1819
+ conditions.push(`startedAt >= {startedAtStart:DateTime64(3)}`);
1820
+ values.startedAtStart = filters.startedAt.start.getTime();
1821
+ }
1822
+ if (filters.startedAt?.end) {
1823
+ conditions.push(`startedAt <= {startedAtEnd:DateTime64(3)}`);
1824
+ values.startedAtEnd = filters.startedAt.end.getTime();
1825
+ }
1826
+ if (filters.endedAt?.start) {
1827
+ conditions.push(`endedAt >= {endedAtStart:DateTime64(3)}`);
1828
+ values.endedAtStart = filters.endedAt.start.getTime();
1829
+ }
1830
+ if (filters.endedAt?.end) {
1831
+ conditions.push(`endedAt <= {endedAtEnd:DateTime64(3)}`);
1832
+ values.endedAtEnd = filters.endedAt.end.getTime();
1833
+ }
1834
+ if (filters.spanType !== void 0) {
1835
+ conditions.push(`spanType = {spanType:String}`);
1836
+ values.spanType = filters.spanType;
1837
+ }
1838
+ if (filters.entityType !== void 0) {
1839
+ conditions.push(`entityType = {entityType:String}`);
1840
+ values.entityType = filters.entityType;
1841
+ }
1842
+ if (filters.entityId !== void 0) {
1843
+ conditions.push(`entityId = {entityId:String}`);
1844
+ values.entityId = filters.entityId;
1845
+ }
1846
+ if (filters.entityName !== void 0) {
1847
+ conditions.push(`entityName = {entityName:String}`);
1848
+ values.entityName = filters.entityName;
1849
+ }
1850
+ if (filters.userId !== void 0) {
1851
+ conditions.push(`userId = {userId:String}`);
1852
+ values.userId = filters.userId;
1853
+ }
1854
+ if (filters.organizationId !== void 0) {
1855
+ conditions.push(`organizationId = {organizationId:String}`);
1856
+ values.organizationId = filters.organizationId;
1857
+ }
1858
+ if (filters.resourceId !== void 0) {
1859
+ conditions.push(`resourceId = {resourceId:String}`);
1860
+ values.resourceId = filters.resourceId;
1861
+ }
1862
+ if (filters.runId !== void 0) {
1863
+ conditions.push(`runId = {runId:String}`);
1864
+ values.runId = filters.runId;
1865
+ }
1866
+ if (filters.sessionId !== void 0) {
1867
+ conditions.push(`sessionId = {sessionId:String}`);
1868
+ values.sessionId = filters.sessionId;
1869
+ }
1870
+ if (filters.threadId !== void 0) {
1871
+ conditions.push(`threadId = {threadId:String}`);
1872
+ values.threadId = filters.threadId;
1873
+ }
1874
+ if (filters.requestId !== void 0) {
1875
+ conditions.push(`requestId = {requestId:String}`);
1876
+ values.requestId = filters.requestId;
1877
+ }
1878
+ if (filters.environment !== void 0) {
1879
+ conditions.push(`environment = {environment:String}`);
1880
+ values.environment = filters.environment;
1881
+ }
1882
+ if (filters.source !== void 0) {
1883
+ conditions.push(`source = {source:String}`);
1884
+ values.source = filters.source;
1885
+ }
1886
+ if (filters.serviceName !== void 0) {
1887
+ conditions.push(`serviceName = {serviceName:String}`);
1888
+ values.serviceName = filters.serviceName;
1889
+ }
1890
+ if (filters.scope != null) {
1891
+ for (const [key, value] of Object.entries(filters.scope)) {
1892
+ if (!/^[a-zA-Z_][a-zA-Z0-9_]*$/.test(key)) {
1893
+ throw new error.MastraError({
1894
+ id: storage.createStorageErrorId("CLICKHOUSE", "LIST_TRACES", "INVALID_FILTER_KEY"),
1895
+ domain: error.ErrorDomain.STORAGE,
1896
+ category: error.ErrorCategory.USER,
1897
+ details: { key }
1898
+ });
1899
+ }
1900
+ const paramName = `scope_${key}_${paramIndex++}`;
1901
+ conditions.push(`JSONExtractString(scope, '${key}') = {${paramName}:String}`);
1902
+ values[paramName] = typeof value === "string" ? value : JSON.stringify(value);
1903
+ }
1904
+ }
1905
+ if (filters.metadata != null) {
1906
+ for (const [key, value] of Object.entries(filters.metadata)) {
1907
+ if (!/^[a-zA-Z_][a-zA-Z0-9_]*$/.test(key)) {
1908
+ throw new error.MastraError({
1909
+ id: storage.createStorageErrorId("CLICKHOUSE", "LIST_TRACES", "INVALID_FILTER_KEY"),
1910
+ domain: error.ErrorDomain.STORAGE,
1911
+ category: error.ErrorCategory.USER,
1912
+ details: { key }
1913
+ });
1914
+ }
1915
+ const paramName = `metadata_${key}_${paramIndex++}`;
1916
+ conditions.push(`JSONExtractString(metadata, '${key}') = {${paramName}:String}`);
1917
+ values[paramName] = typeof value === "string" ? value : JSON.stringify(value);
1918
+ }
1919
+ }
1920
+ if (filters.tags != null && filters.tags.length > 0) {
1921
+ for (const tag of filters.tags) {
1922
+ const paramName = `tag_${paramIndex++}`;
1923
+ conditions.push(`has(JSONExtract(tags, 'Array(String)'), {${paramName}:String})`);
1924
+ values[paramName] = tag;
1925
+ }
1926
+ }
1927
+ if (filters.status !== void 0) {
1928
+ switch (filters.status) {
1929
+ case storage.TraceStatus.ERROR:
1930
+ conditions.push(`(error IS NOT NULL AND error != '')`);
1931
+ break;
1932
+ case storage.TraceStatus.RUNNING:
1933
+ conditions.push(`(endedAt IS NULL OR endedAt = '') AND (error IS NULL OR error = '')`);
1934
+ break;
1935
+ case storage.TraceStatus.SUCCESS:
1936
+ conditions.push(`(endedAt IS NOT NULL AND endedAt != '') AND (error IS NULL OR error = '')`);
1937
+ break;
1938
+ }
1939
+ }
1940
+ if (filters.hasChildError !== void 0) {
1941
+ const engine2 = TABLE_ENGINES[storage.TABLE_SPANS] ?? "MergeTree()";
1942
+ const finalClause2 = engine2.startsWith("ReplacingMergeTree") ? "FINAL" : "";
1943
+ if (filters.hasChildError) {
1944
+ conditions.push(`EXISTS (
1945
+ SELECT 1 FROM ${storage.TABLE_SPANS} ${finalClause2} c
1946
+ WHERE c.traceId = ${storage.TABLE_SPANS}.traceId AND c.error IS NOT NULL AND c.error != ''
1947
+ )`);
1948
+ } else {
1949
+ conditions.push(`NOT EXISTS (
1950
+ SELECT 1 FROM ${storage.TABLE_SPANS} ${finalClause2} c
1951
+ WHERE c.traceId = ${storage.TABLE_SPANS}.traceId AND c.error IS NOT NULL AND c.error != ''
1952
+ )`);
1953
+ }
1954
+ }
1955
+ }
1956
+ const whereClause = conditions.length > 0 ? `WHERE ${conditions.join(" AND ")}` : "";
1957
+ const engine = TABLE_ENGINES[storage.TABLE_SPANS] ?? "MergeTree()";
1958
+ const finalClause = engine.startsWith("ReplacingMergeTree") ? "FINAL" : "";
1959
+ const sortField = orderBy.field;
1960
+ const sortDirection = orderBy.direction;
1961
+ let orderClause;
1962
+ if (sortField === "endedAt") {
1963
+ const nullSortValue = sortDirection === "DESC" ? 0 : 1;
1964
+ const nonNullSortValue = sortDirection === "DESC" ? 1 : 0;
1965
+ orderClause = `ORDER BY CASE WHEN ${sortField} IS NULL OR ${sortField} = '' THEN ${nullSortValue} ELSE ${nonNullSortValue} END, ${sortField} ${sortDirection}`;
1966
+ } else {
1967
+ orderClause = `ORDER BY ${sortField} ${sortDirection}`;
1968
+ }
1969
+ const countResult = await this.client.query({
1970
+ query: `SELECT COUNT(*) as count FROM ${storage.TABLE_SPANS} ${finalClause} ${whereClause}`,
1971
+ query_params: values,
1972
+ format: "JSONEachRow"
1973
+ });
1974
+ const countRows = await countResult.json();
1975
+ const total = Number(countRows[0]?.count ?? 0);
1976
+ if (total === 0) {
1977
+ return {
1978
+ pagination: { total: 0, page, perPage, hasMore: false },
1979
+ spans: []
1980
+ };
1981
+ }
1982
+ const result = await this.client.query({
1983
+ query: `
1984
+ SELECT *
1985
+ FROM ${storage.TABLE_SPANS} ${finalClause}
1986
+ ${whereClause}
1987
+ ${orderClause}
1988
+ LIMIT {limit:UInt32}
1989
+ OFFSET {offset:UInt32}
1990
+ `,
1991
+ query_params: { ...values, limit: perPage, offset: page * perPage },
1992
+ format: "JSONEachRow",
1330
1993
  clickhouse_settings: {
1331
- // Allows to insert serialized JS Dates (such as '2023-12-06T10:54:48.000Z')
1332
1994
  date_time_input_format: "best_effort",
1333
1995
  date_time_output_format: "iso",
1334
1996
  use_client_time_zone: 1,
1335
1997
  output_format_json_quote_64bit_integers: 0
1336
1998
  }
1337
1999
  });
2000
+ const rows = await result.json();
2001
+ const spans = transformRows(rows);
2002
+ return {
2003
+ pagination: {
2004
+ total,
2005
+ page,
2006
+ perPage,
2007
+ hasMore: (page + 1) * perPage < total
2008
+ },
2009
+ spans
2010
+ };
1338
2011
  } catch (error$1) {
2012
+ if (error$1 instanceof error.MastraError) throw error$1;
1339
2013
  throw new error.MastraError(
1340
2014
  {
1341
- id: storage.createStorageErrorId("CLICKHOUSE", "CLEAR_TABLE", "FAILED"),
2015
+ id: storage.createStorageErrorId("CLICKHOUSE", "LIST_TRACES", "FAILED"),
1342
2016
  domain: error.ErrorDomain.STORAGE,
1343
- category: error.ErrorCategory.THIRD_PARTY,
1344
- details: { tableName }
2017
+ category: error.ErrorCategory.USER
1345
2018
  },
1346
2019
  error$1
1347
2020
  );
1348
2021
  }
1349
2022
  }
1350
- async dropTable({ tableName }) {
1351
- await this.client.query({
1352
- query: `DROP TABLE IF EXISTS ${tableName}`
1353
- });
1354
- }
1355
- async insert({ tableName, record }) {
1356
- const createdAt = (record.createdAt || record.created_at || /* @__PURE__ */ new Date()).toISOString();
1357
- const updatedAt = (record.updatedAt || /* @__PURE__ */ new Date()).toISOString();
2023
+ async batchCreateSpans(args) {
1358
2024
  try {
1359
- const result = await this.client.insert({
1360
- table: tableName,
1361
- values: [
1362
- {
1363
- ...record,
1364
- createdAt,
1365
- updatedAt
1366
- }
1367
- ],
1368
- format: "JSONEachRow",
1369
- clickhouse_settings: {
1370
- // Allows to insert serialized JS Dates (such as '2023-12-06T10:54:48.000Z')
1371
- output_format_json_quote_64bit_integers: 0,
1372
- date_time_input_format: "best_effort",
1373
- use_client_time_zone: 1
1374
- }
2025
+ const now = Date.now();
2026
+ await this.#db.batchInsert({
2027
+ tableName: storage.TABLE_SPANS,
2028
+ records: args.records.map((record) => ({
2029
+ ...record,
2030
+ // Convert Date objects to millisecond timestamps for DateTime64(3)
2031
+ startedAt: record.startedAt instanceof Date ? record.startedAt.getTime() : record.startedAt,
2032
+ endedAt: record.endedAt instanceof Date ? record.endedAt.getTime() : record.endedAt,
2033
+ createdAt: now,
2034
+ updatedAt: now
2035
+ }))
1375
2036
  });
1376
- console.info("INSERT RESULT", result);
1377
2037
  } catch (error$1) {
1378
2038
  throw new error.MastraError(
1379
2039
  {
1380
- id: storage.createStorageErrorId("CLICKHOUSE", "INSERT", "FAILED"),
2040
+ id: storage.createStorageErrorId("CLICKHOUSE", "BATCH_CREATE_SPANS", "FAILED"),
1381
2041
  domain: error.ErrorDomain.STORAGE,
1382
- category: error.ErrorCategory.THIRD_PARTY,
1383
- details: { tableName }
2042
+ category: error.ErrorCategory.THIRD_PARTY
1384
2043
  },
1385
2044
  error$1
1386
2045
  );
1387
2046
  }
1388
2047
  }
1389
- async batchInsert({ tableName, records }) {
1390
- const recordsToBeInserted = records.map((record) => ({
1391
- ...Object.fromEntries(
1392
- Object.entries(record).map(([key, value]) => [
1393
- key,
1394
- storage.TABLE_SCHEMAS[tableName]?.[key]?.type === "timestamp" ? new Date(value).toISOString() : value
1395
- ])
1396
- )
1397
- }));
2048
+ async batchUpdateSpans(args) {
1398
2049
  try {
1399
- await this.client.insert({
1400
- table: tableName,
1401
- values: recordsToBeInserted,
1402
- format: "JSONEachRow",
1403
- clickhouse_settings: {
1404
- // Allows to insert serialized JS Dates (such as '2023-12-06T10:54:48.000Z')
1405
- date_time_input_format: "best_effort",
1406
- use_client_time_zone: 1,
1407
- output_format_json_quote_64bit_integers: 0
2050
+ const now = Date.now();
2051
+ for (const record of args.records) {
2052
+ const existing = await this.#db.load({
2053
+ tableName: storage.TABLE_SPANS,
2054
+ keys: { spanId: record.spanId, traceId: record.traceId }
2055
+ });
2056
+ if (existing) {
2057
+ const updates = { ...record.updates };
2058
+ if (updates.startedAt instanceof Date) {
2059
+ updates.startedAt = updates.startedAt.getTime();
2060
+ }
2061
+ if (updates.endedAt instanceof Date) {
2062
+ updates.endedAt = updates.endedAt.getTime();
2063
+ }
2064
+ const updated = {
2065
+ ...existing,
2066
+ ...updates,
2067
+ updatedAt: now
2068
+ };
2069
+ await this.client.insert({
2070
+ table: storage.TABLE_SPANS,
2071
+ values: [updated],
2072
+ format: "JSONEachRow",
2073
+ clickhouse_settings: {
2074
+ date_time_input_format: "best_effort",
2075
+ use_client_time_zone: 1,
2076
+ output_format_json_quote_64bit_integers: 0
2077
+ }
2078
+ });
1408
2079
  }
1409
- });
2080
+ }
1410
2081
  } catch (error$1) {
1411
2082
  throw new error.MastraError(
1412
2083
  {
1413
- id: storage.createStorageErrorId("CLICKHOUSE", "BATCH_INSERT", "FAILED"),
2084
+ id: storage.createStorageErrorId("CLICKHOUSE", "BATCH_UPDATE_SPANS", "FAILED"),
1414
2085
  domain: error.ErrorDomain.STORAGE,
1415
- category: error.ErrorCategory.THIRD_PARTY,
1416
- details: { tableName }
2086
+ category: error.ErrorCategory.THIRD_PARTY
1417
2087
  },
1418
2088
  error$1
1419
2089
  );
1420
2090
  }
1421
2091
  }
1422
- async load({ tableName, keys }) {
2092
+ async batchDeleteTraces(args) {
1423
2093
  try {
1424
- const engine = TABLE_ENGINES[tableName] ?? "MergeTree()";
1425
- const keyEntries = Object.entries(keys);
1426
- const conditions = keyEntries.map(
1427
- ([key]) => `"${key}" = {var_${key}:${COLUMN_TYPES[storage.TABLE_SCHEMAS[tableName]?.[key]?.type ?? "text"]}}`
1428
- ).join(" AND ");
1429
- const values = keyEntries.reduce((acc, [key, value]) => {
1430
- return { ...acc, [`var_${key}`]: value };
1431
- }, {});
1432
- const hasUpdatedAt = storage.TABLE_SCHEMAS[tableName]?.updatedAt;
1433
- const selectClause = `SELECT *, toDateTime64(createdAt, 3) as createdAt${hasUpdatedAt ? ", toDateTime64(updatedAt, 3) as updatedAt" : ""}`;
1434
- const result = await this.client.query({
1435
- query: `${selectClause} FROM ${tableName} ${engine.startsWith("ReplacingMergeTree") ? "FINAL" : ""} WHERE ${conditions} ORDER BY createdAt DESC LIMIT 1`,
1436
- query_params: values,
1437
- clickhouse_settings: {
1438
- // Allows to insert serialized JS Dates (such as '2023-12-06T10:54:48.000Z')
1439
- date_time_input_format: "best_effort",
1440
- date_time_output_format: "iso",
1441
- use_client_time_zone: 1,
1442
- output_format_json_quote_64bit_integers: 0
1443
- }
2094
+ if (args.traceIds.length === 0) return;
2095
+ await this.client.command({
2096
+ query: `DELETE FROM ${storage.TABLE_SPANS} WHERE traceId IN {traceIds:Array(String)}`,
2097
+ query_params: { traceIds: args.traceIds }
1444
2098
  });
1445
- if (!result) {
1446
- return null;
1447
- }
1448
- const rows = await result.json();
1449
- if (tableName === storage.TABLE_WORKFLOW_SNAPSHOT) {
1450
- const snapshot = rows.data[0];
1451
- if (!snapshot) {
1452
- return null;
1453
- }
1454
- if (typeof snapshot.snapshot === "string") {
1455
- snapshot.snapshot = JSON.parse(snapshot.snapshot);
1456
- }
1457
- return transformRow(snapshot);
1458
- }
1459
- const data = transformRow(rows.data[0]);
1460
- return data;
1461
2099
  } catch (error$1) {
1462
2100
  throw new error.MastraError(
1463
2101
  {
1464
- id: storage.createStorageErrorId("CLICKHOUSE", "LOAD", "FAILED"),
2102
+ id: storage.createStorageErrorId("CLICKHOUSE", "BATCH_DELETE_TRACES", "FAILED"),
1465
2103
  domain: error.ErrorDomain.STORAGE,
1466
- category: error.ErrorCategory.THIRD_PARTY,
1467
- details: { tableName }
2104
+ category: error.ErrorCategory.THIRD_PARTY
1468
2105
  },
1469
2106
  error$1
1470
2107
  );
@@ -1473,11 +2110,18 @@ var StoreOperationsClickhouse = class extends storage.StoreOperations {
1473
2110
  };
1474
2111
  var ScoresStorageClickhouse = class extends storage.ScoresStorage {
1475
2112
  client;
1476
- operations;
1477
- constructor({ client, operations }) {
2113
+ #db;
2114
+ constructor(config) {
1478
2115
  super();
2116
+ const { client, ttl } = resolveClickhouseConfig(config);
1479
2117
  this.client = client;
1480
- this.operations = operations;
2118
+ this.#db = new ClickhouseDB({ client, ttl });
2119
+ }
2120
+ async init() {
2121
+ await this.#db.createTable({ tableName: storage.TABLE_SCORERS, schema: storage.TABLE_SCHEMAS[storage.TABLE_SCORERS] });
2122
+ }
2123
+ async dangerouslyClearAll() {
2124
+ await this.#db.clearTable({ tableName: storage.TABLE_SCORERS });
1481
2125
  }
1482
2126
  /**
1483
2127
  * ClickHouse-specific score row transformation.
@@ -1531,7 +2175,7 @@ var ScoresStorageClickhouse = class extends storage.ScoresStorage {
1531
2175
  domain: error.ErrorDomain.STORAGE,
1532
2176
  category: error.ErrorCategory.USER,
1533
2177
  details: {
1534
- scorer: score.scorer?.id ?? "unknown",
2178
+ scorer: typeof score.scorer?.id === "string" ? score.scorer.id : String(score.scorer?.id ?? "unknown"),
1535
2179
  entityId: score.entityId ?? "unknown",
1536
2180
  entityType: score.entityType ?? "unknown",
1537
2181
  traceId: score.traceId ?? "",
@@ -1893,46 +2537,85 @@ var ScoresStorageClickhouse = class extends storage.ScoresStorage {
1893
2537
  };
1894
2538
  var WorkflowsStorageClickhouse = class extends storage.WorkflowsStorage {
1895
2539
  client;
1896
- operations;
1897
- constructor({ client, operations }) {
2540
+ #db;
2541
+ constructor(config) {
1898
2542
  super();
1899
- this.operations = operations;
2543
+ const { client, ttl } = resolveClickhouseConfig(config);
1900
2544
  this.client = client;
2545
+ this.#db = new ClickhouseDB({ client, ttl });
2546
+ }
2547
+ async init() {
2548
+ const schema = storage.TABLE_SCHEMAS[storage.TABLE_WORKFLOW_SNAPSHOT];
2549
+ await this.#db.createTable({ tableName: storage.TABLE_WORKFLOW_SNAPSHOT, schema });
2550
+ await this.#db.alterTable({
2551
+ tableName: storage.TABLE_WORKFLOW_SNAPSHOT,
2552
+ schema,
2553
+ ifNotExists: ["resourceId"]
2554
+ });
2555
+ }
2556
+ async dangerouslyClearAll() {
2557
+ await this.#db.clearTable({ tableName: storage.TABLE_WORKFLOW_SNAPSHOT });
1901
2558
  }
1902
- updateWorkflowResults({
1903
- // workflowName,
1904
- // runId,
1905
- // stepId,
1906
- // result,
1907
- // requestContext,
2559
+ async updateWorkflowResults({
2560
+ workflowName,
2561
+ runId,
2562
+ stepId,
2563
+ result,
2564
+ requestContext
1908
2565
  }) {
1909
- throw new error.MastraError({
1910
- id: storage.createStorageErrorId("CLICKHOUSE", "UPDATE_WORKFLOW_RESULTS", "NOT_IMPLEMENTED"),
1911
- domain: error.ErrorDomain.STORAGE,
1912
- category: error.ErrorCategory.SYSTEM,
1913
- text: "Method not implemented."
1914
- });
2566
+ let snapshot = await this.loadWorkflowSnapshot({ workflowName, runId });
2567
+ if (!snapshot) {
2568
+ snapshot = {
2569
+ context: {},
2570
+ activePaths: [],
2571
+ timestamp: Date.now(),
2572
+ suspendedPaths: {},
2573
+ activeStepsPath: {},
2574
+ resumeLabels: {},
2575
+ serializedStepGraph: [],
2576
+ status: "pending",
2577
+ value: {},
2578
+ waitingPaths: {},
2579
+ runId,
2580
+ requestContext: {}
2581
+ };
2582
+ }
2583
+ snapshot.context[stepId] = result;
2584
+ snapshot.requestContext = { ...snapshot.requestContext, ...requestContext };
2585
+ await this.persistWorkflowSnapshot({ workflowName, runId, snapshot });
2586
+ return snapshot.context;
1915
2587
  }
1916
- updateWorkflowState({
1917
- // workflowName,
1918
- // runId,
1919
- // opts,
2588
+ async updateWorkflowState({
2589
+ workflowName,
2590
+ runId,
2591
+ opts
1920
2592
  }) {
1921
- throw new error.MastraError({
1922
- id: storage.createStorageErrorId("CLICKHOUSE", "UPDATE_WORKFLOW_STATE", "NOT_IMPLEMENTED"),
1923
- domain: error.ErrorDomain.STORAGE,
1924
- category: error.ErrorCategory.SYSTEM,
1925
- text: "Method not implemented."
1926
- });
2593
+ const snapshot = await this.loadWorkflowSnapshot({ workflowName, runId });
2594
+ if (!snapshot) {
2595
+ return void 0;
2596
+ }
2597
+ if (!snapshot.context) {
2598
+ throw new error.MastraError({
2599
+ id: storage.createStorageErrorId("CLICKHOUSE", "UPDATE_WORKFLOW_STATE", "CONTEXT_MISSING"),
2600
+ domain: error.ErrorDomain.STORAGE,
2601
+ category: error.ErrorCategory.SYSTEM,
2602
+ text: `Snapshot context is missing for runId ${runId}`
2603
+ });
2604
+ }
2605
+ const updatedSnapshot = { ...snapshot, ...opts };
2606
+ await this.persistWorkflowSnapshot({ workflowName, runId, snapshot: updatedSnapshot });
2607
+ return updatedSnapshot;
1927
2608
  }
1928
2609
  async persistWorkflowSnapshot({
1929
2610
  workflowName,
1930
2611
  runId,
1931
2612
  resourceId,
1932
- snapshot
2613
+ snapshot,
2614
+ createdAt,
2615
+ updatedAt
1933
2616
  }) {
1934
2617
  try {
1935
- const currentSnapshot = await this.operations.load({
2618
+ const currentSnapshot = await this.#db.load({
1936
2619
  tableName: storage.TABLE_WORKFLOW_SNAPSHOT,
1937
2620
  keys: { workflow_name: workflowName, run_id: runId }
1938
2621
  });
@@ -1941,14 +2624,14 @@ var WorkflowsStorageClickhouse = class extends storage.WorkflowsStorage {
1941
2624
  ...currentSnapshot,
1942
2625
  resourceId,
1943
2626
  snapshot: JSON.stringify(snapshot),
1944
- updatedAt: now.toISOString()
2627
+ updatedAt: (updatedAt ?? now).toISOString()
1945
2628
  } : {
1946
2629
  workflow_name: workflowName,
1947
2630
  run_id: runId,
1948
2631
  resourceId,
1949
2632
  snapshot: JSON.stringify(snapshot),
1950
- createdAt: now.toISOString(),
1951
- updatedAt: now.toISOString()
2633
+ createdAt: (createdAt ?? now).toISOString(),
2634
+ updatedAt: (updatedAt ?? now).toISOString()
1952
2635
  };
1953
2636
  await this.client.insert({
1954
2637
  table: storage.TABLE_WORKFLOW_SNAPSHOT,
@@ -1978,7 +2661,7 @@ var WorkflowsStorageClickhouse = class extends storage.WorkflowsStorage {
1978
2661
  runId
1979
2662
  }) {
1980
2663
  try {
1981
- const result = await this.operations.load({
2664
+ const result = await this.#db.load({
1982
2665
  tableName: storage.TABLE_WORKFLOW_SNAPSHOT,
1983
2666
  keys: {
1984
2667
  workflow_name: workflowName,
@@ -2040,7 +2723,7 @@ var WorkflowsStorageClickhouse = class extends storage.WorkflowsStorage {
2040
2723
  values.var_status = status;
2041
2724
  }
2042
2725
  if (resourceId) {
2043
- const hasResourceId = await this.operations.hasColumn(storage.TABLE_WORKFLOW_SNAPSHOT, "resourceId");
2726
+ const hasResourceId = await this.#db.hasColumn(storage.TABLE_WORKFLOW_SNAPSHOT, "resourceId");
2044
2727
  if (hasResourceId) {
2045
2728
  conditions.push(`resourceId = {var_resourceId:String}`);
2046
2729
  values.var_resourceId = resourceId;
@@ -2182,34 +2865,51 @@ var WorkflowsStorageClickhouse = class extends storage.WorkflowsStorage {
2182
2865
  };
2183
2866
 
2184
2867
  // src/storage/index.ts
2868
+ var isClientConfig = (config) => {
2869
+ return "client" in config;
2870
+ };
2185
2871
  var ClickhouseStore = class extends storage.MastraStorage {
2186
2872
  db;
2187
2873
  ttl = {};
2188
2874
  stores;
2189
2875
  constructor(config) {
2190
2876
  super({ id: config.id, name: "ClickhouseStore", disableInit: config.disableInit });
2191
- this.db = client.createClient({
2192
- url: config.url,
2193
- username: config.username,
2194
- password: config.password,
2195
- clickhouse_settings: {
2196
- date_time_input_format: "best_effort",
2197
- date_time_output_format: "iso",
2198
- // This is crucial
2199
- use_client_time_zone: 1,
2200
- output_format_json_quote_64bit_integers: 0
2877
+ if (isClientConfig(config)) {
2878
+ this.db = config.client;
2879
+ } else {
2880
+ if (!config.url || typeof config.url !== "string" || config.url.trim() === "") {
2881
+ throw new Error("ClickhouseStore: url is required and cannot be empty.");
2201
2882
  }
2202
- });
2883
+ if (typeof config.username !== "string") {
2884
+ throw new Error("ClickhouseStore: username must be a string.");
2885
+ }
2886
+ if (typeof config.password !== "string") {
2887
+ throw new Error("ClickhouseStore: password must be a string.");
2888
+ }
2889
+ this.db = client.createClient({
2890
+ url: config.url,
2891
+ username: config.username,
2892
+ password: config.password,
2893
+ clickhouse_settings: {
2894
+ date_time_input_format: "best_effort",
2895
+ date_time_output_format: "iso",
2896
+ // This is crucial
2897
+ use_client_time_zone: 1,
2898
+ output_format_json_quote_64bit_integers: 0
2899
+ }
2900
+ });
2901
+ }
2203
2902
  this.ttl = config.ttl;
2204
- const operations = new StoreOperationsClickhouse({ client: this.db, ttl: this.ttl });
2205
- const workflows = new WorkflowsStorageClickhouse({ client: this.db, operations });
2206
- const scores = new ScoresStorageClickhouse({ client: this.db, operations });
2207
- const memory = new MemoryStorageClickhouse({ client: this.db, operations });
2903
+ const domainConfig = { client: this.db, ttl: this.ttl };
2904
+ const workflows = new WorkflowsStorageClickhouse(domainConfig);
2905
+ const scores = new ScoresStorageClickhouse(domainConfig);
2906
+ const memory = new MemoryStorageClickhouse(domainConfig);
2907
+ const observability = new ObservabilityStorageClickhouse(domainConfig);
2208
2908
  this.stores = {
2209
- operations,
2210
2909
  workflows,
2211
2910
  scores,
2212
- memory
2911
+ memory,
2912
+ observability
2213
2913
  };
2214
2914
  }
2215
2915
  get supports() {
@@ -2218,13 +2918,13 @@ var ClickhouseStore = class extends storage.MastraStorage {
2218
2918
  resourceWorkingMemory: true,
2219
2919
  hasColumn: true,
2220
2920
  createTable: true,
2221
- deleteMessages: false,
2222
- listScoresBySpan: true
2921
+ deleteMessages: true,
2922
+ observability: true,
2923
+ indexManagement: false,
2924
+ listScoresBySpan: true,
2925
+ agents: false
2223
2926
  };
2224
2927
  }
2225
- async batchInsert({ tableName, records }) {
2226
- await this.stores.operations.batchInsert({ tableName, records });
2227
- }
2228
2928
  async optimizeTable({ tableName }) {
2229
2929
  try {
2230
2930
  await this.db.command({
@@ -2259,145 +2959,25 @@ var ClickhouseStore = class extends storage.MastraStorage {
2259
2959
  );
2260
2960
  }
2261
2961
  }
2262
- async createTable({
2263
- tableName,
2264
- schema
2265
- }) {
2266
- return this.stores.operations.createTable({ tableName, schema });
2267
- }
2268
- async dropTable({ tableName }) {
2269
- return this.stores.operations.dropTable({ tableName });
2270
- }
2271
- async alterTable({
2272
- tableName,
2273
- schema,
2274
- ifNotExists
2275
- }) {
2276
- return this.stores.operations.alterTable({ tableName, schema, ifNotExists });
2277
- }
2278
- async clearTable({ tableName }) {
2279
- return this.stores.operations.clearTable({ tableName });
2280
- }
2281
- async insert({ tableName, record }) {
2282
- return this.stores.operations.insert({ tableName, record });
2283
- }
2284
- async load({ tableName, keys }) {
2285
- return this.stores.operations.load({ tableName, keys });
2286
- }
2287
- async updateWorkflowResults({
2288
- workflowName,
2289
- runId,
2290
- stepId,
2291
- result,
2292
- requestContext
2293
- }) {
2294
- return this.stores.workflows.updateWorkflowResults({ workflowName, runId, stepId, result, requestContext });
2295
- }
2296
- async updateWorkflowState({
2297
- workflowName,
2298
- runId,
2299
- opts
2300
- }) {
2301
- return this.stores.workflows.updateWorkflowState({ workflowName, runId, opts });
2302
- }
2303
- async persistWorkflowSnapshot({
2304
- workflowName,
2305
- runId,
2306
- resourceId,
2307
- snapshot
2308
- }) {
2309
- return this.stores.workflows.persistWorkflowSnapshot({ workflowName, runId, resourceId, snapshot });
2310
- }
2311
- async loadWorkflowSnapshot({
2312
- workflowName,
2313
- runId
2314
- }) {
2315
- return this.stores.workflows.loadWorkflowSnapshot({ workflowName, runId });
2316
- }
2317
- async listWorkflowRuns(args = {}) {
2318
- return this.stores.workflows.listWorkflowRuns(args);
2319
- }
2320
- async getWorkflowRunById({
2321
- runId,
2322
- workflowName
2323
- }) {
2324
- return this.stores.workflows.getWorkflowRunById({ runId, workflowName });
2325
- }
2326
- async deleteWorkflowRunById({ runId, workflowName }) {
2327
- return this.stores.workflows.deleteWorkflowRunById({ runId, workflowName });
2328
- }
2329
- async getThreadById({ threadId }) {
2330
- return this.stores.memory.getThreadById({ threadId });
2331
- }
2332
- async saveThread({ thread }) {
2333
- return this.stores.memory.saveThread({ thread });
2334
- }
2335
- async updateThread({
2336
- id,
2337
- title,
2338
- metadata
2339
- }) {
2340
- return this.stores.memory.updateThread({ id, title, metadata });
2341
- }
2342
- async deleteThread({ threadId }) {
2343
- return this.stores.memory.deleteThread({ threadId });
2344
- }
2345
- async saveMessages(args) {
2346
- return this.stores.memory.saveMessages(args);
2347
- }
2348
- async updateMessages(args) {
2349
- return this.stores.memory.updateMessages(args);
2350
- }
2351
- async getResourceById({ resourceId }) {
2352
- return this.stores.memory.getResourceById({ resourceId });
2353
- }
2354
- async saveResource({ resource }) {
2355
- return this.stores.memory.saveResource({ resource });
2356
- }
2357
- async updateResource({
2358
- resourceId,
2359
- workingMemory,
2360
- metadata
2361
- }) {
2362
- return this.stores.memory.updateResource({ resourceId, workingMemory, metadata });
2363
- }
2364
- async getScoreById({ id }) {
2365
- return this.stores.scores.getScoreById({ id });
2366
- }
2367
- async saveScore(score) {
2368
- return this.stores.scores.saveScore(score);
2369
- }
2370
- async listScoresByRunId({
2371
- runId,
2372
- pagination
2373
- }) {
2374
- return this.stores.scores.listScoresByRunId({ runId, pagination });
2375
- }
2376
- async listScoresByEntityId({
2377
- entityId,
2378
- entityType,
2379
- pagination
2380
- }) {
2381
- return this.stores.scores.listScoresByEntityId({ entityId, entityType, pagination });
2382
- }
2383
- async listScoresByScorerId({
2384
- scorerId,
2385
- pagination,
2386
- entityId,
2387
- entityType,
2388
- source
2389
- }) {
2390
- return this.stores.scores.listScoresByScorerId({ scorerId, pagination, entityId, entityType, source });
2391
- }
2392
- async listScoresBySpan({
2393
- traceId,
2394
- spanId,
2395
- pagination
2396
- }) {
2397
- return this.stores.scores.listScoresBySpan({ traceId, spanId, pagination });
2398
- }
2962
+ /**
2963
+ * Closes the ClickHouse client connection.
2964
+ *
2965
+ * This will close the ClickHouse client, including pre-configured clients.
2966
+ * The store assumes ownership of all clients and manages their lifecycle.
2967
+ */
2399
2968
  async close() {
2400
- await this.db.close();
2969
+ try {
2970
+ await this.db.close();
2971
+ } catch (error$1) {
2972
+ throw new error.MastraError(
2973
+ {
2974
+ id: storage.createStorageErrorId("CLICKHOUSE", "CLOSE", "FAILED"),
2975
+ domain: error.ErrorDomain.STORAGE,
2976
+ category: error.ErrorCategory.THIRD_PARTY
2977
+ },
2978
+ error$1
2979
+ );
2980
+ }
2401
2981
  }
2402
2982
  };
2403
2983