@mastra/duckdb 1.15.0 → 1.16.0-alpha.0

package/dist/{observability-642joQI8.js → observability-Bv3eObB-.js}

@@ -1,4 +1,4 @@
-import { a as DuckDBConnection, i as timeToFirstChunkSupported, o as DuckDBQueryTimeoutError, r as timeToFirstChunkFeatures, t as spanQueryFeatures } from "./features-BI3qUHJ7.js";
+import { a as timeToFirstChunkSupported, i as timeToFirstChunkFeatures, n as spanQueryFeatures, o as DuckDBConnection, s as DuckDBQueryTimeoutError, t as spanAggregateFeatures } from "./features-uOeZunqR.js";
 import { ErrorCategory, ErrorDomain, MastraError } from "@mastra/core/error";
 import * as coreStorage from "@mastra/core/storage";
 import { BRANCH_SPAN_TYPES, METRIC_DISTINCT_COLUMNS, ObservabilityStorage, computeTraceStatus, createStorageErrorId, executeRetentionPrune, feedbackRecordSchema, listBranchesArgsSchema, listFeedbackArgsSchema, listLogsArgsSchema, listMetricsArgsSchema, listScoresArgsSchema, listTracesArgsSchema, resolveRetentionTargets, retentionCutoffMs, toTraceSpans } from "@mastra/core/storage";
