@mastra/clickhouse 0.0.0-vnextWorkflows-20250422142014 → 0.0.0-workflow-deno-20250616115451

package/src/storage/index.ts

@@ -1,7 +1,8 @@
 import type { ClickHouseClient } from '@clickhouse/client';
 import { createClient } from '@clickhouse/client';
+import { MessageList } from '@mastra/core/agent';
 import type { MetricResult, TestInfo } from '@mastra/core/eval';
-import type { MessageType, StorageThreadType } from '@mastra/core/memory';
+import type { MastraMessageV1, MastraMessageV2, StorageThreadType } from '@mastra/core/memory';
 import {
   MastraStorage,
   TABLE_EVALS,
@@ -11,7 +12,17 @@ import {
   TABLE_TRACES,
   TABLE_WORKFLOW_SNAPSHOT,
 } from '@mastra/core/storage';
-import type { EvalRow, StorageColumn, StorageGetMessagesArg, TABLE_NAMES } from '@mastra/core/storage';
+import type {
+  EvalRow,
+  PaginationInfo,
+  StorageColumn,
+  StorageGetMessagesArg,
+  TABLE_NAMES,
+  WorkflowRun,
+  WorkflowRuns,
+  StorageGetTracesArg,
+} from '@mastra/core/storage';
+import type { Trace } from '@mastra/core/telemetry';
 import type { WorkflowRunState } from '@mastra/core/workflows';
 
 function safelyParseJSON(jsonString: string): any {
@@ -203,6 +214,8 @@ export class ClickhouseStore extends MastraStorage {
     perPage,
     attributes,
     filters,
+    fromDate,
+    toDate,
   }: {
     name?: string;
     scope?: string;
@@ -210,6 +223,8 @@ export class ClickhouseStore extends MastraStorage {
     perPage: number;
     attributes?: Record<string, string>;
     filters?: Record<string, any>;
+    fromDate?: Date;
+    toDate?: Date;
   }): Promise<any[]> {
     const limit = perPage;
     const offset = page * perPage;
@@ -241,6 +256,16 @@ export class ClickhouseStore extends MastraStorage {
       });
     }
 
+    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 ')}` : '';
 
     const result = await this.db.query({
@@ -316,7 +341,6 @@ export class ClickhouseStore extends MastraStorage {
           ${['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)
         ${rowTtl ? `TTL toDateTime(${rowTtl.ttlKey ?? 'createdAt'}) + INTERVAL ${rowTtl.interval} ${rowTtl.unit}` : ''}
@@ -327,7 +351,6 @@ export class ClickhouseStore extends MastraStorage {
           ${columns}
         )
         ENGINE = ${TABLE_ENGINES[tableName]}
-        PARTITION BY "createdAt"
         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}` : ''}
@@ -350,6 +373,73 @@ export class ClickhouseStore extends MastraStorage {
     }
   }
 
+  protected getSqlType(type: StorageColumn['type']): string {
+    switch (type) {
+      case 'text':
+        return 'String';
+      case 'timestamp':
+        return 'DateTime64(3)';
+      case 'integer':
+      case 'bigint':
+        return 'Int64';
+      case 'jsonb':
+        return 'String';
+      default:
+        return super.getSqlType(type); // fallback to base implementation
+    }
+  }
+
+  /**
+   * Alters table schema to add columns if they don't exist
+   * @param tableName Name of the table
+   * @param schema Schema of the table
+   * @param ifNotExists Array of column names to add if they don't exist
+   */
+  async alterTable({
+    tableName,
+    schema,
+    ifNotExists,
+  }: {
+    tableName: TABLE_NAMES;
+    schema: Record<string, StorageColumn>;
+    ifNotExists: string[];
+  }): Promise<void> {
+    try {
+      // 1. Get existing columns
+      const describeSql = `DESCRIBE TABLE ${tableName}`;
+      const result = await this.db.query({
+        query: describeSql,
+      });
+      const rows = await result.json();
+      const existingColumnNames = new Set(rows.data.map((row: any) => row.name.toLowerCase()));
+
+      // 2. Add missing columns
+      for (const columnName of ifNotExists) {
+        if (!existingColumnNames.has(columnName.toLowerCase()) && schema[columnName]) {
+          const columnDef = schema[columnName];
+          let sqlType = this.getSqlType(columnDef.type);
+          if (columnDef.nullable !== false) {
+            sqlType = `Nullable(${sqlType})`;
+          }
+          const defaultValue = columnDef.nullable === false ? this.getDefaultValue(columnDef.type) : '';
+          // Use backticks or double quotes as needed for identifiers
+          const alterSql =
+            `ALTER TABLE ${tableName} ADD COLUMN IF NOT EXISTS "${columnName}" ${sqlType} ${defaultValue}`.trim();
+
+          await this.db.query({
+            query: alterSql,
+          });
+          this.logger?.debug?.(`Added column ${columnName} to table ${tableName}`);
+        }
+      }
+    } catch (error) {
+      this.logger?.error?.(
+        `Error altering table ${tableName}: ${error instanceof Error ? error.message : String(error)}`,
+      );
+      throw new Error(`Failed to alter table ${tableName}: ${error}`);
+    }
+  }
+
   async clearTable({ tableName }: { tableName: TABLE_NAMES }): Promise<void> {
     try {
       await this.db.query({
@@ -580,15 +670,18 @@ export class ClickhouseStore extends MastraStorage {
 
       await this.db.insert({
         table: TABLE_THREADS,
+        format: 'JSONEachRow',
         values: [
           {
-            ...updatedThread,
+            id: updatedThread.id,
+            resourceId: updatedThread.resourceId,
+            title: updatedThread.title,
+            metadata: updatedThread.metadata,
+            createdAt: updatedThread.createdAt,
             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,
@@ -606,7 +699,7 @@ export class ClickhouseStore extends MastraStorage {
     try {
       // First delete all messages associated with this thread
       await this.db.command({
-        query: `DELETE FROM "${TABLE_MESSAGES}" WHERE thread_id = '${threadId}';`,
+        query: `DELETE FROM "${TABLE_MESSAGES}" WHERE thread_id = {var_thread_id:String};`,
         query_params: { var_thread_id: threadId },
         clickhouse_settings: {
           output_format_json_quote_64bit_integers: 0,
@@ -627,7 +720,14 @@ export class ClickhouseStore extends MastraStorage {
     }
   }
 
-  async getMessages<T = unknown>({ threadId, selectBy }: StorageGetMessagesArg): Promise<T> {
+  public async getMessages(args: StorageGetMessagesArg & { format?: 'v1' }): Promise<MastraMessageV1[]>;
+  public async getMessages(args: StorageGetMessagesArg & { format: 'v2' }): Promise<MastraMessageV2[]>;
+  public async getMessages({
+    threadId,
+    resourceId,
+    selectBy,
+    format,
+  }: StorageGetMessagesArg & { format?: 'v1' | 'v2' }): Promise<MastraMessageV1[] | MastraMessageV2[]> {
     try {
       const messages: any[] = [];
       const limit = typeof selectBy?.last === `number` ? selectBy.last : 40;
@@ -734,18 +834,26 @@ export class ClickhouseStore extends MastraStorage {
         }
       });
 
-      return messages as T;
+      const list = new MessageList({ threadId, resourceId }).add(messages, 'memory');
+      if (format === `v2`) return list.get.all.v2();
+      return list.get.all.v1();
     } catch (error) {
       console.error('Error getting messages:', error);
       throw error;
     }
   }
 
-  async saveMessages({ messages }: { messages: MessageType[] }): Promise<MessageType[]> {
+  async saveMessages(args: { messages: MastraMessageV1[]; format?: undefined | 'v1' }): Promise<MastraMessageV1[]>;
+  async saveMessages(args: { messages: MastraMessageV2[]; format: 'v2' }): Promise<MastraMessageV2[]>;
+  async saveMessages(
+    args: { messages: MastraMessageV1[]; format?: undefined | 'v1' } | { messages: MastraMessageV2[]; format: 'v2' },
+  ): Promise<MastraMessageV2[] | MastraMessageV1[]> {
+    const { messages, format = 'v1' } = args;
     if (messages.length === 0) return messages;
 
     try {
       const threadId = messages[0]?.threadId;
+      const resourceId = messages[0]?.resourceId;
       if (!threadId) {
         throw new Error('Thread ID is required');
       }
@@ -756,26 +864,52 @@ export class ClickhouseStore extends MastraStorage {
         throw new Error(`Thread ${threadId} not found`);
       }
 
-      await this.db.insert({
-        table: 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,
-        },
-      });
+      // Execute message inserts and thread update in parallel for better performance
+      await Promise.all([
+        // Insert messages
+        this.db.insert({
+          table: 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 || 'v2',
+          })),
+          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,
+          },
+        }),
+        // Update thread's updatedAt timestamp
+        this.db.insert({
+          table: TABLE_THREADS,
+          format: 'JSONEachRow',
+          values: [
+            {
+              id: thread.id,
+              resourceId: thread.resourceId,
+              title: thread.title,
+              metadata: thread.metadata,
+              createdAt: thread.createdAt,
+              updatedAt: new Date().toISOString(),
+            },
+          ],
+          clickhouse_settings: {
+            date_time_input_format: 'best_effort',
+            use_client_time_zone: 1,
+            output_format_json_quote_64bit_integers: 0,
+          },
+        }),
+      ]);
 
-      return messages;
+      const list = new MessageList({ threadId, resourceId }).add(messages, 'memory');
+      if (format === `v2`) return list.get.all.v2();
+      return list.get.all.v1();
     } catch (error) {
       console.error('Error saving messages:', error);
       throw error;
@@ -856,28 +990,42 @@ export class ClickhouseStore extends MastraStorage {
     }
   }
 
+  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;
-  } = {}): Promise<{
-    runs: Array<{
-      workflowName: string;
-      runId: string;
-      snapshot: WorkflowRunState | string;
-      createdAt: Date;
-      updatedAt: Date;
-    }>;
-    total: number;
-  }> {
+    resourceId?: string;
+  } = {}): Promise<WorkflowRuns> {
     try {
       const conditions: string[] = [];
       const values: Record<string, any> = {};
@@ -887,6 +1035,16 @@ export class ClickhouseStore extends MastraStorage {
         values.var_workflow_name = workflowName;
       }
 
+      if (resourceId) {
+        const hasResourceId = await this.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
@@ -921,7 +1079,8 @@ export class ClickhouseStore extends MastraStorage {
             run_id,
             snapshot,
             toDateTime64(createdAt, 3) as createdAt,
-            toDateTime64(updatedAt, 3) as updatedAt
+            toDateTime64(updatedAt, 3) as updatedAt,
+            resourceId
           FROM ${TABLE_WORKFLOW_SNAPSHOT} ${TABLE_ENGINES[TABLE_WORKFLOW_SNAPSHOT].startsWith('ReplacingMergeTree') ? 'FINAL' : ''}
           ${whereClause}
           ORDER BY createdAt DESC
@@ -935,23 +1094,7 @@ export class ClickhouseStore extends MastraStorage {
       const resultJson = await result.json();
       const rows = resultJson as any[];
       const runs = rows.map(row => {
-        let parsedSnapshot: WorkflowRunState | string = row.snapshot;
-        if (typeof parsedSnapshot === 'string') {
-          try {
-            parsedSnapshot = JSON.parse(row.snapshot) 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),
-        };
+        return this.parseWorkflowRun(row);
       });
 
       // Use runs.length as total when not paginating
@@ -962,6 +1105,84 @@ export class ClickhouseStore extends MastraStorage {
     }
   }
 
+  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.db.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) {
+      console.error('Error getting workflow run by ID:', error);
+      throw error;
+    }
+  }
+
+  private async hasColumn(table: string, column: string): Promise<boolean> {
+    const result = await this.db.query({
+      query: `DESCRIBE TABLE ${table}`,
+      format: 'JSONEachRow',
+    });
+    const columns = (await result.json()) as { name: string }[];
+    return columns.some(c => c.name === column);
+  }
+
+  async getTracesPaginated(_args: StorageGetTracesArg): Promise<PaginationInfo & { traces: Trace[] }> {
+    throw new Error('Method not implemented.');
+  }
+
+  async getThreadsByResourceIdPaginated(_args: {
+    resourceId: string;
+    page?: number;
+    perPage?: number;
+  }): Promise<PaginationInfo & { threads: StorageThreadType[] }> {
+    throw new Error('Method not implemented.');
+  }
+
+  async getMessagesPaginated(
+    _args: StorageGetMessagesArg,
+  ): Promise<PaginationInfo & { messages: MastraMessageV1[] | MastraMessageV2[] }> {
+    throw new Error('Method not implemented.');
+  }
+
   async close(): Promise<void> {
     await this.db.close();
   }