@mastra/clickhouse 1.0.0-beta.0 → 1.0.0-beta.2

package/dist/index.js

@@ -1,6 +1,6 @@
 import { createClient } from '@clickhouse/client';
 import { MastraError, ErrorCategory, ErrorDomain } from '@mastra/core/error';
-import { TABLE_SPANS, TABLE_RESOURCES, TABLE_SCORERS, TABLE_THREADS, TABLE_TRACES, TABLE_WORKFLOW_SNAPSHOT, TABLE_MESSAGES, MastraStorage, StoreOperations, TABLE_SCHEMAS, WorkflowsStorage, normalizePerPage, ScoresStorage, safelyParseJSON, calculatePagination, MemoryStorage } from '@mastra/core/storage';
+import { TABLE_SPANS, TABLE_RESOURCES, TABLE_SCORERS, TABLE_THREADS, TABLE_TRACES, TABLE_WORKFLOW_SNAPSHOT, TABLE_MESSAGES, MastraStorage, StoreOperations, TABLE_SCHEMAS, WorkflowsStorage, normalizePerPage, ScoresStorage, transformScoreRow, SCORERS_SCHEMA, calculatePagination, MemoryStorage, safelyParseJSON } from '@mastra/core/storage';
 import { MessageList } from '@mastra/core/agent';
 import { saveScorePayloadSchema } from '@mastra/core/evals';
 
@@ -125,6 +125,8 @@ var MemoryStorageClickhouse = class extends MemoryStorage {
   }
   async listMessages(args) {
     const { threadId, resourceId, include, filter, perPage: perPageInput, page = 0, orderBy } = args;
+    const rawThreadIds = Array.isArray(threadId) ? threadId : [threadId];
+    const threadIds = rawThreadIds.filter((id) => id !== void 0 && id !== null).map((id) => (typeof id === "string" ? id : String(id)).trim()).filter((id) => id.length > 0);
     if (page < 0) {
       throw new MastraError(
         {
@@ -136,20 +138,21 @@ var MemoryStorageClickhouse = class extends MemoryStorage {
         new Error("page must be >= 0")
       );
     }
-    if (!threadId.trim()) {
+    if (threadIds.length === 0) {
       throw new MastraError(
         {
           id: "STORAGE_CLICKHOUSE_LIST_MESSAGES_INVALID_THREAD_ID",
           domain: ErrorDomain.STORAGE,
           category: ErrorCategory.THIRD_PARTY,
-          details: { threadId }
+          details: { threadId: Array.isArray(threadId) ? JSON.stringify(threadId) : String(threadId) }
         },
-        new Error("threadId must be a non-empty string")
+        new Error("threadId must be a non-empty string or array of non-empty strings")
       );
     }
     const perPageForQuery = normalizePerPage(perPageInput, 40);
     const { offset, perPage: perPageForResponse } = calculatePagination(page, perPageInput, perPageForQuery);
     try {
+      const threadCondition = threadIds.length === 1 ? `thread_id = {threadId0:String}` : `thread_id IN (${threadIds.map((_, i) => `{threadId${i}:String}`).join(", ")})`;
       let dataQuery = `
         SELECT 
           id,
@@ -160,9 +163,12 @@ var MemoryStorageClickhouse = class extends MemoryStorage {
           thread_id AS "threadId",
           resourceId
         FROM ${TABLE_MESSAGES}
-        WHERE thread_id = {threadId:String}
+        WHERE ${threadCondition}
       `;
-      const dataParams = { threadId };
+      const dataParams = {};
+      threadIds.forEach((tid, i) => {
+        dataParams[`threadId${i}`] = tid;
+      });
       if (resourceId) {
         dataQuery += ` AND resourceId = {resourceId:String}`;
         dataParams.resourceId = resourceId;
@@ -197,8 +203,11 @@ var MemoryStorageClickhouse = class extends MemoryStorage {
       const rows = await result.json();
       const paginatedMessages = transformRows(rows.data);
       const paginatedCount = paginatedMessages.length;
-      let countQuery = `SELECT count() as total FROM ${TABLE_MESSAGES} WHERE thread_id = {threadId:String}`;
-      const countParams = { threadId };
+      let countQuery = `SELECT count() as total FROM ${TABLE_MESSAGES} WHERE ${threadCondition}`;
+      const countParams = {};
+      threadIds.forEach((tid, i) => {
+        countParams[`threadId${i}`] = tid;
+      });
       if (resourceId) {
         countQuery += ` AND resourceId = {resourceId:String}`;
         countParams.resourceId = resourceId;
@@ -237,12 +246,25 @@ 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 searchId = inc.threadId || threadId;
+          const searchThreadId = inc.threadId ?? threadByMessageId.get(id);
+          if (!searchThreadId) continue;
           unionQueries.push(`
             SELECT * FROM (
               WITH numbered_messages AS (
@@ -264,31 +286,33 @@ var MemoryStorageClickhouse = class extends MemoryStorage {
             ) AS query_${paramIdx}
           `);
           params.push(
-            { [`var_thread_id_${paramIdx}`]: searchId },
+            { [`var_thread_id_${paramIdx}`]: searchThreadId },
             { [`var_include_id_${paramIdx}`]: id },
             { [`var_withPreviousMessages_${paramIdx}`]: withPreviousMessages },
             { [`var_withNextMessages_${paramIdx}`]: withNextMessages }
           );
           paramIdx++;
         }
-        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);
+        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);
+            }
           }
         }
       }
@@ -306,7 +330,10 @@ var MemoryStorageClickhouse = class extends MemoryStorage {
         }
         return direction === "ASC" ? String(aValue).localeCompare(String(bValue)) : String(bValue).localeCompare(String(aValue));
       });
