@mastra/clickhouse 1.2.3 → 1.3.0

package/dist/index.js

@@ -1,6 +1,6 @@
 import { createClient } from '@clickhouse/client';
 import { MastraError, ErrorCategory, ErrorDomain } from '@mastra/core/error';
-import { TABLE_SKILL_BLOBS, TABLE_SKILL_VERSIONS, TABLE_SKILLS, TABLE_WORKSPACE_VERSIONS, TABLE_WORKSPACES, TABLE_MCP_SERVER_VERSIONS, TABLE_MCP_SERVERS, TABLE_MCP_CLIENT_VERSIONS, TABLE_MCP_CLIENTS, TABLE_SCORER_DEFINITION_VERSIONS, TABLE_SCORER_DEFINITIONS, TABLE_PROMPT_BLOCK_VERSIONS, TABLE_PROMPT_BLOCKS, TABLE_EXPERIMENT_RESULTS, TABLE_EXPERIMENTS, TABLE_DATASET_VERSIONS, TABLE_DATASET_ITEMS, TABLE_DATASETS, TABLE_AGENT_VERSIONS, 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, toTraceSpans, ScoresStorage, transformScoreRow, SCORERS_SCHEMA, WorkflowsStorage, MastraCompositeStore, TraceStatus, getSqlType, getDefaultValue, safelyParseJSON } from '@mastra/core/storage';
+import { TABLE_SKILL_BLOBS, TABLE_SKILL_VERSIONS, TABLE_SKILLS, TABLE_WORKSPACE_VERSIONS, TABLE_WORKSPACES, TABLE_MCP_SERVER_VERSIONS, TABLE_MCP_SERVERS, TABLE_MCP_CLIENT_VERSIONS, TABLE_MCP_CLIENTS, TABLE_SCORER_DEFINITION_VERSIONS, TABLE_SCORER_DEFINITIONS, TABLE_PROMPT_BLOCK_VERSIONS, TABLE_PROMPT_BLOCKS, TABLE_EXPERIMENT_RESULTS, TABLE_EXPERIMENTS, TABLE_DATASET_VERSIONS, TABLE_DATASET_ITEMS, TABLE_DATASETS, TABLE_AGENT_VERSIONS, 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, toTraceSpans, ScoresStorage, transformScoreRow, SCORERS_SCHEMA, WorkflowsStorage, MastraCompositeStore, getSqlType, getDefaultValue, safelyParseJSON, TraceStatus } from '@mastra/core/storage';
 import { MessageList } from '@mastra/core/agent';
 import { MastraBase } from '@mastra/core/base';
 import { saveScorePayloadSchema } from '@mastra/core/evals';
