@mastra/clickhouse 1.25.1 → 1.26.0-alpha.1

package/dist/index.js

@@ -3815,9 +3815,27 @@ function readPreview(preview, payload, jsonEncoded = false) {
 }
 /** Older response builders may shorten previews or turn empty strings into null. */
 function buildSpanQueryResponse(...args) {
-	if (canGeneratePreviews) return coreStorage.buildSpanQueryResponse(...args);
-	const [plan, selected, payloads, metrics] = args;
-	const response = coreStorage.buildSpanQueryResponse(plan, selected, [], metrics);
+	const [plan, selected, payloads, metrics = []] = args;
+	const compatibleSelected = coreStorage.SPAN_QUERY_FIELD_REGISTRY ? selected : selected.map((row) => ({
+		organizationId: row.organizationId,
+		resourceId: row.resourceId,
+		traceId: row.traceId,
+		spanId: row.spanId,
+		parentSpanId: row.parentSpanId,
+		name: row.name,
+		spanType: row.spanType,
+		status: row.status,
+		startedAt: row.startedAt,
+		endedAt: row.endedAt,
+		durationMs: row.durationMs,
+		entityType: row.entityType,
+		entityId: row.entityId,
+		entityName: row.entityName,
+		model: row.model,
+		provider: row.provider
+	}));
+	if (canGeneratePreviews) return coreStorage.buildSpanQueryResponse(plan, compatibleSelected, payloads, metrics);
+	const response = coreStorage.buildSpanQueryResponse(plan, compatibleSelected, [], metrics);
 	const byIdentity = new Map(payloads.map((payload) => [coreStorage.spanQueryIdentityKey(payload), payload]));
 	for (const span of response.spans) {
 		const payload = byIdentity.get(coreStorage.spanQueryIdentityKey(span));
@@ -4769,6 +4787,7 @@ function buildMetricsFilterConditions(filters, tableAlias) {
 	assertNoDeprecatedSourceFilter(filters.source, "executionSource", "metrics");
 	addCommonFilterFields(filters, tableAlias, out);
 	addIn(col("traceId"), filters.traceIds, "traceIds", out);
+	addIn(col("experimentId"), filters.experimentIds, "experimentIds", out);
 	addIn(col("name"), filters.name, "metricNames", out);
 	addEq(col("provider"), filters.provider, "provider", "String", out);
 	addEq(col("model"), filters.model, "model", "String", out);
@@ -6962,6 +6981,73 @@ const SPAN_FIELDS = {
 		parameterType: "String"
 	}
 };
+const SPAN_QUERY_FIELDS = {
+	...SPAN_FIELDS,
+	traceId: {
+		sql: "s.traceId",
+		parameterType: "String"
+	},
+	spanId: {
+		sql: "s.spanId",
+		parameterType: "String"
+	},
+	parentSpanId: {
+		sql: "s.parentSpanId",
+		parameterType: "String"
+	},
+	threadId: {
+		sql: "s.threadId",
+		parameterType: "String"
+	},
+	resourceId: {
+		sql: "s.resourceId",
+		parameterType: "String"
+	},
+	requestId: {
+		sql: "s.requestId",
+		parameterType: "String"
+	},
+	serviceName: {
+		sql: "s.serviceName",
+		parameterType: "String"
+	},
+	experimentId: {
+		sql: "s.experimentId",
+		parameterType: "String"
+	},
+	environment: {
+		sql: "s.environment",
+		parameterType: "String"
+	},
+	executionSource: {
+		sql: "s.executionSource",
+		parameterType: "String"
+	},
+	parentEntityType: {
+		sql: "s.parentEntityType",
+		parameterType: "String"
+	},
+	parentEntityId: {
+		sql: "s.parentEntityId",
+		parameterType: "String"
+	},
+	parentEntityName: {
+		sql: "s.parentEntityName",
+		parameterType: "String"
+	},
+	rootEntityType: {
+		sql: "s.rootEntityType",
+		parameterType: "String"
+	},
+	rootEntityId: {
+		sql: "s.rootEntityId",
+		parameterType: "String"
+	},
+	rootEntityName: {
+		sql: "s.rootEntityName",
+		parameterType: "String"
+	}
+};
 const SCORE_FIELDS = {
 	scorerId: {
 		sql: "s.scorerId",
@@ -7833,11 +7919,11 @@ async function queryThreads(client, plan, timeoutMs) {
 		}) : null }
 	});
 }
