@mastra/clickhouse 1.16.1 → 1.17.0-alpha.0

package/dist/index.js

@@ -1,5 +1,6 @@
 import { createClient } from "@clickhouse/client";
 import { ErrorCategory, ErrorDomain, MastraError } from "@mastra/core/error";
+import * as coreStorage from "@mastra/core/storage";
 import { BRANCH_SPAN_TYPES, EntityType, METRIC_DISTINCT_COLUMNS, MastraCompositeStore, MemoryStorage, ObservabilityStorage, SCORERS_SCHEMA, SPAN_SCHEMA, ScoresStorage, TABLE_AGENT_VERSIONS, TABLE_DATASETS, TABLE_DATASET_ITEMS, TABLE_DATASET_VERSIONS, TABLE_EXPERIMENTS, TABLE_EXPERIMENT_RESULTS, TABLE_FAVORITES, TABLE_HARNESS_SESSIONS, TABLE_KNOWLEDGE_ACTIVITY, TABLE_KNOWLEDGE_CURSORS, TABLE_KNOWLEDGE_MENTIONS, TABLE_KNOWLEDGE_NODES, TABLE_KNOWLEDGE_RECORDS, TABLE_KNOWLEDGE_SEMANTIC_OUTBOX, TABLE_MCP_CLIENTS, TABLE_MCP_CLIENT_VERSIONS, TABLE_MCP_SERVERS, TABLE_MCP_SERVER_VERSIONS, TABLE_MESSAGES, TABLE_NOTIFICATIONS, TABLE_PROMPT_BLOCKS, TABLE_PROMPT_BLOCK_VERSIONS, TABLE_RESOURCES, TABLE_SCHEDULES, TABLE_SCHEDULE_TRIGGERS, TABLE_SCHEMAS, TABLE_SCORERS, TABLE_SCORER_DEFINITIONS, TABLE_SCORER_DEFINITION_VERSIONS, TABLE_SKILLS, TABLE_SKILL_BLOBS, TABLE_SKILL_VERSIONS, TABLE_SPANS, TABLE_THREADS, TABLE_THREAD_STATE, TABLE_TOOL_PROVIDER_CONNECTIONS, TABLE_TRACES, TABLE_WORKFLOW_DEFINITIONS, TABLE_WORKFLOW_SNAPSHOT, TABLE_WORKSPACES, TABLE_WORKSPACE_VERSIONS, TraceStatus, WorkflowsStorage, buildInputPreview, calculatePagination, computeTraceStatus, createStorageErrorId, getDefaultValue, getSqlType, listBranchesArgsSchema, listFeedbackArgsSchema, listLogsArgsSchema, listMetricsArgsSchema, listScoresArgsSchema, listTracesArgsSchema, normalizePerPage, safelyParseJSON, toTraceSpans, transformScoreRow, validateStorageMetadataFilter } from "@mastra/core/storage";
 import { MessageList } from "@mastra/core/agent";
 import { MastraBase } from "@mastra/core/base";
@@ -6091,6 +6092,269 @@ async function getScorePercentiles(client, args) {
 	return { series };
 }
 //#endregion
