@mastra/clickhouse 1.0.0-beta.11 → 1.0.0-beta.12

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, MemoryStorage, TABLE_SCHEMAS, createStorageErrorId, normalizePerPage, calculatePagination, ObservabilityStorage, SPAN_SCHEMA, listTracesArgsSchema, ScoresStorage, transformScoreRow, SCORERS_SCHEMA, WorkflowsStorage, MastraStorage, TraceStatus, getSqlType, getDefaultValue, safelyParseJSON } from '@mastra/core/storage';
+import { TABLE_SPANS, TABLE_RESOURCES, TABLE_SCORERS, TABLE_THREADS, TABLE_TRACES, TABLE_WORKFLOW_SNAPSHOT, TABLE_MESSAGES, MemoryStorage, TABLE_SCHEMAS, createStorageErrorId, normalizePerPage, calculatePagination, ObservabilityStorage, SPAN_SCHEMA, listTracesArgsSchema, ScoresStorage, transformScoreRow, SCORERS_SCHEMA, WorkflowsStorage, MastraCompositeStore, TraceStatus, getSqlType, getDefaultValue, safelyParseJSON } from '@mastra/core/storage';
 import { MessageList } from '@mastra/core/agent';
 import { MastraBase } from '@mastra/core/base';
 import { saveScorePayloadSchema } from '@mastra/core/evals';
