@mastra/clickhouse 0.12.0 → 0.12.1

package/src/storage/domains/traces/index.ts

@@ -0,0 +1,275 @@
+import type { ClickHouseClient } from '@clickhouse/client';
+import { ErrorCategory, ErrorDomain, MastraError } from '@mastra/core/error';
+import { safelyParseJSON, TABLE_SCHEMAS, TABLE_TRACES, TracesStorage } from '@mastra/core/storage';
+import type { PaginationInfo, StorageGetTracesPaginatedArg, StorageGetTracesArg } from '@mastra/core/storage';
+import type { Trace } from '@mastra/core/telemetry';
+import type { StoreOperationsClickhouse } from '../operations';
+
+export class TracesStorageClickhouse extends TracesStorage {
+  protected client: ClickHouseClient;
+  protected operations: StoreOperationsClickhouse;
+
+  constructor({ client, operations }: { client: ClickHouseClient; operations: StoreOperationsClickhouse }) {
+    super();
+    this.client = client;
+    this.operations = operations;
+  }
+
+  async getTracesPaginated(args: StorageGetTracesPaginatedArg): Promise<PaginationInfo & { traces: Trace[] }> {
+    const { name, scope, page = 0, perPage = 100, attributes, filters, dateRange } = args;
+    const fromDate = dateRange?.start;
+    const toDate = dateRange?.end;
+    const currentOffset = page * perPage;
+
+    const queryArgs: Record<string, any> = {};
+    const conditions: string[] = [];
+
+    if (name) {
+      conditions.push(`name LIKE CONCAT({var_name:String}, '%')`);
+      queryArgs.var_name = name;
+    }
+    if (scope) {
+      conditions.push(`scope = {var_scope:String}`);
+      queryArgs.var_scope = scope;
+    }
+    if (attributes) {
+      Object.entries(attributes).forEach(([key, value]) => {
+        conditions.push(`JSONExtractString(attributes, '${key}') = {var_attr_${key}:String}`);
+        queryArgs[`var_attr_${key}`] = value;
+      });
+    }
+    if (filters) {
+      Object.entries(filters).forEach(([key, value]) => {
+        conditions.push(`${key} = {var_col_${key}:${TABLE_SCHEMAS.mastra_traces?.[key]?.type ?? 'text'}}`);
+        queryArgs[`var_col_${key}`] = value;
+      });
+    }
+    if (fromDate) {
+      conditions.push(`createdAt >= parseDateTime64BestEffort({var_from_date:String})`);
+      queryArgs.var_from_date = fromDate.toISOString();
+    }
+    if (toDate) {
+      conditions.push(`createdAt <= parseDateTime64BestEffort({var_to_date:String})`);
+      queryArgs.var_to_date = toDate.toISOString();
+    }
+
+    const whereClause = conditions.length > 0 ? `WHERE ${conditions.join(' AND ')}` : '';
+
+    try {
+      // Get total count
+      const countResult = await this.client.query({
+        query: `SELECT COUNT(*) as count FROM ${TABLE_TRACES} ${whereClause}`,
+        query_params: queryArgs,
+        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 countData = await countResult.json();
+      const total = Number((countData.data?.[0] as any)?.count ?? 0);
+
+      if (total === 0) {
+        return {
+          traces: [],
+          total: 0,
+          page,
+          perPage,
+          hasMore: false,
+        };
+      }
+
+      // Get traces with pagination
+      const result = await this.client.query({
+        query: `SELECT *, toDateTime64(createdAt, 3) as createdAt FROM ${TABLE_TRACES} ${whereClause} ORDER BY "createdAt" DESC LIMIT {var_limit:UInt32} OFFSET {var_offset:UInt32}`,
+        query_params: { ...queryArgs, var_limit: perPage, var_offset: currentOffset },
+        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 {
+          traces: [],
+          total,
+          page,
+          perPage,
+          hasMore: false,
+        };
+      }
+
+      const resp = await result.json();
+      const rows: any[] = resp.data;
+      const traces = 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 as string),
+        events: safelyParseJSON(row.events as string),
+        links: safelyParseJSON(row.links as string),
+        attributes: safelyParseJSON(row.attributes as string),
+        startTime: row.startTime,
+        endTime: row.endTime,
+        other: safelyParseJSON(row.other as string),
+        createdAt: row.createdAt,
+      }));
+
+      return {
+        traces,
+        total,
+        page,
+        perPage,
+        hasMore: currentOffset + traces.length < total,
+      };
+    } catch (error: any) {
+      if (error?.message?.includes('no such table') || error?.message?.includes('does not exist')) {
+        return {
+          traces: [],
+          total: 0,
+          page,
+          perPage,
+          hasMore: false,
+        };
+      }
+      throw new MastraError(
+        {
+          id: 'CLICKHOUSE_STORAGE_GET_TRACES_PAGINATED_FAILED',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: {
+            name: name ?? null,
+            scope: scope ?? null,
+            page,
+            perPage,
+            attributes: attributes ? JSON.stringify(attributes) : null,
+            filters: filters ? JSON.stringify(filters) : null,
+            dateRange: dateRange ? JSON.stringify(dateRange) : null,
+          },
+        },
+        error,
+      );
+    }
+  }
+
+  async getTraces({
+    name,
+    scope,
+    page,
+    perPage,
+    attributes,
+    filters,
+    fromDate,
+    toDate,
+  }: StorageGetTracesArg): Promise<any[]> {
+    const limit = perPage;
+    const offset = page * perPage;
+
+    const args: Record<string, any> = {};
+
+    const conditions: string[] = [];
+    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_attr_${key}:String}`);
+        args[`var_attr_${key}`] = value;
+      });
+    }
+
+    if (filters) {
+      Object.entries(filters).forEach(([key, value]) => {
+        conditions.push(`${key} = {var_col_${key}:${TABLE_SCHEMAS.mastra_traces?.[key]?.type ?? 'text'}}`);
+        args[`var_col_${key}`] = value;
+      });
+    }
+
+    if (fromDate) {
+      conditions.push(`createdAt >= {var_from_date:DateTime64(3)}`);
+      args.var_from_date = fromDate.getTime() / 1000; // Convert to Unix timestamp
+    }
+
+    if (toDate) {
+      conditions.push(`createdAt <= {var_to_date:DateTime64(3)}`);
+      args.var_to_date = toDate.getTime() / 1000; // Convert to Unix timestamp
+    }
+
+    const whereClause = conditions.length > 0 ? `WHERE ${conditions.join(' AND ')}` : '';
+
+    try {
+      const result = await this.client.query({
+        query: `SELECT *, toDateTime64(createdAt, 3) as createdAt FROM ${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: any[] = 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 as string),
+        events: safelyParseJSON(row.events as string),
+        links: safelyParseJSON(row.links as string),
+        attributes: safelyParseJSON(row.attributes as string),
+        startTime: row.startTime,
+        endTime: row.endTime,
+        other: safelyParseJSON(row.other as string),
+        createdAt: row.createdAt,
+      }));
+    } catch (error: any) {
+      if (error?.message?.includes('no such table') || error?.message?.includes('does not exist')) {
+        return [];
+      }
+      throw new MastraError(
+        {
+          id: 'CLICKHOUSE_STORAGE_GET_TRACES_FAILED',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: {
+            name: name ?? null,
+            scope: scope ?? null,
+            page,
+            perPage,
+            attributes: attributes ? JSON.stringify(attributes) : null,
+            filters: filters ? JSON.stringify(filters) : null,
+            fromDate: fromDate?.toISOString() ?? null,
+            toDate: toDate?.toISOString() ?? null,
+          },
+        },
+        error,
+      );
+    }
+  }
+
+  async batchTraceInsert(args: { records: Trace[] }): Promise<void> {
+    await this.operations.batchInsert({ tableName: TABLE_TRACES, records: args.records });
+  }
+}