+//#region src/storage/domains/observability/v-next/trace-query.ts
+const TRACE_STATUS_SQL = `if(isNotNull(r.error), 'error', 'success')`;
+const TRACE_FIELDS = {
+	traceId: {
+		sql: "r.traceId",
+		parameterType: "String"
+	},
+	threadId: {
+		sql: "r.threadId",
+		parameterType: "String"
+	},
+	resourceId: {
+		sql: "r.resourceId",
+		parameterType: "String"
+	},
+	startedAt: {
+		sql: "r.startedAt",
+		parameterType: "DateTime64(3, 'UTC')"
+	},
+	endedAt: {
+		sql: "r.endedAt",
+		parameterType: "DateTime64(3, 'UTC')"
+	},
+	entityName: {
+		sql: "r.entityName",
+		parameterType: "String"
+	},
+	entityType: {
+		sql: "r.entityType",
+		parameterType: "String"
+	},
+	environment: {
+		sql: "r.environment",
+		parameterType: "String"
+	},
+	status: {
+		sql: TRACE_STATUS_SQL,
+		parameterType: "String"
+	}
+};
+const SPAN_FIELDS = {
+	spanType: {
+		sql: "s.spanType",
+		parameterType: "String"
+	},
+	error: {
+		sql: "s.error",
+		parameterType: "String"
+	}
+};
+const SCORE_FIELDS = {
+	scorerId: {
+		sql: "s.scorerId",
+		parameterType: "String"
+	},
+	score: {
+		sql: "s.score",
+		parameterType: "Float64"
+	}
+};
+const TRACE_SELECT = `
+  r.traceId AS traceId,
+  r.spanId AS rootSpanId,
+  r.threadId AS threadId,
+  r.resourceId AS resourceId,
+  r.startedAt AS startedAt,
+  r.endedAt AS endedAt,
+  r.entityName AS entityName,
+  r.entityType AS entityType,
+  r.environment AS environment,
+  ${TRACE_STATUS_SQL} AS status`;
+var ParameterBuilder = class {
+	params = {};
+	#next = 1;
+	add(value, type) {
+		const name = `trace_query_${this.#next++}`;
+		this.params[name] = type === "DateTime64(3, 'UTC')" ? new Date(value).toISOString().replace("T", " ").replace(/Z$/, "") : value;
+		return `{${name}:${type}}`;
+	}
+};
+function fieldDefinition(registry, field) {
+	const definition = registry[field];
+	if (definition === void 0) throw new Error(`Unsupported trusted trace-query field: ${field}`);
+	return definition;
+}
+function resolveOrderField(field) {
+	if (field === "startedAt" || field === "endedAt") return field;
+	throw new Error(`Unsupported trusted trace-query field: ${field}`);
+}
+function compileScalarPredicate(predicate, registry, parameters) {
+	if (predicate.type === "boolean") return predicate.args.map((arg) => `(${compileScalarPredicate(arg, registry, parameters)})`).join(predicate.operator === "and" ? " AND " : " OR ");
+	if (predicate.type === "not") return `NOT (${compileScalarPredicate(predicate.arg, registry, parameters)})`;
+	const field = fieldDefinition(registry, predicate.field);
+	if (predicate.type === "presence") return `${predicate.operator === "exists" ? "isNotNull" : "isNull"}(${field.sql})`;
+	if (predicate.type === "membership") {
+		const values = predicate.values.map((value) => parameters.add(value, field.parameterType)).join(", ");
+		return `ifNull(${`${field.sql} ${predicate.operator === "in" ? "IN" : "NOT IN"} (${values})`}, ${predicate.operator === "in" ? "0" : "1"})`;
+	}
+	const parameter = parameters.add(predicate.value, field.parameterType);
+	const operator = {
+		eq: "=",
+		ne: "!=",
+		lt: "<",
+		lte: "<=",
+		gt: ">",
+		gte: ">="
+	}[predicate.operator];
+	if (operator === void 0) throw new Error(`Unsupported trusted trace-query operator: ${predicate.operator}`);
+	return `ifNull(${field.sql} ${operator} ${parameter}, ${predicate.operator === "ne" ? "1" : "0"})`;
+}
+function collectRelationCollections(predicate, collections = /* @__PURE__ */ new Set()) {
+	if (!predicate) return collections;
+	if (predicate.type === "relation") collections.add(predicate.collection);
+	else if (predicate.type === "boolean") for (const arg of predicate.args) collectRelationCollections(arg, collections);
+	else if (predicate.type === "not") collectRelationCollections(predicate.arg, collections);
+	return collections;
+}
+function compilePredicate(predicate, parameters) {
+	if (predicate.type === "relation") {
+		const registry = predicate.collection === "spans" ? SPAN_FIELDS : SCORE_FIELDS;
+		const existence = `EXISTS (
+      SELECT 1 FROM ${predicate.collection === "spans" ? "current_spans" : "current_scores"} s
+      WHERE s.traceId = r.traceId
+        AND (${compileScalarPredicate(predicate.predicate, registry, parameters)})
+    )`;
+		return predicate.quantifier === "some" ? existence : `NOT ${existence}`;
+	}
+	if (predicate.type === "boolean") return predicate.args.map((arg) => `(${compilePredicate(arg, parameters)})`).join(predicate.operator === "and" ? " AND " : " OR ");
+	if (predicate.type === "not") return `NOT (${compilePredicate(predicate.arg, parameters)})`;
+	return compileScalarPredicate(predicate, TRACE_FIELDS, parameters);
+}
+function compileClickHouseTraceQuery(plan) {
+	const parameters = new ParameterBuilder();
+	const from = parameters.add(plan.timeRange.from, "DateTime64(3, 'UTC')");
+	const to = parameters.add(plan.timeRange.to, "DateTime64(3, 'UTC')");
+	const relationCollections = collectRelationCollections(plan.where);
+	const ctes = [`current_roots AS (
+    SELECT * FROM (
+      SELECT *
+      FROM ${TABLE_TRACE_ROOTS}
+      ORDER BY dedupeKey
+      LIMIT 1 BY dedupeKey
+    )
+    ORDER BY traceId, dedupeKey
+    LIMIT 1 BY traceId
+  )`, `root_scope AS (
+    SELECT *
+    FROM current_roots
+    WHERE startedAt >= ${from}
+      AND startedAt < ${to}
+  )`];
+	if (relationCollections.has("spans")) ctes.push(`current_spans AS (
+    SELECT traceId, spanType, error
+    FROM ${TABLE_SPAN_EVENTS}
+    WHERE traceId IN (SELECT traceId FROM root_scope)
+    ORDER BY dedupeKey
+    LIMIT 1 BY dedupeKey
+  )`);
+	if (relationCollections.has("scores")) ctes.push(`current_scores AS (
+    SELECT traceId, scorerId, score
+    FROM ${TABLE_SCORE_EVENTS}
+    WHERE isNotNull(traceId)
+      AND traceId IN (SELECT traceId FROM root_scope)
+  )`);
+	const predicate = plan.where ? compilePredicate(plan.where, parameters) : "1";
+	ctes.push(`candidates AS (
+    SELECT ${TRACE_SELECT}
+    FROM root_scope r
+    WHERE ${predicate}
+  )`);
+	const candidates = `WITH ${ctes.join(",\n")}`;
+	if (plan.result === "groups") return {
+		query: `${candidates}
+SELECT threadId
+FROM candidates
+WHERE isNotNull(threadId) ${plan.cursor ? `AND threadId > ${parameters.add(plan.cursor.threadId, "String")}` : ""}
+GROUP BY threadId
+ORDER BY threadId ASC
+LIMIT ${parameters.add(plan.limit + 1, "UInt64")}`,
+		query_params: parameters.params
+	};
+	const orderField = resolveOrderField(plan.orderBy.field);
+	const direction = plan.orderBy.direction === "asc" ? "ASC" : "DESC";
+	let pageCondition = "";
+	if (plan.cursor) {
+		const comparison = plan.orderBy.direction === "asc" ? ">" : "<";
+		const sortValue = parameters.add(plan.cursor.sortValue, "DateTime64(3, 'UTC')");
+		pageCondition = `WHERE (${orderField} ${comparison} ${sortValue} OR (${orderField} = ${sortValue} AND traceId > ${parameters.add(plan.cursor.traceId, "String")}))`;
+	}
+	const limit = parameters.add(plan.limit + 1, "UInt64");
+	return {
+		query: `${candidates}
+SELECT *
+FROM candidates
+${pageCondition}
+ORDER BY ${orderField} ${direction}, traceId ASC
+LIMIT ${limit}`,
+		query_params: parameters.params
+	};
+}
+function asIsoTimestamp(value) {
+	return new Date(value).toISOString();
+}
+function isClickHouseExecutionTimeout(error) {
+	if (!error || typeof error !== "object") return false;
+	const candidate = error;
+	return String(candidate.code ?? "") === "159" || candidate.type === "TIMEOUT_EXCEEDED";
+}
+async function runWithClickHouseTraceQueryTimeout(client, timeoutMs, compiled, queryId) {
+	const resolvedTimeoutMs = coreStorage.resolveTraceQueryTimeoutMs(timeoutMs);
+	try {
+		return await (await client.query({
+			query: compiled.query,
+			query_params: compiled.query_params,
+			query_id: queryId,
+			format: "JSONEachRow",
+			clickhouse_settings: {
+				...CH_SETTINGS,
+				max_execution_time: resolvedTimeoutMs / 1e3
+			}
+		})).json();
+	} catch (error) {
+		if (isClickHouseExecutionTimeout(error)) throw new coreStorage.TraceQueryExecutionError();
+		throw error;
+	}
+}
+async function queryTraces(client, plan, timeoutMs) {
+	const rows = await runWithClickHouseTraceQueryTimeout(client, timeoutMs, compileClickHouseTraceQuery(plan));
+	const visibleRows = rows.slice(0, plan.limit);
+	if (plan.result === "groups") {
+		const groups = visibleRows.map((row) => ({ threadId: String(row.threadId) }));
+		const last = groups.at(-1);
+		return coreStorage.traceQueryResponseSchema.parse({
+			groups,
+			page: { next: rows.length > plan.limit && last ? coreStorage.encodeTraceQueryCursor(plan, {
+				result: "groups",
+				threadId: last.threadId
+			}) : null }
+		});
+	}
+	const traces = visibleRows.map((row) => ({
+		traceId: String(row.traceId),
+		rootSpanId: String(row.rootSpanId),
+		threadId: row.threadId == null ? null : String(row.threadId),
+		resourceId: row.resourceId == null ? null : String(row.resourceId),
+		startedAt: asIsoTimestamp(row.startedAt),
+		endedAt: asIsoTimestamp(row.endedAt),
+		entityName: row.entityName == null ? null : String(row.entityName),
+		entityType: row.entityType == null ? null : String(row.entityType),
+		environment: row.environment == null ? null : String(row.environment),
+		status: row.status
+	}));
+	const last = traces.at(-1);
+	return coreStorage.traceQueryResponseSchema.parse({
+		traces,
+		page: { next: rows.length > plan.limit && last ? coreStorage.encodeTraceQueryCursor(plan, {
+			result: "traces",
+			sortValue: last[plan.orderBy.field],
+			traceId: last.traceId
+		}) : null }
+	});
+}
+//#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.
@@ -7033,6 +7297,7 @@ var ObservabilityStorageClickhouseVNext = class extends ObservabilityStorage {
 	#retention;
 	#replication;
 	#deltaCursorStrategyOverride;
+	#traceQueryTimeoutMs;
 	#deltaCursorStrategy = "fallback";
 	constructor(config) {
 		super();
@@ -7041,6 +7306,7 @@ var ObservabilityStorageClickhouseVNext = class extends ObservabilityStorage {
 		this.#replication = replication;
 		this.#retention = config.retention;
 		this.#deltaCursorStrategyOverride = config.deltaCursorStrategy;
+		this.#traceQueryTimeoutMs = coreStorage.resolveTraceQueryTimeoutMs(config.traceQueryTimeoutMs);
 	}
 	async init() {
 		const migrationStatus = await checkSignalTablesMigrationStatus(this.#client);
@@ -7144,11 +7410,16 @@ var ObservabilityStorageClickhouseVNext = class extends ObservabilityStorage {
 		};
 	}
 	getFeatures() {
-		if (!deltaPollingSupported(this.#deltaCursorStrategy)) return ["metrics", "logs"];
+		if (!deltaPollingSupported(this.#deltaCursorStrategy)) return [
+			"metrics",
+			"logs",
+			"trace-query"
+		];
 		return [
 			"metrics",
 			"logs",
-			"delta-polling"
+			"delta-polling",
+			"trace-query"
 		];
 	}
 	async createSpan(args) {
@@ -7263,6 +7534,18 @@ var ObservabilityStorageClickhouseVNext = class extends ObservabilityStorage {
 			}, error);
 		}
 	}
+	async queryTraces(plan) {
+		try {
+			return await queryTraces(this.#client, plan, this.#traceQueryTimeoutMs);
+		} catch (error) {
+			if (error instanceof MastraError || error instanceof coreStorage.TraceQueryExecutionError) throw error;
+			throw new MastraError({
+				id: createStorageErrorId("CLICKHOUSE", "QUERY_TRACES", "FAILED"),
+				domain: ErrorDomain.STORAGE,
+				category: ErrorCategory.THIRD_PARTY
+			}, error);
+		}
+	}
 	async listTracesLight(args) {
 		try {
 			return await listTracesLight(this.#client, args, this.#deltaCursorStrategy);
@@ -8539,7 +8822,8 @@ var ClickhouseStoreVNext = class extends ClickhouseStore {
 		this.name = "ClickhouseStoreVNext";
 		const observability = new ObservabilityStorageClickhouseVNext({
 			client: this.db,
-			replication: config.replication
+			replication: config.replication,
+			traceQueryTimeoutMs: config.traceQueryTimeoutMs
 		});
 		this.stores = {
 			...this.stores,