@@ -103,6 +103,8 @@ function resolveClickhouseConfig(config) {
 var ClickhouseDB = class extends MastraBase {
   ttl;
   client;
+  /** Cache of actual table columns: tableName -> Promise<Set<columnName>> (stores in-flight promise to coalesce concurrent calls) */
+  tableColumnsCache = /* @__PURE__ */ new Map();
   constructor({ client, ttl }) {
     super({
       name: "CLICKHOUSE_DB"
@@ -110,6 +112,48 @@ var ClickhouseDB = class extends MastraBase {
     this.ttl = ttl;
     this.client = client;
   }
+  /**
+   * Gets the set of column names that actually exist in the database table.
+   * Results are cached; the cache is invalidated when alterTable() adds new columns.
+   */
+  async getTableColumns(tableName) {
+    const cached = this.tableColumnsCache.get(tableName);
+    if (cached) return cached;
+    const promise = (async () => {
+      try {
+        const result = await this.client.query({
+          query: `DESCRIBE TABLE ${tableName}`,
+          format: "JSONEachRow"
+        });
+        const rows = await result.json();
+        const columns = new Set(rows.map((r) => r.name));
+        if (columns.size === 0) {
+          this.tableColumnsCache.delete(tableName);
+        }
+        return columns;
+      } catch {
+        this.tableColumnsCache.delete(tableName);
+        return /* @__PURE__ */ new Set();
+      }
+    })();
+    this.tableColumnsCache.set(tableName, promise);
+    return promise;
+  }
+  /**
+   * Filters a record to only include columns that exist in the actual database table.
+   * Unknown columns are silently dropped to ensure forward compatibility.
+   */
+  async filterRecordToKnownColumns(tableName, record) {
+    const knownColumns = await this.getTableColumns(tableName);
+    if (knownColumns.size === 0) return record;
+    const filtered = {};
+    for (const [key, value] of Object.entries(record)) {
+      if (knownColumns.has(key)) {
+        filtered[key] = value;
+      }
+    }
+    return filtered;
+  }
   async hasColumn(table, column) {
     const result = await this.client.query({
       query: `DESCRIBE TABLE ${table}`,
@@ -421,6 +465,8 @@ var ClickhouseDB = class extends MastraBase {
         },
         error
       );
+    } finally {
+      this.tableColumnsCache.delete(tableName);
     }
   }
   async alterTable({
@@ -460,6 +506,8 @@ var ClickhouseDB = class extends MastraBase {
         },
         error
       );
+    } finally {
+      this.tableColumnsCache.delete(tableName);
     }
   }
   async clearTable({ tableName }) {
@@ -487,25 +535,39 @@ var ClickhouseDB = class extends MastraBase {
     }
   }
   async dropTable({ tableName }) {
-    await this.client.query({
-      query: `DROP TABLE IF EXISTS ${tableName}`
-    });
+    try {
+      await this.client.query({
+        query: `DROP TABLE IF EXISTS ${tableName}`
+      });
+    } catch (error) {
+      throw new MastraError(
+        {
+          id: createStorageErrorId("CLICKHOUSE", "DROP_TABLE", "FAILED"),
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: { tableName }
+        },
+        error
+      );
+    } finally {
+      this.tableColumnsCache.delete(tableName);
+    }
   }
   async insert({ tableName, record }) {
-    const rawCreatedAt = record.createdAt || record.created_at || /* @__PURE__ */ new Date();
-    const rawUpdatedAt = record.updatedAt || /* @__PURE__ */ new Date();
-    const createdAt = rawCreatedAt instanceof Date ? rawCreatedAt.toISOString() : rawCreatedAt;
-    const updatedAt = rawUpdatedAt instanceof Date ? rawUpdatedAt.toISOString() : rawUpdatedAt;
     try {
+      const filteredRecord = await this.filterRecordToKnownColumns(tableName, record);
+      if (Object.keys(filteredRecord).length === 0) return;
+      const rawCreatedAt = filteredRecord.createdAt || filteredRecord.created_at || /* @__PURE__ */ new Date();
+      const rawUpdatedAt = filteredRecord.updatedAt || /* @__PURE__ */ new Date();
+      if ("createdAt" in filteredRecord || (await this.getTableColumns(tableName)).has("createdAt")) {
+        filteredRecord.createdAt = rawCreatedAt instanceof Date ? rawCreatedAt.toISOString() : rawCreatedAt;
+      }
+      if ("updatedAt" in filteredRecord || (await this.getTableColumns(tableName)).has("updatedAt")) {
+        filteredRecord.updatedAt = rawUpdatedAt instanceof Date ? rawUpdatedAt.toISOString() : rawUpdatedAt;
+      }
       await this.client.insert({
         table: tableName,
-        values: [
-          {
-            ...record,
-            createdAt,
-            updatedAt
-          }
-        ],
+        values: [filteredRecord],
         format: "JSONEachRow",
         clickhouse_settings: {
           // Allows to insert serialized JS Dates (such as '2023-12-06T10:54:48.000Z')
@@ -527,7 +589,7 @@ var ClickhouseDB = class extends MastraBase {
     }
   }
   async batchInsert({ tableName, records }) {
-    const recordsToBeInserted = records.map((record) => ({
+    const processedRecords = records.map((record) => ({
       ...Object.fromEntries(
         Object.entries(record).map(([key, value]) => [
           key,
@@ -537,10 +599,15 @@ var ClickhouseDB = class extends MastraBase {
         ])
       )
     }));
+    const recordsToBeInserted = await Promise.all(
+      processedRecords.map((r) => this.filterRecordToKnownColumns(tableName, r))
+    );
+    const nonEmptyRecords = recordsToBeInserted.filter((r) => Object.keys(r).length > 0);
+    if (nonEmptyRecords.length === 0) return;
     try {
       await this.client.insert({
         table: tableName,
-        values: recordsToBeInserted,
+        values: nonEmptyRecords,
         format: "JSONEachRow",
         clickhouse_settings: {
           // Allows to insert serialized JS Dates (such as '2023-12-06T10:54:48.000Z')
@@ -806,6 +873,20 @@ var MemoryStorageClickhouse = class extends MemoryStorage {
       }
       const { field, direction } = this.parseOrderBy(orderBy, "ASC");
       dataQuery += ` ORDER BY "${field}" ${direction}`;
+      if (perPageForQuery === 0 && (!include || include.length === 0)) {
+        return { messages: [], total: 0, page, perPage: perPageForResponse, hasMore: false };
+      }
+      if (perPageForQuery === 0 && include && include.length > 0) {
+        const includeResult = await this._getIncludedMessages({ include });
+        const list2 = new MessageList().add(includeResult, "memory");
+        return {
+          messages: this._sortMessages(list2.get.all.db(), field, direction),
+          total: 0,
+          page,
+          perPage: perPageForResponse,
+          hasMore: false
+        };
+      }
       if (perPageForResponse === false) ; else {
         dataQuery += ` LIMIT {limit:Int64} OFFSET {offset:Int64}`;
         dataParams.limit = perPageForQuery;
@@ -867,92 +948,17 @@ var MemoryStorageClickhouse = class extends MemoryStorage {
         };
       }
       const messageIds = new Set(paginatedMessages.map((m) => m.id));
-      let includeMessages = [];
       if (include && include.length > 0) {
-        const includesNeedingThread = include.filter((inc) => !inc.threadId);
-        const threadByMessageId = /* @__PURE__ */ new Map();
-        if (includesNeedingThread.length > 0) {
-          const { messages: includeLookup } = await this.listMessagesById({
-            messageIds: includesNeedingThread.map((inc) => inc.id)
-          });
-          for (const msg of includeLookup) {
-            if (msg.threadId) {
-              threadByMessageId.set(msg.id, msg.threadId);
-            }
-          }
-        }
-        const unionQueries = [];
-        const params = [];
-        let paramIdx = 1;
-        for (const inc of include) {
-          const { id, withPreviousMessages = 0, withNextMessages = 0 } = inc;
-          const searchThreadId = inc.threadId ?? threadByMessageId.get(id);
-          if (!searchThreadId) continue;
-          unionQueries.push(`
-            SELECT * FROM (
-              WITH numbered_messages AS (
-                SELECT
-                  id, content, role, type, "createdAt", thread_id, "resourceId",
-                  ROW_NUMBER() OVER (ORDER BY "createdAt" ASC) as row_num
-                FROM "${TABLE_MESSAGES}"
-                WHERE thread_id = {var_thread_id_${paramIdx}:String}
-              ),
-              target_positions AS (
-                SELECT row_num as target_pos
-                FROM numbered_messages
-                WHERE id = {var_include_id_${paramIdx}:String}
-              )
-              SELECT DISTINCT m.id, m.content, m.role, m.type, m."createdAt", m.thread_id AS "threadId", m."resourceId"
-              FROM numbered_messages m
-              CROSS JOIN target_positions t
-              WHERE m.row_num BETWEEN (t.target_pos - {var_withPreviousMessages_${paramIdx}:Int64}) AND (t.target_pos + {var_withNextMessages_${paramIdx}:Int64})
-            ) AS query_${paramIdx}
-          `);
-          params.push(
-            { [`var_thread_id_${paramIdx}`]: searchThreadId },
-            { [`var_include_id_${paramIdx}`]: id },
-            { [`var_withPreviousMessages_${paramIdx}`]: withPreviousMessages },
-            { [`var_withNextMessages_${paramIdx}`]: withNextMessages }
-          );
-          paramIdx++;
-        }
-        if (unionQueries.length > 0) {
-          const finalQuery = unionQueries.join(" UNION ALL ") + ' ORDER BY "createdAt" ASC';
-          const mergedParams = params.reduce((acc, paramObj) => ({ ...acc, ...paramObj }), {});
-          const includeResult = await this.client.query({
-            query: finalQuery,
-            query_params: mergedParams,
-            clickhouse_settings: {
-              date_time_input_format: "best_effort",
-              date_time_output_format: "iso",
-              use_client_time_zone: 1,
-              output_format_json_quote_64bit_integers: 0
-            }
-          });
-          const includeRows = await includeResult.json();
-          includeMessages = transformRows(includeRows.data);
-          for (const includeMsg of includeMessages) {
-            if (!messageIds.has(includeMsg.id)) {
-              paginatedMessages.push(includeMsg);
-              messageIds.add(includeMsg.id);
-            }
+        const includeMessages = await this._getIncludedMessages({ include });
+        for (const includeMsg of includeMessages) {
+          if (!messageIds.has(includeMsg.id)) {
+            paginatedMessages.push(includeMsg);
+            messageIds.add(includeMsg.id);
           }
         }
       }
       const list = new MessageList().add(paginatedMessages, "memory");
-      let finalMessages = list.get.all.db();
-      finalMessages = finalMessages.sort((a, b) => {
-        const isDateField = field === "createdAt" || field === "updatedAt";
-        const aValue = isDateField ? new Date(a[field]).getTime() : a[field];
-        const bValue = isDateField ? new Date(b[field]).getTime() : b[field];
-        if (aValue === bValue) {
-          return a.id.localeCompare(b.id);
-        }
-        if (typeof aValue === "number" && typeof bValue === "number") {
-          return direction === "ASC" ? aValue - bValue : bValue - aValue;
-        }
-        return direction === "ASC" ? String(aValue).localeCompare(String(bValue)) : String(bValue).localeCompare(String(aValue));
-      });
+      const finalMessages = this._sortMessages(list.get.all.db(), field, direction);
       const threadIdSet = new Set(threadIds);
       const returnedThreadMessageIds = new Set(
         finalMessages.filter((m) => m.threadId && threadIdSet.has(m.threadId)).map((m) => m.id)
@@ -990,6 +996,95 @@ var MemoryStorageClickhouse = class extends MemoryStorage {
       };
     }
   }
+  _sortMessages(messages, field, direction) {
+    return messages.sort((a, b) => {
+      const isDateField = field === "createdAt" || field === "updatedAt";
+      const aValue = isDateField ? new Date(a[field]).getTime() : a[field];
+      const bValue = isDateField ? new Date(b[field]).getTime() : b[field];
+      if (aValue === bValue) {
+        return a.id.localeCompare(b.id);
+      }
+      if (typeof aValue === "number" && typeof bValue === "number") {
+        return direction === "ASC" ? aValue - bValue : bValue - aValue;
+      }
+      return direction === "ASC" ? String(aValue).localeCompare(String(bValue)) : String(bValue).localeCompare(String(aValue));
+    });
+  }
+  async _getIncludedMessages({
+    include
+  }) {
+    if (!include || include.length === 0) return [];
+    const targetIds = include.map((inc) => inc.id).filter(Boolean);
+    if (targetIds.length === 0) return [];
+    const { messages: targetDocs } = await this.listMessagesById({ messageIds: targetIds });
+    const targetMap = new Map(
+      targetDocs.map((msg) => [msg.id, { threadId: msg.threadId, createdAt: msg.createdAt }])
+    );
+    const unionQueries = [];
+    const params = {};
+    let paramIdx = 1;
+    for (const inc of include) {
+      const { id, withPreviousMessages = 0, withNextMessages = 0 } = inc;
+      const target = targetMap.get(id);
+      if (!target) continue;
+      const threadParam = `var_thread_${paramIdx}`;
+      const createdAtParam = `var_createdAt_${paramIdx}`;
+      const limitParam = `var_limit_${paramIdx}`;
+      unionQueries.push(`
+        SELECT id, content, role, type, "createdAt", thread_id AS "threadId", "resourceId"
+        FROM "${TABLE_MESSAGES}"
+        WHERE thread_id = {${threadParam}:String}
+          AND createdAt <= parseDateTime64BestEffort({${createdAtParam}:String}, 3)
+        ORDER BY createdAt DESC, id DESC
+        LIMIT {${limitParam}:Int64}
+      `);
+      params[threadParam] = target.threadId;
+      params[createdAtParam] = target.createdAt;
+      params[limitParam] = withPreviousMessages + 1;
+      paramIdx++;
+      if (withNextMessages > 0) {
+        const threadParam2 = `var_thread_${paramIdx}`;
+        const createdAtParam2 = `var_createdAt_${paramIdx}`;
+        const limitParam2 = `var_limit_${paramIdx}`;
+        unionQueries.push(`
+          SELECT id, content, role, type, "createdAt", thread_id AS "threadId", "resourceId"
+          FROM "${TABLE_MESSAGES}"
+          WHERE thread_id = {${threadParam2}:String}
+            AND createdAt > parseDateTime64BestEffort({${createdAtParam2}:String}, 3)
+          ORDER BY createdAt ASC, id ASC
+          LIMIT {${limitParam2}:Int64}
+        `);
+        params[threadParam2] = target.threadId;
+        params[createdAtParam2] = target.createdAt;
+        params[limitParam2] = withNextMessages;
+        paramIdx++;
+      }
+    }
+    if (unionQueries.length === 0) return [];
+    let finalQuery;
+    if (unionQueries.length === 1) {
+      finalQuery = unionQueries[0];
+    } else {
+      finalQuery = `SELECT * FROM (${unionQueries.join(" UNION ALL ")}) ORDER BY "createdAt" ASC, id ASC`;
+    }
+    const includeResult = await this.client.query({
+      query: finalQuery,
+      query_params: params,
+      clickhouse_settings: {
+        date_time_input_format: "best_effort",
+        date_time_output_format: "iso",
+        use_client_time_zone: 1,
+        output_format_json_quote_64bit_integers: 0
+      }
+    });
+    const includeRows = await includeResult.json();
+    const seen = /* @__PURE__ */ new Set();
+    return transformRows(includeRows.data).filter((row) => {
+      if (seen.has(row.id)) return false;
+      seen.add(row.id);
+      return true;
+    });
+  }
   async saveMessages(args) {
     const { messages } = args;
     if (messages.length === 0) return { messages };
@@ -1910,6 +2005,11 @@ time for large tables. Please ensure you have a backup before proceeding.
       });
     }
     await this.#db.createTable({ tableName: TABLE_SPANS, schema: SPAN_SCHEMA });
+    await this.#db.alterTable({
+      tableName: TABLE_SPANS,
+      schema: SPAN_SCHEMA,
+      ifNotExists: ["requestContext"]
+    });
   }
   async dangerouslyClearAll() {
     await this.#db.clearTable({ tableName: TABLE_SPANS });
@@ -2182,7 +2282,8 @@ time for large tables. Please ensure you have a backup before proceeding.
   }
   async listTraces(args) {
     const { filters, pagination, orderBy } = listTracesArgsSchema.parse(args);
-    const { page, perPage } = pagination;
+    const page = pagination?.page ?? 0;
+    const perPage = pagination?.perPage ?? 10;
     try {
       const conditions = [`(parentSpanId IS NULL OR parentSpanId = '')`];
       const values = {};
@@ -2329,8 +2430,8 @@ time for large tables. Please ensure you have a backup before proceeding.
       const whereClause = conditions.length > 0 ? `WHERE ${conditions.join(" AND ")}` : "";
       const engine = TABLE_ENGINES[TABLE_SPANS] ?? "MergeTree()";
       const finalClause = engine.startsWith("ReplacingMergeTree") ? "FINAL" : "";
-      const sortField = orderBy.field;
-      const sortDirection = orderBy.direction;
+      const sortField = orderBy?.field ?? "startedAt";
+      const sortDirection = orderBy?.direction ?? "DESC";
       let orderClause;
       if (sortField === "endedAt") {
         const nullSortValue = sortDirection === "DESC" ? 0 : 1;