@mastra/clickhouse 1.20.0-alpha.3 → 1.20.0

package/dist/index.js

@@ -7,6 +7,7 @@ import { MastraBase } from "@mastra/core/base";
 import { randomUUID } from "crypto";
 import { parseFieldKey } from "@mastra/core/utils";
 import { coreFeatures } from "@mastra/core/features";
+import { z } from "zod/v4";
 import { saveScorePayloadSchema } from "@mastra/core/evals";
 //#region src/storage/db/replication.ts
 const DEFAULT_ZOOKEEPER_PATH = "/clickhouse/tables/{shard}/{database}/{table}";
@@ -6849,7 +6850,18 @@ function compileClickHouseTraceScope(selection, relationCollections, parameters)
   )`);
 	return ctes;
 }
-function compileClickHouseTraceQuery(plan) {
+const deltaWatermarkSchema = z.object({
+	cursorId: z.string().regex(/^\d+$/).refine((value) => BigInt(value) <= 18446744073709551615n),
+	traceId: z.string()
+}).strict();
+function parseDeltaWatermark(value) {
+	try {
+		return deltaWatermarkSchema.parse(JSON.parse(value));
+	} catch {
+		throw new coreStorage.TraceQueryCursorError("TRACE_QUERY_CURSOR_MALFORMED");
+	}
+}
+function compileClickHouseTraceQuery(plan, deltaHead) {
 	const parameters = new ParameterBuilder();
 	const ctes = compileClickHouseTraceScope(plan, collectRelationCollections(plan.where), parameters);
 	const predicate = plan.where ? compilePredicate(plan.where, parameters) : "1";
@@ -6869,6 +6881,28 @@ ORDER BY threadId ASC
 LIMIT ${parameters.add(plan.limit + 1, "UInt64")}`,
 		query_params: parameters.params
 	};
+	if (plan.paginationMode === "delta") {
+		const watermark = coreStorage.getTraceQueryDeltaWatermark(plan, "clickhouse");
+		const after = watermark ? parseDeltaWatermark(watermark) : {
+			cursorId: "0",
+			traceId: ""
+		};
+		return {
+			query: `${candidates}, delta_candidates AS (
+  SELECT traceId, max(cursorId) AS latestCursorId
+  FROM ${TABLE_TRACE_ROOTS_DELTA}
+  WHERE tuple(cursorId, traceId) > ${`tuple(${parameters.add(after.cursorId, "UInt64")}, ${parameters.add(after.traceId, "String")})`} ${deltaHead ? `AND tuple(cursorId, traceId) <= tuple(${parameters.add(deltaHead.cursorId, "UInt64")}, ${parameters.add(deltaHead.traceId, "String")})` : ""}
+  GROUP BY traceId
+)
+SELECT c.*, toString(d.latestCursorId) AS __delta_cursor
+FROM candidates c
+INNER JOIN delta_candidates d ON c.traceId = d.traceId
+ORDER BY d.latestCursorId ASC, c.traceId ASC
+LIMIT ${parameters.add(plan.limit + 1, "UInt64")}`,
+			query_params: parameters.params,
+			sharedSnapshot: true
+		};
+	}
 	const orderField = resolveOrderField(plan.orderBy.field);
 	const direction = plan.orderBy.direction === "asc" ? "ASC" : "DESC";
 	if (plan.paginationMode === "page") return {
@@ -7088,9 +7122,39 @@ async function getTraceQueryValues(client, plan, limits) {
 		valuesTruncated: rows.length > plan.limit
 	});
 }