-/** Compile a span-row filter with exactly the same rules as trace span predicates. */
+/** Compile a direct span-query filter using its dedicated field mapping. */
 function compileSpanQueryPredicate(predicate) {
 	const parameters = new ParameterBuilder();
 	return {
-		sql: compileScalarPredicate(predicate, SPAN_FIELDS, parameters),
+		sql: compileScalarPredicate(predicate, SPAN_QUERY_FIELDS, parameters),
 		params: parameters.params
 	};
 }
@@ -7851,7 +7937,6 @@ const identityFields = [
 ];
 const identities = identityFields.join(", ");
 const identitySetColumns = "isNull(organizationId), ifNull(organizationId, ''), isNull(resourceId), ifNull(resourceId, ''), traceId, spanId";
-const identitySetKey = `tuple(${identitySetColumns})`;
 const displayFields = [
 	"parentSpanId",
 	"name",
@@ -7859,6 +7944,24 @@ const displayFields = [
 	"entityType",
 	"entityId",
 	"entityName",
+	"entityVersionId",
+	"parentEntityType",
+	"parentEntityId",
+	"parentEntityName",
+	"parentEntityVersionId",
+	"rootEntityType",
+	"rootEntityId",
+	"rootEntityName",
+	"rootEntityVersionId",
+	"threadId",
+	"runId",
+	"sessionId",
+	"userId",
+	"requestId",
+	"serviceName",
+	"experimentId",
+	"environment",
+	"executionSource",
 	"startedAt",
 	"endedAt",
 	"model",
@@ -7937,15 +8040,11 @@ function compileClickHouseSpanQuery(plan) {
 	const expressions = selectionFields.map(fieldExpression);
 	const after = afterCursor(plan, parameters);
 	return {
-		query: `WITH candidate_ids AS (
-    SELECT ${identitySetKey} AS identity FROM ${TABLE_SPAN_EVENTS}
-    PREWHERE endedAt >= ${from}
-    WHERE startedAt >= ${from} AND startedAt < ${to} ${scope.length ? `AND ${scope.join(" AND ")}` : ""}
-  ), candidates AS (
+		query: `WITH candidates AS (
     SELECT ${identities}, ${selectionFields.map((field, index) => `${expressions[index]} AS ${field}`).join(", ")}
     FROM ${TABLE_SPAN_EVENTS} e
     PREWHERE e.endedAt >= ${from}
-    WHERE ${identitySetKey} IN (SELECT identity FROM candidate_ids) ${scope.length ? `AND ${scope.join(" AND ")}` : ""}
+    WHERE e.startedAt >= ${from} AND e.startedAt < ${to} ${scope.length ? `AND ${scope.join(" AND ")}` : ""}
     ORDER BY ${identities}, e.endedAt DESC
     LIMIT 1 BY ${identities}
   )
@@ -7963,6 +8062,7 @@ function compileClickHouseSpanQuery(plan) {
 		}
 	};
 }
+const LEGACY_PAYLOAD_BATCH_ROWS = 32;
 function iso(value) {
 	return new Date(value.includes("T") ? value : `${value.replace(" ", "T")}Z`).toISOString();
 }
@@ -7996,34 +8096,25 @@ async function querySpans(client, plan, timeoutMs) {
 	});
 	const missingKeys = new Set(rowsWithPreviews.filter(needsPreviewPayloads).map(coreStorage.spanQueryIdentityKey));
 	const missing = page.filter((row) => missingKeys.has(coreStorage.spanQueryIdentityKey(row)));
-	const previewPayloads = /* @__PURE__ */ new Map();
-	if (missing.length) {
+	const storedByIdentity = new Map(rowsWithPreviews.map((row) => [coreStorage.spanQueryIdentityKey(row), row]));
+	for (let start = 0; start < missing.length; start += LEGACY_PAYLOAD_BATCH_ROWS) {
 		const rows = await run({
 			query: `SELECT ${identities}, data.1 AS input, data.2 AS output, data.3 AS error
       FROM (SELECT ${identities}, any(tuple(input, output, error)) AS data
         FROM ${TABLE_SPAN_EVENTS} e WHERE ${payloadWhere} GROUP BY ${identities})`,
-			query_params: { page: pageParameter(missing) }
+			query_params: { page: pageParameter(missing.slice(start, start + LEGACY_PAYLOAD_BATCH_ROWS)) }
 		});
-		for (const row of rows) previewPayloads.set(coreStorage.spanQueryIdentityKey(row), row);
+		for (const row of rows) {
+			const stored = storedByIdentity.get(coreStorage.spanQueryIdentityKey(row));
+			if (stored) Object.assign(stored, resolveStoredPreviews(stored, row));
+		}
 	}
 	const payloads = rowsWithPreviews.map((row) => ({
 		...row,
-		...resolveStoredPreviews(row, previewPayloads.get(coreStorage.spanQueryIdentityKey(row)))
+		...resolveStoredPreviews(row)
 	}));
-	const metricWhere = `${identitySetKey} IN {page:Array(Tuple(UInt8, String, UInt8, String, String, String))}`;
-	const metrics = await run({
-		query: `SELECT ${identities}, data.1 AS name, data.2 AS estimatedCost,
-      data.3 AS costUnit, data.4 AS allocation, data.5 AS costError
-    FROM (SELECT ${identities}, metricId, argMax(tuple(name, estimatedCost, costUnit,
-        JSONExtractString(costMetadata, 'allocation'), JSONHas(costMetadata, 'error')), timestamp) AS data
-      FROM ${TABLE_METRIC_EVENTS} WHERE (${metricWhere}) AND name LIKE 'mastra_model_%'
-      GROUP BY ${identities}, metricId)
-    LIMIT ${coreStorage.SPAN_QUERY_MAX_COST_METRICS + 1}`,
-		query_params: { page: page.map((row) => new TupleParam(identityValues(row))) }
-	});
-	if (metrics.length > coreStorage.SPAN_QUERY_MAX_COST_METRICS) throw new coreStorage.TraceQueryResourceLimitError();
 	const byIdentity = new Map(payloads.map((row) => [coreStorage.spanQueryIdentityKey(row), row]));
-	const result = buildSpanQueryResponse(plan, page.flatMap((identity) => {
+	const hydrated = page.flatMap((identity) => {
 		const row = byIdentity.get(coreStorage.spanQueryIdentityKey(identity));
 		if (!row) return [];
 		const { inputPreview: _input, outputPreview: _output, errorPreview: _error, ...metadata } = row;
@@ -8032,7 +8123,22 @@ async function querySpans(client, plan, timeoutMs) {
 			...identity,
 			durationMs: Number(row.durationMs)
 		}];
-	}), payloads, metrics);
+	});
+	let metrics = [];
+	if (!coreStorage.SPAN_QUERY_FIELD_REGISTRY) {
+		metrics = await run({
+			query: `SELECT ${identities}, data.1 AS name, data.2 AS estimatedCost,
+        data.3 AS costUnit, data.4 AS allocation, data.5 AS costError
+      FROM (SELECT ${identities}, metricId, argMax(tuple(name, estimatedCost, costUnit,
+          JSONExtractString(costMetadata, 'allocation'), JSONHas(costMetadata, 'error')), timestamp) AS data
+        FROM ${TABLE_METRIC_EVENTS} WHERE (${`tuple(${identitySetColumns}) IN {page:Array(Tuple(UInt8, String, UInt8, String, String, String))}`}) AND name LIKE 'mastra_model_%'
+        GROUP BY ${identities}, metricId)
+      LIMIT ${coreStorage.SPAN_QUERY_MAX_COST_METRICS + 1}`,
+			query_params: { page: page.map((row) => new TupleParam(identityValues(row))) }
+		});
+		if (metrics.length > coreStorage.SPAN_QUERY_MAX_COST_METRICS) throw new coreStorage.TraceQueryResourceLimitError();
+	}
+	const result = buildSpanQueryResponse(plan, hydrated, payloads, metrics);
 	const last = page.at(-1);
 	result.page.next = selected.length > plan.limit ? coreStorage.encodeSpanQueryCursor(plan, {
 		organizationId: last.organizationId,
@@ -9191,7 +9297,7 @@ function buildBranchCursor(row) {
 }
 //#endregion
 //#region src/storage/domains/observability/v-next/index.ts
-const spanQueryFeatures = typeof coreStorage.planSpanQuery === "function" ? ["span-query"] : [];
+const spanQueryFeatures = typeof coreStorage.planSpanQuery === "function" ? ["span-query", "span-query-context-fields"] : [];
 function buildSignalMigrationRequiredMessage(args) {
 	const tableList = args.tables.map((table) => `  - ${table.table} (${table.engine})`).join("\n");
 	return `