@@ -13,8 +13,10 @@ var TABLE_ENGINES = {
   [TABLE_THREADS]: `ReplacingMergeTree()`,
   [TABLE_SCORERS]: `MergeTree()`,
   [TABLE_RESOURCES]: `ReplacingMergeTree()`,
-  // TODO: verify this is the correct engine for Spans when implementing clickhouse storage
-  [TABLE_SPANS]: `ReplacingMergeTree()`,
+  // ReplacingMergeTree(updatedAt) deduplicates rows with the same (traceId, spanId) sorting key,
+  // keeping the row with the highest updatedAt value. Combined with ORDER BY (traceId, spanId),
+  // this provides eventual uniqueness for the (traceId, spanId) composite key.
+  [TABLE_SPANS]: `ReplacingMergeTree(updatedAt)`,
   mastra_agents: `ReplacingMergeTree()`
 };
 var COLUMN_TYPES = {
@@ -97,6 +99,213 @@ var ClickhouseDB = class extends MastraBase {
     const columns = await result.json();
     return columns.some((c) => c.name === column);
   }
+  /**
+   * Checks if a table exists in the database.
+   */
+  async tableExists(tableName) {
+    try {
+      const result = await this.client.query({
+        query: `EXISTS TABLE ${tableName}`,
+        format: "JSONEachRow"
+      });
+      const rows = await result.json();
+      return rows[0]?.result === 1;
+    } catch {
+      return false;
+    }
+  }
+  /**
+   * Gets the sorting key (ORDER BY columns) for a table.
+   * Returns null if the table doesn't exist.
+   */
+  async getTableSortingKey(tableName) {
+    try {
+      const result = await this.client.query({
+        query: `SELECT sorting_key FROM system.tables WHERE database = currentDatabase() AND name = {tableName:String}`,
+        query_params: { tableName },
+        format: "JSONEachRow"
+      });
+      const rows = await result.json();
+      return rows[0]?.sorting_key ?? null;
+    } catch {
+      return null;
+    }
+  }
+  /**
+   * Checks if migration is needed for the spans table.
+   * Returns information about the current state.
+   */
+  async checkSpansMigrationStatus(tableName) {
+    const exists = await this.tableExists(tableName);
+    if (!exists) {
+      return { needsMigration: false, currentSortingKey: null };
+    }
+    const currentSortingKey = await this.getTableSortingKey(tableName);
+    if (!currentSortingKey) {
+      return { needsMigration: false, currentSortingKey: null };
+    }
+    const needsMigration = currentSortingKey.toLowerCase().startsWith("createdat");
+    return { needsMigration, currentSortingKey };
+  }
+  /**
+   * Checks for duplicate (traceId, spanId) combinations in the spans table.
+   * Returns information about duplicates for logging/CLI purposes.
+   */
+  async checkForDuplicateSpans(tableName) {
+    try {
+      const result = await this.client.query({
+        query: `
+          SELECT count() as duplicate_count
+          FROM (
+            SELECT traceId, spanId
+            FROM ${tableName}
+            GROUP BY traceId, spanId
+            HAVING count() > 1
+          )
+        `,
+        format: "JSONEachRow"
+      });
+      const rows = await result.json();
+      const duplicateCount = parseInt(rows[0]?.duplicate_count ?? "0", 10);
+      return {
+        hasDuplicates: duplicateCount > 0,
+        duplicateCount
+      };
+    } catch (error) {
+      this.logger?.debug?.(`Could not check for duplicates: ${error}`);
+      return { hasDuplicates: false, duplicateCount: 0 };
+    }
+  }
+  /**
+   * Migrates the spans table from the old sorting key (createdAt, traceId, spanId)
+   * to the new sorting key (traceId, spanId) for proper uniqueness enforcement.
+   *
+   * This migration:
+   * 1. Renames the old table to a backup
+   * 2. Creates a new table with the correct sorting key
+   * 3. Copies all data from the backup to the new table, deduplicating by (traceId, spanId)
+   *    using priority-based selection:
+   *    - First, prefer completed spans (those with endedAt set)
+   *    - Then prefer the most recently updated span (highest updatedAt)
+   *    - Finally use creation time as tiebreaker (highest createdAt)
+   * 4. Drops the backup table
+   *
+   * The deduplication strategy matches the PostgreSQL migration (PR #12073) to ensure
+   * consistent behavior across storage backends.
+   *
+   * The migration is idempotent - it only runs if the old sorting key is detected.
+   *
+   * @returns true if migration was performed, false if not needed
+   */
+  async migrateSpansTableSortingKey({
+    tableName,
+    schema
+  }) {
+    if (tableName !== TABLE_SPANS) {
+      return false;
+    }
+    const exists = await this.tableExists(tableName);
+    if (!exists) {
+      return false;
+    }
+    const currentSortingKey = await this.getTableSortingKey(tableName);
+    if (!currentSortingKey) {
+      return false;
+    }
+    const needsMigration = currentSortingKey.toLowerCase().startsWith("createdat");
+    if (!needsMigration) {
+      this.logger?.debug?.(`Spans table already has correct sorting key: ${currentSortingKey}`);
+      return false;
+    }
+    this.logger?.info?.(`Migrating spans table from sorting key "${currentSortingKey}" to "(traceId, spanId)"`);
+    const backupTableName = `${tableName}_backup_${Date.now()}`;
+    const rowTtl = this.ttl?.[tableName]?.row;
+    try {
+      await this.client.command({
+        query: `RENAME TABLE ${tableName} TO ${backupTableName}`
+      });
+      const columns = Object.entries(schema).map(([name, def]) => {
+        let sqlType = this.getSqlType(def.type);
+        let isNullable = def.nullable === true;
+        if (tableName === TABLE_SPANS && name === "updatedAt") {
+          isNullable = false;
+        }
+        if (isNullable) {
+          sqlType = `Nullable(${sqlType})`;
+        }
+        const constraints = [];
+        if (name === "metadata" && (def.type === "text" || def.type === "jsonb") && isNullable) {
+          constraints.push("DEFAULT '{}'");
+        }
+        const columnTtl = this.ttl?.[tableName]?.columns?.[name];
+        return `"${name}" ${sqlType} ${constraints.join(" ")} ${columnTtl ? `TTL toDateTime(${columnTtl.ttlKey ?? "createdAt"}) + INTERVAL ${columnTtl.interval} ${columnTtl.unit}` : ""}`;
+      }).join(",\n");
+      const createSql = `
+        CREATE TABLE ${tableName} (
+          ${columns}
+        )
+        ENGINE = ${TABLE_ENGINES[tableName] ?? "MergeTree()"}
+        PRIMARY KEY (traceId, spanId)
+        ORDER BY (traceId, spanId)
+        ${rowTtl ? `TTL toDateTime(${rowTtl.ttlKey ?? "createdAt"}) + INTERVAL ${rowTtl.interval} ${rowTtl.unit}` : ""}
+        SETTINGS index_granularity = 8192
+      `;
+      await this.client.command({
+        query: createSql
+      });
+      const describeResult = await this.client.query({
+        query: `DESCRIBE TABLE ${backupTableName}`,
+        format: "JSONEachRow"
+      });
+      const backupColumns = await describeResult.json();
+      const backupColumnNames = new Set(backupColumns.map((c) => c.name));
+      const columnsToInsert = Object.keys(schema).filter((col) => backupColumnNames.has(col));
+      const columnList = columnsToInsert.map((c) => `"${c}"`).join(", ");
+      const selectExpressions = columnsToInsert.map((c) => c === "updatedAt" ? `COALESCE("updatedAt", "createdAt") as "updatedAt"` : `"${c}"`).join(", ");
+      await this.client.command({
+        query: `INSERT INTO ${tableName} (${columnList})
+                SELECT ${selectExpressions}
+                FROM ${backupTableName}
+                ORDER BY traceId, spanId,
+                         (endedAt IS NOT NULL AND endedAt != '') DESC,
+                         COALESCE(updatedAt, createdAt) DESC,
+                         createdAt DESC
+                LIMIT 1 BY traceId, spanId`
+      });
+      await this.client.command({
+        query: `DROP TABLE ${backupTableName}`
+      });
+      this.logger?.info?.(`Successfully migrated spans table to new sorting key`);
+      return true;
+    } catch (error) {
+      this.logger?.error?.(`Migration failed: ${error.message}`);
+      try {
+        const originalExists = await this.tableExists(tableName);
+        const backupExists = await this.tableExists(backupTableName);
+        if (!originalExists && backupExists) {
+          this.logger?.info?.(`Restoring spans table from backup`);
+          await this.client.command({
+            query: `RENAME TABLE ${backupTableName} TO ${tableName}`
+          });
+        } else if (originalExists && backupExists) {
+          await this.client.command({
+            query: `DROP TABLE IF EXISTS ${backupTableName}`
+          });
+        }
+      } catch (restoreError) {
+        this.logger?.error?.(`Failed to restore from backup: ${restoreError}`);
+      }
+      throw new MastraError(
+        {
+          id: createStorageErrorId("CLICKHOUSE", "MIGRATE_SPANS_SORTING_KEY", "FAILED"),
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: { tableName, currentSortingKey }
+        },
+        error
+      );
+    }
+  }
   getSqlType(type) {
     switch (type) {
       case "text":
@@ -123,7 +332,10 @@ var ClickhouseDB = class extends MastraBase {
     try {
       const columns = Object.entries(schema).map(([name, def]) => {
         let sqlType = this.getSqlType(def.type);
-        const isNullable = def.nullable === true;
+        let isNullable = def.nullable === true;
+        if (tableName === TABLE_SPANS && name === "updatedAt") {
+          isNullable = false;
+        }
         if (isNullable) {
           sqlType = `Nullable(${sqlType})`;
         }
@@ -153,8 +365,8 @@ var ClickhouseDB = class extends MastraBase {
               ${columns}
             )
             ENGINE = ${TABLE_ENGINES[tableName] ?? "MergeTree()"}
-            PRIMARY KEY (createdAt, traceId, spanId)
-            ORDER BY (createdAt, traceId, spanId)
+            PRIMARY KEY (traceId, spanId)
+            ORDER BY (traceId, spanId)
             ${rowTtl ? `TTL toDateTime(${rowTtl.ttlKey ?? "createdAt"}) + INTERVAL ${rowTtl.interval} ${rowTtl.unit}` : ""}
             SETTINGS index_granularity = 8192
           `;
@@ -300,7 +512,9 @@ var ClickhouseDB = class extends MastraBase {
       ...Object.fromEntries(
         Object.entries(record).map(([key, value]) => [
           key,
-          TABLE_SCHEMAS[tableName]?.[key]?.type === "timestamp" ? new Date(value).toISOString() : value
+          // Only convert to Date if it's a timestamp column AND value is not null/undefined
+          // new Date(null) returns epoch date, not null, so we must check first
+          TABLE_SCHEMAS[tableName]?.[key]?.type === "timestamp" && value != null ? new Date(value).toISOString() : value
         ])
       )
     }));
@@ -1642,11 +1856,101 @@ var ObservabilityStorageClickhouse = class extends ObservabilityStorage {
     this.#db = new ClickhouseDB({ client, ttl });
   }
   async init() {
+    const migrationStatus = await this.#db.checkSpansMigrationStatus(TABLE_SPANS);
+    if (migrationStatus.needsMigration) {
+      const duplicateInfo = await this.#db.checkForDuplicateSpans(TABLE_SPANS);
+      const duplicateMessage = duplicateInfo.hasDuplicates ? `
+Found ${duplicateInfo.duplicateCount} duplicate (traceId, spanId) combinations that will be removed.
+` : "";
+      const errorMessage = `
+===========================================================================
+MIGRATION REQUIRED: ClickHouse spans table needs sorting key update
+===========================================================================
+
+The spans table structure has changed. ClickHouse requires a table recreation
+to update the sorting key from (traceId) to (traceId, spanId).
+` + duplicateMessage + `
+To fix this, run the manual migration command:
+
+  npx mastra migrate
+
+This command will:
+  1. Create a new table with the correct sorting key
+  2. Copy data from the old table (deduplicating if needed)
+  3. Replace the old table with the new one
+
+WARNING: This migration involves table recreation and may take significant
+time for large tables. Please ensure you have a backup before proceeding.
+===========================================================================
+`;
+      throw new MastraError({
+        id: createStorageErrorId("CLICKHOUSE", "MIGRATION_REQUIRED", "SORTING_KEY_CHANGE"),
+        domain: ErrorDomain.STORAGE,
+        category: ErrorCategory.USER,
+        text: errorMessage
+      });
+    }
     await this.#db.createTable({ tableName: TABLE_SPANS, schema: SPAN_SCHEMA });
   }
   async dangerouslyClearAll() {
     await this.#db.clearTable({ tableName: TABLE_SPANS });
   }
+  /**
+   * Manually run the spans migration to deduplicate and update the sorting key.
+   * This is intended to be called from the CLI when duplicates are detected.
+   *
+   * @returns Migration result with status and details
+   */
+  async migrateSpans() {
+    const migrationStatus = await this.#db.checkSpansMigrationStatus(TABLE_SPANS);
+    if (!migrationStatus.needsMigration) {
+      return {
+        success: true,
+        alreadyMigrated: true,
+        duplicatesRemoved: 0,
+        message: `Migration already complete. Spans table has correct sorting key.`
+      };
+    }
+    const duplicateInfo = await this.#db.checkForDuplicateSpans(TABLE_SPANS);
+    if (duplicateInfo.hasDuplicates) {
+      this.logger?.info?.(
+        `Found ${duplicateInfo.duplicateCount} duplicate (traceId, spanId) combinations. Starting migration with deduplication...`
+      );
+    } else {
+      this.logger?.info?.(`No duplicate spans found. Starting sorting key migration...`);
+    }
+    await this.#db.migrateSpansTableSortingKey({ tableName: TABLE_SPANS, schema: SPAN_SCHEMA });
+    return {
+      success: true,
+      alreadyMigrated: false,
+      duplicatesRemoved: duplicateInfo.duplicateCount,
+      message: duplicateInfo.hasDuplicates ? `Migration complete. Removed duplicates and updated sorting key for ${TABLE_SPANS}.` : `Migration complete. Updated sorting key for ${TABLE_SPANS}.`
+    };
+  }
+  /**
+   * Check migration status for the spans table.
+   * Returns information about whether migration is needed.
+   */
+  async checkSpansMigrationStatus() {
+    const migrationStatus = await this.#db.checkSpansMigrationStatus(TABLE_SPANS);
+    if (!migrationStatus.needsMigration) {
+      return {
+        needsMigration: false,
+        hasDuplicates: false,
+        duplicateCount: 0,
+        constraintExists: true,
+        tableName: TABLE_SPANS
+      };
+    }
+    const duplicateInfo = await this.#db.checkForDuplicateSpans(TABLE_SPANS);
+    return {
+      needsMigration: true,
+      hasDuplicates: duplicateInfo.hasDuplicates,
+      duplicateCount: duplicateInfo.duplicateCount,
+      constraintExists: false,
+      tableName: TABLE_SPANS
+    };
+  }
   get tracingStrategy() {
     return {
       preferred: "insert-only",
@@ -1980,10 +2284,10 @@ var ObservabilityStorageClickhouse = class extends ObservabilityStorage {
               conditions.push(`(error IS NOT NULL AND error != '')`);
               break;
             case TraceStatus.RUNNING:
-              conditions.push(`(endedAt IS NULL OR endedAt = '') AND (error IS NULL OR error = '')`);
+              conditions.push(`endedAt IS NULL AND (error IS NULL OR error = '')`);
               break;
             case TraceStatus.SUCCESS:
-              conditions.push(`(endedAt IS NOT NULL AND endedAt != '') AND (error IS NULL OR error = '')`);
+              conditions.push(`endedAt IS NOT NULL AND (error IS NULL OR error = '')`);
               break;
           }
         }
@@ -2012,7 +2316,7 @@ var ObservabilityStorageClickhouse = class extends ObservabilityStorage {
       if (sortField === "endedAt") {
         const nullSortValue = sortDirection === "DESC" ? 0 : 1;
         const nonNullSortValue = sortDirection === "DESC" ? 1 : 0;
-        orderClause = `ORDER BY CASE WHEN ${sortField} IS NULL OR ${sortField} = '' THEN ${nullSortValue} ELSE ${nonNullSortValue} END, ${sortField} ${sortDirection}`;
+        orderClause = `ORDER BY CASE WHEN ${sortField} IS NULL THEN ${nullSortValue} ELSE ${nonNullSortValue} END, ${sortField} ${sortDirection}`;
       } else {
         orderClause = `ORDER BY ${sortField} ${sortDirection}`;
       }
@@ -2918,7 +3222,7 @@ var WorkflowsStorageClickhouse = class extends WorkflowsStorage {
 var isClientConfig = (config) => {
   return "client" in config;
 };
-var ClickhouseStore = class extends MastraStorage {
+var ClickhouseStore = class extends MastraCompositeStore {
   db;
   ttl = {};
   stores;