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