@mastra/pg 1.30.0-alpha.0 → 1.30.0-alpha.2

package/dist/index.js

@@ -18256,6 +18256,166 @@ async function queryThreads(client, schema, plan, timeoutMs) {
 		}) : null }
 	});
 }
+/** Compile a span-row filter with exactly the same rules as trace span predicates. */
+function compileSpanQueryPredicate(predicate, offset) {
+	return compileScalarPredicate(predicate, SPAN_FIELDS, offset);
+}
+//#endregion
+//#region src/storage/domains/observability/v-next/span-query.ts
+const identityFields = [
+	"organizationId",
+	"resourceId",
+	"traceId",
+	"spanId"
+];
+const fields = [
+	...identityFields,
+	"parentSpanId",
+	"name",
+	"spanType",
+	"entityType",
+	"entityId",
+	"entityName",
+	"entityVersionId",
+	"parentEntityVersionId",
+	"rootEntityVersionId",
+	"runId",
+	"sessionId",
+	"userId"
+];
+const identities = identityFields.map((field) => `"${field}"`).join(", ");
+const identityOrder = identityFields.map((field) => `"${field}" COLLATE "C" ASC NULLS LAST`).join(", ");
+function selectedKeys(rows, values) {
+	return rows.map((row) => `(${identityFields.map((field) => {
+		values.push(row[field]);
+		return `"${field}" ${field === "traceId" || field === "spanId" ? "=" : "IS NOT DISTINCT FROM"} $${values.length}::text`;
+	}).join(" AND ")})`).join(" OR ");
+}
+function afterCursor(plan, values) {
+	if (!plan.cursor) return "TRUE";
+	const cursor = plan.cursor;
+	values.push(cursor.sortValue);
+	const timestamp = `$${values.length}::timestamptz`;
+	const field = `s."${plan.orderBy.field}"`;
+	const terms = [`${field} ${plan.orderBy.direction === "asc" ? ">" : "<"} ${timestamp}`];
+	const equal = [`${field} = ${timestamp}`];
+	for (const key of identityFields) {
+		const ref = `s."${key}" COLLATE "C"`;
+		const value = cursor[key];
+		if (value === null) equal.push(`${ref} IS NULL`);
+		else {
+			values.push(value);
+			const parameter = `$${values.length}::text COLLATE "C"`;
+			terms.push(`(${equal.join(" AND ")} AND (${ref} > ${parameter} OR ${ref} IS NULL))`);
+			equal.push(`${ref} = ${parameter}`);
+		}
+	}
+	return `(${terms.join(" OR ")})`;
+}
+function compilePostgresSpanQuery(schema, plan) {
+	const values = [plan.timeRange.from, plan.timeRange.to];
+	const scope = [];
+	if (plan.scope) {
+		values.push(plan.scope.organizationId);
+		scope.push(`"organizationId" = $${values.length}`);
+		if (plan.scope.resourceId !== void 0) {
+			values.push(plan.scope.resourceId);
+			scope.push(`"resourceId" = $${values.length}`);
+		}
+	}
+	const predicate = plan.where ? compileSpanQueryPredicate(plan.where, values.length + 1) : {
+		sql: "TRUE",
+		values: []
+	};
+	values.push(...predicate.values);
+	const after = afterCursor(plan, values);
+	values.push(plan.limit + 1);
+	const table = qualifiedTable(schema, TABLE_SPAN_EVENTS);
+	return {
+		text: `WITH 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 (
+    SELECT DISTINCT ON (${identities})
+      ${fields.map((field) => `"${field}"`).join(", ")}, "startedAt", "endedAt", "cursorId"::text AS version,
+      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,
+      (EXTRACT(EPOCH FROM ("endedAt" - "startedAt")) * 1000)::float8 AS "durationMs",
+      CASE WHEN error IS NULL THEN 'success' ELSE 'error' END AS status,
+      error IS NOT NULL AS "hasError"
+    FROM ${table} e
+    WHERE NOT "isPending" AND "endedAt" >= $1 AND (
+      "organizationId" IS NULL, COALESCE("organizationId", ''), "resourceId" IS NULL, COALESCE("resourceId", ''), "traceId", "spanId"
+    ) IN (
+      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 (
+    SELECT *, CASE WHEN "hasError" THEN 'error' ELSE NULL END AS error FROM current_spans
+  )
+  SELECT ${[
+			...identityFields,
+			"parentSpanId",
+			"name",
+			"spanType",
+			"entityType",
+			"entityId",
+			"entityName",
+			"model",
+			"provider",
+			"startedAt",
+			"endedAt",
+			"durationMs",
+			"status",
+			"version"
+		].map((field) => `s."${field}"`).join(", ")}
+  FROM candidates s WHERE s."startedAt" >= $1 AND s."startedAt" < $2 AND (${predicate.sql}) AND ${after}
+  ORDER BY s."${plan.orderBy.field}" ${plan.orderBy.direction === "asc" ? "ASC" : "DESC"}, ${identityOrder}
+  LIMIT $${values.length}`,
+		values
+	};
+}
+async function querySpans(client, schema, plan, timeoutMs) {
+	const deadline = performance.now() + coreStorage.resolveTraceQueryTimeoutMs(timeoutMs);
+	return runWithPostgresTraceQueryTimeout(client, timeoutMs, async (tx) => {
+		const query = compilePostgresSpanQuery(schema, plan);
+		const rows = await tx.any(query.text, query.values);
+		const selected = rows.map(({ version: _version, startedAt, endedAt, ...row }) => ({
+			...row,
+			startedAt: startedAt.toISOString(),
+			endedAt: endedAt.toISOString()
+		}));
+		const page = rows.slice(0, plan.limit);
+		if (!page.length) return coreStorage.buildSpanQueryResponse(plan, selected, [], []);
+		const remaining = async () => {
+			const ms = Math.floor(deadline - performance.now());
+			if (ms <= 0) throw new coreStorage.TraceQueryExecutionError();
+			await tx.query(`SELECT set_config('statement_timeout', $1, true)`, [`${ms}ms`]);
+		};
+		await remaining();
+		const payloadValues = [];
+		const payloadWhere = page.map((row) => {
+			const key = selectedKeys([row], payloadValues);
+			payloadValues.push(row.version);
+			return `(${key} AND "cursorId" = $${payloadValues.length}::bigint)`;
+		}).join(" OR ");
+		const payloads = await tx.any(`SELECT ${identities},
+      left(input::text, ${coreStorage.SPAN_QUERY_MAX_PREVIEW_CHARACTERS + 1}) AS "inputPreview", left(output::text, ${coreStorage.SPAN_QUERY_MAX_PREVIEW_CHARACTERS + 1}) AS "outputPreview"
+      FROM ${qualifiedTable(schema, TABLE_SPAN_EVENTS)} WHERE ${payloadWhere}`, payloadValues);
+		await remaining();
+		const metricValues = [];
+		const metricWhere = selectedKeys(page, metricValues);
+		const metrics = await tx.any(`SELECT ${identities}, name, "estimatedCost", "costUnit",
+      "costMetadata" ->> 'allocation' AS allocation, COALESCE("costMetadata" ? 'error', false) AS "costError"
+      FROM (SELECT DISTINCT ON (${identities}, "metricId") ${identities}, name, "estimatedCost", "costUnit", "costMetadata" FROM ${qualifiedTable(schema, TABLE_METRIC_EVENTS)}
+        WHERE (${metricWhere}) AND name LIKE 'mastra_model_%'
+        ORDER BY ${identities}, "metricId", "cursorId" DESC) m
+      LIMIT ${coreStorage.SPAN_QUERY_MAX_COST_METRICS + 1}`, metricValues);
+		if (metrics.length > coreStorage.SPAN_QUERY_MAX_COST_METRICS) throw new coreStorage.TraceQueryResourceLimitError();
+		return coreStorage.buildSpanQueryResponse(plan, selected, payloads, metrics);
+	}, { repeatableRead: true });
+}
 //#endregion
 //#region src/storage/domains/observability/v-next/trace-aggregate.ts
 /** Trace-root column expressions (alias `r`) for each groupable dimension; same values `where` sees. */
