@mastra/pg 1.25.0 → 1.26.0-alpha.1

package/dist/index.js

@@ -2460,6 +2460,7 @@ async function loadSchemaSnapshot(client, schemaName) {
 }
 //#endregion
 //#region src/storage/db/index.ts
+const POSTGRES_MAX_BIND_PARAMETERS$1 = 65535;
 /**
 * Resolves PgDomainConfig to a database client and schema.
 * Handles creating a new pool if config is provided.
@@ -3013,6 +3014,90 @@ var PgDB = class extends MastraBase {
              ON CONFLICT ("traceId", "spanId") DO NOTHING`, values);
 		} else await client.none(`INSERT INTO ${fullTableName} (${columnList}) VALUES (${placeholders})`, values);
 	}
+	getChunkRowLimit(columnCount) {
+		if (columnCount === 0) return 0;
+		return Math.max(1, Math.floor(POSTGRES_MAX_BIND_PARAMETERS$1 / columnCount));
+	}
+	getSpanConflictIdentifier(record) {
+		const traceId = record.traceId;
+		const spanId = record.spanId;
+		if (traceId === void 0 || spanId === void 0) return;
+		return `${String(traceId)}|${String(spanId)}`;
+	}
+	async normalizeForInsert(tableName, record) {
+		this.addTimestampZColumns(record);
+		const filteredRecord = await this.filterRecordToKnownColumns(tableName, record);
+		return {
+			columns: Object.keys(filteredRecord).map((column) => parseSqlIdentifier(column, "column name")),
+			values: this.prepareValuesForInsert(filteredRecord, tableName),
+			conflictKey: tableName === TABLE_SPANS ? this.getSpanConflictIdentifier(filteredRecord) : void 0
+		};
+	}
+	buildMultiRowInsertStatement({ tableName, columns, rows }) {
+		const fullTableName = getTableName$6({
+			indexName: tableName,
+			schemaName: getSchemaName$6(this.schemaName)
+		});
+		const columnList = columns.map((column) => `"${column}"`).join(", ");
+		const bindParams = [];
+		const values = [];
+		let bindIndex = 1;
+		for (const rowValues of rows) {
+			const placeholders = rowValues.map(() => `$${bindIndex++}`);
+			bindParams.push(`(${placeholders.join(", ")})`);
+			values.push(...rowValues);
+		}
+		let query = `INSERT INTO ${fullTableName} (${columnList}) VALUES ${bindParams.join(", ")}`;
+		if (tableName === TABLE_SPANS) {
+			const updateColumns = columns.filter((column) => column !== "traceId" && column !== "spanId");
+			if (updateColumns.length > 0) {
+				const updateClause = updateColumns.map((column) => `"${column}" = EXCLUDED."${column}"`).join(", ");
+				query += ` ON CONFLICT ("traceId", "spanId") DO UPDATE SET ${updateClause}`;
+			} else query += ` ON CONFLICT ("traceId", "spanId") DO NOTHING`;
+		}
+		return {
+			query,
+			values
+		};
+	}
+	async executeBatchInsert(client, { tableName, records }) {
+		const preparedRecords = [];
+		for (const record of records) preparedRecords.push(await this.normalizeForInsert(tableName, record));
+		let pendingColumns;
+		let pendingConflictKeys = /* @__PURE__ */ new Set();
+		let pendingRows = [];
+		let pendingLimit = 0;
+		const flush = async () => {
+			if (!pendingColumns || pendingRows.length === 0) return;
+			const statement = this.buildMultiRowInsertStatement({
+				tableName,
+				columns: pendingColumns,
+				rows: pendingRows
+			});
+			await client.none(statement.query, statement.values);
+			pendingColumns = void 0;
+			pendingRows = [];
+			pendingConflictKeys = /* @__PURE__ */ new Set();
+			pendingLimit = 0;
+		};
+		for (const { columns, values, conflictKey } of preparedRecords) {
+			if (columns.length === 0) continue;
+			const columnsSignature = columns.join("\0");
+			const currentPendingColumns = pendingColumns;
+			const isSpans = tableName === TABLE_SPANS;
+			const conflictDuplicate = isSpans && conflictKey !== void 0 && pendingConflictKeys.has(conflictKey);
+			const exceedsLimit = pendingRows.length >= pendingLimit;
+			if (currentPendingColumns === void 0 || columnsSignature !== currentPendingColumns.join("\0") || conflictDuplicate || exceedsLimit) {
+				await flush();
+				pendingColumns = columns;
+				pendingLimit = this.getChunkRowLimit(columns.length);
+				pendingRows = [values];
+				pendingConflictKeys = /* @__PURE__ */ new Set();
+			} else pendingRows.push(values);
+			if (isSpans && conflictKey !== void 0) pendingConflictKeys.add(conflictKey);
+		}
+		await flush();
+	}
 	async insert({ tableName, record }) {
 		try {
 			await this.executeInsert(this.client, {
@@ -3461,9 +3546,9 @@ MIGRATION REQUIRED: Duplicate spans detected in ${duplicateInfo.tableName}\n====
 	async batchInsert({ tableName, records }) {
 		try {
 			await this.client.tx(async (tx) => {
-				for (const record of records) await this.executeInsert(tx, {
+				await this.executeBatchInsert(tx, {
 					tableName,
-					record
+					records
 				});
 			});
 		} catch (error) {
@@ -17607,6 +17692,75 @@ LIMIT $${values.length}`,
 		values
 	};
 }
+function discoveryRegistry(scope) {
+	if (scope === "trace") return TRACE_FIELDS;
+	if (scope === "spans") return SPAN_FIELDS;
+	if (scope === "scores") return SCORE_FIELDS;
+	return FEEDBACK_FIELDS;
+}
+function discoverySource(scope) {
+	if (scope === "trace") return "root_scope r";
+	if (scope === "spans") return "current_spans s";
+	if (scope === "scores") return "current_scores s";
+	return "current_feedback s";
+}
+function discoveryCollections(scope) {
+	return scope === "trace" ? /* @__PURE__ */ new Set() : /* @__PURE__ */ new Set([scope]);
+}
+function compilePostgresTraceQueryObservedFields(schema, plan) {
+	const { ctes, values } = compilePostgresTraceScope(schema, plan, /* @__PURE__ */ new Set());
+	const searchParameter = values.length + 1;
+	const search = plan.search ? `AND strpos(lower('metadata.' || entry.key), lower($${searchParameter})) > 0` : "";
+	if (plan.search) values.push(plan.search);
+	values.push(plan.limit + 1);
+	return {
+		text: `WITH ${ctes.join(",\n")}
+SELECT 'metadata.' || entry.key AS path, count(*)::bigint AS occurrences
+FROM root_scope r
+CROSS JOIN LATERAL jsonb_each(CASE WHEN jsonb_typeof(r."metadataRaw") = 'object' THEN r."metadataRaw" ELSE '{}'::jsonb END) entry
+WHERE jsonb_typeof(entry.value) = 'string'
+  AND btrim(entry.value #>> '{}') <> ''
+  AND entry.key <> ''
+  AND strpos(entry.key, '.') = 0
+  AND octet_length('metadata.' || entry.key) <= ${coreStorage.TRACE_QUERY_MAX_PATH_BYTES}
+  AND octet_length(entry.value #>> '{}') <= ${coreStorage.TRACE_QUERY_MAX_STRING_BYTES}
+  ${search}
+GROUP BY entry.key
+ORDER BY occurrences DESC, ('metadata.' || entry.key) COLLATE "C" ASC
+LIMIT $${values.length}`,
+		values
+	};
+}
+function compilePostgresTraceQueryValues(schema, plan) {
+	const { ctes, values } = compilePostgresTraceScope(schema, plan, discoveryCollections(plan.predicateScope));
+	let field;
+	if (plan.predicateScope === "trace" && plan.path.startsWith("metadata.")) {
+		const keyParameter = `$${values.length + 1}`;
+		field = `COALESCE(
+      CASE WHEN jsonb_typeof(r."metadataSearch" -> ${keyParameter}) = 'string' THEN r."metadataSearch" ->> ${keyParameter} END,
+      CASE WHEN jsonb_typeof(r."metadataRaw" -> ${keyParameter}) = 'string' THEN NULLIF(btrim(r."metadataRaw" ->> ${keyParameter}), '') END
+    )`;
+		values.push(plan.path.slice(9));
+	} else field = fieldSql(discoveryRegistry(plan.predicateScope), plan.path);
+	const searchParameter = values.length + 1;
+	const search = plan.search ? `AND strpos(lower(value), lower($${searchParameter})) > 0` : "";
+	if (plan.search) values.push(plan.search);
+	values.push(plan.limit + 1);
+	return {
+		text: `WITH ${ctes.join(",\n")}, extracted AS (
+  SELECT ${field}::text AS value FROM ${discoverySource(plan.predicateScope)}
+)
+SELECT value, count(*)::bigint AS count
+FROM extracted
+WHERE value IS NOT NULL
+  AND octet_length(value) <= ${coreStorage.TRACE_QUERY_MAX_STRING_BYTES}
+  ${search}
+GROUP BY value
+ORDER BY count DESC, value COLLATE "C" ASC
+LIMIT $${values.length}`,
+		values
+	};
+}
 function asIsoTimestamp$1(value) {
 	if (value === null || value === void 0) throw new Error("Trace query returned a null timestamp");
 	return value instanceof Date ? value.toISOString() : new Date(value).toISOString();
@@ -17628,6 +17782,29 @@ async function runWithPostgresTraceQueryTimeout(client, timeoutMs, execute) {
 		throw error;
 	}
 }
+async function getTraceQueryObservedFields(client, schema, plan, timeoutMs) {
+	if (plan.predicateScope !== "trace") return {
+		observedFields: [],
+		observedFieldsTruncated: false
+	};
+	const query = compilePostgresTraceQueryObservedFields(schema, plan);
+	const rows = await runWithPostgresTraceQueryTimeout(client, timeoutMs, (transaction) => transaction.any(query.text, query.values));
+	return {
+		observedFields: rows.slice(0, plan.limit).map((row) => coreStorage.createTraceQueryObservedFieldDescriptor(String(row.path), Number(row.occurrences))),
+		observedFieldsTruncated: rows.length > plan.limit
+	};
+}
+async function getTraceQueryValues(client, schema, plan, timeoutMs) {
+	const query = compilePostgresTraceQueryValues(schema, plan);
+	const rows = await runWithPostgresTraceQueryTimeout(client, timeoutMs, (transaction) => transaction.any(query.text, query.values));
+	return coreStorage.getTraceQueryValuesResponseSchema.parse({
+		values: rows.slice(0, plan.limit).map((row) => ({
+			value: String(row.value),
+			count: Number(row.count)
+		})),
+		valuesTruncated: rows.length > plan.limit
+	});
+}
 async function queryTraces(client, schema, plan, timeoutMs) {
 	const query = compilePostgresTraceQuery(schema, plan);
 	const rows = await runWithPostgresTraceQueryTimeout(client, timeoutMs, (transaction) => transaction.any(query.text, query.values));
@@ -18574,6 +18751,7 @@ var ObservabilityStoragePostgresVNext = class ObservabilityStoragePostgresVNext
 			"metrics",
 			"logs",
 			"trace-query",
+			"trace-query-discovery",
 			"thread-query"
 		];
 		return [
@@ -18581,6 +18759,7 @@ var ObservabilityStoragePostgresVNext = class ObservabilityStoragePostgresVNext
 			"logs",
 			"delta-polling",
 			"trace-query",
+			"trace-query-discovery",
 			"thread-query"
 		];
 	}
@@ -18627,6 +18806,12 @@ var ObservabilityStoragePostgresVNext = class ObservabilityStoragePostgresVNext
 	async queryTraces(plan) {
 		return this.#run("QUERY_TRACES", () => queryTraces(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));
+	}
+	async getTraceQueryValues(plan) {
+		return this.#run("GET_TRACE_QUERY_VALUES", () => getTraceQueryValues(this.#readClient, this.#schema, plan, this.#traceQueryTimeoutMs));
+	}
 	async queryThreads(plan) {
 		return this.#run("QUERY_THREADS", () => queryThreads(this.#readClient, this.#schema, plan, this.#traceQueryTimeoutMs));
 	}
@@ -23406,7 +23591,12 @@ var PgFactoryStorageOps = class {
 		const assignments = columns.map((column, i) => `"${column}" = $${i + 1}`).join(", ");
 		const filter = this.#buildWhere(schema, where, columns.length + 1);
 		const args = [...columns.map((column) => this.#serialize(this.#column(schema, column), set[column])), ...filter.args];
-		return (await queryable.query(`UPDATE "${schema.name}" SET ${assignments} WHERE ${filter.sql}`, args)).rowCount ?? 0;
+		try {
+			return (await queryable.query(`UPDATE "${schema.name}" SET ${assignments} WHERE ${filter.sql}`, args)).rowCount ?? 0;
+		} catch (error) {
+			if (isUniqueViolation(error)) throw new UniqueViolationError(collection, { cause: error });
+			throw error;
+		}
 	}
 	async deleteMany(collection, where) {
 		const schema = this.#schema(collection);