@mastra/clickhouse 1.11.1 → 1.12.0-alpha.0

package/dist/index.js

@@ -5947,6 +5947,11 @@ async function getTableColumns(client, table) {
   const rows = await result.json();
   return rows.map((r) => r.name);
 }
+async function getTableColumnTypes(client, table) {
+  const result = await client.query({ query: `DESCRIBE TABLE ${table}`, format: "JSONEachRow" });
+  const rows = await result.json();
+  return new Map(rows.map((r) => [r.name, r.type]));
+}
 function buildTemporaryTableDDL(createDDL, table, tempTable) {
   return createDDL.replace(`CREATE TABLE IF NOT EXISTS ${table}`, `CREATE TABLE ${tempTable}`);
 }
@@ -5980,6 +5985,218 @@ async function checkSignalTablesMigrationStatus(client) {
     tables
   };
 }
+var TABLE_LEGACY_SPANS = "mastra_ai_spans";
+var PROMOTED_METADATA_KEYS = [
+  "experimentId",
+  "entityType",
+  "entityId",
+  "entityName",
+  "entityVersionId",
+  "parentEntityVersionId",
+  "rootEntityVersionId",
+  "userId",
+  "organizationId",
+  "resourceId",
+  "runId",
+  "sessionId",
+  "threadId",
+  "requestId",
+  "environment",
+  "executionSource",
+  "serviceName"
+];
+var TABLE_LEGACY_MIGRATION_DONE = "mastra_legacy_span_migration_done";
+async function hasMarkerRow(client) {
+  const engine = await getTableEngine(client, TABLE_LEGACY_MIGRATION_DONE);
+  if (!engine) return false;
+  const result = await client.query({
+    query: `SELECT count() as cnt FROM ${TABLE_LEGACY_MIGRATION_DONE}`,
+    format: "JSONEachRow"
+  });
+  const rows = await result.json();
+  return Number(rows[0]?.cnt ?? 0) > 0;
+}
+async function checkLegacySpanMigrationStatus(client) {
+  if (await hasMarkerRow(client)) {
+    return { needsMigration: false };
+  }
+  const engine = await getTableEngine(client, TABLE_LEGACY_SPANS);
+  if (!engine) {
+    return { needsMigration: false };
+  }
+  const countResult = await client.query({
+    query: `SELECT count() as cnt FROM ${TABLE_LEGACY_SPANS}`,
+    format: "JSONEachRow"
+  });
+  const countRows = await countResult.json();
+  const legacyRowCount = Number(countRows[0]?.cnt ?? 0);
+  if (legacyRowCount === 0) {
+    return { needsMigration: false };
+  }
+  return { needsMigration: true };
+}
+var VNEXT_SPAN_COLUMNS = [
+  "dedupeKey",
+  "traceId",
+  "spanId",
+  "parentSpanId",
+  "experimentId",
+  "entityType",
+  "entityId",
+  "entityName",
+  "entityVersionId",
+  "parentEntityVersionId",
+  "parentEntityType",
+  "parentEntityId",
+  "parentEntityName",
+  "rootEntityVersionId",
+  "rootEntityType",
+  "rootEntityId",
+  "rootEntityName",
+  "userId",
+  "organizationId",
+  "resourceId",
+  "runId",
+  "sessionId",
+  "threadId",
+  "requestId",
+  "environment",
+  "executionSource",
+  "serviceName",
+  "name",
+  "spanType",
+  "isEvent",
+  "startedAt",
+  "endedAt",
+  "tags",
+  "metadataSearch",
+  "attributes",
+  "scope",
+  "links",
+  "input",
+  "output",
+  "error",
+  "metadataRaw",
+  "requestContext"
+];
+function buildLegacySpanSelectExprs(legacyColumnTypes) {
+  const has = (col) => legacyColumnTypes.has(col);
+  const colOrNull = (legacy) => has(legacy) ? `"${legacy}"` : "NULL";
+  const promotedIn = PROMOTED_METADATA_KEYS.map((k) => `'${k}'`).join(", ");
+  const metadataSearchExpr = has("metadata") ? `CAST(arrayFilter(x -> x.1 NOT IN (${promotedIn}) AND trim(BOTH ' ' FROM x.2) != '', arrayMap(x -> (x.1, trim(BOTH ' ' FROM x.2)), JSONExtractKeysAndValues(COALESCE("metadata", '{}'), 'String'))), 'Map(String, String)')` : `CAST(map(), 'Map(String, String)')`;
+  let tagsExpr = "[]";
+  if (has("tags")) {
+    const tagsType = legacyColumnTypes.get("tags");
+    if (tagsType.includes("Array")) {
+      tagsExpr = `arrayFilter(x -> x != '', arrayDistinct(arrayMap(x -> trim(BOTH ' ' FROM x), "tags")))`;
+    } else {
+      tagsExpr = `arrayFilter(x -> x != '', arrayDistinct(arrayMap(x -> trim(BOTH ' ' FROM x), JSONExtract(COALESCE("tags", '[]'), 'Array(String)'))))`;
+    }
+  }
+  const exprs = {
+    dedupeKey: `concat("traceId", ':', "spanId")`,
+    traceId: `"traceId"`,
+    spanId: `"spanId"`,
+    // Legacy stores '' for root spans; VNext MV filters on IS NULL.
+    parentSpanId: has("parentSpanId") ? `nullIf("parentSpanId", '')` : "NULL",
+    experimentId: colOrNull("experimentId"),
+    entityType: colOrNull("entityType"),
+    entityId: colOrNull("entityId"),
+    entityName: colOrNull("entityName"),
+    entityVersionId: colOrNull("entityVersionId"),
+    parentEntityVersionId: colOrNull("parentEntityVersionId"),
+    parentEntityType: colOrNull("parentEntityType"),
+    parentEntityId: colOrNull("parentEntityId"),
+    parentEntityName: colOrNull("parentEntityName"),
+    rootEntityVersionId: colOrNull("rootEntityVersionId"),
+    rootEntityType: colOrNull("rootEntityType"),
+    rootEntityId: colOrNull("rootEntityId"),
+    rootEntityName: colOrNull("rootEntityName"),
+    userId: colOrNull("userId"),
+    organizationId: colOrNull("organizationId"),
+    resourceId: colOrNull("resourceId"),
+    runId: colOrNull("runId"),
+    sessionId: colOrNull("sessionId"),
+    threadId: colOrNull("threadId"),
+    requestId: colOrNull("requestId"),
+    environment: colOrNull("environment"),
+    executionSource: has("source") ? `"source"` : "NULL",
+    serviceName: colOrNull("serviceName"),
+    name: `"name"`,
+    spanType: `"spanType"`,
+    isEvent: `"isEvent"`,
+    startedAt: `"startedAt"`,
+    // VNext requires non-null endedAt; events always use startedAt.
+    endedAt: `if("isEvent", "startedAt", COALESCE("endedAt", "startedAt"))`,
+    tags: tagsExpr,
+    metadataSearch: metadataSearchExpr,
+    attributes: colOrNull("attributes"),
+    scope: colOrNull("scope"),
+    links: colOrNull("links"),
+    input: colOrNull("input"),
+    output: colOrNull("output"),
+    error: colOrNull("error"),
+    metadataRaw: has("metadata") ? `"metadata"` : "NULL",
+    requestContext: colOrNull("requestContext")
+  };
+  return VNEXT_SPAN_COLUMNS.map((c) => `${exprs[c]} AS "${c}"`).join(",\n    ");
+}
+async function migrateLegacySpans(client, logger) {
+  if (await hasMarkerRow(client)) {
+    return { migratedRows: 0, batches: 0 };
+  }
+  const engine = await getTableEngine(client, TABLE_LEGACY_SPANS);
+  if (!engine) {
+    return { migratedRows: 0, batches: 0 };
+  }
+  const legacyColumnTypes = await getTableColumnTypes(client, TABLE_LEGACY_SPANS);
+  const hasUpdatedAt = legacyColumnTypes.has("updatedAt");
+  await client.command({ query: SPAN_EVENTS_DDL });
+  await client.command({ query: TRACE_ROOTS_DDL });
+  await client.command({ query: TRACE_BRANCHES_DDL });
+  await client.command({ query: TRACE_ROOTS_MV_DDL });
+  await client.command({ query: TRACE_BRANCHES_MV_DDL });
+  const selectExprs = buildLegacySpanSelectExprs(legacyColumnTypes);
+  const columnList = VNEXT_SPAN_COLUMNS.map((c) => `"${c}"`).join(", ");
+  const dedupOrder = hasUpdatedAt ? `ORDER BY "traceId", "spanId", COALESCE("updatedAt", "createdAt") DESC` : `ORDER BY "traceId", "spanId", "createdAt" DESC`;
+  const daysResult = await client.query({
+    query: `SELECT
+      toString(toDate(COALESCE(endedAt, startedAt))) as day,
+      count() as cnt
+    FROM ${TABLE_LEGACY_SPANS}
+    GROUP BY day
+    ORDER BY day`,
+    format: "JSONEachRow"
+  });
+  const days = await daysResult.json();
+  let migratedRows = 0;
+  let batches = 0;
+  for (const { day, cnt } of days) {
+    const dayCount = Number(cnt);
+    await client.command({
+      query: `INSERT INTO ${TABLE_SPAN_EVENTS} (${columnList})
+        SELECT ${selectExprs}
+        FROM ${TABLE_LEGACY_SPANS}
+        WHERE toDate(COALESCE(endedAt, startedAt)) = {batchDate:Date}
+        ${dedupOrder}
+        LIMIT 1 BY "traceId", "spanId"`,
+      query_params: { batchDate: day }
+    });
+    migratedRows += dayCount;
+    batches++;
+    logger?.info?.(`Migrated batch ${day}: ${dayCount} rows`);
+  }
+  await client.command({
+    query: `CREATE TABLE IF NOT EXISTS ${TABLE_LEGACY_MIGRATION_DONE} (completedAt DateTime64(3, 'UTC')) ENGINE = MergeTree ORDER BY completedAt`
+  });
+  await client.command({
+    query: `INSERT INTO ${TABLE_LEGACY_MIGRATION_DONE} VALUES (now64(3))`
+  });
+  logger?.info?.(
+    `Legacy span migration complete: ${migratedRows} rows in ${batches} batches. The legacy table '${TABLE_LEGACY_SPANS}' has been preserved. Drop it manually after verifying the migration.`
+  );
+  return { migratedRows, batches };
+}
 async function migrateSignalTables(client, logger) {
   for (const { table, createDDL, idColumn } of SIGNAL_MIGRATIONS) {
     const engine = await getTableEngine(client, table);
@@ -7294,6 +7511,15 @@ var ObservabilityStorageClickhouseVNext = class extends ObservabilityStorage {
         })
       });
     }