@@ -80,6 +80,12 @@ CREATE TABLE IF NOT EXISTS span_events (
   entityVersionId VARCHAR,
   parentEntityVersionId VARCHAR,
   rootEntityVersionId VARCHAR,
+  parentEntityType VARCHAR,
+  parentEntityId VARCHAR,
+  parentEntityName VARCHAR,
+  rootEntityType VARCHAR,
+  rootEntityId VARCHAR,
+  rootEntityName VARCHAR,
 
   -- Context
   userId VARCHAR,
@@ -372,6 +378,12 @@ const ALL_MIGRATIONS = [
 	`ALTER TABLE span_events ADD COLUMN IF NOT EXISTS hasOutputPayload BOOLEAN`,
 	`ALTER TABLE span_events ADD COLUMN IF NOT EXISTS hasErrorPayload BOOLEAN`,
 	`ALTER TABLE span_events ADD COLUMN IF NOT EXISTS timeToFirstChunkMs DOUBLE`,
+	`ALTER TABLE span_events ADD COLUMN IF NOT EXISTS parentEntityType VARCHAR`,
+	`ALTER TABLE span_events ADD COLUMN IF NOT EXISTS parentEntityId VARCHAR`,
+	`ALTER TABLE span_events ADD COLUMN IF NOT EXISTS parentEntityName VARCHAR`,
+	`ALTER TABLE span_events ADD COLUMN IF NOT EXISTS rootEntityType VARCHAR`,
+	`ALTER TABLE span_events ADD COLUMN IF NOT EXISTS rootEntityId VARCHAR`,
+	`ALTER TABLE span_events ADD COLUMN IF NOT EXISTS rootEntityName VARCHAR`,
 	`ALTER TABLE metric_events ADD COLUMN IF NOT EXISTS cursorId BIGINT`,
 	`ALTER TABLE metric_events ADD COLUMN IF NOT EXISTS entityVersionId VARCHAR`,
 	`ALTER TABLE metric_events ADD COLUMN IF NOT EXISTS parentEntityVersionId VARCHAR`,
@@ -3956,15 +3968,8 @@ const identityFields = [
 	"traceId",
 	"spanId"
 ];
-const missingFields = /* @__PURE__ */ new Set([
-	"parentEntityType",
-	"parentEntityId",
-	"parentEntityName",
-	"rootEntityType",
-	"rootEntityId",
-	"rootEntityName",
-	"executionSource"
-]);
+/** Span rows store the execution source as `source`. */
+const columnSql = (field) => field === "executionSource" ? "e.source AS executionSource" : `e.${field}`;
 const identities = identityFields.join(", ");
 const fields = [
 	...identityFields,
@@ -4028,7 +4033,11 @@ function afterCursor(plan) {
 		values
 	};
 }
-function compileDuckDBSpanQuery(plan) {
+/**
+* The spans `querySpans()` selects, as CTEs ending in `selected`: the current completed copy of
+* each span that started in the window and matches `where`. Values are in placeholder order.
+*/
+function compileDuckDBSpanSelection(plan) {
 	const values = [plan.timeRange.from, plan.timeRange.to];
 	const scope = [];
 	if (plan.scope) {
@@ -4043,28 +4052,44 @@ function compileDuckDBSpanQuery(plan) {
 		sql: "TRUE",
 		values: []
 	};
-	const after = afterCursor(plan);
-	values.push(plan.timeRange.from, plan.timeRange.to, ...predicate.values, ...after.values, plan.limit + 1);
+	values.push(plan.timeRange.from, plan.timeRange.to, ...predicate.values);
 	const join = identityFields.map((field) => `e.${field} IS NOT DISTINCT FROM c.${field}`).join(" AND ");
+	const partition = identityFields.map((field) => `e.${field}`).join(", ");
 	const legacyModelFields = coreStorage.SPAN_QUERY_FIELD_REGISTRY ? "" : `CASE WHEN json_type(e.attributes, '$.model') = 'VARCHAR' THEN json_extract_string(e.attributes, '$.model') END AS model,
        CASE WHEN json_type(e.attributes, '$.provider') = 'VARCHAR' THEN json_extract_string(e.attributes, '$.provider') END AS provider,`;
 	return {
-		sql: `WITH candidates AS MATERIALIZED (
+		ctes: [
+			`candidates AS MATERIALIZED (
     SELECT ${identities} FROM span_events
     WHERE eventType = 'start' AND timestamp >= CAST(? AS TIMESTAMP) AND timestamp < CAST(? AS TIMESTAMP)
       ${scope.length ? `AND ${scope.join(" AND ")}` : ""}
     GROUP BY ${identities}
-  ), ranked AS (
-    SELECT ${fields.map((field) => missingFields.has(field) ? `CAST(NULL AS VARCHAR) AS ${field}` : `e.${field}`).join(", ")},
-      min(e.timestamp) FILTER (WHERE e.eventType = 'start') OVER (PARTITION BY ${identityFields.map((field) => `e.${field}`).join(", ")}) AS startedAt, e.endedAt, e.cursorId,
+  )`,
+			`span_versions AS (
+    SELECT ${fields.map(columnSql).join(", ")},
+      min(e.timestamp) FILTER (WHERE e.eventType = 'start') OVER (PARTITION BY ${partition}) AS startedAt, e.endedAt, e.cursorId,
       ${legacyModelFields}
-      date_diff('millisecond', min(e.timestamp) FILTER (WHERE e.eventType = 'start') OVER (PARTITION BY ${identityFields.map((field) => `e.${field}`).join(", ")}), e.endedAt) AS durationMs,
+      date_diff('millisecond', min(e.timestamp) FILTER (WHERE e.eventType = 'start') OVER (PARTITION BY ${partition}), e.endedAt) AS durationMs,
       CASE WHEN e.error IS NULL THEN 'success' ELSE 'error' END AS status,
       CASE WHEN e.error IS NULL THEN NULL ELSE 'error' END AS error,
-      row_number() OVER (PARTITION BY ${identityFields.map((field) => `e.${field}`).join(", ")}
+      row_number() OVER (PARTITION BY ${partition}
         ORDER BY CASE WHEN e.endedAt IS NULL THEN 1 ELSE 0 END, e.cursorId DESC NULLS LAST, e.endedAt DESC) AS rank
     FROM span_events e INNER JOIN candidates c ON ${join}
-  )
+  )`,
+			`selected AS (
+    SELECT * FROM span_versions s
+    WHERE rank = 1 AND endedAt IS NOT NULL AND startedAt >= CAST(? AS TIMESTAMP) AND startedAt < CAST(? AS TIMESTAMP) AND (${predicate.sql})
+  )`
+		],
+		values
+	};
+}
+function compileDuckDBSpanQuery(plan) {
+	const { ctes, values } = compileDuckDBSpanSelection(plan);
+	const after = afterCursor(plan);
+	values.push(...after.values, plan.limit + 1);
+	return {
+		sql: `WITH ${ctes.join(", ")}
   SELECT ${[
 			...fields,
 			"startedAt",
@@ -4073,7 +4098,7 @@ function compileDuckDBSpanQuery(plan) {
 			"status",
 			"cursorId"
 		].map((field) => `s.${field}`).join(", ")}
-  FROM ranked s WHERE rank = 1 AND endedAt IS NOT NULL AND startedAt >= CAST(? AS TIMESTAMP) AND startedAt < CAST(? AS TIMESTAMP) AND (${predicate.sql}) AND ${after.sql}
+  FROM selected s WHERE ${after.sql}
   ORDER BY s.${plan.orderBy.field} ${plan.orderBy.direction === "asc" ? "ASC" : "DESC"} NULLS LAST, ${identityFields.map((field) => `s.${field} ASC NULLS LAST`).join(", ")}
   LIMIT ?`,
 		values
@@ -4416,7 +4441,7 @@ const DIMENSION_SQL = {
 	environment: "r.environment",
 	status: TRACE_STATUS_SQL,
 	serviceName: "r.serviceName",
-	executionSource: "r.executionSource",
+	executionSource: "r.source",
 	threadId: "r.threadId",
 	resourceId: "r.resourceId",
 	userId: "r.userId",
@@ -4424,6 +4449,28 @@ const DIMENSION_SQL = {
 	organizationId: "r.organizationId",
 	experimentId: "r.experimentId"
 };
+/** Span column expressions (alias `r`, a `selected` span row) for each groupable span dimension. */
+const SPAN_DIMENSION_SQL = Object.fromEntries([
+	"name",
+	"spanType",
+	"entityType",
+	"entityId",
+	"entityName",
+	"parentEntityType",
+	"parentEntityName",
+	"rootEntityType",
+	"rootEntityName",
+	"environment",
+	"status",
+	"serviceName",
+	"executionSource",
+	"threadId",
+	"resourceId",
+	"userId",
+	"sessionId",
+	"organizationId",
+	"experimentId"
+].map((field) => [field, `r.${field}`]));
 const PERCENTILES = {
 	"duration.p50": .5,
 	"duration.p90": .9,
@@ -4440,18 +4487,47 @@ const COMPARISON_SQL = {
 };
 function compileDuckDBTraceAggregate(plan, options = {}) {
 	const { ctes, values } = compileDuckDBTraceCandidates(plan, "r.*");
-	const dimensionSql = (field) => {
-		if (field.startsWith("metadata.")) return {
-			sql: TRACE_METADATA_VALUE_SQL,
-			values: traceMetadataPathValues(field.slice(9))
-		};
-		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,
+	return compileDuckDBAggregate(plan, options, {
+		ctes,
+		values,
+		from: "candidates",
+		dimensionSql: (field) => {
+			if (field.startsWith("metadata.")) return {
+				sql: TRACE_METADATA_VALUE_SQL,
+				values: traceMetadataPathValues(field.slice(9))
+			};
+			return {
+				sql: lookupDimensionSql(DIMENSION_SQL, field),
+				values: []
+			};
+		},
+		usageKey: ["traceId"]
+	});
+}
+/**
+* Compiles a trusted span aggregate plan over the spans `querySpans()` selects. Usage is
+* attributed by `(traceId, spanId)`, so only the span that recorded a metric row carries it.
+*/
+function compileDuckDBSpanAggregate(plan, options = {}) {
+	const { ctes, values } = compileDuckDBSpanSelection(plan);
+	return compileDuckDBAggregate(plan, options, {
+		ctes,
+		values,
+		from: "selected",
+		dimensionSql: (field) => ({
+			sql: lookupDimensionSql(SPAN_DIMENSION_SQL, field),
 			values: []
-		};
-	};
+		}),
+		usageKey: ["traceId", "spanId"]
+	});
+}
+function lookupDimensionSql(columns, field) {
+	const sql = Object.hasOwn(columns, field) ? columns[field] : void 0;
+	if (sql === void 0) throw new Error(`Unsupported trusted aggregate field: ${field}`);
+	return sql;
+}
+function compileDuckDBAggregate(plan, options, source) {
+	const { ctes, values, dimensionSql } = source;
 	const distinctSql = (field) => field === "traceId" ? {
 		sql: "r.traceId",
 		values: []
@@ -4482,16 +4558,18 @@ function compileDuckDBTraceAggregate(plan, options = {}) {
 	const hasCost = measureKinds.has("cost");
 	if (!hasCost && !measureKinds.has("tokens")) ctes.push(`facts AS (
     SELECT ${factColumns.join(",\n      ")}
-    FROM candidates r
+    FROM ${source.from} r
   )`);
 	else {
+		const keyColumns = source.usageKey.map((key) => `r.${key}`).join(", ");
 		ctes.push(`base_facts AS (
-    SELECT r.traceId, ${factColumns.join(",\n      ")}
-    FROM candidates r
+    SELECT ${keyColumns}, ${factColumns.join(",\n      ")}
+    FROM ${source.from} r
   )`);
 		ctes.push(...compileUsageCtes(plan, values, {
 			hasCost,
-			dedupe: !options.metricIdsUnique
+			dedupe: !options.metricIdsUnique,
+			key: source.usageKey
 		}));
 		const usageColumns = USAGE_METRIC_NAMES.map((_, index) => `CASE WHEN u.traceId IS NOT NULL THEN COALESCE(u.t${index}, 0) END AS t${index}`);
 		usageColumns.push(`u.traceId IS NOT NULL AS usageBearing`);
@@ -4500,7 +4578,7 @@ function compileDuckDBTraceAggregate(plan, options = {}) {
     SELECT b.*,
       ${usageColumns.join(",\n      ")}
     FROM base_facts b
-    LEFT JOIN usage u ON u.traceId = b.traceId
+    LEFT JOIN usage u ON ${source.usageKey.map((key) => `u.${key} = b.${key}`).join(" AND ")}
   )`);
 	}
 	const measureSql = (name, alias = "") => {
@@ -4622,8 +4700,8 @@ function costColumns(alias) {
 	return [`CAST(count_if(${alias}covered) AS DOUBLE) / NULLIF(count_if(${alias}usageBearing), 0) AS costCoverage`, `CASE WHEN ${mixedUnitsSql(alias)} THEN '${mixed}' ELSE min(${alias}unitMin) END AS costUnit`];
 }
 /**
-* Builds `usage_bounds`, `usage_rows`, and `usage`: one row per candidate trace that has at least
-* one token metric row, holding per-name token sums (`t<i>`, indexed by
+* Builds `usage_bounds`, `usage_rows`, and `usage`: one row per population row (trace, or span)
+* that has at least one token metric row, keyed by the source's usage key, holding per-name token sums (`t<i>`, indexed by
 * `TRACE_AGGREGATE_USAGE_METRIC_NAMES`) and, for cost requests, the trace's pricing state. Metric
 * rows are matched by traceId, usage metric name, tenant scope, and `timestamp >= from` only: a
 * trace that starts in the window can emit metrics after `to`, and the semi-join bounds the scan.
@@ -4637,7 +4715,7 @@ function costColumns(alias) {
 *
 * Placeholders are bound in text order: `param` pushes each value as its placeholder is emitted.
 */
-function compileUsageCtes(plan, values, { hasCost, dedupe }) {
+function compileUsageCtes(plan, values, { hasCost, dedupe, key }) {
 	const param = (value) => {
 		values.push(value);
 		return "?";
@@ -4664,9 +4742,10 @@ function compileUsageCtes(plan, values, { hasCost, dedupe }) {
         CASE WHEN m.name IN (${list(costNames)})
           THEN COALESCE(json_type(m.costMetadata, '$.error'), 'NULL') <> 'NULL'
           ELSE FALSE
-        END AS hasError` : "";
+        END AS hasError,
+        COALESCE(json_extract_string(m.costMetadata, '$.error') = 'partial_cost', FALSE) AS partialCost` : "";
 	const rows = `usage_rows AS (
-      SELECT ${dedupe ? "DISTINCT ON (m.metricId) " : ""}m.traceId, m.name, m.value${costSelect}
+      SELECT ${dedupe ? "DISTINCT ON (m.metricId) " : ""}${key.map((column) => `m.${column}`).join(", ")}, m.name, m.value${costSelect}
       FROM metric_events m
       WHERE m.timestamp >= (SELECT minTs FROM usage_bounds)
         AND m.timestamp <= (SELECT maxTs FROM usage_bounds)
@@ -4675,17 +4754,17 @@ function compileUsageCtes(plan, values, { hasCost, dedupe }) {
     )`;
 	const columns = USAGE_METRIC_NAMES.map((name) => `sum(value) FILTER (WHERE name = ${param(name)})`).map((sql, index) => `${sql} AS t${index}`);
 	if (hasCost) {
-		const priced = "costRow AND estimatedCost IS NOT NULL AND costUnit IS NOT NULL AND NOT hasError";
+		const priced = "costRow AND estimatedCost IS NOT NULL AND costUnit IS NOT NULL AND (NOT hasError OR partialCost)";
 		columns.push(`sum(estimatedCost) FILTER (WHERE ${priced}) AS cost`, `bool_or(${priced}) AS priced`, `bool_or(costRow AND (hasError OR (estimatedCost IS NOT NULL AND costUnit IS NULL))) AS pricingFailure`, `min(costUnit) FILTER (WHERE ${priced}) AS unitMin`, `max(costUnit) FILTER (WHERE ${priced}) AS unitMax`);
 	}
 	return [
 		bounds,
 		rows,
 		`usage AS (
-      SELECT traceId,
+      SELECT ${key.join(", ")},
         ${columns.join(",\n        ")}
       FROM usage_rows
-      GROUP BY traceId
+      GROUP BY ${key.join(", ")}
     )`
 	];
 }
@@ -4693,7 +4772,12 @@ function nullableNumber(value) {
 	return value === null || value === void 0 ? null : Number(value);
 }
 async function aggregateTraces(db, plan, options = {}) {
-	const compiled = compileDuckDBTraceAggregate(plan, options);
+	return runAggregate(db, plan, compileDuckDBTraceAggregate(plan, options));
+}
+async function aggregateSpans(db, plan, options = {}) {
+	return runAggregate(db, plan, compileDuckDBSpanAggregate(plan, options));
+}
+async function runAggregate(db, plan, compiled) {
 	let result;
 	try {
 		result = await db.query(compiled.sql, compiled.values);
@@ -4748,6 +4832,12 @@ const COLUMNS_SQL = [
 	"entityVersionId",
 	"parentEntityVersionId",
 	"rootEntityVersionId",
+	"parentEntityType",
+	"parentEntityId",
+	"parentEntityName",
+	"rootEntityType",
+	"rootEntityId",
+	"rootEntityName",
 	"userId",
 	"organizationId",
 	"resourceId",
@@ -4826,6 +4916,12 @@ const SPAN_RECONSTRUCT_SELECT = `
     ${argMaxNonNull("entityVersionId")},
     ${argMaxNonNull("parentEntityVersionId")},
     ${argMaxNonNull("rootEntityVersionId")},
+    ${argMaxNonNull("parentEntityType")},
+    ${argMaxNonNull("parentEntityId")},
+    ${argMaxNonNull("parentEntityName")},
+    ${argMaxNonNull("rootEntityType")},
+    ${argMaxNonNull("rootEntityId")},
+    ${argMaxNonNull("rootEntityName")},
     ${argMaxNonNull("userId")},
     ${argMaxNonNull("organizationId")},
     ${argMaxNonNull("resourceId")},
@@ -5056,6 +5152,12 @@ function rowToSpanRecord(row) {
 		entityVersionId: row.entityVersionId ?? null,
 		parentEntityVersionId: row.parentEntityVersionId ?? null,
 		rootEntityVersionId: row.rootEntityVersionId ?? null,
+		parentEntityType: row.parentEntityType ?? null,
+		parentEntityId: row.parentEntityId ?? null,
+		parentEntityName: row.parentEntityName ?? null,
+		rootEntityType: row.rootEntityType ?? null,
+		rootEntityId: row.rootEntityId ?? null,
+		rootEntityName: row.rootEntityName ?? null,
 		userId: row.userId ?? null,
 		organizationId: row.organizationId ?? null,
 		resourceId: row.resourceId ?? null,
@@ -5236,6 +5338,12 @@ function toValuesTuple(row) {
 		v(row.entityVersionId),
 		v(row.parentEntityVersionId),
 		v(row.rootEntityVersionId),
+		v(row.parentEntityType),
+		v(row.parentEntityId),
+		v(row.parentEntityName),
+		v(row.rootEntityType),
+		v(row.rootEntityId),
+		v(row.rootEntityName),
 		v(row.userId),
 		v(row.organizationId),
 		v(row.resourceId),
@@ -5289,6 +5397,12 @@ function createStartSpanRow(s) {
 		entityVersionId: s.entityVersionId ?? null,
 		parentEntityVersionId: s.parentEntityVersionId ?? null,
 		rootEntityVersionId: s.rootEntityVersionId ?? null,
+		parentEntityType: s.parentEntityType ?? null,
+		parentEntityId: s.parentEntityId ?? null,
+		parentEntityName: s.parentEntityName ?? null,
+		rootEntityType: s.rootEntityType ?? null,
+		rootEntityId: s.rootEntityId ?? null,
+		rootEntityName: s.rootEntityName ?? null,
 		userId: s.userId ?? null,
 		organizationId: s.organizationId ?? null,
 		resourceId: s.resourceId ?? null,
@@ -5338,6 +5452,12 @@ function createEndSpanRow(s) {
 		entityVersionId: s.entityVersionId ?? null,
 		parentEntityVersionId: s.parentEntityVersionId ?? null,
 		rootEntityVersionId: s.rootEntityVersionId ?? null,
+		parentEntityType: s.parentEntityType ?? null,
+		parentEntityId: s.parentEntityId ?? null,
+		parentEntityName: s.parentEntityName ?? null,
+		rootEntityType: s.rootEntityType ?? null,
+		rootEntityId: s.rootEntityId ?? null,
+		rootEntityName: s.rootEntityName ?? null,
 		userId: s.userId ?? null,
 		organizationId: s.organizationId ?? null,
 		resourceId: s.resourceId ?? null,
@@ -6181,6 +6301,7 @@ var ObservabilityStorageDuckDB = class ObservabilityStorageDuckDB extends Observ
 			"trace-query",
 			"trace-aggregate",
 			...spanQueryFeatures,
+			...spanAggregateFeatures,
 			...typeof coreStorage.planThreadRowsQuery === "function" ? ["thread-query-rows"] : [],
 			"trace-query-root-duration",
 			"trace-query-discovery",
@@ -6205,6 +6326,7 @@ var ObservabilityStorageDuckDB = class ObservabilityStorageDuckDB extends Observ
 			"trace-query",
 			"trace-aggregate",
 			...spanQueryFeatures,
+			...spanAggregateFeatures,
 			...typeof coreStorage.planThreadRowsQuery === "function" ? ["thread-query-rows"] : [],
 			"trace-query-root-duration",
 			"trace-query-discovery",
@@ -6253,6 +6375,9 @@ var ObservabilityStorageDuckDB = class ObservabilityStorageDuckDB extends Observ
 	async aggregateTraces(plan) {
 		return aggregateTraces(this.db, plan, { metricIdsUnique: this.metricIdsUnique });
 	}
+	async aggregateSpans(plan) {
+		return aggregateSpans(this.db, plan, { metricIdsUnique: this.metricIdsUnique });
+	}
 	async getTraceQueryObservedFields(plan) {
 		return getTraceQueryObservedFields(this.db, plan);
 	}
@@ -6376,4 +6501,4 @@ var ObservabilityStorageDuckDB = class ObservabilityStorageDuckDB extends Observ
 //#endregion
 export { ObservabilityStorageDuckDB };
 
-//# sourceMappingURL=observability-642joQI8.js.map
+//# sourceMappingURL=observability-Bv3eObB-.js.map