package/src/storage/domains/utils.ts

@@ -0,0 +1,86 @@
+import type { TABLE_NAMES, TABLE_SCHEMAS, StorageColumn } from '@mastra/core/storage';
+import {
+  TABLE_MESSAGES,
+  TABLE_RESOURCES,
+  TABLE_EVALS,
+  TABLE_SCORERS,
+  TABLE_THREADS,
+  TABLE_TRACES,
+  TABLE_WORKFLOW_SNAPSHOT,
+  safelyParseJSON,
+} from '@mastra/core/storage';
+
+export const TABLE_ENGINES: Record<TABLE_NAMES, string> = {
+  [TABLE_MESSAGES]: `MergeTree()`,
+  [TABLE_WORKFLOW_SNAPSHOT]: `ReplacingMergeTree()`,
+  [TABLE_TRACES]: `MergeTree()`,
+  [TABLE_THREADS]: `ReplacingMergeTree()`,
+  [TABLE_EVALS]: `MergeTree()`,
+  [TABLE_SCORERS]: `MergeTree()`,
+  [TABLE_RESOURCES]: `ReplacingMergeTree()`,
+};
+
+export const COLUMN_TYPES: Record<StorageColumn['type'], string> = {
+  text: 'String',
+  timestamp: 'DateTime64(3)',
+  uuid: 'String',
+  jsonb: 'String',
+  integer: 'Int64',
+  float: 'Float64',
+  bigint: 'Int64',
+};
+
+export type IntervalUnit =
+  | 'NANOSECOND'
+  | 'MICROSECOND'
+  | 'MILLISECOND'
+  | 'SECOND'
+  | 'MINUTE'
+  | 'HOUR'
+  | 'DAY'
+  | 'WEEK'
+  | 'MONTH'
+  | 'QUARTER'
+  | 'YEAR';
+
+export type ClickhouseConfig = {
+  url: string;
+  username: string;
+  password: string;
+  ttl?: {
+    [TableKey in TABLE_NAMES]?: {
+      row?: { interval: number; unit: IntervalUnit; ttlKey?: string };
+      columns?: Partial<{
+        [ColumnKey in keyof (typeof TABLE_SCHEMAS)[TableKey]]: {
+          interval: number;
+          unit: IntervalUnit;
+          ttlKey?: string;
+        };
+      }>;
+    };
+  };
+};
+
+export function transformRow<R>(row: any): R {
+  if (!row) {
+    return row;
+  }
+
+  if (row.createdAt) {
+    row.createdAt = new Date(row.createdAt);
+  }
+  if (row.updatedAt) {
+    row.updatedAt = new Date(row.updatedAt);
+  }
+
+  // Parse content field if it's a JSON string
+  if (row.content && typeof row.content === 'string') {
+    row.content = safelyParseJSON(row.content);
+  }
+
+  return row;
+}
+
+export function transformRows<R>(rows: any[]): R[] {
+  return rows.map((row: any) => transformRow<R>(row));
+}