-async function queryTraces(client, plan, timeoutMs) {
+async function queryTraces(client, plan, timeoutMs, strategy = null) {
+	const deadline = performance.now() + coreStorage.resolveTraceQueryTimeoutMs(timeoutMs);
+	const remaining = () => {
+		const value = Math.floor(deadline - performance.now());
+		if (value <= 0) throw new coreStorage.TraceQueryExecutionError();
+		return value;
+	};
+	let deltaHead;
+	if (plan.paginationMode === "delta") {
+		assertDeltaPollingSupported(strategy);
+		const watermark = coreStorage.getTraceQueryDeltaWatermark(plan, "clickhouse");
+		if (watermark) parseDeltaWatermark(watermark);
+	}
+	if (plan.paginationMode === "delta" || plan.paginationMode === "page" && deltaPollingSupported(strategy) && typeof coreStorage.encodeTraceQueryDeltaCursor === "function") {
+		const head = await runWithClickHouseTraceQueryTimeout(client, { timeoutMs: remaining() }, {
+			query: `SELECT toString(cursorId) AS cursorId, traceId FROM ${TABLE_TRACE_ROOTS_DELTA} ORDER BY cursorId DESC, traceId DESC LIMIT 1`,
+			query_params: {}
+		});
+		deltaHead = head[0] ? parseDeltaWatermark(JSON.stringify(head[0])) : {
+			cursorId: "0",
+			traceId: ""
+		};
+		if (plan.paginationMode === "delta" && !plan.deltaCursor) return coreStorage.traceQueryResponseSchema.parse({
+			traces: [],
+			delta: {
+				limit: plan.limit,
+				hasMore: false
+			},
+			deltaCursor: coreStorage.encodeTraceQueryDeltaCursor(plan, "clickhouse", JSON.stringify(deltaHead))
+		});
+	}
 	if (plan.paginationMode === "page") {
-		const rows = await runWithClickHouseTraceQueryTimeout(client, { timeoutMs }, compileClickHouseTraceQuery(plan));
+		const rows = await runWithClickHouseTraceQueryTimeout(client, { timeoutMs: deltaHead ? remaining() : timeoutMs }, compileClickHouseTraceQuery(plan, deltaHead));
 		const total = Number(rows.at(-1)?.total ?? 0);
 		const traces = rows.filter((row) => Number(row.__metadata) === 0).map((row) => ({
 			traceId: String(row.traceId),
@@ -7112,6 +7176,7 @@ async function queryTraces(client, plan, timeoutMs) {
 		}));
 		return coreStorage.traceQueryResponseSchema.parse({
 			traces,
+			...deltaHead ? { deltaCursor: coreStorage.encodeTraceQueryDeltaCursor(plan, "clickhouse", JSON.stringify(deltaHead)) } : {},
 			pagination: {
 				total,
 				page: plan.page,
@@ -7120,7 +7185,7 @@ async function queryTraces(client, plan, timeoutMs) {
 			}
 		});
 	}
-	const rows = await runWithClickHouseTraceQueryTimeout(client, { timeoutMs }, compileClickHouseTraceQuery(plan));
+	const rows = await runWithClickHouseTraceQueryTimeout(client, { timeoutMs: deltaHead ? remaining() : timeoutMs }, compileClickHouseTraceQuery(plan, deltaHead));
 	const visibleRows = rows.slice(0, plan.limit);
 	if (plan.result === "groups") {
 		const groups = visibleRows.map((row) => ({ threadId: String(row.threadId) }));
@@ -7152,6 +7217,23 @@ async function queryTraces(client, plan, timeoutMs) {
 		status: row.status
 	}));
 	const last = traces.at(-1);
+	if (plan.paginationMode === "delta") {
+		const lastRow = visibleRows.at(-1);
+		let watermark = lastRow ? {
+			cursorId: String(lastRow.__delta_cursor),
+			traceId: String(lastRow.traceId)
+		} : deltaHead;
+		const previous = parseDeltaWatermark(coreStorage.getTraceQueryDeltaWatermark(plan, "clickhouse"));
+		if (BigInt(previous.cursorId) > BigInt(watermark.cursorId) || previous.cursorId === watermark.cursorId && previous.traceId > watermark.traceId) watermark = previous;
+		return coreStorage.traceQueryResponseSchema.parse({
+			traces,
+			delta: {
+				limit: plan.limit,
+				hasMore: rows.length > plan.limit
+			},
+			deltaCursor: coreStorage.encodeTraceQueryDeltaCursor(plan, "clickhouse", JSON.stringify(watermark))
+		});
+	}
 	return coreStorage.traceQueryResponseSchema.parse({
 		traces,
 		page: { next: rows.length > plan.limit && last ? coreStorage.encodeTraceQueryCursor(plan, {
@@ -8510,9 +8592,9 @@ var ObservabilityStorageClickhouseVNext = class extends ObservabilityStorage {
 	}
 	async queryTraces(plan) {
 		try {
-			return await queryTraces(this.#client, plan, this.#traceQueryTimeoutMs);
+			return await queryTraces(this.#client, plan, this.#traceQueryTimeoutMs, this.#deltaCursorStrategy);
 		} catch (error) {
-			if (error instanceof MastraError || error instanceof coreStorage.TraceQueryExecutionError || error instanceof coreStorage.TraceQueryResourceLimitError) throw error;
+			if (error instanceof MastraError || error instanceof coreStorage.TraceQueryExecutionError || error instanceof coreStorage.TraceQueryCursorError || error instanceof coreStorage.TraceQueryResourceLimitError) throw error;
 			throw new MastraError({
 				id: createStorageErrorId("CLICKHOUSE", "QUERY_TRACES", "FAILED"),
 				domain: ErrorDomain.STORAGE,