@mastra/clickhouse 1.18.0 → 1.19.0

package/dist/index.js

@@ -6578,6 +6578,13 @@ function collectRelationCollections(predicate, collections = /* @__PURE__ */ new
 	else if (predicate.type === "not") collectRelationCollections(predicate.arg, collections);
 	return collections;
 }
+function collectThreadRelationCollections(predicate, collections) {
+	if (!predicate) return collections;
+	if (predicate.type === "relation") collectRelationCollections(predicate.predicate, collections);
+	else if (predicate.type === "boolean") for (const arg of predicate.args) collectThreadRelationCollections(arg, collections);
+	else collectThreadRelationCollections(predicate.arg, collections);
+	return collections;
+}
 function compilePredicate(predicate, parameters) {
 	if (predicate.type === "relation") {
 		const existence = `EXISTS (
@@ -6592,11 +6599,21 @@ function compilePredicate(predicate, parameters) {
 	if (predicate.type === "not") return `NOT (${compilePredicate(predicate.arg, parameters)})`;
 	return compileScalarPredicate(predicate, TRACE_FIELDS, parameters, true);
 }
-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);
+function compileThreadPredicate(predicate, parameters) {
+	if (predicate.type === "relation") {
+		const existence = `EXISTS (
+      SELECT 1 FROM eligible_roots r
+      WHERE r.threadId = t.threadId
+        AND (${compilePredicate(predicate.predicate, parameters)})
+    )`;
+		return predicate.quantifier === "some" ? existence : `NOT ${existence}`;
+	}
+	if (predicate.type === "boolean") return predicate.args.map((arg) => `(${compileThreadPredicate(arg, parameters)})`).join(predicate.operator === "and" ? " AND " : " OR ");
+	return `NOT (${compileThreadPredicate(predicate.arg, parameters)})`;
+}
+function compileClickHouseTraceScope(selection, relationCollections, parameters) {
+	const from = parameters.add(selection.timeRange.from, "DateTime64(3, 'UTC')");
+	const to = parameters.add(selection.timeRange.to, "DateTime64(3, 'UTC')");
 	const ctes = [`current_roots AS (
     SELECT * FROM (
       SELECT *
@@ -6677,6 +6694,11 @@ function compileClickHouseTraceQuery(plan) {
     WHERE isNotNull(traceId)
       AND traceId IN (SELECT traceId FROM root_scope)
   )`);
+	return ctes;
+}
+function compileClickHouseTraceQuery(plan) {
+	const parameters = new ParameterBuilder();
+	const ctes = compileClickHouseTraceScope(plan, collectRelationCollections(plan.where), parameters);
 	const predicate = plan.where ? compilePredicate(plan.where, parameters) : "1";
 	ctes.push(`candidates AS (
     SELECT ${TRACE_SELECT}
@@ -6713,6 +6735,41 @@ LIMIT ${limit}`,
 		query_params: parameters.params
 	};
 }
+function compileClickHouseThreadQuery(plan) {
+	const parameters = new ParameterBuilder();
+	const relationCollections = collectRelationCollections(plan.traces.where);
+	collectThreadRelationCollections(plan.where, relationCollections);
+	const ctes = compileClickHouseTraceScope(plan.traces, relationCollections, parameters);
+	const eligibility = plan.traces.where ? compilePredicate(plan.traces.where, parameters) : "1";
+	ctes.push(`eligible_roots AS (
+    SELECT *
+    FROM root_scope r
+    WHERE ${eligibility}
+  )`);
+	ctes.push(`thread_ids AS (
+    SELECT threadId
+    FROM eligible_roots
+    WHERE isNotNull(threadId)
+    GROUP BY threadId
+  )`);
+	const threadPredicate = plan.where ? compileThreadPredicate(plan.where, parameters) : "1";
+	ctes.push(`qualified_threads AS (
+    SELECT t.threadId
+    FROM thread_ids t
+    WHERE ${threadPredicate}
+  )`);
+	const pageCondition = plan.cursor ? `WHERE threadId > ${parameters.add(plan.cursor.threadId, "String")}` : "";
+	const limit = parameters.add(plan.limit + 1, "UInt64");
+	return {
+		query: `WITH ${ctes.join(",\n")}
+SELECT threadId
+FROM qualified_threads
+${pageCondition}
+ORDER BY threadId ASC
+LIMIT ${limit}`,
+		query_params: parameters.params
+	};
+}
 function asIsoTimestamp(value) {
 	return new Date(value).toISOString();
 }
@@ -6775,6 +6832,18 @@ async function queryTraces(client, plan, timeoutMs) {
 		}) : null }
 	});
 }
+async function queryThreads(client, plan, timeoutMs) {
+	const rows = await runWithClickHouseTraceQueryTimeout(client, timeoutMs, compileClickHouseThreadQuery(plan));
+	const threads = rows.slice(0, plan.limit).map((row) => ({ threadId: String(row.threadId) }));
+	const last = threads.at(-1);
+	return coreStorage.queryThreadsResultSchema.parse({
+		threads,
+		page: { next: rows.length > plan.limit && last ? coreStorage.encodeTraceQueryCursor(plan, {
+			result: "threads",
+			threadId: last.threadId
+		}) : null }
+	});
+}
 //#endregion
 //#region src/storage/domains/observability/v-next/trace-roots.ts
 /**
@@ -7891,13 +7960,15 @@ var ObservabilityStorageClickhouseVNext = class extends ObservabilityStorage {
 		if (!deltaPollingSupported(this.#deltaCursorStrategy)) return [
 			"metrics",
 			"logs",
-			"trace-query"
+			"trace-query",
+			"thread-query"
 		];
 		return [
 			"metrics",
 			"logs",
 			"delta-polling",
-			"trace-query"
+			"trace-query",
+			"thread-query"
 		];
 	}
 	async createSpan(args) {
@@ -8024,6 +8095,18 @@ var ObservabilityStorageClickhouseVNext = class extends ObservabilityStorage {
 			}, error);
 		}
 	}
+	async queryThreads(plan) {
+		try {
+			return await queryThreads(this.#client, plan, this.#traceQueryTimeoutMs);
+		} catch (error) {
+			if (error instanceof MastraError || error instanceof coreStorage.TraceQueryExecutionError) throw error;
+			throw new MastraError({
+				id: createStorageErrorId("CLICKHOUSE", "QUERY_THREADS", "FAILED"),
+				domain: ErrorDomain.STORAGE,
+				category: ErrorCategory.THIRD_PARTY
+			}, error);
+		}
+	}
 	async listTracesLight(args) {
 		try {
 			return await listTracesLight(this.#client, args, this.#deltaCursorStrategy);