+    try {
+      const legacyStatus = await checkLegacySpanMigrationStatus(this.#client);
+      if (legacyStatus.needsMigration) {
+        this.logger?.warn?.(
+          `Legacy span table 'mastra_ai_spans' detected. Run 'npx mastra migrate' to migrate historical spans to the v-next schema.`
+        );
+      }
+    } catch {
+    }
     try {
       await assertExistingTablesCompatibleWithReplication(this.#client, this.#replication);
       const existingStrategy = await detectExistingDeltaCursorStrategy(this.#client);
@@ -7368,34 +7594,44 @@ var ObservabilityStorageClickhouseVNext = class extends ObservabilityStorage {
     }
   }
   /**
-   * Manually migrate legacy signal tables to the signal-ID ReplacingMergeTree schema.
-   * The public method name is historical; the CLI still calls `migrateSpans()`
-   * for observability migrations even though this now also migrates signal tables.
+   * Manually migrate legacy tables to the v-next schema.
+   * Handles both signal table migrations (MergeTree → ReplacingMergeTree)
+   * and legacy span migration (mastra_ai_spans → mastra_span_events).
    */
   async migrateSpans() {
-    const migrationStatus = await checkSignalTablesMigrationStatus(this.#client);
-    if (!migrationStatus.needsMigration) {
-      return {
-        success: true,
-        alreadyMigrated: true,
-        duplicatesRemoved: 0,
-        message: "Migration already complete. Signal tables already use signal-ID dedupe keys."
-      };
+    const messages = [];
+    const signalStatus = await checkSignalTablesMigrationStatus(this.#client);
+    if (signalStatus.needsMigration) {
+      if (isReplicationConfigured(this.#replication)) {
+        throw new MastraError({
+          id: createStorageErrorId("CLICKHOUSE", "REPLICATION", "SIGNAL_TABLES_MIGRATION_UNSUPPORTED"),
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.USER,
+          text: "ClickHouse replication is enabled, so Mastra will not run copy-and-swap signal table migrations automatically. Migrate existing local signal tables manually before enabling replication."
+        });
+      }
+      await migrateSignalTables(this.#client, this.logger);
+      messages.push(`Migrated signal tables: ${signalStatus.tables.map((t) => t.table).join(", ")}.`);
     }
-    if (isReplicationConfigured(this.#replication)) {
-      throw new MastraError({
-        id: createStorageErrorId("CLICKHOUSE", "REPLICATION", "SIGNAL_TABLES_MIGRATION_UNSUPPORTED"),
-        domain: ErrorDomain.STORAGE,
-        category: ErrorCategory.USER,
-        text: "ClickHouse replication is enabled, so Mastra will not run copy-and-swap signal table migrations automatically. Migrate existing local signal tables manually before enabling replication."
-      });
+    const legacyStatus = await checkLegacySpanMigrationStatus(this.#client);
+    if (legacyStatus.needsMigration) {
+      if (isReplicationConfigured(this.#replication)) {
+        throw new MastraError({
+          id: createStorageErrorId("CLICKHOUSE", "REPLICATION", "LEGACY_SPAN_MIGRATION_UNSUPPORTED"),
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.USER,
+          text: "ClickHouse replication is enabled. Migrate legacy mastra_ai_spans manually before enabling replication."
+        });
+      }
+      const result = await migrateLegacySpans(this.#client, this.logger);
+      messages.push(`Migrated ${result.migratedRows} legacy spans in ${result.batches} batches.`);
     }
-    await migrateSignalTables(this.#client, this.logger);
+    const alreadyMigrated = !signalStatus.needsMigration && !legacyStatus.needsMigration;
     return {
       success: true,
-      alreadyMigrated: false,
+      alreadyMigrated,
       duplicatesRemoved: 0,
-      message: `Migration complete. Migrated signal tables: ${migrationStatus.tables.map((t) => t.table).join(", ")}.`
+      message: alreadyMigrated ? "Migration already complete." : `Migration complete. ${messages.join(" ")}`
     };
   }
   // -------------------------------------------------------------------------