@mastra/clickhouse 1.0.0-beta.6 → 1.0.0-beta.8

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,278 +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
- case "uuid":
1220
- case "jsonb":
1221
- return "String";
1222
- case "timestamp":
1223
- return "DateTime64(3)";
1224
- case "integer":
1225
- case "bigint":
1226
- return "Int64";
1227
- case "float":
1228
- return "Float64";
1229
- case "boolean":
1230
- return "Bool";
1231
- default:
1232
- 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
+ );
1233
1634
  }
1234
1635
  }
1235
- async createTable({
1236
- tableName,
1237
- schema
1238
- }) {
1636
+ async getSpan(args) {
1637
+ const { traceId, spanId } = args;
1239
1638
  try {
1240
- const columns = Object.entries(schema).map(([name, def]) => {
1241
- const constraints = [];
1242
- if (!def.nullable) constraints.push("NOT NULL");
1243
- if (name === "metadata" && def.type === "text" && def.nullable) {
1244
- constraints.push("DEFAULT '{}'");
1245
- }
1246
- const columnTtl = this.ttl?.[tableName]?.columns?.[name];
1247
- return `"${name}" ${this.getSqlType(def.type)} ${constraints.join(" ")} ${columnTtl ? `TTL toDateTime(${columnTtl.ttlKey ?? "createdAt"}) + INTERVAL ${columnTtl.interval} ${columnTtl.unit}` : ""}`;
1248
- }).join(",\n");
1249
- const rowTtl = this.ttl?.[tableName]?.row;
1250
- const sql = tableName === storage.TABLE_WORKFLOW_SNAPSHOT ? `
1251
- CREATE TABLE IF NOT EXISTS ${tableName} (
1252
- ${["id String"].concat(columns)}
1253
- )
1254
- ENGINE = ${TABLE_ENGINES[tableName] ?? "MergeTree()"}
1255
- PRIMARY KEY (createdAt, run_id, workflow_name)
1256
- ORDER BY (createdAt, run_id, workflow_name)
1257
- ${rowTtl ? `TTL toDateTime(${rowTtl.ttlKey ?? "createdAt"}) + INTERVAL ${rowTtl.interval} ${rowTtl.unit}` : ""}
1258
- SETTINGS index_granularity = 8192
1259
- ` : `
1260
- CREATE TABLE IF NOT EXISTS ${tableName} (
1261
- ${columns}
1262
- )
1263
- ENGINE = ${TABLE_ENGINES[tableName] ?? "MergeTree()"}
1264
- PRIMARY KEY (createdAt, ${"id"})
1265
- ORDER BY (createdAt, ${"id"})
1266
- ${this.ttl?.[tableName]?.row ? `TTL toDateTime(createdAt) + INTERVAL ${this.ttl[tableName].row.interval} ${this.ttl[tableName].row.unit}` : ""}
1267
- SETTINGS index_granularity = 8192
1268
- `;
1269
- await this.client.query({
1270
- 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",
1271
1649
  clickhouse_settings: {
1272
- // Allows to insert serialized JS Dates (such as '2023-12-06T10:54:48.000Z')
1273
1650
  date_time_input_format: "best_effort",
1274
1651
  date_time_output_format: "iso",
1275
1652
  use_client_time_zone: 1,
1276
1653
  output_format_json_quote_64bit_integers: 0
1277
1654
  }
1278
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 };
1279
1666
  } catch (error$1) {
1280
1667
  throw new error.MastraError(
1281
1668
  {
1282
- id: storage.createStorageErrorId("CLICKHOUSE", "CREATE_TABLE", "FAILED"),
1669
+ id: storage.createStorageErrorId("CLICKHOUSE", "GET_SPAN", "FAILED"),
1283
1670
  domain: error.ErrorDomain.STORAGE,
1284
1671
  category: error.ErrorCategory.THIRD_PARTY,
1285
- details: { tableName }
1672
+ details: { traceId, spanId }
1286
1673
  },
1287
1674
  error$1
1288
1675
  );
1289
1676
  }
1290
1677
  }
1291
- async alterTable({
1292
- tableName,
1293
- schema,
1294
- ifNotExists
1295
- }) {
1678
+ async getRootSpan(args) {
1679
+ const { traceId } = args;
1296
1680
  try {
1297
- const describeSql = `DESCRIBE TABLE ${tableName}`;
1681
+ const engine = TABLE_ENGINES[storage.TABLE_SPANS] ?? "MergeTree()";
1298
1682
  const result = await this.client.query({
1299
- 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
+ }
1300
1697
  });
1301
1698
  const rows = await result.json();
1302
- const existingColumnNames = new Set(rows.data.map((row) => row.name.toLowerCase()));
1303
- for (const columnName of ifNotExists) {
1304
- if (!existingColumnNames.has(columnName.toLowerCase()) && schema[columnName]) {
1305
- const columnDef = schema[columnName];
1306
- let sqlType = this.getSqlType(columnDef.type);
1307
- if (columnDef.nullable !== false) {
1308
- sqlType = `Nullable(${sqlType})`;
1309
- }
1310
- const defaultValue = columnDef.nullable === false ? this.getDefaultValue(columnDef.type) : "";
1311
- const alterSql = `ALTER TABLE ${tableName} ADD COLUMN IF NOT EXISTS "${columnName}" ${sqlType} ${defaultValue}`.trim();
1312
- await this.client.query({
1313
- query: alterSql
1314
- });
1315
- this.logger?.debug?.(`Added column ${columnName} to table ${tableName}`);
1316
- }
1699
+ if (!rows || rows.length === 0) {
1700
+ return null;
1317
1701
  }
1702
+ const spans = transformRows(rows);
1703
+ const span = spans[0];
1704
+ if (!span) {
1705
+ return null;
1706
+ }
1707
+ return { span };
1318
1708
  } catch (error$1) {
1319
1709
  throw new error.MastraError(
1320
1710
  {
1321
- id: storage.createStorageErrorId("CLICKHOUSE", "ALTER_TABLE", "FAILED"),
1711
+ id: storage.createStorageErrorId("CLICKHOUSE", "GET_ROOT_SPAN", "FAILED"),
1322
1712
  domain: error.ErrorDomain.STORAGE,
1323
1713
  category: error.ErrorCategory.THIRD_PARTY,
1324
- details: { tableName }
1714
+ details: { traceId }
1325
1715
  },
1326
1716
  error$1
1327
1717
  );
1328
1718
  }
1329
1719
  }
1330
- async clearTable({ tableName }) {
1720
+ async getTrace(args) {
1721
+ const { traceId } = args;
1331
1722
  try {
1332
- await this.client.query({
1333
- query: `TRUNCATE TABLE ${tableName}`,
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",
1334
1733
  clickhouse_settings: {
1335
- // Allows to insert serialized JS Dates (such as '2023-12-06T10:54:48.000Z')
1336
1734
  date_time_input_format: "best_effort",
1337
1735
  date_time_output_format: "iso",
1338
1736
  use_client_time_zone: 1,
1339
1737
  output_format_json_quote_64bit_integers: 0
1340
1738
  }
1341
1739
  });
1740
+ const rows = await result.json();
1741
+ if (!rows || rows.length === 0) {
1742
+ return null;
1743
+ }
1744
+ return {
1745
+ traceId,
1746
+ spans: transformRows(rows)
1747
+ };
1342
1748
  } catch (error$1) {
1343
1749
  throw new error.MastraError(
1344
1750
  {
1345
- id: storage.createStorageErrorId("CLICKHOUSE", "CLEAR_TABLE", "FAILED"),
1751
+ id: storage.createStorageErrorId("CLICKHOUSE", "GET_TRACE", "FAILED"),
1346
1752
  domain: error.ErrorDomain.STORAGE,
1347
1753
  category: error.ErrorCategory.THIRD_PARTY,
1348
- details: { tableName }
1754
+ details: { traceId }
1349
1755
  },
1350
1756
  error$1
1351
1757
  );
1352
1758
  }
1353
1759
  }
1354
- async dropTable({ tableName }) {
1355
- await this.client.query({
1356
- query: `DROP TABLE IF EXISTS ${tableName}`
1357
- });
1760
+ async updateSpan(args) {
1761
+ const { traceId, spanId, updates } = args;
1762
+ try {
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
+ }
1358
1809
  }
1359
- async insert({ tableName, record }) {
1360
- const createdAt = (record.createdAt || record.created_at || /* @__PURE__ */ new Date()).toISOString();
1361
- const updatedAt = (record.updatedAt || /* @__PURE__ */ new Date()).toISOString();
1810
+ async listTraces(args) {
1811
+ const { filters, pagination, orderBy } = storage.listTracesArgsSchema.parse(args);
1812
+ const { page, perPage } = pagination;
1362
1813
  try {
1363
- const result = await this.client.insert({
1364
- table: tableName,
1365
- values: [
1366
- {
1367
- ...record,
1368
- createdAt,
1369
- updatedAt
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);
1370
1903
  }
1371
- ],
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 },
1372
1992
  format: "JSONEachRow",
1373
1993
  clickhouse_settings: {
1374
- // Allows to insert serialized JS Dates (such as '2023-12-06T10:54:48.000Z')
1375
- output_format_json_quote_64bit_integers: 0,
1376
1994
  date_time_input_format: "best_effort",
1377
- use_client_time_zone: 1
1995
+ date_time_output_format: "iso",
1996
+ use_client_time_zone: 1,
1997
+ output_format_json_quote_64bit_integers: 0
1378
1998
  }
1379
1999
  });
1380
- console.info("INSERT RESULT", result);
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
+ };
1381
2011
  } catch (error$1) {
2012
+ if (error$1 instanceof error.MastraError) throw error$1;
1382
2013
  throw new error.MastraError(
1383
2014
  {
1384
- id: storage.createStorageErrorId("CLICKHOUSE", "INSERT", "FAILED"),
2015
+ id: storage.createStorageErrorId("CLICKHOUSE", "LIST_TRACES", "FAILED"),
1385
2016
  domain: error.ErrorDomain.STORAGE,
1386
- category: error.ErrorCategory.THIRD_PARTY,
1387
- details: { tableName }
2017
+ category: error.ErrorCategory.USER
1388
2018
  },
1389
2019
  error$1
1390
2020
  );
1391
2021
  }
1392
2022
  }
1393
- async batchInsert({ tableName, records }) {
1394
- const recordsToBeInserted = records.map((record) => ({
1395
- ...Object.fromEntries(
1396
- Object.entries(record).map(([key, value]) => [
1397
- key,
1398
- storage.TABLE_SCHEMAS[tableName]?.[key]?.type === "timestamp" ? new Date(value).toISOString() : value
1399
- ])
1400
- )
1401
- }));
2023
+ async batchCreateSpans(args) {
1402
2024
  try {
1403
- await this.client.insert({
1404
- table: tableName,
1405
- values: recordsToBeInserted,
1406
- format: "JSONEachRow",
1407
- clickhouse_settings: {
1408
- // Allows to insert serialized JS Dates (such as '2023-12-06T10:54:48.000Z')
1409
- date_time_input_format: "best_effort",
1410
- use_client_time_zone: 1,
1411
- output_format_json_quote_64bit_integers: 0
1412
- }
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
+ }))
1413
2036
  });
1414
2037
  } catch (error$1) {
1415
2038
  throw new error.MastraError(
1416
2039
  {
1417
- id: storage.createStorageErrorId("CLICKHOUSE", "BATCH_INSERT", "FAILED"),
2040
+ id: storage.createStorageErrorId("CLICKHOUSE", "BATCH_CREATE_SPANS", "FAILED"),
1418
2041
  domain: error.ErrorDomain.STORAGE,
1419
- category: error.ErrorCategory.THIRD_PARTY,
1420
- details: { tableName }
2042
+ category: error.ErrorCategory.THIRD_PARTY
1421
2043
  },
1422
2044
  error$1
1423
2045
  );
1424
2046
  }
1425
2047
  }
1426
- async load({ tableName, keys }) {
2048
+ async batchUpdateSpans(args) {
1427
2049
  try {
1428
- const engine = TABLE_ENGINES[tableName] ?? "MergeTree()";
1429
- const keyEntries = Object.entries(keys);
1430
- const conditions = keyEntries.map(
1431
- ([key]) => `"${key}" = {var_${key}:${this.getSqlType(storage.TABLE_SCHEMAS[tableName]?.[key]?.type ?? "text")}}`
1432
- ).join(" AND ");
1433
- const values = keyEntries.reduce((acc, [key, value]) => {
1434
- return { ...acc, [`var_${key}`]: value };
1435
- }, {});
1436
- const hasUpdatedAt = storage.TABLE_SCHEMAS[tableName]?.updatedAt;
1437
- const selectClause = `SELECT *, toDateTime64(createdAt, 3) as createdAt${hasUpdatedAt ? ", toDateTime64(updatedAt, 3) as updatedAt" : ""}`;
1438
- const result = await this.client.query({
1439
- query: `${selectClause} FROM ${tableName} ${engine.startsWith("ReplacingMergeTree") ? "FINAL" : ""} WHERE ${conditions} ORDER BY createdAt DESC LIMIT 1`,
1440
- query_params: values,
1441
- clickhouse_settings: {
1442
- // Allows to insert serialized JS Dates (such as '2023-12-06T10:54:48.000Z')
1443
- date_time_input_format: "best_effort",
1444
- date_time_output_format: "iso",
1445
- use_client_time_zone: 1,
1446
- output_format_json_quote_64bit_integers: 0
1447
- }
1448
- });
1449
- if (!result) {
1450
- return null;
1451
- }
1452
- const rows = await result.json();
1453
- if (tableName === storage.TABLE_WORKFLOW_SNAPSHOT) {
1454
- const snapshot = rows.data[0];
1455
- if (!snapshot) {
1456
- return null;
1457
- }
1458
- if (typeof snapshot.snapshot === "string") {
1459
- snapshot.snapshot = JSON.parse(snapshot.snapshot);
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
+ });
1460
2079
  }
1461
- return transformRow(snapshot);
1462
2080
  }
1463
- const data = transformRow(rows.data[0]);
1464
- return data;
1465
2081
  } catch (error$1) {
1466
2082
  throw new error.MastraError(
1467
2083
  {
1468
- id: storage.createStorageErrorId("CLICKHOUSE", "LOAD", "FAILED"),
2084
+ id: storage.createStorageErrorId("CLICKHOUSE", "BATCH_UPDATE_SPANS", "FAILED"),
1469
2085
  domain: error.ErrorDomain.STORAGE,
1470
- category: error.ErrorCategory.THIRD_PARTY,
1471
- details: { tableName }
2086
+ category: error.ErrorCategory.THIRD_PARTY
2087
+ },
2088
+ error$1
2089
+ );
2090
+ }
2091
+ }
2092
+ async batchDeleteTraces(args) {
2093
+ try {
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 }
2098
+ });
2099
+ } catch (error$1) {
2100
+ throw new error.MastraError(
2101
+ {
2102
+ id: storage.createStorageErrorId("CLICKHOUSE", "BATCH_DELETE_TRACES", "FAILED"),
2103
+ domain: error.ErrorDomain.STORAGE,
2104
+ category: error.ErrorCategory.THIRD_PARTY
1472
2105
  },
1473
2106
  error$1
1474
2107
  );
@@ -1477,11 +2110,18 @@ var StoreOperationsClickhouse = class extends storage.StoreOperations {
1477
2110
  };
1478
2111
  var ScoresStorageClickhouse = class extends storage.ScoresStorage {
1479
2112
  client;
1480
- operations;
1481
- constructor({ client, operations }) {
2113
+ #db;
2114
+ constructor(config) {
1482
2115
  super();
2116
+ const { client, ttl } = resolveClickhouseConfig(config);
1483
2117
  this.client = client;
1484
- 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 });
1485
2125
  }
1486
2126
  /**
1487
2127
  * ClickHouse-specific score row transformation.
@@ -1535,7 +2175,7 @@ var ScoresStorageClickhouse = class extends storage.ScoresStorage {
1535
2175
  domain: error.ErrorDomain.STORAGE,
1536
2176
  category: error.ErrorCategory.USER,
1537
2177
  details: {
1538
- scorer: score.scorer?.id ?? "unknown",
2178
+ scorer: typeof score.scorer?.id === "string" ? score.scorer.id : String(score.scorer?.id ?? "unknown"),
1539
2179
  entityId: score.entityId ?? "unknown",
1540
2180
  entityType: score.entityType ?? "unknown",
1541
2181
  traceId: score.traceId ?? "",
@@ -1897,46 +2537,85 @@ var ScoresStorageClickhouse = class extends storage.ScoresStorage {
1897
2537
  };
1898
2538
  var WorkflowsStorageClickhouse = class extends storage.WorkflowsStorage {
1899
2539
  client;
1900
- operations;
1901
- constructor({ client, operations }) {
2540
+ #db;
2541
+ constructor(config) {
1902
2542
  super();
1903
- this.operations = operations;
2543
+ const { client, ttl } = resolveClickhouseConfig(config);
1904
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
+ });
1905
2555
  }
1906
- updateWorkflowResults({
1907
- // workflowName,
1908
- // runId,
1909
- // stepId,
1910
- // result,
1911
- // requestContext,
2556
+ async dangerouslyClearAll() {
2557
+ await this.#db.clearTable({ tableName: storage.TABLE_WORKFLOW_SNAPSHOT });
2558
+ }
2559
+ async updateWorkflowResults({
2560
+ workflowName,
2561
+ runId,
2562
+ stepId,
2563
+ result,
2564
+ requestContext
1912
2565
  }) {
1913
- throw new error.MastraError({
1914
- id: storage.createStorageErrorId("CLICKHOUSE", "UPDATE_WORKFLOW_RESULTS", "NOT_IMPLEMENTED"),
1915
- domain: error.ErrorDomain.STORAGE,
1916
- category: error.ErrorCategory.SYSTEM,
1917
- text: "Method not implemented."
1918
- });
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;
1919
2587
  }
1920
- updateWorkflowState({
1921
- // workflowName,
1922
- // runId,
1923
- // opts,
2588
+ async updateWorkflowState({
2589
+ workflowName,
2590
+ runId,
2591
+ opts
1924
2592
  }) {
1925
- throw new error.MastraError({
1926
- id: storage.createStorageErrorId("CLICKHOUSE", "UPDATE_WORKFLOW_STATE", "NOT_IMPLEMENTED"),
1927
- domain: error.ErrorDomain.STORAGE,
1928
- category: error.ErrorCategory.SYSTEM,
1929
- text: "Method not implemented."
1930
- });
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;
1931
2608
  }
1932
2609
  async persistWorkflowSnapshot({
1933
2610
  workflowName,
1934
2611
  runId,
1935
2612
  resourceId,
1936
- snapshot
2613
+ snapshot,
2614
+ createdAt,
2615
+ updatedAt
1937
2616
  }) {
1938
2617
  try {
1939
- const currentSnapshot = await this.operations.load({
2618
+ const currentSnapshot = await this.#db.load({
1940
2619
  tableName: storage.TABLE_WORKFLOW_SNAPSHOT,
1941
2620
  keys: { workflow_name: workflowName, run_id: runId }
1942
2621
  });
@@ -1945,14 +2624,14 @@ var WorkflowsStorageClickhouse = class extends storage.WorkflowsStorage {
1945
2624
  ...currentSnapshot,
1946
2625
  resourceId,
1947
2626
  snapshot: JSON.stringify(snapshot),
1948
- updatedAt: now.toISOString()
2627
+ updatedAt: (updatedAt ?? now).toISOString()
1949
2628
  } : {
1950
2629
  workflow_name: workflowName,
1951
2630
  run_id: runId,
1952
2631
  resourceId,
1953
2632
  snapshot: JSON.stringify(snapshot),
1954
- createdAt: now.toISOString(),
1955
- updatedAt: now.toISOString()
2633
+ createdAt: (createdAt ?? now).toISOString(),
2634
+ updatedAt: (updatedAt ?? now).toISOString()
1956
2635
  };
1957
2636
  await this.client.insert({
1958
2637
  table: storage.TABLE_WORKFLOW_SNAPSHOT,
@@ -1982,7 +2661,7 @@ var WorkflowsStorageClickhouse = class extends storage.WorkflowsStorage {
1982
2661
  runId
1983
2662
  }) {
1984
2663
  try {
1985
- const result = await this.operations.load({
2664
+ const result = await this.#db.load({
1986
2665
  tableName: storage.TABLE_WORKFLOW_SNAPSHOT,
1987
2666
  keys: {
1988
2667
  workflow_name: workflowName,
@@ -2044,7 +2723,7 @@ var WorkflowsStorageClickhouse = class extends storage.WorkflowsStorage {
2044
2723
  values.var_status = status;
2045
2724
  }
2046
2725
  if (resourceId) {
2047
- const hasResourceId = await this.operations.hasColumn(storage.TABLE_WORKFLOW_SNAPSHOT, "resourceId");
2726
+ const hasResourceId = await this.#db.hasColumn(storage.TABLE_WORKFLOW_SNAPSHOT, "resourceId");
2048
2727
  if (hasResourceId) {
2049
2728
  conditions.push(`resourceId = {var_resourceId:String}`);
2050
2729
  values.var_resourceId = resourceId;
@@ -2186,49 +2865,53 @@ var WorkflowsStorageClickhouse = class extends storage.WorkflowsStorage {
2186
2865
  };
2187
2866
 
2188
2867
  // src/storage/index.ts
2868
+ var isClientConfig = (config) => {
2869
+ return "client" in config;
2870
+ };
2189
2871
  var ClickhouseStore = class extends storage.MastraStorage {
2190
2872
  db;
2191
2873
  ttl = {};
2192
2874
  stores;
2193
2875
  constructor(config) {
2194
2876
  super({ id: config.id, name: "ClickhouseStore", disableInit: config.disableInit });
2195
- this.db = client.createClient({
2196
- url: config.url,
2197
- username: config.username,
2198
- password: config.password,
2199
- clickhouse_settings: {
2200
- date_time_input_format: "best_effort",
2201
- date_time_output_format: "iso",
2202
- // This is crucial
2203
- use_client_time_zone: 1,
2204
- 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.");
2205
2882
  }
2206
- });
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
+ }
2207
2902
  this.ttl = config.ttl;
2208
- const operations = new StoreOperationsClickhouse({ client: this.db, ttl: this.ttl });
2209
- const workflows = new WorkflowsStorageClickhouse({ client: this.db, operations });
2210
- const scores = new ScoresStorageClickhouse({ client: this.db, operations });
2211
- 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);
2212
2908
  this.stores = {
2213
- operations,
2214
2909
  workflows,
2215
2910
  scores,
2216
- memory
2217
- };
2218
- }
2219
- get supports() {
2220
- return {
2221
- selectByIncludeResourceScope: true,
2222
- resourceWorkingMemory: true,
2223
- hasColumn: true,
2224
- createTable: true,
2225
- deleteMessages: false,
2226
- listScoresBySpan: true
2911
+ memory,
2912
+ observability
2227
2913
  };
2228
2914
  }
2229
- async batchInsert({ tableName, records }) {
2230
- await this.stores.operations.batchInsert({ tableName, records });
2231
- }
2232
2915
  async optimizeTable({ tableName }) {
2233
2916
  try {
2234
2917
  await this.db.command({
@@ -2263,150 +2946,34 @@ var ClickhouseStore = class extends storage.MastraStorage {
2263
2946
  );
2264
2947
  }
2265
2948
  }
2266
- async createTable({
2267
- tableName,
2268
- schema
2269
- }) {
2270
- return this.stores.operations.createTable({ tableName, schema });
2271
- }
2272
- async dropTable({ tableName }) {
2273
- return this.stores.operations.dropTable({ tableName });
2274
- }
2275
- async alterTable({
2276
- tableName,
2277
- schema,
2278
- ifNotExists
2279
- }) {
2280
- return this.stores.operations.alterTable({ tableName, schema, ifNotExists });
2281
- }
2282
- async clearTable({ tableName }) {
2283
- return this.stores.operations.clearTable({ tableName });
2284
- }
2285
- async insert({ tableName, record }) {
2286
- return this.stores.operations.insert({ tableName, record });
2287
- }
2288
- async load({ tableName, keys }) {
2289
- return this.stores.operations.load({ tableName, keys });
2290
- }
2291
- async updateWorkflowResults({
2292
- workflowName,
2293
- runId,
2294
- stepId,
2295
- result,
2296
- requestContext
2297
- }) {
2298
- return this.stores.workflows.updateWorkflowResults({ workflowName, runId, stepId, result, requestContext });
2299
- }
2300
- async updateWorkflowState({
2301
- workflowName,
2302
- runId,
2303
- opts
2304
- }) {
2305
- return this.stores.workflows.updateWorkflowState({ workflowName, runId, opts });
2306
- }
2307
- async persistWorkflowSnapshot({
2308
- workflowName,
2309
- runId,
2310
- resourceId,
2311
- snapshot
2312
- }) {
2313
- return this.stores.workflows.persistWorkflowSnapshot({ workflowName, runId, resourceId, snapshot });
2314
- }
2315
- async loadWorkflowSnapshot({
2316
- workflowName,
2317
- runId
2318
- }) {
2319
- return this.stores.workflows.loadWorkflowSnapshot({ workflowName, runId });
2320
- }
2321
- async listWorkflowRuns(args = {}) {
2322
- return this.stores.workflows.listWorkflowRuns(args);
2323
- }
2324
- async getWorkflowRunById({
2325
- runId,
2326
- workflowName
2327
- }) {
2328
- return this.stores.workflows.getWorkflowRunById({ runId, workflowName });
2329
- }
2330
- async deleteWorkflowRunById({ runId, workflowName }) {
2331
- return this.stores.workflows.deleteWorkflowRunById({ runId, workflowName });
2332
- }
2333
- async getThreadById({ threadId }) {
2334
- return this.stores.memory.getThreadById({ threadId });
2335
- }
2336
- async saveThread({ thread }) {
2337
- return this.stores.memory.saveThread({ thread });
2338
- }
2339
- async updateThread({
2340
- id,
2341
- title,
2342
- metadata
2343
- }) {
2344
- return this.stores.memory.updateThread({ id, title, metadata });
2345
- }
2346
- async deleteThread({ threadId }) {
2347
- return this.stores.memory.deleteThread({ threadId });
2348
- }
2349
- async saveMessages(args) {
2350
- return this.stores.memory.saveMessages(args);
2351
- }
2352
- async updateMessages(args) {
2353
- return this.stores.memory.updateMessages(args);
2354
- }
2355
- async getResourceById({ resourceId }) {
2356
- return this.stores.memory.getResourceById({ resourceId });
2357
- }
2358
- async saveResource({ resource }) {
2359
- return this.stores.memory.saveResource({ resource });
2360
- }
2361
- async updateResource({
2362
- resourceId,
2363
- workingMemory,
2364
- metadata
2365
- }) {
2366
- return this.stores.memory.updateResource({ resourceId, workingMemory, metadata });
2367
- }
2368
- async getScoreById({ id }) {
2369
- return this.stores.scores.getScoreById({ id });
2370
- }
2371
- async saveScore(score) {
2372
- return this.stores.scores.saveScore(score);
2373
- }
2374
- async listScoresByRunId({
2375
- runId,
2376
- pagination
2377
- }) {
2378
- return this.stores.scores.listScoresByRunId({ runId, pagination });
2379
- }
2380
- async listScoresByEntityId({
2381
- entityId,
2382
- entityType,
2383
- pagination
2384
- }) {
2385
- return this.stores.scores.listScoresByEntityId({ entityId, entityType, pagination });
2386
- }
2387
- async listScoresByScorerId({
2388
- scorerId,
2389
- pagination,
2390
- entityId,
2391
- entityType,
2392
- source
2393
- }) {
2394
- return this.stores.scores.listScoresByScorerId({ scorerId, pagination, entityId, entityType, source });
2395
- }
2396
- async listScoresBySpan({
2397
- traceId,
2398
- spanId,
2399
- pagination
2400
- }) {
2401
- return this.stores.scores.listScoresBySpan({ traceId, spanId, pagination });
2402
- }
2949
+ /**
2950
+ * Closes the ClickHouse client connection.
2951
+ *
2952
+ * This will close the ClickHouse client, including pre-configured clients.
2953
+ * The store assumes ownership of all clients and manages their lifecycle.
2954
+ */
2403
2955
  async close() {
2404
- await this.db.close();
2956
+ try {
2957
+ await this.db.close();
2958
+ } catch (error$1) {
2959
+ throw new error.MastraError(
2960
+ {
2961
+ id: storage.createStorageErrorId("CLICKHOUSE", "CLOSE", "FAILED"),
2962
+ domain: error.ErrorDomain.STORAGE,
2963
+ category: error.ErrorCategory.THIRD_PARTY
2964
+ },
2965
+ error$1
2966
+ );
2967
+ }
2405
2968
  }
2406
2969
  };
2407
2970
 
2408
2971
  exports.COLUMN_TYPES = COLUMN_TYPES;
2409
2972
  exports.ClickhouseStore = ClickhouseStore;
2973
+ exports.MemoryStorageClickhouse = MemoryStorageClickhouse;
2974
+ exports.ObservabilityStorageClickhouse = ObservabilityStorageClickhouse;
2975
+ exports.ScoresStorageClickhouse = ScoresStorageClickhouse;
2410
2976
  exports.TABLE_ENGINES = TABLE_ENGINES;
2977
+ exports.WorkflowsStorageClickhouse = WorkflowsStorageClickhouse;
2411
2978
  //# sourceMappingURL=index.cjs.map
2412
2979
  //# sourceMappingURL=index.cjs.map