@mastra/duckdb 1.6.4 → 1.7.0-alpha.0

package/dist/{observability-CT5eidI9.js → observability-Cqe4mNHq.js}

@@ -1,5 +1,6 @@
 import { t as DuckDBConnection } from "./db-Dl0fY488.js";
 import { ErrorCategory, ErrorDomain, MastraError } from "@mastra/core/error";
+import * as coreStorage from "@mastra/core/storage";
 import { BRANCH_SPAN_TYPES, METRIC_DISTINCT_COLUMNS, ObservabilityStorage, buildInputPreview, computeTraceStatus, createStorageErrorId, feedbackRecordSchema, listBranchesArgsSchema, listFeedbackArgsSchema, listLogsArgsSchema, listMetricsArgsSchema, listScoresArgsSchema, listTracesArgsSchema, toTraceSpans } from "@mastra/core/storage";
 import { coreFeatures } from "@mastra/core/features";
 import { EntityType } from "@mastra/core/observability";
@@ -2233,7 +2234,7 @@ function getComparisonDateRange(comparePeriod, timestamp) {
 async function createScore(db, args) {
 	const s = args.score;
 	const scoreSource = s.scoreSource ?? s.source ?? null;
-	await db.execute(`INSERT INTO score_events (
+	await db.execute(`INSERT OR REPLACE INTO score_events (
       scoreId, timestamp, cursorId, traceId, spanId, experimentId, scoreTraceId,
       entityType, entityId, entityName, entityVersionId, parentEntityVersionId, parentEntityType, parentEntityId, parentEntityName, rootEntityVersionId, rootEntityType, rootEntityId, rootEntityName,
       userId, organizationId, resourceId, runId, sessionId, threadId, requestId, environment, executionSource, serviceName,
@@ -2277,8 +2278,7 @@ async function createScore(db, args) {
 		jsonV(s.tags ?? null),
 		jsonV(s.metadata),
 		jsonV(s.scope ?? null)
-	].join(", ")})
-     ON CONFLICT DO NOTHING`);
+	].join(", ")})`);
 }
 /** Insert multiple score events in a single statement. */
 async function batchCreateScores(db, args) {
@@ -2326,14 +2326,13 @@ async function batchCreateScores(db, args) {
 			jsonV(legacyScore.scope ?? null)
 		].join(", ")})`;
 	});
-	await db.execute(`INSERT INTO score_events (
+	await db.execute(`INSERT OR REPLACE INTO score_events (
       scoreId, timestamp, cursorId, traceId, spanId, experimentId, scoreTraceId,
       entityType, entityId, entityName, entityVersionId, parentEntityVersionId, parentEntityType, parentEntityId, parentEntityName, rootEntityVersionId, rootEntityType, rootEntityId, rootEntityName,
       userId, organizationId, resourceId, runId, sessionId, threadId, requestId, environment, executionSource, serviceName,
       scorerId, scorerVersion, scoreSource, score, reason, tags, metadata, scope
     )
-     VALUES ${tuples.join(",\n       ")}
-     ON CONFLICT DO NOTHING`);
+     VALUES ${tuples.join(",\n       ")}`);
 }
 /** Query score events with filtering, ordering, and pagination. */
 async function listScores(db, args) {
@@ -2519,6 +2518,323 @@ async function getScorePercentiles(db, args) {
 	return { series };
 }
 //#endregion
+//#region src/storage/domains/observability/trace-query.ts
+const TRACE_STATUS_SQL = `CASE WHEN r.error IS NOT NULL THEN 'error' ELSE 'success' END`;
+const TRACE_FIELDS = {
+	traceId: {
+		sql: "r.traceId",
+		parameterType: "scalar"
+	},
+	threadId: {
+		sql: "r.threadId",
+		parameterType: "scalar"
+	},
+	resourceId: {
+		sql: "r.resourceId",
+		parameterType: "scalar"
+	},
+	startedAt: {
+		sql: "r.startedAt",
+		parameterType: "timestamp"
+	},
+	endedAt: {
+		sql: "r.endedAt",
+		parameterType: "timestamp"
+	},
+	entityName: {
+		sql: "r.entityName",
+		parameterType: "scalar"
+	},
+	entityType: {
+		sql: "r.entityType",
+		parameterType: "scalar"
+	},
+	environment: {
+		sql: "r.environment",
+		parameterType: "scalar"
+	},
+	status: {
+		sql: TRACE_STATUS_SQL,
+		parameterType: "scalar"
+	}
+};
+const SPAN_FIELDS = {
+	spanType: {
+		sql: "s.spanType",
+		parameterType: "scalar"
+	},
+	error: {
+		sql: "s.error",
+		parameterType: "scalar"
+	}
+};
+const SCORE_FIELDS = {
+	scorerId: {
+		sql: "s.scorerId",
+		parameterType: "scalar"
+	},
+	score: {
+		sql: "s.score",
+		parameterType: "scalar"
+	}
+};
+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`;
+function fieldDefinition(registry, field) {
+	const definition = registry[field];
+	if (definition === void 0) throw new Error(`Unsupported trusted trace-query field: ${field}`);
+	return definition;
+}
+function parameterSql(type) {
+	return type === "timestamp" ? "CAST(? AS TIMESTAMP)" : "?";
+}
+function compileScalarPredicate(predicate, registry) {
+	if (predicate.type === "boolean") {
+		const values = [];
+		return {
+			sql: predicate.args.map((arg) => {
+				const compiled = compileScalarPredicate(arg, registry);
+				values.push(...compiled.values);
+				return `(${compiled.sql})`;
+			}).join(predicate.operator === "and" ? " AND " : " OR "),
+			values
+		};
+	}
+	if (predicate.type === "not") {
+		const compiled = compileScalarPredicate(predicate.arg, registry);
+		return {
+			sql: `NOT (${compiled.sql})`,
+			values: compiled.values
+		};
+	}
+	const field = fieldDefinition(registry, predicate.field);
+	if (predicate.type === "presence") return {
+		sql: `${field.sql} IS ${predicate.operator === "exists" ? "NOT " : ""}NULL`,
+		values: []
+	};
+	if (predicate.type === "membership") {
+		const list = predicate.values.map(() => parameterSql(field.parameterType)).join(", ");
+		if (predicate.operator === "in") return {
+			sql: `${field.sql} IS NOT NULL AND ${field.sql} IN (${list})`,
+			values: predicate.values
+		};
+		return {
+			sql: `${field.sql} IS NULL OR ${field.sql} NOT IN (${list})`,
+			values: predicate.values
+		};
+	}
+	const parameter = parameterSql(field.parameterType);
+	const operators = {
+		lt: "<",
+		lte: "<=",
+		gt: ">",
+		gte: ">="
+	};
+	if (predicate.operator === "eq") return {
+		sql: `${field.sql} IS NOT DISTINCT FROM ${parameter}`,
+		values: [predicate.value]
+	};
+	if (predicate.operator === "ne") return {
+		sql: `${field.sql} IS DISTINCT FROM ${parameter}`,
+		values: [predicate.value]
+	};
+	const operator = operators[predicate.operator];
+	if (operator === void 0) throw new Error(`Unsupported trusted trace-query operator: ${predicate.operator}`);
+	return {
+		sql: `${field.sql} IS NOT NULL AND ${field.sql} ${operator} ${parameter}`,
+		values: [predicate.value]
+	};
+}
+function compilePredicate(predicate) {
+	if (predicate.type === "relation") {
+		const registry = predicate.collection === "spans" ? SPAN_FIELDS : SCORE_FIELDS;
+		const compiled = compileScalarPredicate(predicate.predicate, registry);
+		const existence = `EXISTS (
+      SELECT 1 FROM ${predicate.collection === "spans" ? "current_spans" : "current_scores"} s
+      WHERE s.traceId IS NOT NULL
+        AND s.traceId = r.traceId
+        AND (${compiled.sql})
+    )`;
+		return {
+			sql: predicate.quantifier === "some" ? existence : `NOT ${existence}`,
+			values: compiled.values
+		};
+	}
+	if (predicate.type === "boolean") {
+		const values = [];
+		return {
+			sql: predicate.args.map((arg) => {
+				const compiled = compilePredicate(arg);
+				values.push(...compiled.values);
+				return `(${compiled.sql})`;
+			}).join(predicate.operator === "and" ? " AND " : " OR "),
+			values
+		};
+	}
+	if (predicate.type === "not") {
+		const compiled = compilePredicate(predicate.arg);
+		return {
+			sql: `NOT (${compiled.sql})`,
+			values: compiled.values
+		};
+	}
+	return compileScalarPredicate(predicate, TRACE_FIELDS);
+}
+function collectRelatedCollections(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) collectRelatedCollections(arg, collections);
+	else if (predicate.type === "not") collectRelatedCollections(predicate.arg, collections);
+	return collections;
+}
+function compileDuckDBTraceQuery(plan) {
+	const values = [plan.timeRange.from, plan.timeRange.to];
+	const conditions = [
+		`r.endedAt IS NOT NULL`,
+		`r.startedAt >= CAST(? AS TIMESTAMP)`,
+		`r.startedAt < CAST(? AS TIMESTAMP)`
+	];
+	if (plan.where) {
+		const predicate = compilePredicate(plan.where);
+		conditions.push(`(${predicate.sql})`);
+		values.push(...predicate.values);
+	}
+	const relatedCollections = collectRelatedCollections(plan.where);
+	const ctes = [
+		`root_events AS (
+      SELECT
+        *,
+        CASE
+          WHEN eventType = 'start' THEN timestamp
+          ELSE lag(timestamp) OVER (PARTITION BY traceId, spanId ORDER BY cursorId)
+        END AS startedAt
+      FROM span_events
+      WHERE parentSpanId IS NULL
+    )`,
+		`current_roots AS (
+      SELECT * EXCLUDE (rootRank)
+      FROM (
+        SELECT *, row_number() OVER (PARTITION BY traceId ORDER BY cursorId DESC) AS rootRank
+        FROM root_events
+      )
+      WHERE rootRank = 1
+    )`,
+		`root_scope AS (
+      SELECT *
+      FROM current_roots r
+      WHERE r.endedAt IS NOT NULL
+        AND r.startedAt >= CAST(? AS TIMESTAMP)
+        AND r.startedAt < CAST(? AS TIMESTAMP)
+    )`
+	];
+	if (relatedCollections.has("spans")) ctes.push(`current_spans AS (
+      SELECT * EXCLUDE (currentRank)
+      FROM (
+        SELECT
+          e.*,
+          row_number() OVER (
+            PARTITION BY e.traceId, e.spanId
+            ORDER BY CASE WHEN e.endedAt IS NULL THEN 1 ELSE 0 END ASC, e.cursorId DESC
+          ) AS currentRank
+        FROM span_events e
+        INNER JOIN root_scope roots ON roots.traceId = e.traceId
+      )
+      WHERE currentRank = 1
+    )`);
+	if (relatedCollections.has("scores")) ctes.push(`current_scores AS (
+      SELECT s.*
+      FROM score_events s
+      INNER JOIN root_scope roots ON roots.traceId = s.traceId
+    )`);
+	ctes.push(`candidates AS (
+    SELECT ${TRACE_SELECT}
+    FROM root_scope r
+    WHERE ${conditions.slice(3).join("\n      AND ") || "TRUE"}
+  )`);
+	const candidates = `WITH ${ctes.join(",\n  ")}`;
+	if (plan.result === "groups") {
+		const pageCondition = plan.cursor ? `AND threadId > ?` : "";
+		if (plan.cursor) values.push(plan.cursor.threadId);
+		values.push(plan.limit + 1);
+		return {
+			sql: `${candidates}
+SELECT threadId
+FROM candidates
+WHERE threadId IS NOT NULL ${pageCondition}
+GROUP BY threadId
+ORDER BY threadId ASC
+LIMIT ?`,
+			values
+		};
+	}
+	const orderField = plan.orderBy.field;
+	const direction = plan.orderBy.direction === "asc" ? "ASC" : "DESC";
+	let pageCondition = "";
+	if (plan.cursor) {
+		pageCondition = `WHERE (${orderField} ${plan.orderBy.direction === "asc" ? ">" : "<"} CAST(? AS TIMESTAMP) OR (${orderField} = CAST(? AS TIMESTAMP) AND traceId > ?))`;
+		values.push(plan.cursor.sortValue, plan.cursor.sortValue, plan.cursor.traceId);
+	}
+	values.push(plan.limit + 1);
+	return {
+		sql: `${candidates}
+SELECT *
+FROM candidates
+${pageCondition}
+ORDER BY ${orderField} ${direction}, traceId ASC
+LIMIT ?`,
+		values
+	};
+}
+function asIsoTimestamp(value) {
+	return value instanceof Date ? value.toISOString() : new Date(value).toISOString();
+}
+async function queryTraces(db, plan) {
+	const query = compileDuckDBTraceQuery(plan);
+	const rows = await db.query(query.sql, query.values);
+	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/tracing.ts
 const COLUMNS_SQL = [
 	"eventType",
@@ -3717,11 +4033,16 @@ var ObservabilityStorageDuckDB = class extends ObservabilityStorage {
 		};
 	}
 	getFeatures() {
-		if (!deltaPollingFeatureEnabled()) return ["metrics", "logs"];
+		if (!deltaPollingFeatureEnabled()) return [
+			"metrics",
+			"logs",
+			"trace-query"
+		];
 		return [
 			"metrics",
 			"logs",
-			"delta-polling"
+			"delta-polling",
+			"trace-query"
 		];
 	}
 	async createSpan(args) {
@@ -3751,6 +4072,9 @@ var ObservabilityStorageDuckDB = class extends ObservabilityStorage {
 	async listTraces(args) {
 		return listTraces(this.db, args);
 	}
+	async queryTraces(plan) {
+		return queryTraces(this.db, plan);
+	}
 	async listTracesLight(args) {
 		if (args.mode === "delta") return super.listTracesLight(args);
 		return listTracesLight(this.db, args);
@@ -3858,4 +4182,4 @@ var ObservabilityStorageDuckDB = class extends ObservabilityStorage {
 //#endregion
 export { ObservabilityStorageDuckDB };
 
-//# sourceMappingURL=observability-CT5eidI9.js.map
+//# sourceMappingURL=observability-Cqe4mNHq.js.map