@mastra/clickhouse 1.22.1 → 1.23.0-alpha.0

package/dist/index.js

@@ -6557,6 +6557,10 @@ const TRACE_STATUS_SQL = `if(isNotNull(r.error), 'error', 'success')`;
 function durationMsSql(startedAt, endedAt) {
 	return `dateDiff('millisecond', ${startedAt}, ${endedAt})`;
 }
+/** Top-level metadata string value on a trace root (alias `r`); `key` is a bound parameter placeholder. */
+function traceMetadataValueSql(key) {
+	return `coalesce(if(mapContains(r.metadataSearch, ${key}), r.metadataSearch[${key}], NULL), nullIf(trim(JSONExtractString(r.metadataRaw, ${key})), ''))`;
+}
 const TRACE_FIELDS = {
 	traceId: {
 		sql: "r.traceId",
@@ -6815,9 +6819,8 @@ function compileScalarPredicate(predicate, registry, parameters, allowMetadata =
 	if (predicate.type === "not") return `NOT (${compileScalarPredicate(predicate.arg, registry, parameters, allowMetadata)})`;
 	const field = isMetadataField(predicate.field) ? (() => {
 		if (!allowMetadata) throw new Error(`Unsupported trusted trace-query field: ${predicate.field}`);
-		const key = parameters.add(predicate.field.slice(9), "String");
 		return {
-			sql: `coalesce(if(mapContains(r.metadataSearch, ${key}), r.metadataSearch[${key}], NULL), nullIf(trim(JSONExtractString(r.metadataRaw, ${key})), ''))`,
+			sql: traceMetadataValueSql(parameters.add(predicate.field.slice(9), "String")),
 			parameterType: "String"
 		};
 	})() : fieldDefinition(registry, predicate.field);
@@ -6910,7 +6913,7 @@ function compileTenantScope(scope, parameters) {
 	if (scope.resourceId !== void 0) sql += `\n      AND resourceId = ${parameters.add(scope.resourceId, "String")}`;
 	return sql;
 }
-function compileClickHouseTraceScope(selection, relationCollections, parameters, scope) {
+function compileClickHouseTraceScope(selection, relationCollections, parameters, scope, seedConjuncts = []) {
 	const from = parameters.add(selection.timeRange.from, "DateTime64(3, 'UTC')");
 	const to = parameters.add(selection.timeRange.to, "DateTime64(3, 'UTC')");
 	const tenant = compileTenantScope(scope, parameters);
@@ -6920,9 +6923,9 @@ function compileClickHouseTraceScope(selection, relationCollections, parameters,
       FROM ${TABLE_TRACE_ROOTS}
       WHERE traceId IN (
         SELECT traceId
-        FROM ${TABLE_TRACE_ROOTS}
+        FROM ${TABLE_TRACE_ROOTS} r
         WHERE startedAt >= ${from}
-          AND startedAt < ${to}${tenant}
+          AND startedAt < ${to}${tenant}${seedConjuncts.map((conjunct) => `\n          AND (${compilePredicate(conjunct, parameters)})`).join("")}
       )
       ORDER BY dedupeKey
       LIMIT 1 BY dedupeKey
@@ -7015,16 +7018,25 @@ function parseDeltaWatermark(value) {
 		throw new coreStorage.TraceQueryCursorError("TRACE_QUERY_CURSOR_MALFORMED");
 	}
 }
-function compileClickHouseTraceQuery(plan, deltaHead) {
-	const parameters = new ParameterBuilder();
-	const ctes = compileClickHouseTraceScope(plan, collectRelationCollections(plan.where), parameters, plan.scope);
-	const predicate = plan.where ? compilePredicate(plan.where, parameters) : "1";
+/**
+* Builds the CTE chain ending in `candidates`: the completed, current trace roots in the
+* selection's time range and tenant scope that match its `where` predicate. Trace queries and
+* trace aggregates both select from this CTE so they always see the same population.
+*/
+function compileClickHouseTraceCandidates(selection, columns, parameters) {
+	const rootConjuncts = selection.where ? (selection.where.type === "boolean" && selection.where.operator === "and" ? selection.where.args : [selection.where]).filter((conjunct) => collectRelationCollections(conjunct).size === 0) : [];
+	const ctes = compileClickHouseTraceScope(selection, collectRelationCollections(selection.where), parameters, selection.scope, rootConjuncts);
+	const predicate = selection.where ? compilePredicate(selection.where, parameters) : "1";
 	ctes.push(`candidates AS (
-    SELECT ${TRACE_SELECT}
+    SELECT ${columns}
     FROM root_scope r
     WHERE ${predicate}
   )`);
-	const candidates = `WITH ${ctes.join(",\n")}`;
+	return ctes;
+}
+function compileClickHouseTraceQuery(plan, deltaHead) {
+	const parameters = new ParameterBuilder();
+	const candidates = `WITH ${compileClickHouseTraceCandidates(plan, TRACE_SELECT, parameters).join(",\n")}`;
 	if (plan.result === "groups") return {
 		query: `${candidates}
 SELECT threadId
@@ -7200,10 +7212,8 @@ function compileClickHouseTraceQueryValues(plan) {
 	const parameters = new ParameterBuilder();
 	const ctes = compileClickHouseTraceScope(plan, discoveryCollections(plan.predicateScope), parameters, plan.scope);
 	let field;
-	if (plan.predicateScope === "trace" && plan.path.startsWith("metadata.")) {
-		const key = parameters.add(plan.path.slice(9), "String");
-		field = `coalesce(if(mapContains(r.metadataSearch, ${key}), r.metadataSearch[${key}], NULL), nullIf(trim(JSONExtractString(r.metadataRaw, ${key})), ''))`;
-	} else if (plan.predicateScope === "trace" && plan.path === "tags") field = `arrayJoin(arrayDistinct(${TRACE_FIELDS.tags.sql}))`;
+	if (plan.predicateScope === "trace" && plan.path.startsWith("metadata.")) field = traceMetadataValueSql(parameters.add(plan.path.slice(9), "String"));
+	else if (plan.predicateScope === "trace" && plan.path === "tags") field = `arrayJoin(arrayDistinct(${TRACE_FIELDS.tags.sql}))`;
 	else field = fieldDefinition(discoveryRegistry(plan.predicateScope), plan.path).sql;
 	const search = plan.search ? `AND positionCaseInsensitiveUTF8(value, ${parameters.add(plan.search, "String")}) > 0` : "";
 	const limit = parameters.add(plan.limit + 1, "UInt64");
@@ -7411,6 +7421,185 @@ async function queryThreads(client, plan, timeoutMs) {
 	});
 }
 //#endregion
+//#region src/storage/domains/observability/v-next/trace-aggregate.ts
+/** Trace-root column expressions (alias `r`) for each groupable dimension; same values `where` sees. */
+const DIMENSION_SQL = {
+	entityType: "r.entityType",
+	entityName: "r.entityName",
+	environment: "r.environment",
+	status: TRACE_STATUS_SQL,
+	serviceName: "r.serviceName",
+	executionSource: "r.executionSource",
+	threadId: "r.threadId",
+	resourceId: "r.resourceId",
+	userId: "r.userId",
+	sessionId: "r.sessionId",
+	organizationId: "r.organizationId",
+	experimentId: "r.experimentId"
+};
+const PERCENTILES = {
+	"duration.p50": .5,
+	"duration.p90": .9,
+	"duration.p95": .95,
+	"duration.p99": .99
+};
+const COMPARISON_SQL = {
+	eq: "=",
+	ne: "!=",
+	lt: "<",
+	lte: "<=",
+	gt: ">",
+	gte: ">="
+};
+/**
+* Compiles a trusted aggregate plan. The plan's caps are enforced by the planner and trusted here.
+*
+* Every candidate root is first reduced to a `facts` row (dimension values, duration, error flag,
+* countDistinct inputs, bucket). Without `interval`, groups are filtered, ordered, and limited in
+* one grouped select. With `interval`, one `GROUPING SETS` pass computes both the whole-window
+* group rows and the per-bucket rows: group rows are ranked (having → orderBy), and each bucket
+* row inherits its group's rank. ClickHouse re-executes a CTE at every reference, so this keeps
+* the candidate scan to a single pass instead of joining bucket rows back to a ranked CTE.
+*/
+function compileClickHouseTraceAggregate(plan) {
+	const parameters = new ParameterBuilder();
+	const ctes = compileClickHouseTraceCandidates(plan, "*", parameters);
+	const dimensionSql = (field) => {
+		if (field.startsWith("metadata.")) return traceMetadataValueSql(parameters.add(field.slice(9), "String"));
+		const sql = Object.hasOwn(DIMENSION_SQL, field) ? DIMENSION_SQL[field] : void 0;
+		if (sql === void 0) throw new Error(`Unsupported trusted trace-aggregate field: ${field}`);
+		return sql;
+	};
+	const distinctSql = (field) => field === "traceId" ? "r.traceId" : dimensionSql(field);
+	const dimensionColumns = plan.dimensions.map((_, index) => `d${index}`);
+	const distinctColumns = /* @__PURE__ */ new Map();
+	const factColumns = plan.dimensions.map((dimension, index) => `CAST(${dimensionSql(dimension)} AS Nullable(String)) AS d${index}`);
+	for (const measure of plan.measures) {
+		if (measure.type !== "countDistinct" || distinctColumns.has(measure.field)) continue;
+		const column = `cd${distinctColumns.size}`;
+		distinctColumns.set(measure.field, column);
+		factColumns.push(`CAST(${distinctSql(measure.field)} AS Nullable(String)) AS ${column}`);
+	}
+	factColumns.push(`toFloat64(${durationMsSql("r.startedAt", "r.endedAt")}) AS durationMs`, `isNotNull(r.error) AS isError`, `cityHash64(r.traceId) AS traceSeed`);
+	if (plan.interval !== void 0) {
+		const intervalMs = coreStorage.TRACE_AGGREGATE_INTERVAL_MS[plan.interval];
+		if (intervalMs === void 0) throw new Error(`Unsupported trusted trace-aggregate interval: ${plan.interval}`);
+		factColumns.push(`fromUnixTimestamp64Milli(intDiv(toUnixTimestamp64Milli(r.startedAt), ${intervalMs}) * ${intervalMs}, 'UTC') AS bucket`);
+	}
+	ctes.push(`facts AS (
+    SELECT ${factColumns.join(",\n      ")}
+    FROM candidates r
+  )`);
+	const measureSql = (name) => {
+		if (name === "count") return "toFloat64(count())";
+		if (name === "errorCount") return "toFloat64(countIf(isError))";
+		if (name === "errorRate") return "countIf(isError) / count()";
+		if (name === "duration.avg") return "avg(durationMs)";
+		if (name === "duration.min") return "min(durationMs)";
+		if (name === "duration.max") return "max(durationMs)";
+		const percentile = PERCENTILES[name];
+		if (percentile !== void 0) return `quantileDeterministic(${percentile})(durationMs, traceSeed)`;
+		const measure = plan.measures.find((candidate) => candidate.type === "countDistinct" && candidate.name === name);
+		const distinctColumn = measure && distinctColumns.get(measure.field);
+		if (distinctColumn === void 0) throw new Error(`Unsupported trusted trace-aggregate measure: ${name}`);
+		return `toFloat64(uniqExact(${distinctColumn}))`;
+	};
+	const compileHaving = (predicate) => {
+		if (predicate.type === "boolean") return predicate.args.map((arg) => `(${compileHaving(arg)})`).join(predicate.operator === "and" ? " AND " : " OR ");
+		if (predicate.type === "not") return `NOT (${compileHaving(predicate.arg)})`;
+		const measure = measureSql(predicate.measure);
+		if (predicate.type === "membership") {
+			if (predicate.values.length === 0) return predicate.operator === "in" ? "0" : "1";
+			const list = predicate.values.map((value) => parameters.add(value, "Float64")).join(", ");
+			return `${measure} ${predicate.operator === "in" ? "IN" : "NOT IN"} (${list})`;
+		}
+		const operator = COMPARISON_SQL[predicate.operator];
+		if (operator === void 0) throw new Error(`Unsupported trusted trace-aggregate operator: ${predicate.operator}`);
+		return `${measure} ${operator} ${parameters.add(predicate.value, "Float64")}`;
+	};
+	const havingSql = ["count() > 0", ...plan.having ? [`(${compileHaving(plan.having)})`] : []].join(" AND ");
+	const direction = plan.orderBy.direction === "asc" ? "ASC" : "DESC";
+	let orderTarget;
+	if (plan.orderBy.target === "measure") orderTarget = measureSql(plan.orderBy.measure);
+	else {
+		const index = plan.dimensions.indexOf(plan.orderBy.dimension);
+		if (index === -1) throw new Error(`Unsupported trusted trace-aggregate order dimension: ${plan.orderBy.dimension}`);
+		orderTarget = dimensionColumns[index];
+	}
+	const tiebreakSql = dimensionColumns.map((column) => `${column} ASC NULLS LAST`);
+	const measureColumns = plan.measures.map((measure, index) => `${measureSql(measure.name)} AS m${index}`);
+	const measureNames = plan.measures.map((_, index) => `m${index}`);
+	if (plan.interval === void 0) {
+		const limit = parameters.add(plan.limit + 1, "UInt64");
+		return {
+			query: `WITH ${ctes.join(",\n")}
+SELECT ${[...dimensionColumns, ...measureColumns].join(", ")}
+FROM facts
+${dimensionColumns.length > 0 ? `GROUP BY ${dimensionColumns.join(", ")}\n` : ""}HAVING ${havingSql}
+ORDER BY ${[`${orderTarget} ${direction} NULLS LAST`, ...tiebreakSql].join(", ")}
+LIMIT ${limit}`,
+			query_params: parameters.params
+		};
+	}
+	const maxRank = parameters.add(plan.limit, "UInt64");
+	const groupOrderSql = [`__order ${direction} NULLS LAST`, ...tiebreakSql].join(", ");
+	const groupPartition = dimensionColumns.length > 0 ? `PARTITION BY ${dimensionColumns.join(", ")}` : "";
+	ctes.push(`grouped AS (
+    SELECT ${[
+		...dimensionColumns,
+		"bucket",
+		"grouping(bucket) = 1 AS __isGroup",
+		...measureColumns,
+		`(${havingSql}) AS __keep`,
+		`${orderTarget} AS __order`
+	].join(", ")}
+    FROM facts
+    GROUP BY GROUPING SETS ((${[...dimensionColumns, "bucket"].join(", ")}), (${dimensionColumns.join(", ")}))
+  )`, `ranked AS (
+    SELECT *,
+      if(__isGroup AND __keep, row_number() OVER (PARTITION BY __isGroup, __keep ORDER BY ${groupOrderSql}), 0) AS __groupRank,
+      countIf(__isGroup AND __keep) OVER () AS __survivingGroups
+    FROM grouped
+  )`, `expanded AS (
+    SELECT *, maxIf(__groupRank, __isGroup) OVER (${groupPartition}) AS __rank
+    FROM ranked
+  )`);
+	return {
+		query: `WITH ${ctes.join(",\n")}
+SELECT ${[
+			...dimensionColumns,
+			"bucket",
+			...measureNames,
+			`toUInt8(__survivingGroups > ${maxRank}) AS truncated`
+		].join(", ")}
+FROM expanded
+WHERE NOT __isGroup AND __rank BETWEEN 1 AND ${maxRank}
+ORDER BY __rank ASC, bucket ASC`,
+		query_params: parameters.params
+	};
+}
+async function aggregateTraces(client, plan, timeoutMs) {
+	const result = await runWithClickHouseTraceQueryTimeout(client, { timeoutMs }, compileClickHouseTraceAggregate(plan));
+	let truncated;
+	let resultRows;
+	if (plan.interval === void 0) {
+		truncated = result.length > plan.limit;
+		resultRows = result.slice(0, plan.limit);
+	} else {
+		truncated = Number(result[0]?.truncated ?? 0) === 1;
+		resultRows = result;
+	}
+	return {
+		rows: resultRows.map((row) => {
+			const shaped = { measures: Object.fromEntries(plan.measures.map((measure, index) => [measure.name, Number(row[`m${index}`])])) };
+			if (plan.dimensions.length > 0) shaped.dimensions = Object.fromEntries(plan.dimensions.map((dimension, index) => [dimension, row[`d${index}`] ?? null]));
+			if (plan.interval !== void 0) shaped.bucket = new Date(row.bucket).toISOString();
+			return shaped;
+		}),
+		truncated
+	};
+}
+//#endregion
 //#region src/storage/domains/observability/v-next/trace-roots.ts
 /**
 * Get the root span for a trace, reading from trace_roots as compatibility path.
@@ -8658,7 +8847,8 @@ var ObservabilityStorageClickhouseVNext = class extends ObservabilityStorage {
 			"thread-query",
 			"trace-query-tenant-scope",
 			"feedback",
-			"trace-query-context-ids"
+			"trace-query-context-ids",
+			"trace-aggregate"
 		];
 		return [
 			"metrics",
@@ -8676,7 +8866,8 @@ var ObservabilityStorageClickhouseVNext = class extends ObservabilityStorage {
 			"thread-query",
 			"trace-query-tenant-scope",
 			"feedback",
-			"trace-query-context-ids"
+			"trace-query-context-ids",
+			"trace-aggregate"
 		];
 	}
 	async createSpan(args) {
@@ -8803,6 +8994,18 @@ var ObservabilityStorageClickhouseVNext = class extends ObservabilityStorage {
 			}, error);
 		}
 	}
+	async aggregateTraces(plan) {
+		try {
+			return await aggregateTraces(this.#client, plan, this.#traceQueryTimeoutMs);
+		} catch (error) {
+			if (error instanceof MastraError || error instanceof coreStorage.TraceQueryExecutionError || error instanceof coreStorage.TraceQueryResourceLimitError) throw error;
+			throw new MastraError({
+				id: createStorageErrorId("CLICKHOUSE", "AGGREGATE_TRACES", "FAILED"),
+				domain: ErrorDomain.STORAGE,
+				category: ErrorCategory.THIRD_PARTY
+			}, error);
+		}
+	}
 	async getTraceQueryObservedFields(plan) {
 		try {
 			return await getTraceQueryObservedFields(this.#client, plan, this.#traceQueryDiscoveryLimits);