@mastra/pg 1.32.0-alpha.6 → 1.33.0-alpha.0

package/dist/index.js

@@ -20256,7 +20256,12 @@ function afterCursor(plan, values) {
 	}
 	return `(${terms.join(" OR ")})`;
 }
-function compilePostgresSpanQuery(schema, plan) {
+/**
+* The spans `querySpans()` pages through, as a `selected` CTE: the latest completed copy of each
+* span that starts in the window and matches `where`. `aggregateSpans()` groups the same CTE.
+* `values` holds the CTEs' parameters; callers append their own.
+*/
+function compilePostgresSpanSelection(schema, plan) {
 	const values = [plan.timeRange.from, plan.timeRange.to];
 	const scope = [];
 	if (plan.scope) {
@@ -20272,17 +20277,17 @@ function compilePostgresSpanQuery(schema, plan) {
 		values: []
 	};
 	values.push(...predicate.values);
-	const after = afterCursor(plan, values);
-	values.push(plan.limit + 1);
 	const table = qualifiedTable(schema, TABLE_SPAN_EVENTS);
 	const legacyModelFields = coreStorage.SPAN_QUERY_FIELD_REGISTRY ? "" : `CASE WHEN jsonb_typeof(attributes -> 'model') = 'string' THEN attributes ->> 'model' END AS model,
        CASE WHEN jsonb_typeof(attributes -> 'provider') = 'string' THEN attributes ->> 'provider' END AS provider,`;
 	return {
-		text: `WITH candidate_ids AS MATERIALIZED (
+		ctes: [
+			`candidate_ids AS MATERIALIZED (
     SELECT DISTINCT ${identities} FROM ${table}
     WHERE "startedAt" >= $1 AND "startedAt" < $2 AND "endedAt" >= $1 AND NOT "isPending"
       ${scope.length ? `AND ${scope.join(" AND ")}` : ""}
-  ), current_spans AS MATERIALIZED (
+  )`,
+			`current_spans AS MATERIALIZED (
     SELECT DISTINCT ON (${identities})
       ${fields.map((field) => `"${field}"`).join(", ")}, "startedAt", "endedAt", "cursorId"::text AS version,
       ${legacyModelFields}
@@ -20296,9 +20301,23 @@ function compilePostgresSpanQuery(schema, plan) {
       SELECT "organizationId" IS NULL, COALESCE("organizationId", ''), "resourceId" IS NULL, COALESCE("resourceId", ''), "traceId", "spanId" FROM candidate_ids
     ) ${scope.length ? `AND ${scope.join(" AND ")}` : ""}
     ORDER BY ${identities}, "endedAt" DESC, "cursorId" DESC
-  ), candidates AS (
+  )`,
+			`candidates AS (
     SELECT *, CASE WHEN "hasError" THEN 'error' ELSE NULL END AS error FROM current_spans
-  )
+  )`,
+			`selected AS (
+    SELECT * FROM candidates s WHERE s."startedAt" >= $1 AND s."startedAt" < $2 AND (${predicate.sql})
+  )`
+		],
+		values
+	};
+}
+function compilePostgresSpanQuery(schema, plan) {
+	const { ctes, values } = compilePostgresSpanSelection(schema, plan);
+	const after = afterCursor(plan, values);
+	values.push(plan.limit + 1);
+	return {
+		text: `WITH ${ctes.join(", ")}
   SELECT ${[
 			...fields,
 			"startedAt",
@@ -20307,7 +20326,7 @@ function compilePostgresSpanQuery(schema, plan) {
 			"status",
 			"version"
 		].map((field) => `s."${field}"`).join(", ")}
-  FROM candidates s WHERE s."startedAt" >= $1 AND s."startedAt" < $2 AND (${predicate.sql}) AND ${after}
+  FROM selected s WHERE ${after}
   ORDER BY s."${plan.orderBy.field}" ${plan.orderBy.direction === "asc" ? "ASC" : "DESC"} NULLS LAST, ${identityOrder}
   LIMIT $${values.length}`,
 		values
@@ -20641,6 +20660,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,
@@ -20666,16 +20707,43 @@ const COMPARISON_SQL = {
 */
 function compilePostgresTraceAggregate(schema, plan) {
 	const { ctes, values } = compilePostgresTraceCandidates(schema, plan, "r.*");
+	return compilePostgresAggregate(schema, plan, {
+		ctes,
+		values,
+		from: "candidates",
+		dimensionSql: (field, param) => {
+			if (field.startsWith("metadata.")) return traceMetadataValueSql(param(field.slice(9)));
+			return lookupDimensionSql(DIMENSION_SQL, field);
+		},
+		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 compilePostgresSpanAggregate(schema, plan) {
+	const { ctes, values } = compilePostgresSpanSelection(schema, plan);
+	return compilePostgresAggregate(schema, plan, {
+		ctes,
+		values,
+		from: "selected",
+		dimensionSql: (field) => lookupDimensionSql(SPAN_DIMENSION_SQL, field),
+		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 compilePostgresAggregate(schema, plan, source) {
+	const { ctes, values } = source;
 	const param = (value) => {
 		values.push(value);
 		return `$${values.length}`;
 	};
-	const dimensionSql = (field) => {
-		if (field.startsWith("metadata.")) return traceMetadataValueSql(param(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;
-	};
+	const dimensionSql = (field) => source.dimensionSql(field, param);
 	const distinctSql = (field) => field === "traceId" ? "r.\"traceId\"" : dimensionSql(field);
 	const dimensionColumns = plan.dimensions.map((_, index) => `"d${index}"`);
 	const distinctColumns = /* @__PURE__ */ new Map();
@@ -20696,14 +20764,14 @@ function compilePostgresTraceAggregate(schema, plan) {
 	const hasCost = measureKinds.has("cost");
 	let usageJoin = "";
 	if (hasCost || measureKinds.has("tokens")) {
-		ctes.push(...compileUsageCtes(schema, plan, param));
-		usageJoin = "\n    LEFT JOIN usage u ON u.\"traceId\" = r.\"traceId\"";
+		ctes.push(...compileUsageCtes(schema, plan, source, param));
+		usageJoin = `\n    LEFT JOIN usage u ON ${source.usageKey.map((key) => `u."${key}" = r."${key}"`).join(" AND ")}`;
 		USAGE_METRIC_NAMES.forEach((_, index) => factColumns.push(`CASE WHEN u."traceId" IS NOT NULL THEN COALESCE(u."t${index}", 0) END AS "t${index}"`));
 		factColumns.push(`u."traceId" IS NOT NULL AS "usageBearing"`, `CASE WHEN u."priced" THEN u."cost" END AS "traceCost"`, `COALESCE(u."priced" AND NOT u."pricingFailure", FALSE) AS "covered"`, `u."unitMin"`, `u."unitMax"`);
 	}
 	ctes.push(`facts AS (
     SELECT ${factColumns.join(",\n      ")}
-    FROM candidates r${usageJoin}
+    FROM ${source.from} r${usageJoin}
   )`);
 	const mixedUnitsSql = (alias) => `MIN(${alias}"unitMin" COLLATE "C") <> MAX(${alias}"unitMax" COLLATE "C")`;
 	const measureSql = (name, alias = "") => {
@@ -20819,36 +20887,39 @@ function usageMetricIndex(name) {
 	return index;
 }
 /**
-* Builds `usage`: one row per candidate trace that has at least one token metric row, holding
-* per-name token sums (`t<i>`, indexed by `TRACE_AGGREGATE_USAGE_METRIC_NAMES`) and 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 candidates semi-join already bounds the scan.
+* Builds `usage`: one row per population row (trace, or span) that has at least one token metric
+* row, holding per-name token sums (`t<i>`, indexed by `TRACE_AGGREGATE_USAGE_METRIC_NAMES`) and
+* its pricing state. Metric rows are matched by the source's usage key, usage metric name, tenant
+* scope, and `timestamp >= from` only: a trace or span that starts in the window can emit metrics
+* after `to`, and the population semi-join already bounds the scan.
 *
 * Rows are deduplicated on `metricId` alone. The primary key is (`metricId`, `timestamp`)
 * because the table is partitioned by timestamp, so retried copies with different timestamps
 * are both stored.
 */
-function compileUsageCtes(schema, plan, param) {
+function compileUsageCtes(schema, plan, source, param) {
+	const key = source.usageKey.map((column) => `"${column}"`).join(", ");
+	const metricKey = source.usageKey.map((column) => `m."${column}"`).join(", ");
 	const scopeConditions = plan.scope ? [`m."organizationId" = ${param(plan.scope.organizationId)}`, ...plan.scope.resourceId === void 0 ? [] : [`m."resourceId" = ${param(plan.scope.resourceId)}`]] : [];
 	const rowConditions = [
-		`m."traceId" IN (SELECT "traceId" FROM candidates)`,
+		`${source.usageKey.length === 1 ? metricKey : `(${metricKey})`} IN (SELECT ${key} FROM ${source.from})`,
 		`m."name" = ANY(${param(USAGE_METRIC_NAMES)}::text[])`,
 		`m."timestamp" >= ${param(plan.timeRange.from)}::timestamptz`,
 		...scopeConditions
 	];
 	const costRow = `"name" = ANY(${param(coreStorage.TRACE_AGGREGATE_COST_METRIC_NAMES)}::text[])`;
-	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")`;
 	const failed = `${costRow} AND ("hasError" OR ("estimatedCost" IS NOT NULL AND "costUnit" IS NULL))`;
 	const tokenColumns = USAGE_METRIC_NAMES.map((name, index) => `SUM("value") FILTER (WHERE "name" = ${param(name)}) AS "t${index}"`);
 	return [`usage_rows AS (
-    SELECT DISTINCT ON (m."metricId") m."traceId", m."name", m."value", m."estimatedCost", m."costUnit",
-      COALESCE(jsonb_typeof(m."costMetadata" -> 'error'), 'null') <> 'null' AS "hasError"
+    SELECT DISTINCT ON (m."metricId") ${metricKey}, m."name", m."value", m."estimatedCost", m."costUnit",
+      COALESCE(jsonb_typeof(m."costMetadata" -> 'error'), 'null') <> 'null' AS "hasError",
+      COALESCE(m."costMetadata" ->> 'error' = 'partial_cost', FALSE) AS "partialCost"
     FROM ${qualifiedTable(schema, TABLE_METRIC_EVENTS)} m
     WHERE ${rowConditions.join("\n      AND ")}
     ORDER BY m."metricId", m."timestamp"
   )`, `usage AS (
-    SELECT "traceId",
+    SELECT ${key},
       ${tokenColumns.join(",\n      ")},
       SUM("estimatedCost") FILTER (WHERE ${priced}) AS "cost",
       bool_or(${priced}) AS "priced",
@@ -20856,7 +20927,7 @@ function compileUsageCtes(schema, plan, param) {
       MIN("costUnit" COLLATE "C") FILTER (WHERE ${priced}) AS "unitMin",
       MAX("costUnit" COLLATE "C") FILTER (WHERE ${priced}) AS "unitMax"
     FROM usage_rows
-    GROUP BY "traceId"
+    GROUP BY ${key}
   )`];
 }
 function nullableNumber(value) {
@@ -20866,7 +20937,12 @@ function toIsoTimestamp(value) {
 	return value instanceof Date ? value.toISOString() : new Date(value).toISOString();
 }
 async function aggregateTraces(client, schema, plan, timeoutMs) {
-	const compiled = compilePostgresTraceAggregate(schema, plan);
+	return runAggregate(client, plan, compilePostgresTraceAggregate(schema, plan), timeoutMs);
+}
+async function aggregateSpans(client, schema, plan, timeoutMs) {
+	return runAggregate(client, plan, compilePostgresSpanAggregate(schema, plan), timeoutMs);
+}
+async function runAggregate(client, plan, compiled, timeoutMs) {
 	const result = await runWithPostgresTraceQueryTimeout(client, timeoutMs, (transaction) => transaction.any(compiled.text, compiled.values));
 	let truncated;
 	let resultRows;
@@ -21662,6 +21738,7 @@ async function dangerouslyClearTracing(client, schema) {
 * Use it through `MastraCompositeStore` with a dedicated Postgres connection.
 */
 const spanQueryFeatures = typeof coreStorage.planSpanQuery === "function" ? ["span-query", "span-query-context-fields"] : [];
+const spanAggregateFeatures = typeof coreStorage.planSpanAggregate === "function" ? ["span-aggregate"] : [];
 function wrapError(op, error, details) {
 	if (error instanceof MastraError || error instanceof coreStorage.TraceQueryCursorError || error instanceof coreStorage.TraceQueryExecutionError || error instanceof coreStorage.TraceQueryResourceLimitError) throw error;
 	throw new MastraError({
@@ -21852,6 +21929,7 @@ var ObservabilityStoragePostgresVNext = class ObservabilityStoragePostgresVNext
 			"trace-query",
 			"trace-aggregate",
 			...spanQueryFeatures,
+			...spanAggregateFeatures,
 			...typeof coreStorage.planThreadRowsQuery === "function" ? ["thread-query-rows"] : [],
 			"trace-query-root-duration",
 			"trace-query-discovery",
@@ -21876,6 +21954,7 @@ var ObservabilityStoragePostgresVNext = class ObservabilityStoragePostgresVNext
 			"trace-query",
 			"trace-aggregate",
 			...spanQueryFeatures,
+			...spanAggregateFeatures,
 			...typeof coreStorage.planThreadRowsQuery === "function" ? ["thread-query-rows"] : [],
 			"trace-query-root-duration",
 			"trace-query-discovery",
@@ -21940,6 +22019,9 @@ var ObservabilityStoragePostgresVNext = class ObservabilityStoragePostgresVNext
 	async aggregateTraces(plan) {
 		return this.#run("AGGREGATE_TRACES", () => aggregateTraces(this.#readClient, this.#schema, plan, this.#traceQueryTimeoutMs));
 	}
+	async aggregateSpans(plan) {
+		return this.#run("AGGREGATE_SPANS", () => aggregateSpans(this.#readClient, this.#schema, plan, this.#traceQueryTimeoutMs));
+	}
 	async getTraceQueryObservedFields(plan) {
 		return this.#run("GET_TRACE_QUERY_OBSERVED_FIELDS", () => getTraceQueryObservedFields(this.#readClient, this.#schema, plan, this.#traceQueryTimeoutMs));
 	}