@mastra/clickhouse 0.2.7-alpha.1 → 0.2.7-alpha.3

package/dist/index.js

@@ -1,5 +1,5 @@
-import { MastraStorage, TABLE_SCHEMAS, TABLE_TRACES, TABLE_WORKFLOW_SNAPSHOT, TABLE_THREADS, TABLE_MESSAGES, TABLE_EVALS } from '@mastra/core/storage';
 import { createClient } from '@clickhouse/client';
+import { MastraStorage, TABLE_EVALS, TABLE_SCHEMAS, TABLE_TRACES, TABLE_WORKFLOW_SNAPSHOT, TABLE_THREADS, TABLE_MESSAGES } from '@mastra/core/storage';
 
 // src/storage/index.ts
 function safelyParseJSON(jsonString) {
@@ -41,6 +41,7 @@ function transformRow(row) {
 }
 var ClickhouseStore = class extends MastraStorage {
   db;
+  ttl = {};
   constructor(config) {
     super({ name: "ClickhouseStore" });
     this.db = createClient({
@@ -55,9 +56,54 @@ var ClickhouseStore = class extends MastraStorage {
         output_format_json_quote_64bit_integers: 0
       }
     });
+    this.ttl = config.ttl;
   }
-  getEvalsByAgentName(_agentName, _type) {
-    throw new Error("Method not implemented.");
+  transformEvalRow(row) {
+    row = transformRow(row);
+    const resultValue = JSON.parse(row.result);
+    const testInfoValue = row.test_info ? JSON.parse(row.test_info) : void 0;
+    if (!resultValue || typeof resultValue !== "object" || !("score" in resultValue)) {
+      throw new Error(`Invalid MetricResult format: ${JSON.stringify(resultValue)}`);
+    }
+    return {
+      input: row.input,
+      output: row.output,
+      result: resultValue,
+      agentName: row.agent_name,
+      metricName: row.metric_name,
+      instructions: row.instructions,
+      testInfo: testInfoValue,
+      globalRunId: row.global_run_id,
+      runId: row.run_id,
+      createdAt: row.created_at
+    };
+  }
+  async getEvalsByAgentName(agentName, type) {
+    try {
+      const baseQuery = `SELECT *, toDateTime64(createdAt, 3) as createdAt FROM ${TABLE_EVALS} WHERE agent_name = {var_agent_name:String}`;
+      const typeCondition = type === "test" ? " AND test_info IS NOT NULL AND JSONExtractString(test_info, 'testPath') IS NOT NULL" : type === "live" ? " AND (test_info IS NULL OR JSONExtractString(test_info, 'testPath') IS NULL)" : "";
+      const result = await this.db.query({
+        query: `${baseQuery}${typeCondition} ORDER BY createdAt DESC`,
+        query_params: { var_agent_name: agentName },
+        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
+        }
+      });
+      if (!result) {
+        return [];
+      }
+      const rows = await result.json();
+      return rows.data.map((row) => this.transformEvalRow(row));
+    } catch (error) {
+      if (error instanceof Error && error.message.includes("no such table")) {
+        return [];
+      }
+      this.logger.error("Failed to get evals for the specified agent: " + error?.message);
+      throw error;
+    }
   }
   async batchInsert({ tableName, records }) {
     try {
@@ -89,7 +135,8 @@ var ClickhouseStore = class extends MastraStorage {
     scope,
     page,
     perPage,
-    attributes
+    attributes,
+    filters
   }) {
     const limit = perPage;
     const offset = page * perPage;
@@ -105,8 +152,16 @@ var ClickhouseStore = class extends MastraStorage {
     }
     if (attributes) {
       Object.entries(attributes).forEach(([key, value]) => {
-        conditions.push(`JSONExtractString(attributes, '${key}') = {var_${key}:String}`);
-        args[`var_${key}`] = value;
+        conditions.push(`JSONExtractString(attributes, '${key}') = {var_attr_${key}:String}`);
+        args[`var_attr_${key}`] = value;
+      });
+    }
+    if (filters) {
+      Object.entries(filters).forEach(([key, value]) => {
+        conditions.push(
+          `${key} = {var_col_${key}:${COLUMN_TYPES[TABLE_SCHEMAS.mastra_traces?.[key]?.type ?? "text"]}}`
+        );
+        args[`var_col_${key}`] = value;
       });
     }
     const whereClause = conditions.length > 0 ? `WHERE ${conditions.join(" AND ")}` : "";
@@ -143,6 +198,16 @@ var ClickhouseStore = class extends MastraStorage {
       createdAt: row.createdAt
     }));
   }
+  async optimizeTable({ tableName }) {
+    await this.db.command({
+      query: `OPTIMIZE TABLE ${tableName} FINAL`
+    });
+  }
+  async materializeTtl({ tableName }) {
+    await this.db.command({
+      query: `ALTER TABLE ${tableName} MATERIALIZE TTL;`
+    });
+  }
   async createTable({
     tableName,
     schema
@@ -151,8 +216,10 @@ var ClickhouseStore = class extends MastraStorage {
       const columns = Object.entries(schema).map(([name, def]) => {
         const constraints = [];
         if (!def.nullable) constraints.push("NOT NULL");
-        return `"${name}" ${COLUMN_TYPES[def.type]} ${constraints.join(" ")}`;
+        const columnTtl = this.ttl?.[tableName]?.columns?.[name];
+        return `"${name}" ${COLUMN_TYPES[def.type]} ${constraints.join(" ")} ${columnTtl ? `TTL toDateTime(${columnTtl.ttlKey ?? "createdAt"}) + INTERVAL ${columnTtl.interval} ${columnTtl.unit}` : ""}`;
       }).join(",\n");
+      const rowTtl = this.ttl?.[tableName]?.row;
       const sql = tableName === TABLE_WORKFLOW_SNAPSHOT ? `
         CREATE TABLE IF NOT EXISTS ${tableName} (
           ${["id String"].concat(columns)}
@@ -161,16 +228,18 @@ var ClickhouseStore = class extends MastraStorage {
         PARTITION BY "createdAt"
         PRIMARY KEY (createdAt, run_id, workflow_name)
         ORDER BY (createdAt, run_id, workflow_name)
-        SETTINGS index_granularity = 8192;
+        ${rowTtl ? `TTL toDateTime(${rowTtl.ttlKey ?? "createdAt"}) + INTERVAL ${rowTtl.interval} ${rowTtl.unit}` : ""}
+        SETTINGS index_granularity = 8192
           ` : `
         CREATE TABLE IF NOT EXISTS ${tableName} (
           ${columns}
         )
         ENGINE = ${TABLE_ENGINES[tableName]}
         PARTITION BY "createdAt"
-        PRIMARY KEY (createdAt, id)
-        ORDER BY (createdAt, id)
-        SETTINGS index_granularity = 8192;
+        PRIMARY KEY (createdAt, ${tableName === TABLE_EVALS ? "run_id" : "id"})
+        ORDER BY (createdAt, ${tableName === TABLE_EVALS ? "run_id" : "id"})
+        ${this.ttl?.[tableName]?.row ? `TTL toDateTime(createdAt) + INTERVAL ${this.ttl[tableName].row.interval} ${this.ttl[tableName].row.unit}` : ""}
+        SETTINGS index_granularity = 8192
       `;
       await this.db.query({
         query: sql,
@@ -625,6 +694,83 @@ var ClickhouseStore = class extends MastraStorage {
       throw error;
     }
   }
+  async getWorkflowRuns({
+    workflowName,
+    fromDate,
+    toDate,
+    limit,
+    offset
+  } = {}) {
+    try {
+      const conditions = [];
+      const values = {};
+      if (workflowName) {
+        conditions.push(`workflow_name = {var_workflow_name:String}`);
+        values.var_workflow_name = workflowName;
+      }
+      if (fromDate) {
+        conditions.push(`createdAt >= {var_from_date:DateTime64(3)}`);
+        values.var_from_date = fromDate.getTime() / 1e3;
+      }
+      if (toDate) {
+        conditions.push(`createdAt <= {var_to_date:DateTime64(3)}`);
+        values.var_to_date = toDate.getTime() / 1e3;
+      }
+      const whereClause = conditions.length > 0 ? `WHERE ${conditions.join(" AND ")}` : "";
+      const limitClause = limit !== void 0 ? `LIMIT ${limit}` : "";
+      const offsetClause = offset !== void 0 ? `OFFSET ${offset}` : "";
+      let total = 0;
+      if (limit !== void 0 && offset !== void 0) {
+        const countResult = await this.db.query({
+          query: `SELECT COUNT(*) as count FROM ${TABLE_WORKFLOW_SNAPSHOT} ${TABLE_ENGINES[TABLE_WORKFLOW_SNAPSHOT].startsWith("ReplacingMergeTree") ? "FINAL" : ""} ${whereClause}`,
+          query_params: values,
+          format: "JSONEachRow"
+        });
+        const countRows = await countResult.json();
+        total = Number(countRows[0]?.count ?? 0);
+      }
+      const result = await this.db.query({
+        query: `
+          SELECT 
+            workflow_name,
+            run_id,
+            snapshot,
+            toDateTime64(createdAt, 3) as createdAt,
+            toDateTime64(updatedAt, 3) as updatedAt
+          FROM ${TABLE_WORKFLOW_SNAPSHOT} ${TABLE_ENGINES[TABLE_WORKFLOW_SNAPSHOT].startsWith("ReplacingMergeTree") ? "FINAL" : ""}
+          ${whereClause}
+          ORDER BY createdAt DESC
+          ${limitClause}
+          ${offsetClause}
+        `,
+        query_params: values,
+        format: "JSONEachRow"
+      });
+      const resultJson = await result.json();
+      const rows = resultJson;
+      const runs = rows.map((row) => {
+        let parsedSnapshot = row.snapshot;
+        if (typeof parsedSnapshot === "string") {
+          try {
+            parsedSnapshot = JSON.parse(row.snapshot);
+          } catch (e) {
+            console.warn(`Failed to parse snapshot for workflow ${row.workflow_name}: ${e}`);
+          }
+        }
+        return {
+          workflowName: row.workflow_name,
+          runId: row.run_id,
+          snapshot: parsedSnapshot,
+          createdAt: new Date(row.createdAt),
+          updatedAt: new Date(row.updatedAt)
+        };
+      });
+      return { runs, total: total || runs.length };
+    } catch (error) {
+      console.error("Error getting workflow runs:", error);
+      throw error;
+    }
+  }
   async close() {
     await this.db.close();
   }

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@mastra/clickhouse",
-  "version": "0.2.7-alpha.1",
+  "version": "0.2.7-alpha.3",
   "description": "Clickhouse provider for Mastra - includes db storage capabilities",
   "type": "module",
   "main": "dist/index.js",
@@ -20,7 +20,7 @@
   },
   "dependencies": {
     "@clickhouse/client": "^1.11.0",
-    "@mastra/core": "^0.8.0-alpha.2"
+    "@mastra/core": "^0.8.0-alpha.4"
   },
   "devDependencies": {
     "@microsoft/api-extractor": "^7.52.1",

package/src/storage/index.test.ts

@@ -1,4 +1,5 @@
 import { randomUUID } from 'crypto';
+import { TABLE_THREADS, TABLE_MESSAGES, TABLE_WORKFLOW_SNAPSHOT } from '@mastra/core/storage';
 import type { WorkflowRunState } from '@mastra/core/workflows';
 import { describe, it, expect, beforeAll, beforeEach, afterAll } from 'vitest';
 
@@ -9,6 +10,16 @@ const TEST_CONFIG: ClickhouseConfig = {
   url: process.env.CLICKHOUSE_URL || 'http://localhost:8123',
   username: process.env.CLICKHOUSE_USERNAME || 'default',
   password: process.env.CLICKHOUSE_PASSWORD || 'password',
+  ttl: {
+    mastra_traces: {
+      row: { interval: 10, unit: 'SECOND' },
+    },
+    mastra_evals: {
+      columns: {
+        result: { interval: 10, unit: 'SECOND' },
+      },
+    },
+  },
 };
 
 // Sample test data factory functions
@@ -31,6 +42,46 @@ const createSampleMessage = (threadId: string, createdAt: Date = new Date()) =>
     createdAt,
   }) as any;
 
+const createSampleTrace = () => ({
+  id: `trace-${randomUUID()}`,
+  name: 'Test Trace',
+  createdAt: new Date(),
+  updatedAt: new Date(),
+  metadata: { key: 'value' },
+});
+
+const createSampleEval = () => ({
+  agent_name: 'test-agent',
+  run_id: 'test-run-1',
+  result: '{ "score": 1 }',
+  createdAt: new Date(),
+});
+
+const createSampleWorkflowSnapshot = (status: string, createdAt?: Date) => {
+  const runId = `run-${randomUUID()}`;
+  const stepId = `step-${randomUUID()}`;
+  const timestamp = createdAt || new Date();
+  const snapshot = {
+    result: { success: true },
+    value: {},
+    context: {
+      steps: {
+        [stepId]: {
+          status,
+          payload: {},
+          error: undefined,
+        },
+      },
+      triggerData: {},
+      attempts: {},
+    },
+    activePaths: [],
+    runId,
+    timestamp: timestamp.getTime(),
+  } as WorkflowRunState;
+  return { snapshot, runId, stepId };
+};
+
 describe('ClickhouseStore', () => {
   let store: ClickhouseStore;
 
@@ -41,15 +92,14 @@ describe('ClickhouseStore', () => {
 
   beforeEach(async () => {
     // Clear tables before each test
-    await store.clearTable({ tableName: 'mastra_workflow_snapshot' });
-    await store.clearTable({ tableName: 'mastra_messages' });
-    await store.clearTable({ tableName: 'mastra_threads' });
+    await store.clearTable({ tableName: TABLE_THREADS });
+    await store.clearTable({ tableName: TABLE_MESSAGES });
+    await store.clearTable({ tableName: TABLE_WORKFLOW_SNAPSHOT });
   });
 
   describe('Thread Operations', () => {
     it('should create and retrieve a thread', async () => {
       const thread = createSampleThread();
-      console.log('Saving thread:', thread);
 
       // Save thread
       const savedThread = await store.__saveThread({ thread });
@@ -135,8 +185,6 @@ describe('ClickhouseStore', () => {
       // Retrieve messages
       const retrievedMessages = await store.__getMessages({ threadId: thread.id });
       expect(retrievedMessages).toHaveLength(2);
-      console.log('Messages:', messages);
-      console.log('Retrieved messages:', retrievedMessages);
       expect(retrievedMessages).toEqual(expect.arrayContaining(messages));
     }, 10e3);
 
@@ -192,6 +240,59 @@ describe('ClickhouseStore', () => {
     // });
   });
 
+  describe('Traces and TTL', () => {
+    it('should create and retrieve a trace, but not when row level ttl expires', async () => {
+      const trace = createSampleTrace();
+      await store.__batchInsert({
+        tableName: 'mastra_traces',
+        records: [trace],
+      });
+      let traces = await store.__getTraces({
+        page: 0,
+        perPage: 10,
+      });
+
+      expect(traces).toHaveLength(1);
+      expect(traces[0]!.id).toBe(trace.id);
+
+      await new Promise(resolve => setTimeout(resolve, 10e3));
+      await store.optimizeTable({ tableName: 'mastra_traces' });
+
+      traces = await store.__getTraces({
+        page: 0,
+        perPage: 10,
+      });
+
+      expect(traces).toHaveLength(0);
+    }, 60e3);
+
+    // NOTE: unable to clear column level TTLs for the test case nicely, but it does seem to get applied correctly
+    it.skip('should create and retrieve a trace, but not expired columns when column level ttl expires', async () => {
+      await store.clearTable({ tableName: 'mastra_evals' });
+      const ev = createSampleEval();
+      await store.__batchInsert({
+        tableName: 'mastra_evals',
+        records: [ev],
+      });
+      let evals = await store.__getEvalsByAgentName('test-agent');
+      console.log(evals);
+
+      expect(evals).toHaveLength(1);
+      expect(evals[0]!.agentName).toBe('test-agent');
+      expect(evals[0]!.runId).toBe('test-run-1');
+
+      await new Promise(resolve => setTimeout(resolve, 12e3));
+      await store.materializeTtl({ tableName: 'mastra_evals' });
+      await store.optimizeTable({ tableName: 'mastra_evals' });
+
+      evals = await store.__getEvalsByAgentName('test-agent');
+
+      expect(evals).toHaveLength(1);
+      expect(evals[0]!.agentName).toBe('test-agent');
+      expect(evals[0]!.runId).toBeNull();
+    }, 60e3);
+  });
+
   describe('Edge Cases and Error Handling', () => {
     it('should handle large metadata objects', async () => {
       const thread = createSampleThread();
@@ -385,6 +486,186 @@ describe('ClickhouseStore', () => {
     }, 10e3);
   });
 
+  describe('getWorkflowRuns', () => {
+    beforeEach(async () => {
+      await store.clearTable({ tableName: TABLE_WORKFLOW_SNAPSHOT });
+    });
+    it('returns empty array when no workflows exist', async () => {
+      const { runs, total } = await store.__getWorkflowRuns();
+      expect(runs).toEqual([]);
+      expect(total).toBe(0);
+    });
+
+    it('returns all workflows by default', async () => {
+      const workflowName1 = 'default_test_1';
+      const workflowName2 = 'default_test_2';
+
+      const { snapshot: workflow1, runId: runId1, stepId: stepId1 } = createSampleWorkflowSnapshot('completed');
+      const { snapshot: workflow2, runId: runId2, stepId: stepId2 } = createSampleWorkflowSnapshot('running');
+
+      await store.persistWorkflowSnapshot({
+        workflowName: workflowName1,
+        runId: runId1,
+        snapshot: workflow1,
+      });
+      await new Promise(resolve => setTimeout(resolve, 10)); // Small delay to ensure different timestamps
+      await store.persistWorkflowSnapshot({
+        workflowName: workflowName2,
+        runId: runId2,
+        snapshot: workflow2,
+      });
+
+      const { runs, total } = await store.__getWorkflowRuns();
+      expect(runs).toHaveLength(2);
+      expect(total).toBe(2);
+      expect(runs[0]!.workflowName).toBe(workflowName2); // Most recent first
+      expect(runs[1]!.workflowName).toBe(workflowName1);
+      const firstSnapshot = runs[0]!.snapshot as WorkflowRunState;
+      const secondSnapshot = runs[1]!.snapshot as WorkflowRunState;
+      expect(firstSnapshot.context?.steps[stepId2]?.status).toBe('running');
+      expect(secondSnapshot.context?.steps[stepId1]?.status).toBe('completed');
+    });
+
+    it('filters by workflow name', async () => {
+      const workflowName1 = 'filter_test_1';
+      const workflowName2 = 'filter_test_2';
+
+      const { snapshot: workflow1, runId: runId1, stepId: stepId1 } = createSampleWorkflowSnapshot('completed');
+      const { snapshot: workflow2, runId: runId2 } = createSampleWorkflowSnapshot('failed');
+
+      await store.persistWorkflowSnapshot({
+        workflowName: workflowName1,
+        runId: runId1,
+        snapshot: workflow1,
+      });
+      await new Promise(resolve => setTimeout(resolve, 10)); // Small delay to ensure different timestamps
+      await store.persistWorkflowSnapshot({
+        workflowName: workflowName2,
+        runId: runId2,
+        snapshot: workflow2,
+      });
+
+      const { runs, total } = await store.__getWorkflowRuns({
+        workflowName: workflowName1,
+      });
+      expect(runs).toHaveLength(1);
+      expect(total).toBe(1);
+      expect(runs[0]!.workflowName).toBe(workflowName1);
+      const snapshot = runs[0]!.snapshot as WorkflowRunState;
+      expect(snapshot.context?.steps[stepId1]?.status).toBe('completed');
+    });
+
+    it('filters by date range', async () => {
+      const now = new Date();
+      const yesterday = new Date(now.getTime() - 24 * 60 * 60 * 1000);
+      const twoDaysAgo = new Date(now.getTime() - 2 * 24 * 60 * 60 * 1000);
+      const workflowName1 = 'date_test_1';
+      const workflowName2 = 'date_test_2';
+      const workflowName3 = 'date_test_3';
+
+      const { snapshot: workflow1, runId: runId1 } = createSampleWorkflowSnapshot('completed');
+      const { snapshot: workflow2, runId: runId2, stepId: stepId2 } = createSampleWorkflowSnapshot('running');
+      const { snapshot: workflow3, runId: runId3, stepId: stepId3 } = createSampleWorkflowSnapshot('waiting');
+
+      await store.insert({
+        tableName: TABLE_WORKFLOW_SNAPSHOT,
+        record: {
+          workflow_name: workflowName1,
+          run_id: runId1,
+          snapshot: workflow1,
+          createdAt: twoDaysAgo,
+          updatedAt: twoDaysAgo,
+        },
+      });
+      await store.insert({
+        tableName: TABLE_WORKFLOW_SNAPSHOT,
+        record: {
+          workflow_name: workflowName2,
+          run_id: runId2,
+          snapshot: workflow2,
+          createdAt: yesterday,
+          updatedAt: yesterday,
+        },
+      });
+      await store.insert({
+        tableName: TABLE_WORKFLOW_SNAPSHOT,
+        record: {
+          workflow_name: workflowName3,
+          run_id: runId3,
+          snapshot: workflow3,
+          createdAt: now,
+          updatedAt: now,
+        },
+      });
+
+      const { runs } = await store.__getWorkflowRuns({
+        fromDate: yesterday,
+        toDate: now,
+      });
+
+      expect(runs).toHaveLength(2);
+      expect(runs[0]!.workflowName).toBe(workflowName3);
+      expect(runs[1]!.workflowName).toBe(workflowName2);
+      const firstSnapshot = runs[0]!.snapshot as WorkflowRunState;
+      const secondSnapshot = runs[1]!.snapshot as WorkflowRunState;
+      expect(firstSnapshot.context?.steps[stepId3]?.status).toBe('waiting');
+      expect(secondSnapshot.context?.steps[stepId2]?.status).toBe('running');
+    });
+
+    it('handles pagination', async () => {
+      const workflowName1 = 'page_test_1';
+      const workflowName2 = 'page_test_2';
+      const workflowName3 = 'page_test_3';
+
+      const { snapshot: workflow1, runId: runId1, stepId: stepId1 } = createSampleWorkflowSnapshot('completed');
+      const { snapshot: workflow2, runId: runId2, stepId: stepId2 } = createSampleWorkflowSnapshot('running');
+      const { snapshot: workflow3, runId: runId3, stepId: stepId3 } = createSampleWorkflowSnapshot('waiting');
+
+      await store.persistWorkflowSnapshot({
+        workflowName: workflowName1,
+        runId: runId1,
+        snapshot: workflow1,
+      });
+      await new Promise(resolve => setTimeout(resolve, 10)); // Small delay to ensure different timestamps
+      await store.persistWorkflowSnapshot({
+        workflowName: workflowName2,
+        runId: runId2,
+        snapshot: workflow2,
+      });
+      await new Promise(resolve => setTimeout(resolve, 10)); // Small delay to ensure different timestamps
+      await store.persistWorkflowSnapshot({
+        workflowName: workflowName3,
+        runId: runId3,
+        snapshot: workflow3,
+      });
+
+      // Get first page
+      const page1 = await store.__getWorkflowRuns({
+        limit: 2,
+        offset: 0,
+      });
+      expect(page1.runs).toHaveLength(2);
+      expect(page1.total).toBe(3); // Total count of all records
+      expect(page1.runs[0]!.workflowName).toBe(workflowName3);
+      expect(page1.runs[1]!.workflowName).toBe(workflowName2);
+      const firstSnapshot = page1.runs[0]!.snapshot as WorkflowRunState;
+      const secondSnapshot = page1.runs[1]!.snapshot as WorkflowRunState;
+      expect(firstSnapshot.context?.steps[stepId3]?.status).toBe('waiting');
+      expect(secondSnapshot.context?.steps[stepId2]?.status).toBe('running');
+
+      // Get second page
+      const page2 = await store.__getWorkflowRuns({
+        limit: 2,
+        offset: 2,
+      });
+      expect(page2.runs).toHaveLength(1);
+      expect(page2.total).toBe(3);
+      expect(page2.runs[0]!.workflowName).toBe(workflowName1);
+      const snapshot = page2.runs[0]!.snapshot as WorkflowRunState;
+      expect(snapshot.context?.steps[stepId1]?.status).toBe('completed');
+    }, 10e3);
+  });
+
   afterAll(async () => {
     await store.close();
   });