@mastra/clickhouse 1.4.2-alpha.0 → 1.5.0

package/dist/index.js

@@ -2428,6 +2428,49 @@ time for large tables. Please ensure you have a backup before proceeding.
       );
     }
   }
+  async getTraceLight(args) {
+    const { traceId } = args;
+    try {
+      const engine = TABLE_ENGINES[TABLE_SPANS] ?? "MergeTree()";
+      const result = await this.client.query({
+        query: `
+          SELECT traceId, spanId, parentSpanId, name,
+            entityType, entityId, entityName,
+            spanType, error, isEvent,
+            startedAt, endedAt, createdAt, updatedAt
+          FROM ${TABLE_SPANS} ${engine.startsWith("ReplacingMergeTree") ? "FINAL" : ""}
+          WHERE traceId = {traceId:String}
+          ORDER BY startedAt ASC
+        `,
+        query_params: { traceId },
+        format: "JSONEachRow",
+        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 rows = await result.json();
+      if (!rows || rows.length === 0) {
+        return null;
+      }
+      return {
+        traceId,
+        spans: transformRows(rows)
+      };
+    } catch (error) {
+      throw new MastraError(
+        {
+          id: createStorageErrorId("CLICKHOUSE", "GET_TRACE_LIGHT", "FAILED"),
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: { traceId }
+        },
+        error
+      );
+    }
+  }
   async updateSpan(args) {
     const { traceId, spanId, updates } = args;
     try {
@@ -2942,6 +2985,7 @@ CREATE TABLE IF NOT EXISTS ${TABLE_METRIC_EVENTS} (
   timestamp          DateTime64(3, 'UTC'),
 
   -- IDs
+  metricId           String,
   traceId            Nullable(String),
   spanId             Nullable(String),
   experimentId       Nullable(String),
@@ -2989,9 +3033,9 @@ CREATE TABLE IF NOT EXISTS ${TABLE_METRIC_EVENTS} (
   metadata           Nullable(String),
   scope              Nullable(String)
 )
-ENGINE = MergeTree
+ENGINE = ReplacingMergeTree
 PARTITION BY toDate(timestamp)
-ORDER BY (name, timestamp)
+ORDER BY (name, timestamp, metricId)
 `;
 var LOG_EVENTS_DDL = `
 CREATE TABLE IF NOT EXISTS ${TABLE_LOG_EVENTS} (
@@ -2999,6 +3043,7 @@ CREATE TABLE IF NOT EXISTS ${TABLE_LOG_EVENTS} (
   timestamp          DateTime64(3, 'UTC'),
 
   -- IDs
+  logId              String,
   traceId            Nullable(String),
   spanId             Nullable(String),
   experimentId       Nullable(String),
@@ -3041,10 +3086,9 @@ CREATE TABLE IF NOT EXISTS ${TABLE_LOG_EVENTS} (
   metadata           Nullable(String),
   scope              Nullable(String)
 )
-ENGINE = MergeTree
+ENGINE = ReplacingMergeTree
 PARTITION BY toDate(timestamp)
-ORDER BY (timestamp, traceId)
-SETTINGS allow_nullable_key = 1
+ORDER BY (timestamp, logId)
 `;
 var SCORE_EVENTS_DDL = `
 CREATE TABLE IF NOT EXISTS ${TABLE_SCORE_EVENTS} (
@@ -3052,6 +3096,7 @@ CREATE TABLE IF NOT EXISTS ${TABLE_SCORE_EVENTS} (
   timestamp          DateTime64(3, 'UTC'),
 
   -- IDs
+  scoreId            String,
   traceId            Nullable(String),
   spanId             Nullable(String),
   experimentId       Nullable(String),
@@ -3101,9 +3146,9 @@ CREATE TABLE IF NOT EXISTS ${TABLE_SCORE_EVENTS} (
   metadata           Nullable(String),
   scope              Nullable(String)
 )
-ENGINE = MergeTree
+ENGINE = ReplacingMergeTree
 PARTITION BY toDate(timestamp)
-ORDER BY (traceId, timestamp)
+ORDER BY (traceId, timestamp, scoreId)
 SETTINGS allow_nullable_key = 1
 `;
 var FEEDBACK_EVENTS_DDL = `
@@ -3112,6 +3157,7 @@ CREATE TABLE IF NOT EXISTS ${TABLE_FEEDBACK_EVENTS} (
   timestamp          DateTime64(3, 'UTC'),
 
   -- IDs
+  feedbackId         String,
   traceId            Nullable(String),
   spanId             Nullable(String),
   experimentId       Nullable(String),
@@ -3164,9 +3210,9 @@ CREATE TABLE IF NOT EXISTS ${TABLE_FEEDBACK_EVENTS} (
   metadata           Nullable(String),
   scope              Nullable(String)
 )
-ENGINE = MergeTree
+ENGINE = ReplacingMergeTree
 PARTITION BY toDate(timestamp)
-ORDER BY (traceId, timestamp)
+ORDER BY (traceId, timestamp, feedbackId)
 SETTINGS allow_nullable_key = 1
 `;
 var DISCOVERY_VALUES_DDL = `
@@ -3521,6 +3567,7 @@ function spanRecordToRow(span) {
 }
 function rowToLogRecord(row) {
   return {
+    logId: row.logId,
     timestamp: toDate(row.timestamp),
     level: row.level,
     message: row.message,
@@ -3557,6 +3604,7 @@ function rowToLogRecord(row) {
 }
 function logRecordToRow(log) {
   return {
+    logId: log.logId,
     timestamp: toISOString(log.timestamp),
     level: log.level,
     message: log.message,
@@ -3593,6 +3641,7 @@ function logRecordToRow(log) {
 }
 function rowToMetricRecord(row) {
   return {
+    metricId: row.metricId,
     timestamp: toDate(row.timestamp),
     name: row.name,
     value: Number(row.value),
@@ -3634,6 +3683,7 @@ function rowToMetricRecord(row) {
 }
 function metricRecordToRow(metric) {
   return {
+    metricId: metric.metricId,
     timestamp: toISOString(metric.timestamp),
     name: metric.name,
     value: metric.value,
@@ -3675,6 +3725,7 @@ function metricRecordToRow(metric) {
 }
 function rowToScoreRecord(row) {
   return {
+    scoreId: row.scoreId,
     timestamp: toDate(row.timestamp),
     // Core score/feedback shapes still type traceId as required for now.
     traceId: nullableString(row.traceId),
@@ -3717,6 +3768,7 @@ function scoreRecordToRow(score) {
   const metadata = score.metadata ?? null;
   const scoreSource = score.scoreSource ?? score.source ?? null;
   return {
+    scoreId: score.scoreId,
     timestamp: toISOString(score.timestamp),
     traceId: score.traceId ?? null,
     spanId: score.spanId ?? null,
@@ -3759,6 +3811,7 @@ function rowToFeedbackRecord(row) {
   const feedbackSource = nullableString(row.feedbackSource);
   const feedbackUserId = nullableString(row.feedbackUserId) ?? nullableString(row.userId);
   return {
+    feedbackId: row.feedbackId,
     timestamp: toDate(row.timestamp),
     // Core score/feedback shapes still type traceId as required for now.
     traceId: nullableString(row.traceId),
@@ -3802,6 +3855,7 @@ function feedbackRecordToRow(feedback) {
   const feedbackSource = feedback.feedbackSource ?? feedback.source ?? "";
   const feedbackUserId = feedback.feedbackUserId ?? feedback.userId ?? null;
   return {
+    feedbackId: feedback.feedbackId,
     timestamp: toISOString(feedback.timestamp),
     traceId: feedback.traceId ?? null,
     spanId: feedback.spanId ?? null,
@@ -4921,6 +4975,93 @@ async function getMetricLabelValues(client, args) {
   );
   return { values: rows.map((r) => r.value) };
 }
+var SIGNAL_MIGRATIONS = [
+  { table: TABLE_METRIC_EVENTS, createDDL: METRIC_EVENTS_DDL, idColumn: "metricId" },
+  { table: TABLE_LOG_EVENTS, createDDL: LOG_EVENTS_DDL, idColumn: "logId" },
+  { table: TABLE_SCORE_EVENTS, createDDL: SCORE_EVENTS_DDL, idColumn: "scoreId" },
+  { table: TABLE_FEEDBACK_EVENTS, createDDL: FEEDBACK_EVENTS_DDL, idColumn: "feedbackId" }
+];
+async function getTableEngine(client, table) {
+  const result = await client.query({
+    query: `SELECT engine FROM system.tables WHERE database = currentDatabase() AND name = {table:String}`,
+    query_params: { table },
+    format: "JSONEachRow"
+  });
+  const rows = await result.json();
+  return rows[0]?.engine ?? null;
+}
+async function getTableColumns(client, table) {
+  const result = await client.query({ query: `DESCRIBE TABLE ${table}`, format: "JSONEachRow" });
+  const rows = await result.json();
+  return rows.map((r) => r.name);
+}
+function buildTemporaryTableDDL(createDDL, table, tempTable) {
+  return createDDL.replace(`CREATE TABLE IF NOT EXISTS ${table}`, `CREATE TABLE ${tempTable}`);
+}
+async function dropTableIfExists(client, table) {
+  if (await getTableEngine(client, table) !== null) {
+    await client.command({ query: `DROP TABLE ${table}` });
+  }
+}
+function createMigrationError(args, error) {
+  return new MastraError(
+    {
+      id: createStorageErrorId("CLICKHOUSE", "MIGRATE_SIGNAL_TABLES", "FAILED"),
+      domain: ErrorDomain.STORAGE,
+      category: ErrorCategory.THIRD_PARTY,
+      details: args
+    },
+    error
+  );
+}
+async function checkSignalTablesMigrationStatus(client) {
+  const tables = [];
+  for (const { table, idColumn } of SIGNAL_MIGRATIONS) {
+    const engine = await getTableEngine(client, table);
+    if (!engine || engine === "ReplacingMergeTree") {
+      continue;
+    }
+    tables.push({ table, engine, idColumn });
+  }
+  return {
+    needsMigration: tables.length > 0,
+    tables
+  };
+}
+async function migrateSignalTables(client, logger) {
+  for (const { table, createDDL, idColumn } of SIGNAL_MIGRATIONS) {
+    const engine = await getTableEngine(client, table);
+    if (!engine || engine === "ReplacingMergeTree") continue;
+    logger?.info?.(`Migrating ${table} from ${engine} to ReplacingMergeTree with ${idColumn} column`);
+    const temp = `${table}_migrating_${Date.now()}`;
+    try {
+      await client.command({ query: buildTemporaryTableDDL(createDDL, table, temp) });
+      const newColumns = await getTableColumns(client, temp);
+      const currentColumns = new Set(await getTableColumns(client, table));
+      const columnList = newColumns.map((c) => `"${c}"`).join(", ");
+      const selectExprs = newColumns.map((c) => {
+        if (c === idColumn) {
+          return currentColumns.has(c) ? `COALESCE(nullIf("${c}", ''), toString(generateUUIDv4())) AS "${c}"` : `toString(generateUUIDv4()) AS "${c}"`;
+        }
+        return currentColumns.has(c) ? `"${c}"` : `NULL AS "${c}"`;
+      }).join(", ");
+      await client.command({
+        query: `INSERT INTO ${temp} (${columnList}) SELECT ${selectExprs} FROM ${table}`
+      });
+      await client.command({ query: `EXCHANGE TABLES ${temp} AND ${table}` });
+      await client.command({ query: `DROP TABLE ${temp}` });
+      logger?.info?.(`Successfully migrated ${table}`);
+    } catch (error) {
+      logger?.error?.(`Migration of ${table} failed: ${error.message}`);
+      try {
+        await dropTableIfExists(client, temp);
+      } catch (restoreError) {
+        logger?.error?.(`Failed to clean up temporary table ${temp}: ${restoreError.message}`);
+      }
+      throw createMigrationError({ table, idColumn }, error);
+    }
+  }
+}
 var SCORE_TYPED_COLUMNS = /* @__PURE__ */ new Set([
   "timestamp",
   "traceId",
@@ -5384,6 +5525,31 @@ async function getTrace(client, args) {
   const spans = rows.map(rowToSpanRecord);
   return { traceId: args.traceId, spans };
 }
+async function getTraceLight(client, args) {
+  const result = await client.query({
+    query: `
+      SELECT traceId, spanId, parentSpanId, name,
+        entityType, entityId, entityName,
+        spanType, error, isEvent,
+        startedAt, endedAt
+      FROM (
+        SELECT *
+        FROM ${TABLE_SPAN_EVENTS}
+        WHERE traceId = {traceId:String}
+        ORDER BY dedupeKey, endedAt DESC
+        LIMIT 1 BY dedupeKey
+      )
+      ORDER BY startedAt ASC
+    `,
+    query_params: { traceId: args.traceId },
+    format: "JSONEachRow",
+    clickhouse_settings: CH_SETTINGS
+  });
+  const rows = await result.json();
+  if (!rows || rows.length === 0) return null;
+  const spans = rows.map(rowToSpanRecord);
+  return { traceId: args.traceId, spans };
+}
 async function batchDeleteTraces(client, args) {
   if (args.traceIds.length === 0) return;
   const params = {};
@@ -5412,6 +5578,32 @@ async function batchDeleteTraces(client, args) {
 }
 
 // src/storage/domains/observability/v-next/index.ts
+function buildSignalMigrationRequiredMessage(args) {
+  const tableList = args.tables.map((table) => `  - ${table.table} (${table.engine})`).join("\n");
+  return `
+===========================================================================
+MIGRATION REQUIRED: ${args.store} observability signal tables need signal IDs
+===========================================================================
+
+The following signal tables still use the legacy schema and must be migrated
+before observability storage can initialize:
+
+${tableList}
+
+To fix this, run the manual migration command:
+
+  npx mastra migrate
+
+This command will:
+  1. Create replacement signal tables with signal-ID dedupe keys
+  2. Backfill missing signal IDs for legacy rows
+  3. Swap the migrated tables into place
+
+WARNING: This migration recreates the signal tables and may take significant
+time for large databases. Please ensure you have a backup before proceeding.
+===========================================================================
+`;
+}
 var ObservabilityStorageClickhouseVNext = class extends ObservabilityStorage {
   #client;
   #retention;
@@ -5425,6 +5617,18 @@ var ObservabilityStorageClickhouseVNext = class extends ObservabilityStorage {
   // Initialization
   // -------------------------------------------------------------------------
   async init() {
+    const migrationStatus = await checkSignalTablesMigrationStatus(this.#client);
+    if (migrationStatus.needsMigration) {
+      throw new MastraError({
+        id: createStorageErrorId("CLICKHOUSE", "MIGRATION_REQUIRED", "SIGNAL_TABLES"),
+        domain: ErrorDomain.STORAGE,
+        category: ErrorCategory.USER,
+        text: buildSignalMigrationRequiredMessage({
+          store: "ClickHouse",
+          tables: migrationStatus.tables.map(({ table, engine }) => ({ table, engine }))
+        })
+      });
+    }
     try {
       for (const ddl of [...ALL_TABLE_DDL, ...ALL_MV_DDL]) {
         await this.#client.command({ query: ddl });
@@ -5439,6 +5643,9 @@ var ObservabilityStorageClickhouseVNext = class extends ObservabilityStorage {
         }
       }
     } catch (error) {
+      if (error instanceof MastraError) {
+        throw error;
+      }
       throw new MastraError(
         {
           id: createStorageErrorId("CLICKHOUSE", "VNEXT_INIT", "FAILED"),
@@ -5460,6 +5667,29 @@ var ObservabilityStorageClickhouseVNext = class extends ObservabilityStorage {
     } catch {
     }
   }
+  /**
+   * 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.
+   */
+  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."
+      };
+    }
+    await migrateSignalTables(this.#client, this.logger);
+    return {
+      success: true,
+      alreadyMigrated: false,
+      duplicatesRemoved: 0,
+      message: `Migration complete. Migrated signal tables: ${migrationStatus.tables.map((t) => t.table).join(", ")}.`
+    };
+  }
   // -------------------------------------------------------------------------
   // Strategy
   // -------------------------------------------------------------------------
@@ -5555,6 +5785,22 @@ var ObservabilityStorageClickhouseVNext = class extends ObservabilityStorage {
       );
     }
   }
+  async getTraceLight(args) {
+    try {
+      return await getTraceLight(this.#client, args);
+    } catch (error) {
+      if (error instanceof MastraError) throw error;
+      throw new MastraError(
+        {
+          id: createStorageErrorId("CLICKHOUSE", "GET_TRACE_LIGHT", "FAILED"),
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: { traceId: args.traceId }
+        },
+        error
+      );
+    }
+  }
   async listTraces(args) {
     try {
       return await listTraces(this.#client, args);