@@ -19151,6 +19311,7 @@ async function dangerouslyClearTracing(client, schema) {
 * application database — observability writes will degrade app performance.
 * Use it through `MastraCompositeStore` with a dedicated Postgres connection.
 */
+const spanQueryFeatures = typeof coreStorage.planSpanQuery === "function" ? ["span-query"] : [];
 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({
@@ -19340,6 +19501,7 @@ var ObservabilityStoragePostgresVNext = class ObservabilityStoragePostgresVNext
 			"metric-discovery",
 			"trace-query",
 			"trace-aggregate",
+			...spanQueryFeatures,
 			"trace-query-root-duration",
 			"trace-query-discovery",
 			"thread-query",
@@ -19359,6 +19521,7 @@ var ObservabilityStoragePostgresVNext = class ObservabilityStoragePostgresVNext
 			"delta-polling",
 			"trace-query",
 			"trace-aggregate",
+			...spanQueryFeatures,
 			"trace-query-root-duration",
 			"trace-query-discovery",
 			"thread-query",
@@ -19407,6 +19570,9 @@ var ObservabilityStoragePostgresVNext = class ObservabilityStoragePostgresVNext
 	async listTraces(args) {
 		return this.#run("LIST_TRACES", () => listTraces(this.#readClient, this.#schema, args));
 	}
+	async querySpans(plan) {
+		return this.#run("QUERY_SPANS", () => querySpans(this.#readClient, this.#schema, plan, this.#traceQueryTimeoutMs));
+	}
 	async queryTraces(plan) {
 		return this.#run("QUERY_TRACES", () => queryTraces(this.#readClient, this.#schema, plan, this.#traceQueryTimeoutMs));
 	}
@@ -20424,9 +20590,9 @@ var SchedulesPG = class SchedulesPG extends SchedulesStorage {
 		if (!updated) throw new Error(`Schedule ${id} not found`);
 		return updated;
 	}
-	async updateScheduleNextFire(id, expectedNextFireAt, newNextFireAt, lastFireAt, lastRunId) {
+	async updateScheduleNextFire(id, expectedNextFireAt, newNextFireAt, lastFireAt, lastRunId, newStatus) {
 		return ((await this.#client.query(`UPDATE ${this.#table(TABLE_SCHEDULES)}
-       SET next_fire_at = $1, last_fire_at = $2, last_run_id = $3, updated_at = $4
+       SET next_fire_at = $1, last_fire_at = $2, last_run_id = $3, updated_at = $4, status = COALESCE($8, status)
        WHERE id = $5 AND next_fire_at = $6 AND status = $7`, [
 			newNextFireAt,
 			lastFireAt,
@@ -20434,7 +20600,8 @@ var SchedulesPG = class SchedulesPG extends SchedulesStorage {
 			Date.now(),
 			id,
 			expectedNextFireAt,
-			"active"
+			"active",
+			newStatus ?? null
 		])).rowCount ?? 0) > 0;
 	}
 	async deleteSchedule(id) {