-      const returnedThreadMessageIds = new Set(finalMessages.filter((m) => m.threadId === threadId).map((m) => m.id));
+      const threadIdSet = new Set(threadIds);
+      const returnedThreadMessageIds = new Set(
+        finalMessages.filter((m) => m.threadId && threadIdSet.has(m.threadId)).map((m) => m.id)
+      );
       const allThreadMessagesReturned = returnedThreadMessageIds.size >= total;
       const hasMore = perPageForResponse === false ? false : allThreadMessagesReturned ? false : offset + paginatedCount < total;
       return {
@@ -323,7 +350,7 @@ var MemoryStorageClickhouse = class extends MemoryStorage {
           domain: ErrorDomain.STORAGE,
           category: ErrorCategory.THIRD_PARTY,
           details: {
-            threadId,
+            threadId: Array.isArray(threadId) ? threadId.join(",") : threadId,
             resourceId: resourceId ?? ""
           }
         },
@@ -1449,30 +1476,15 @@ var ScoresStorageClickhouse = class extends ScoresStorage {
     this.client = client;
     this.operations = operations;
   }
+  /**
+   * ClickHouse-specific score row transformation.
+   * Converts timestamps to Date objects and filters out '_null_' values.
+   */
   transformScoreRow(row) {
-    const scorer = safelyParseJSON(row.scorer);
-    const preprocessStepResult = safelyParseJSON(row.preprocessStepResult);
-    const analyzeStepResult = safelyParseJSON(row.analyzeStepResult);
-    const metadata = safelyParseJSON(row.metadata);
-    const input = safelyParseJSON(row.input);
-    const output = safelyParseJSON(row.output);
-    const additionalContext = safelyParseJSON(row.additionalContext);
-    const requestContext = safelyParseJSON(row.requestContext);
-    const entity = safelyParseJSON(row.entity);
-    return {
-      ...row,
-      scorer,
-      preprocessStepResult,
-      analyzeStepResult,
-      metadata,
-      input,
-      output,
-      additionalContext,
-      requestContext,
-      entity,
-      createdAt: new Date(row.createdAt),
-      updatedAt: new Date(row.updatedAt)
-    };
+    return transformScoreRow(row, {
+      convertTimestamps: true,
+      nullValuePattern: "_null_"
+    });
   }
   async getScoreById({ id }) {
     try {
@@ -1521,9 +1533,15 @@ var ScoresStorageClickhouse = class extends ScoresStorage {
       );
     }
     try {
-      const record = {
-        ...parsedScore
-      };
+      const record = {};
+      for (const key of Object.keys(SCORERS_SCHEMA)) {
+        const value = parsedScore[key];
+        if (key === "createdAt" || key === "updatedAt") {
+          record[key] = (/* @__PURE__ */ new Date()).toISOString();
+          continue;
+        }
+        record[key] = value === void 0 || value === null ? "_null_" : value;
+      }
       await this.client.insert({
         table: TABLE_SCORERS,
         values: [record],
@@ -1980,7 +1998,8 @@ var WorkflowsStorageClickhouse = class extends WorkflowsStorage {
     toDate,
     page,
     perPage,
-    resourceId
+    resourceId,
+    status
   } = {}) {
     try {
       const conditions = [];
@@ -1989,6 +2008,10 @@ var WorkflowsStorageClickhouse = class extends WorkflowsStorage {
         conditions.push(`workflow_name = {var_workflow_name:String}`);
         values.var_workflow_name = workflowName;
       }
+      if (status) {
+        conditions.push(`JSONExtractString(snapshot, 'status') = {var_status:String}`);
+        values.var_status = status;
+      }
       if (resourceId) {
         const hasResourceId = await this.operations.hasColumn(TABLE_WORKFLOW_SNAPSHOT, "resourceId");
         if (hasResourceId) {
@@ -2242,15 +2265,8 @@ var ClickhouseStore = class extends MastraStorage {
   }) {
     return this.stores.workflows.loadWorkflowSnapshot({ workflowName, runId });
   }
-  async listWorkflowRuns({
-    workflowName,
-    fromDate,
-    toDate,
-    perPage,
-    page,
-    resourceId
-  } = {}) {
-    return this.stores.workflows.listWorkflowRuns({ workflowName, fromDate, toDate, perPage, page, resourceId });
+  async listWorkflowRuns(args = {}) {
+    return this.stores.workflows.listWorkflowRuns(args);
   }
   async getWorkflowRunById({
     runId,