package/src/storage/domains/workflows/index.ts

@@ -0,0 +1,285 @@
+import type { ClickHouseClient } from '@clickhouse/client';
+import { ErrorCategory, ErrorDomain, MastraError } from '@mastra/core/error';
+import { TABLE_WORKFLOW_SNAPSHOT, WorkflowsStorage } from '@mastra/core/storage';
+import type { WorkflowRun, WorkflowRuns } from '@mastra/core/storage';
+import type { WorkflowRunState } from '@mastra/core/workflows';
+import type { StoreOperationsClickhouse } from '../operations';
+import { TABLE_ENGINES } from '../utils';
+
+export class WorkflowsStorageClickhouse extends WorkflowsStorage {
+  protected client: ClickHouseClient;
+  protected operations: StoreOperationsClickhouse;
+  constructor({ client, operations }: { client: ClickHouseClient; operations: StoreOperationsClickhouse }) {
+    super();
+    this.operations = operations;
+    this.client = client;
+  }
+
+  async persistWorkflowSnapshot({
+    workflowName,
+    runId,
+    snapshot,
+  }: {
+    workflowName: string;
+    runId: string;
+    snapshot: WorkflowRunState;
+  }): Promise<void> {
+    try {
+      const currentSnapshot = await this.operations.load({
+        tableName: TABLE_WORKFLOW_SNAPSHOT,
+        keys: { workflow_name: workflowName, run_id: runId },
+      });
+
+      const now = 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.client.insert({
+        table: 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: any) {
+      throw new MastraError(
+        {
+          id: 'CLICKHOUSE_STORAGE_PERSIST_WORKFLOW_SNAPSHOT_FAILED',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: { workflowName, runId },
+        },
+        error,
+      );
+    }
+  }
+
+  async loadWorkflowSnapshot({
+    workflowName,
+    runId,
+  }: {
+    workflowName: string;
+    runId: string;
+  }): Promise<WorkflowRunState | null> {
+    try {
+      const result = await this.operations.load({
+        tableName: TABLE_WORKFLOW_SNAPSHOT,
+        keys: {
+          workflow_name: workflowName,
+          run_id: runId,
+        },
+      });
+
+      if (!result) {
+        return null;
+      }
+
+      return (result as any).snapshot;
+    } catch (error: any) {
+      throw new MastraError(
+        {
+          id: 'CLICKHOUSE_STORAGE_LOAD_WORKFLOW_SNAPSHOT_FAILED',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: { workflowName, runId },
+        },
+        error,
+      );
+    }
+  }
+
+  private parseWorkflowRun(row: any): WorkflowRun {
+    let parsedSnapshot: WorkflowRunState | string = row.snapshot as string;
+    if (typeof parsedSnapshot === 'string') {
+      try {
+        parsedSnapshot = JSON.parse(row.snapshot as string) as WorkflowRunState;
+      } catch (e) {
+        // If parsing fails, return the raw snapshot string
+        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),
+      resourceId: row.resourceId,
+    };
+  }
+
+  async getWorkflowRuns({
+    workflowName,
+    fromDate,
+    toDate,
+    limit,
+    offset,
+    resourceId,
+  }: {
+    workflowName?: string;
+    fromDate?: Date;
+    toDate?: Date;
+    limit?: number;
+    offset?: number;
+    resourceId?: string;
+  } = {}): Promise<WorkflowRuns> {
+    try {
+      const conditions: string[] = [];
+      const values: Record<string, any> = {};
+
+      if (workflowName) {
+        conditions.push(`workflow_name = {var_workflow_name:String}`);
+        values.var_workflow_name = workflowName;
+      }
+
+      if (resourceId) {
+        const hasResourceId = await this.operations.hasColumn(TABLE_WORKFLOW_SNAPSHOT, 'resourceId');
+        if (hasResourceId) {
+          conditions.push(`resourceId = {var_resourceId:String}`);
+          values.var_resourceId = resourceId;
+        } else {
+          console.warn(`[${TABLE_WORKFLOW_SNAPSHOT}] resourceId column not found. Skipping resourceId filter.`);
+        }
+      }
+
+      if (fromDate) {
+        conditions.push(`createdAt >= {var_from_date:DateTime64(3)}`);
+        values.var_from_date = fromDate.getTime() / 1000; // Convert to Unix timestamp
+      }
+
+      if (toDate) {
+        conditions.push(`createdAt <= {var_to_date:DateTime64(3)}`);
+        values.var_to_date = toDate.getTime() / 1000; // Convert to Unix timestamp
+      }
+
+      const whereClause = conditions.length > 0 ? `WHERE ${conditions.join(' AND ')}` : '';
+      const limitClause = limit !== undefined ? `LIMIT ${limit}` : '';
+      const offsetClause = offset !== undefined ? `OFFSET ${offset}` : '';
+
+      let total = 0;
+      // Only get total count when using pagination
+      if (limit !== undefined && offset !== undefined) {
+        const countResult = await this.client.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 as Array<{ count: string | number }>)[0]?.count ?? 0);
+      }
+
+      // Get results
+      const result = await this.client.query({
+        query: `
+              SELECT 
+                workflow_name,
+                run_id,
+                snapshot,
+                toDateTime64(createdAt, 3) as createdAt,
+                toDateTime64(updatedAt, 3) as updatedAt,
+                resourceId
+              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 as any[];
+      const runs = rows.map(row => {
+        return this.parseWorkflowRun(row);
+      });
+
+      // Use runs.length as total when not paginating
+      return { runs, total: total || runs.length };
+    } catch (error: any) {
+      throw new MastraError(
+        {
+          id: 'CLICKHOUSE_STORAGE_GET_WORKFLOW_RUNS_FAILED',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: { workflowName: workflowName ?? '', resourceId: resourceId ?? '' },
+        },
+        error,
+      );
+    }
+  }
+
+  async getWorkflowRunById({
+    runId,
+    workflowName,
+  }: {
+    runId: string;
+    workflowName?: string;
+  }): Promise<WorkflowRun | null> {
+    try {
+      const conditions: string[] = [];
+      const values: Record<string, any> = {};
+
+      if (runId) {
+        conditions.push(`run_id = {var_runId:String}`);
+        values.var_runId = runId;
+      }
+
+      if (workflowName) {
+        conditions.push(`workflow_name = {var_workflow_name:String}`);
+        values.var_workflow_name = workflowName;
+      }
+
+      const whereClause = conditions.length > 0 ? `WHERE ${conditions.join(' AND ')}` : '';
+
+      // Get results
+      const result = await this.client.query({
+        query: `
+              SELECT 
+                workflow_name,
+                run_id,
+                snapshot,
+                toDateTime64(createdAt, 3) as createdAt,
+                toDateTime64(updatedAt, 3) as updatedAt,
+                resourceId
+              FROM ${TABLE_WORKFLOW_SNAPSHOT} ${TABLE_ENGINES[TABLE_WORKFLOW_SNAPSHOT].startsWith('ReplacingMergeTree') ? 'FINAL' : ''}
+              ${whereClause}
+            `,
+        query_params: values,
+        format: 'JSONEachRow',
+      });
+
+      const resultJson = await result.json();
+      if (!Array.isArray(resultJson) || resultJson.length === 0) {
+        return null;
+      }
+      return this.parseWorkflowRun(resultJson[0]);
+    } catch (error: any) {
+      throw new MastraError(
+        {
+          id: 'CLICKHOUSE_STORAGE_GET_WORKFLOW_RUN_BY_ID_FAILED',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: { runId: runId ?? '', workflowName: workflowName ?? '' },
+        },
+        error,
+      );
+    }
+  }
+}