@mastra/clickhouse 0.2.7-alpha.1

package/dist/index.cjs

@@ -0,0 +1,635 @@
+'use strict';
+
+var storage = require('@mastra/core/storage');
+var client = require('@clickhouse/client');
+
+// src/storage/index.ts
+function safelyParseJSON(jsonString) {
+  try {
+    return JSON.parse(jsonString);
+  } catch {
+    return {};
+  }
+}
+var TABLE_ENGINES = {
+  [storage.TABLE_MESSAGES]: `MergeTree()`,
+  [storage.TABLE_WORKFLOW_SNAPSHOT]: `ReplacingMergeTree()`,
+  [storage.TABLE_TRACES]: `MergeTree()`,
+  [storage.TABLE_THREADS]: `ReplacingMergeTree()`,
+  [storage.TABLE_EVALS]: `MergeTree()`
+};
+var COLUMN_TYPES = {
+  text: "String",
+  timestamp: "DateTime64(3)",
+  uuid: "String",
+  jsonb: "String",
+  integer: "Int64",
+  bigint: "Int64"
+};
+function transformRows(rows) {
+  return rows.map((row) => transformRow(row));
+}
+function transformRow(row) {
+  if (!row) {
+    return row;
+  }
+  if (row.createdAt) {
+    row.createdAt = new Date(row.createdAt);
+  }
+  if (row.updatedAt) {
+    row.updatedAt = new Date(row.updatedAt);
+  }
+  return row;
+}
+var ClickhouseStore = class extends storage.MastraStorage {
+  db;
+  constructor(config) {
+    super({ name: "ClickhouseStore" });
+    this.db = client.createClient({
+      url: config.url,
+      username: config.username,
+      password: config.password,
+      clickhouse_settings: {
+        date_time_input_format: "best_effort",
+        date_time_output_format: "iso",
+        // This is crucial
+        use_client_time_zone: 1,
+        output_format_json_quote_64bit_integers: 0
+      }
+    });
+  }
+  getEvalsByAgentName(_agentName, _type) {
+    throw new Error("Method not implemented.");
+  }
+  async batchInsert({ tableName, records }) {
+    try {
+      await this.db.insert({
+        table: tableName,
+        values: records.map((record) => ({
+          ...Object.fromEntries(
+            Object.entries(record).map(([key, value]) => [
+              key,
+              storage.TABLE_SCHEMAS[tableName]?.[key]?.type === "timestamp" ? new Date(value).toISOString() : value
+            ])
+          )
+        })),
+        format: "JSONEachRow",
+        clickhouse_settings: {
+          // Allows to insert serialized JS Dates (such as '2023-12-06T10:54:48.000Z')
+          date_time_input_format: "best_effort",
+          use_client_time_zone: 1,
+          output_format_json_quote_64bit_integers: 0
+        }
+      });
+    } catch (error) {
+      console.error(`Error inserting into ${tableName}:`, error);
+      throw error;
+    }
+  }
+  async getTraces({
+    name,
+    scope,
+    page,
+    perPage,
+    attributes
+  }) {
+    const limit = perPage;
+    const offset = page * perPage;
+    const args = {};
+    const conditions = [];
+    if (name) {
+      conditions.push(`name LIKE CONCAT({var_name:String}, '%')`);
+      args.var_name = name;
+    }
+    if (scope) {
+      conditions.push(`scope = {var_scope:String}`);
+      args.var_scope = scope;
+    }
+    if (attributes) {
+      Object.entries(attributes).forEach(([key, value]) => {
+        conditions.push(`JSONExtractString(attributes, '${key}') = {var_${key}:String}`);
+        args[`var_${key}`] = value;
+      });
+    }
+    const whereClause = conditions.length > 0 ? `WHERE ${conditions.join(" AND ")}` : "";
+    const result = await this.db.query({
+      query: `SELECT *, toDateTime64(createdAt, 3) as createdAt FROM ${storage.TABLE_TRACES} ${whereClause} ORDER BY "createdAt" DESC LIMIT ${limit} OFFSET ${offset}`,
+      query_params: args,
+      clickhouse_settings: {
+        // Allows to insert serialized JS Dates (such as '2023-12-06T10:54:48.000Z')
+        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 resp = await result.json();
+    const rows = resp.data;
+    return rows.map((row) => ({
+      id: row.id,
+      parentSpanId: row.parentSpanId,
+      traceId: row.traceId,
+      name: row.name,
+      scope: row.scope,
+      kind: row.kind,
+      status: safelyParseJSON(row.status),
+      events: safelyParseJSON(row.events),
+      links: safelyParseJSON(row.links),
+      attributes: safelyParseJSON(row.attributes),
+      startTime: row.startTime,
+      endTime: row.endTime,
+      other: safelyParseJSON(row.other),
+      createdAt: row.createdAt
+    }));
+  }
+  async createTable({
+    tableName,
+    schema
+  }) {
+    try {
+      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(" ")}`;
+      }).join(",\n");
+      const sql = tableName === storage.TABLE_WORKFLOW_SNAPSHOT ? `
+        CREATE TABLE IF NOT EXISTS ${tableName} (
+          ${["id String"].concat(columns)}
+        )
+        ENGINE = ${TABLE_ENGINES[tableName]}
+        PARTITION BY "createdAt"
+        PRIMARY KEY (createdAt, run_id, workflow_name)
+        ORDER BY (createdAt, run_id, workflow_name)
+        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;
+      `;
+      await this.db.query({
+        query: sql,
+        clickhouse_settings: {
+          // Allows to insert serialized JS Dates (such as '2023-12-06T10:54:48.000Z')
+          date_time_input_format: "best_effort",
+          date_time_output_format: "iso",
+          use_client_time_zone: 1,
+          output_format_json_quote_64bit_integers: 0
+        }
+      });
+    } catch (error) {
+      console.error(`Error creating table ${tableName}:`, error);
+      throw error;
+    }
+  }
+  async clearTable({ tableName }) {
+    try {
+      await this.db.query({
+        query: `TRUNCATE TABLE ${tableName}`,
+        clickhouse_settings: {
+          // Allows to insert serialized JS Dates (such as '2023-12-06T10:54:48.000Z')
+          date_time_input_format: "best_effort",
+          date_time_output_format: "iso",
+          use_client_time_zone: 1,
+          output_format_json_quote_64bit_integers: 0
+        }
+      });
+    } catch (error) {
+      console.error(`Error clearing table ${tableName}:`, error);
+      throw error;
+    }
+  }
+  async insert({ tableName, record }) {
+    try {
+      await this.db.insert({
+        table: tableName,
+        values: [
+          {
+            ...record,
+            createdAt: record.createdAt.toISOString(),
+            updatedAt: record.updatedAt.toISOString()
+          }
+        ],
+        format: "JSONEachRow",
+        clickhouse_settings: {
+          // Allows to insert serialized JS Dates (such as '2023-12-06T10:54:48.000Z')
+          output_format_json_quote_64bit_integers: 0,
+          date_time_input_format: "best_effort",
+          use_client_time_zone: 1
+        }
+      });
+    } catch (error) {
+      console.error(`Error inserting into ${tableName}:`, error);
+      throw error;
+    }
+  }
+  async load({ tableName, keys }) {
+    try {
+      const keyEntries = Object.entries(keys);
+      const conditions = keyEntries.map(
+        ([key], index) => `"${key}" = {var_${key}:${COLUMN_TYPES[storage.TABLE_SCHEMAS[tableName]?.[key]?.type ?? "text"]}}`
+      ).join(" AND ");
+      const values = keyEntries.reduce((acc, [key, value]) => {
+        return { ...acc, [`var_${key}`]: value };
+      }, {});
+      const result = await this.db.query({
+        query: `SELECT *, toDateTime64(createdAt, 3) as createdAt, toDateTime64(updatedAt, 3) as updatedAt FROM ${tableName} ${TABLE_ENGINES[tableName].startsWith("ReplacingMergeTree") ? "FINAL" : ""} WHERE ${conditions}`,
+        query_params: values,
+        clickhouse_settings: {
+          // Allows to insert serialized JS Dates (such as '2023-12-06T10:54:48.000Z')
+          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 null;
+      }
+      const rows = await result.json();
+      if (tableName === storage.TABLE_WORKFLOW_SNAPSHOT) {
+        const snapshot = rows.data[0];
+        if (!snapshot) {
+          return null;
+        }
+        if (typeof snapshot.snapshot === "string") {
+          snapshot.snapshot = JSON.parse(snapshot.snapshot);
+        }
+        return transformRow(snapshot);
+      }
+      const data = transformRow(rows.data[0]);
+      return data;
+    } catch (error) {
+      console.error(`Error loading from ${tableName}:`, error);
+      throw error;
+    }
+  }
+  async getThreadById({ threadId }) {
+    try {
+      const result = await this.db.query({
+        query: `SELECT 
+          id,
+          "resourceId",
+          title,
+          metadata,
+          toDateTime64(createdAt, 3) as createdAt,
+          toDateTime64(updatedAt, 3) as updatedAt
+        FROM "${storage.TABLE_THREADS}"
+        FINAL
+        WHERE id = {var_id:String}`,
+        query_params: { var_id: threadId },
+        clickhouse_settings: {
+          // Allows to insert serialized JS Dates (such as '2023-12-06T10:54:48.000Z')
+          date_time_input_format: "best_effort",
+          date_time_output_format: "iso",
+          use_client_time_zone: 1,
+          output_format_json_quote_64bit_integers: 0
+        }
+      });
+      const rows = await result.json();
+      const thread = transformRow(rows.data[0]);
+      if (!thread) {
+        return null;
+      }
+      return {
+        ...thread,
+        metadata: typeof thread.metadata === "string" ? JSON.parse(thread.metadata) : thread.metadata,
+        createdAt: thread.createdAt,
+        updatedAt: thread.updatedAt
+      };
+    } catch (error) {
+      console.error(`Error getting thread ${threadId}:`, error);
+      throw error;
+    }
+  }
+  async getThreadsByResourceId({ resourceId }) {
+    try {
+      const result = await this.db.query({
+        query: `SELECT 
+          id,
+          "resourceId",
+          title,
+          metadata,
+          toDateTime64(createdAt, 3) as createdAt,
+          toDateTime64(updatedAt, 3) as updatedAt
+        FROM "${storage.TABLE_THREADS}"
+        WHERE "resourceId" = {var_resourceId:String}`,
+        query_params: { var_resourceId: resourceId },
+        clickhouse_settings: {
+          // Allows to insert serialized JS Dates (such as '2023-12-06T10:54:48.000Z')
+          date_time_input_format: "best_effort",
+          date_time_output_format: "iso",
+          use_client_time_zone: 1,
+          output_format_json_quote_64bit_integers: 0
+        }
+      });
+      const rows = await result.json();
+      const threads = transformRows(rows.data);
+      return threads.map((thread) => ({
+        ...thread,
+        metadata: typeof thread.metadata === "string" ? JSON.parse(thread.metadata) : thread.metadata,
+        createdAt: thread.createdAt,
+        updatedAt: thread.updatedAt
+      }));
+    } catch (error) {
+      console.error(`Error getting threads for resource ${resourceId}:`, error);
+      throw error;
+    }
+  }
+  async saveThread({ thread }) {
+    try {
+      await this.db.insert({
+        table: storage.TABLE_THREADS,
+        values: [
+          {
+            ...thread,
+            createdAt: thread.createdAt.toISOString(),
+            updatedAt: thread.updatedAt.toISOString()
+          }
+        ],
+        format: "JSONEachRow",
+        clickhouse_settings: {
+          // Allows to insert serialized JS Dates (such as '2023-12-06T10:54:48.000Z')
+          date_time_input_format: "best_effort",
+          use_client_time_zone: 1,
+          output_format_json_quote_64bit_integers: 0
+        }
+      });
+      return thread;
+    } catch (error) {
+      console.error("Error saving thread:", error);
+      throw error;
+    }
+  }
+  async updateThread({
+    id,
+    title,
+    metadata
+  }) {
+    try {
+      const existingThread = await this.getThreadById({ threadId: id });
+      if (!existingThread) {
+        throw new Error(`Thread ${id} not found`);
+      }
+      const mergedMetadata = {
+        ...existingThread.metadata,
+        ...metadata
+      };
+      const updatedThread = {
+        ...existingThread,
+        title,
+        metadata: mergedMetadata,
+        updatedAt: /* @__PURE__ */ new Date()
+      };
+      await this.db.insert({
+        table: storage.TABLE_THREADS,
+        values: [
+          {
+            ...updatedThread,
+            updatedAt: updatedThread.updatedAt.toISOString()
+          }
+        ],
+        format: "JSONEachRow",
+        clickhouse_settings: {
+          // Allows to insert serialized JS Dates (such as '2023-12-06T10:54:48.000Z')
+          date_time_input_format: "best_effort",
+          use_client_time_zone: 1,
+          output_format_json_quote_64bit_integers: 0
+        }
+      });
+      return updatedThread;
+    } catch (error) {
+      console.error("Error updating thread:", error);
+      throw error;
+    }
+  }
+  async deleteThread({ threadId }) {
+    try {
+      await this.db.command({
+        query: `DELETE FROM "${storage.TABLE_MESSAGES}" WHERE thread_id = '${threadId}';`,
+        query_params: { var_thread_id: threadId },
+        clickhouse_settings: {
+          output_format_json_quote_64bit_integers: 0
+        }
+      });
+      await this.db.command({
+        query: `DELETE FROM "${storage.TABLE_THREADS}" WHERE id = {var_id:String};`,
+        query_params: { var_id: threadId },
+        clickhouse_settings: {
+          output_format_json_quote_64bit_integers: 0
+        }
+      });
+    } catch (error) {
+      console.error("Error deleting thread:", error);
+      throw error;
+    }
+  }
+  async getMessages({ threadId, selectBy }) {
+    try {
+      const messages = [];
+      const limit = typeof selectBy?.last === `number` ? selectBy.last : 40;
+      const include = selectBy?.include || [];
+      if (include.length) {
+        const includeResult = await this.db.query({
+          query: `
+          WITH ordered_messages AS (
+            SELECT 
+              *,
+              toDateTime64(createdAt, 3) as createdAt,
+              toDateTime64(updatedAt, 3) as updatedAt,
+              ROW_NUMBER() OVER (ORDER BY "createdAt" DESC) as row_num
+            FROM "${storage.TABLE_MESSAGES}"
+            WHERE thread_id = {var_thread_id:String}
+          )
+          SELECT
+            m.id AS id, 
+            m.content as content, 
+            m.role as role, 
+            m.type as type,
+            m.createdAt as createdAt, 
+            m.updatedAt as updatedAt,
+            m.thread_id AS "threadId"
+          FROM ordered_messages m
+          WHERE m.id = ANY({var_include:Array(String)})
+          OR EXISTS (
+            SELECT 1 FROM ordered_messages target
+            WHERE target.id = ANY({var_include:Array(String)})
+            AND (
+              -- Get previous messages based on the max withPreviousMessages
+              (m.row_num <= target.row_num + {var_withPreviousMessages:Int64} AND m.row_num > target.row_num)
+              OR
+              -- Get next messages based on the max withNextMessages
+              (m.row_num >= target.row_num - {var_withNextMessages:Int64} AND m.row_num < target.row_num)
+            )
+          )
+          ORDER BY m."createdAt" DESC
+          `,
+          query_params: {
+            var_thread_id: threadId,
+            var_include: include.map((i) => i.id),
+            var_withPreviousMessages: Math.max(...include.map((i) => i.withPreviousMessages || 0)),
+            var_withNextMessages: Math.max(...include.map((i) => i.withNextMessages || 0))
+          },
+          clickhouse_settings: {
+            // Allows to insert serialized JS Dates (such as '2023-12-06T10:54:48.000Z')
+            date_time_input_format: "best_effort",
+            date_time_output_format: "iso",
+            use_client_time_zone: 1,
+            output_format_json_quote_64bit_integers: 0
+          }
+        });
+        const rows2 = await includeResult.json();
+        messages.push(...transformRows(rows2.data));
+      }
+      const result = await this.db.query({
+        query: `
+        SELECT 
+            id, 
+            content, 
+            role, 
+            type,
+            toDateTime64(createdAt, 3) as createdAt,
+            thread_id AS "threadId"
+        FROM "${storage.TABLE_MESSAGES}"
+        WHERE thread_id = {threadId:String}
+        AND id NOT IN ({exclude:Array(String)})
+        ORDER BY "createdAt" DESC
+        LIMIT {limit:Int64}
+        `,
+        query_params: {
+          threadId,
+          exclude: messages.map((m) => m.id),
+          limit
+        },
+        clickhouse_settings: {
+          // Allows to insert serialized JS Dates (such as '2023-12-06T10:54:48.000Z')
+          date_time_input_format: "best_effort",
+          date_time_output_format: "iso",
+          use_client_time_zone: 1,
+          output_format_json_quote_64bit_integers: 0
+        }
+      });
+      const rows = await result.json();
+      messages.push(...transformRows(rows.data));
+      messages.sort((a, b) => new Date(a.createdAt).getTime() - new Date(b.createdAt).getTime());
+      messages.forEach((message) => {
+        if (typeof message.content === "string") {
+          try {
+            message.content = JSON.parse(message.content);
+          } catch {
+          }
+        }
+      });
+      return messages;
+    } catch (error) {
+      console.error("Error getting messages:", error);
+      throw error;
+    }
+  }
+  async saveMessages({ messages }) {
+    if (messages.length === 0) return messages;
+    try {
+      const threadId = messages[0]?.threadId;
+      if (!threadId) {
+        throw new Error("Thread ID is required");
+      }
+      const thread = await this.getThreadById({ threadId });
+      if (!thread) {
+        throw new Error(`Thread ${threadId} not found`);
+      }
+      await this.db.insert({
+        table: storage.TABLE_MESSAGES,
+        format: "JSONEachRow",
+        values: messages.map((message) => ({
+          id: message.id,
+          thread_id: threadId,
+          content: typeof message.content === "string" ? message.content : JSON.stringify(message.content),
+          createdAt: message.createdAt.toISOString(),
+          role: message.role,
+          type: message.type
+        })),
+        clickhouse_settings: {
+          // Allows to insert serialized JS Dates (such as '2023-12-06T10:54:48.000Z')
+          date_time_input_format: "best_effort",
+          use_client_time_zone: 1,
+          output_format_json_quote_64bit_integers: 0
+        }
+      });
+      return messages;
+    } catch (error) {
+      console.error("Error saving messages:", error);
+      throw error;
+    }
+  }
+  async persistWorkflowSnapshot({
+    workflowName,
+    runId,
+    snapshot
+  }) {
+    try {
+      const currentSnapshot = await this.load({
+        tableName: storage.TABLE_WORKFLOW_SNAPSHOT,
+        keys: { workflow_name: workflowName, run_id: runId }
+      });
+      const now = /* @__PURE__ */ new Date();
+      const persisting = currentSnapshot ? {
+        ...currentSnapshot,
+        snapshot: JSON.stringify(snapshot),
+        updatedAt: now.toISOString()
+      } : {
+        workflow_name: workflowName,
+        run_id: runId,
+        snapshot: JSON.stringify(snapshot),
+        createdAt: now.toISOString(),
+        updatedAt: now.toISOString()
+      };
+      await this.db.insert({
+        table: storage.TABLE_WORKFLOW_SNAPSHOT,
+        format: "JSONEachRow",
+        values: [persisting],
+        clickhouse_settings: {
+          // Allows to insert serialized JS Dates (such as '2023-12-06T10:54:48.000Z')
+          date_time_input_format: "best_effort",
+          use_client_time_zone: 1,
+          output_format_json_quote_64bit_integers: 0
+        }
+      });
+    } catch (error) {
+      console.error("Error persisting workflow snapshot:", error);
+      throw error;
+    }
+  }
+  async loadWorkflowSnapshot({
+    workflowName,
+    runId
+  }) {
+    try {
+      const result = await this.load({
+        tableName: storage.TABLE_WORKFLOW_SNAPSHOT,
+        keys: {
+          workflow_name: workflowName,
+          run_id: runId
+        }
+      });
+      if (!result) {
+        return null;
+      }
+      return result.snapshot;
+    } catch (error) {
+      console.error("Error loading workflow snapshot:", error);
+      throw error;
+    }
+  }
+  async close() {
+    await this.db.close();
+  }
+};
+
+exports.ClickhouseStore = ClickhouseStore;

package/dist/index.d.cts

@@ -0,0 +1,2 @@
+export { ClickhouseConfig } from './_tsup-dts-rollup.cjs';
+export { ClickhouseStore } from './_tsup-dts-rollup.cjs';

package/dist/index.d.ts

@@ -0,0 +1,2 @@
+export { ClickhouseConfig } from './_tsup-dts-rollup.js';
+export { ClickhouseStore } from './_tsup-dts-rollup.js';