@mastra/pg 1.36.0-alpha.1 → 1.36.0-alpha.3

package/dist/index.js

@@ -2696,7 +2696,8 @@ async function loadSchemaSnapshot(client, schemaName) {
 		notNullColumns,
 		indexes,
 		replicaIdentityIndexes,
-		primaryKeyIndexes
+		primaryKeyIndexes,
+		createdTables: /* @__PURE__ */ new Set()
 	};
 }
 // src/storage/db/index.ts
@@ -2856,26 +2857,30 @@ function generateTableSQL({ tableName, schema, schemaName, compositePrimaryKey,
             ` : ""}
           `;
 }
+function indexColumnsSQL(columns, opclass) {
+	return columns.map((col) => {
+		const [colName, ...modifiers] = col.includes(" DESC") || col.includes(" ASC") ? col.split(" ") : [col];
+		if (!colName) throw new Error(`Invalid column specification: ${col}`);
+		return [
+			`"${parseSqlIdentifier(colName, "column name")}"`,
+			opclass,
+			...modifiers
+		].filter(Boolean).join(" ");
+	}).join(", ");
+}
 /**
 * Generates a CREATE INDEX SQL statement from index options.
 * Used by exportSchemas to produce index DDL without a database connection.
 */
 function generateIndexSQL(options, schemaName) {
-	const { name, table, columns, unique = false, where, method = "btree" } = options;
+	const { name, table, columns, unique = false, where, method = "btree", opclass } = options;
 	const fullTableName = getTableName$6({
 		indexName: table,
 		schemaName: getSchemaName$6(schemaName)
 	});
 	const uniqueStr = unique ? "UNIQUE " : "";
 	const methodStr = method !== "btree" ? `USING ${method} ` : "";
-	const columnsStr = columns.map((col) => {
-		if (col.includes(" DESC") || col.includes(" ASC")) {
-			const [colName, ...modifiers] = col.split(" ");
-			if (!colName) throw new Error(`Invalid column specification: ${col}`);
-			return `"${parseSqlIdentifier(colName, "column name")}" ${modifiers.join(" ")}`;
-		}
-		return `"${parseSqlIdentifier(col, "column name")}"`;
-	}).join(", ");
+	const columnsStr = indexColumnsSQL(columns, opclass);
 	const whereStr = where ? ` WHERE ${where}` : "";
 	return `CREATE ${uniqueStr}INDEX IF NOT EXISTS ${`"${parseSqlIdentifier(name, "index name")}"`} ON ${fullTableName} ${methodStr}(${columnsStr})${whereStr};`;
 }
@@ -3386,6 +3391,7 @@ var PgDB = class extends MastraBase {
 			if (this.schemaName) await this.setupSchema();
 			const snapshot = this.schemaSnapshot;
 			if (!snapshot || !this.snapshotShowsTableConverged(snapshot, tableName)) {
+				const isNewTable = snapshot ? !snapshot.tables.has(tableName) : false;
 				const sql = generateTableSQL({
 					tableName,
 					schema,
@@ -3399,6 +3405,7 @@ var PgDB = class extends MastraBase {
 				}
 				if (snapshot) {
 					snapshot.tables.add(tableName);
+					if (isNewTable) snapshot.createdTables.add(tableName);
 					const created = this.snapshotColumns(snapshot, tableName);
 					for (const [name, def] of Object.entries(schema)) {
 						const parsedName = parseSqlIdentifier(name, "column name");
@@ -3867,18 +3874,9 @@ MIGRATION REQUIRED: Duplicate spans detected in ${duplicateInfo.tableName}\n====
 			const snapshot = this.schemaSnapshot;
 			if (snapshot?.indexes.has(name)) return;
 			const uniqueStr = unique ? "UNIQUE " : "";
-			const concurrentStr = concurrent ? "CONCURRENTLY " : "";
+			const concurrentStr = concurrent && !this.isTableCreatedInThisInit(table) ? "CONCURRENTLY " : "";
 			const methodStr = method !== "btree" ? `USING ${method} ` : "";
-			const columnsStr = columns.map((col) => {
-				if (col.includes(" DESC") || col.includes(" ASC")) {
-					const [colName, ...modifiers] = col.split(" ");
-					if (!colName) throw new Error(`Invalid column specification: ${col}`);
-					const quotedCol = `"${parseSqlIdentifier(colName, "column name")}" ${modifiers.join(" ")}`;
-					return opclass ? `${quotedCol} ${opclass}` : quotedCol;
-				}
-				const quotedCol = `"${parseSqlIdentifier(col, "column name")}"`;
-				return opclass ? `${quotedCol} ${opclass}` : quotedCol;
-			}).join(", ");
+			const columnsStr = indexColumnsSQL(columns, opclass);
 			const whereStr = where ? ` WHERE ${where}` : "";
 			const tablespaceStr = tablespace ? ` TABLESPACE ${tablespace}` : "";
 			let withStr = "";
@@ -3911,6 +3909,13 @@ MIGRATION REQUIRED: Duplicate spans detected in ${duplicateInfo.tableName}\n====
 		if (snapshot?.indexes.has(indexName)) return;
 		if (await buildIndex(this.client, this.schemaName || "public", indexName, () => this.timedIndexBuild(indexName, sql), this.logger)) snapshot?.indexes.add(indexName);
 	}
+	/**
+	* Whether `tableName` was absent when this init's snapshot was read and this
+	* init created it, so its indexes can be built without CONCURRENTLY.
+	*/
+	isTableCreatedInThisInit(tableName) {
+		return this.schemaSnapshot?.createdTables.has(tableName) ?? false;
+	}
 	async timedIndexBuild(indexName, sql) {
 		const startedAt = Date.now();
 		this.logger.info(`Creating index ${indexName}`);
@@ -11650,6 +11655,13 @@ function dedupeMessagesForSave(messages) {
 	}
 	return Array.from(deduped.values());
 }
+/** Shallow merge; keys patched to `undefined` are dropped when serialized. Non-object metadata is replaced. */
+function mergeMetadata(existing, patch) {
+	return {
+		...existing && typeof existing === "object" && !Array.isArray(existing) ? existing : {},
+		...patch
+	};
+}
 var MemoryPG = class MemoryPG extends MemoryStorage {
 	supportsPartialThreadUpdate = true;
 	supportsObservationalMemory = true;
@@ -12216,7 +12228,7 @@ var MemoryPG = class MemoryPG extends MemoryStorage {
 			const createdAt = toUtcISOString(thread.createdAt);
 			const updatedAt = toUtcISOString(thread.updatedAt);
 			const metadataJson = thread.metadata ? toPgJson(thread.metadata) : null;
-			await this.#fenced(fence, "saveThread", (q) => q.none(`INSERT INTO ${tableName} (
+			const saved = await this.#fenced(fence, "saveThread", (q) => q.oneOrNone(`INSERT INTO ${tableName} AS t (
           id,
           "resourceId",
           title,
@@ -12227,13 +12239,12 @@ var MemoryPG = class MemoryPG extends MemoryStorage {
           "updatedAtZ"
         ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8)
         ON CONFLICT (id) DO UPDATE SET
-          "resourceId" = EXCLUDED."resourceId",
           title = EXCLUDED.title,
           metadata = EXCLUDED.metadata,
-          "createdAt" = EXCLUDED."createdAt",
-          "createdAtZ" = EXCLUDED."createdAtZ",
           "updatedAt" = EXCLUDED."updatedAt",
-          "updatedAtZ" = EXCLUDED."updatedAtZ"`, [
+          "updatedAtZ" = EXCLUDED."updatedAtZ"
+        WHERE t."resourceId" = EXCLUDED."resourceId"
+        RETURNING *`, [
 				thread.id,
 				thread.resourceId,
 				thread.title,
@@ -12243,12 +12254,26 @@ var MemoryPG = class MemoryPG extends MemoryStorage {
 				updatedAt,
 				updatedAt
 			]));
+			if (!saved) throw new MastraError({
+				id: createStorageErrorId("PG", "SAVE_THREAD", "RESOURCE_ID_MISMATCH"),
+				domain: ErrorDomain.STORAGE,
+				category: ErrorCategory.USER,
+				text: `Thread ${thread.id} already exists and belongs to a different resource. Use updateThreadResourceId to move a thread and its messages to another resource.`,
+				details: {
+					threadId: thread.id,
+					resourceId: thread.resourceId
+				}
+			});
 			return {
-				...thread,
-				metadata: metadataJson ? JSON.parse(metadataJson) : thread.metadata
+				id: saved.id,
+				resourceId: saved.resourceId,
+				title: saved.title,
+				metadata: saved.metadata == null ? void 0 : typeof saved.metadata === "string" ? JSON.parse(saved.metadata) : saved.metadata,
+				createdAt: saved.createdAtZ || saved.createdAt,
+				updatedAt: saved.updatedAtZ || saved.updatedAt
 			};
 		} catch (error) {
-			if (isRunFenceConflictError(error)) throw error;
+			if (error instanceof MastraError || isRunFenceConflictError(error)) throw error;
 			throw new MastraError({
 				id: createStorageErrorId("PG", "SAVE_THREAD", "FAILED"),
 				domain: ErrorDomain.STORAGE,
@@ -12257,52 +12282,125 @@ var MemoryPG = class MemoryPG extends MemoryStorage {
 			}, error);
 		}
 	}
-	async updateThread({ id, title, metadata, fence }) {
+	#toThread(row) {
+		return {
+			id: row.id,
+			resourceId: row.resourceId,
+			title: row.title,
+			metadata: typeof row.metadata === "string" ? JSON.parse(row.metadata) : row.metadata,
+			createdAt: row.createdAtZ || row.createdAt,
+			updatedAt: row.updatedAtZ || row.updatedAt
+		};
+	}
+	async #lockThread(t, tableName, { id, resourceId }) {
+		const params = [id];
+		let query = `SELECT * FROM ${tableName} WHERE id = $1`;
+		if (resourceId !== void 0) {
+			query += ` AND "resourceId" = $2`;
+			params.push(resourceId);
+		}
+		const row = await t.oneOrNone(`${query} FOR UPDATE`, params);
+		return row ? this.#toThread(row) : null;
+	}
+	/**
+	* Writes `title` (when given) and the full `metadata` value (when given).
+	* Callers that merge metadata must hold the row lock from `#lockThread`.
+	*/
+	async #writeThread(q, tableName, { id, title, metadata }) {
+		const nowStr = toUtcISOString(/* @__PURE__ */ new Date());
+		const params = [
+			title ?? null,
+			nowStr,
+			nowStr,
+			id
+		];
+		let metadataSet = "";
+		if (metadata !== void 0) {
+			params.push(toPgJson(metadata));
+			metadataSet = "metadata = $5,";
+		}
+		const row = await q.oneOrNone(`UPDATE ${tableName}
+       SET
+         title = COALESCE($1, title),
+         ${metadataSet}
+         "updatedAt" = $2,
+         "updatedAtZ" = $3
+       WHERE id = $4
+       RETURNING *`, params);
+		return row ? this.#toThread(row) : null;
+	}
+	/**
+	* Cross-process atomic metadata read-modify-write: the thread row stays locked
+	* (`SELECT ... FOR UPDATE`) from the read through the write, so concurrent
+	* writers from any process serialize instead of overwriting each other.
+	*/
+	async updateThreadMetadata({ id, resourceId, update, fence: explicitFence }) {
+		const fence = resolveRunFence(this, explicitFence);
+		const threadsTable = getTableName$3({
+			indexName: TABLE_THREADS,
+			schemaName: getSchemaName$3(this.#schema)
+		});
+		try {
+			return await this.#db.client.tx(async (t) => {
+				if (fence) await assertRunFence(t, this.#runFencesTable(), fence, "updateThreadMetadata");
+				const thread = await this.#lockThread(t, threadsTable, {
+					id,
+					resourceId
+				});
+				if (!thread) return null;
+				const metadata = update(thread);
+				if (!metadata) return thread;
+				return this.#writeThread(t, threadsTable, {
+					id,
+					metadata: mergeMetadata(thread.metadata, metadata)
+				});
+			});
+		} catch (error) {
+			if (error instanceof MastraError || isRunFenceConflictError(error)) throw error;
+			throw new MastraError({
+				id: createStorageErrorId("PG", "UPDATE_THREAD_METADATA", "FAILED"),
+				domain: ErrorDomain.STORAGE,
+				category: ErrorCategory.THIRD_PARTY,
+				details: { threadId: id }
+			}, error);
+		}
+	}
+	async updateThread({ id, title, metadata, fence: explicitFence }) {
 		const threadTableName = getTableName$3({
 			indexName: TABLE_THREADS,
 			schemaName: getSchemaName$3(this.#schema)
 		});
+		const notFound = () => new MastraError({
+			id: createStorageErrorId("PG", "UPDATE_THREAD", "FAILED"),
+			domain: ErrorDomain.STORAGE,
+			category: ErrorCategory.USER,
+			text: `Thread ${id} not found`,
+			details: {
+				threadId: id,
+				title: title ?? null
+			}
+		});
 		try {
-			return await this.#fenced(fence, "updateThread", async (q) => {
-				const existingThread = await this.#getThreadById(q, { threadId: id });
-				if (!existingThread) throw new MastraError({
-					id: createStorageErrorId("PG", "UPDATE_THREAD", "FAILED"),
-					domain: ErrorDomain.STORAGE,
-					category: ErrorCategory.USER,
-					text: `Thread ${id} not found`,
-					details: {
-						threadId: id,
-						title: title ?? null
-					}
+			if (metadata === void 0) return await this.#fenced(explicitFence, "updateThread", async (q) => {
+				const thread = await this.#writeThread(q, threadTableName, {
+					id,
+					title
 				});
-				const mergedMetadata = {
-					...existingThread.metadata,
-					...metadata
-				};
-				const nowStr = toUtcISOString(/* @__PURE__ */ new Date());
-				const thread = await q.one(`UPDATE ${threadTableName}
-                    SET
-                        title = COALESCE($1, title),
-                        metadata = $2,
-                        "updatedAt" = $3,
-                        "updatedAtZ" = $4
-                    WHERE id = $5
-                    RETURNING *
-                `, [
-					title ?? null,
-					mergedMetadata,
-					nowStr,
-					nowStr,
-					id
-				]);
-				return {
-					id: thread.id,
-					resourceId: thread.resourceId,
-					title: thread.title,
-					metadata: typeof thread.metadata === "string" ? JSON.parse(thread.metadata) : thread.metadata,
-					createdAt: thread.createdAtZ || thread.createdAt,
-					updatedAt: thread.updatedAtZ || thread.updatedAt
-				};
+				if (!thread) throw notFound();
+				return thread;
+			});
+			const fence = resolveRunFence(this, explicitFence);
+			return await this.#db.client.tx(async (t) => {
+				if (fence) await assertRunFence(t, this.#runFencesTable(), fence, "updateThread");
+				const existing = await this.#lockThread(t, threadTableName, { id });
+				if (!existing) throw notFound();
+				const thread = await this.#writeThread(t, threadTableName, {
+					id,
+					title,
+					metadata: mergeMetadata(existing.metadata, metadata)
+				});
+				if (!thread) throw notFound();
+				return thread;
 			});
 		} catch (error) {
 			if (error instanceof MastraError || isRunFenceConflictError(error)) throw error;
@@ -17358,8 +17456,12 @@ ${onConflict}`,
 async function insertRows(client, schema, table, records, partitioning, onConflict) {
 	const insert = buildInsert(schema, table, records, onConflict);
 	if (!insert) return;
+	await withPartitionRetry(client, schema, table, records, partitioning, () => client.query(insert.text, insert.values));
+}
+/** Retry a complete write after creating a missing native partition, outside any failed transaction. */
+async function withPartitionRetry(client, schema, table, records, partitioning, write) {
 	try {
-		await client.query(insert.text, insert.values);
+		return await write();
 	} catch (error) {
 		if (!isMissingPartitionError(error) || await resolveMode(client, partitioning) !== "native") throw error;
 		const days = withinPartitionWindow(recordDays(records, SIGNAL_TIME_COLUMN[table]), partitioning);
@@ -17372,7 +17474,7 @@ async function insertRows(client, schema, table, records, partitioning, onConfli
 				cause
 			}), { cause });
 		});
-		await client.query(insert.text, insert.values);
+		return await write();
 	}
 }
 /** `42501`: the role may not create the child, or does not own the parent it must attach to. */
@@ -18306,13 +18408,47 @@ function pushFeedbackIdentity(acc, feedbackType, feedbackSource) {
 	acc.conditions.push(`"valueNumber" IS NOT NULL`);
 }
 async function createFeedback(client, schema, args, partitioning) {
-	const row = feedbackRecordToRow(args.feedback);
-	await insertRows(client, schema, TABLE_FEEDBACK_EVENTS, [row], partitioning, FEEDBACK_UPSERT_CLAUSE);
+	await writeFeedback(client, schema, [args.feedback], args.onConflict, partitioning);
+}
+async function replaceFeedback(client, schema, args, partitioning) {
+	await writeFeedback(client, schema, [args.feedback], "replace", partitioning, true);
 }
 async function batchCreateFeedback(client, schema, args, partitioning) {
-	if (args.feedbacks.length === 0) return;
-	const rows = [...new Map(args.feedbacks.map((feedback) => [feedback.feedbackId, feedback])).values()].map(feedbackRecordToRow);
-	await insertRows(client, schema, TABLE_FEEDBACK_EVENTS, rows, partitioning, FEEDBACK_UPSERT_CLAUSE);
+	await writeFeedback(client, schema, args.feedbacks, args.onConflict, partitioning);
+}
+/** Serialize feedback writes by ID, including absent IDs and retries with a different timestamp. */
+async function writeFeedback(client, schema, feedbacks, onConflict, partitioning, preserveIdentity = false) {
+	if (!feedbacks.length) return;
+	const currentFeedback = new Map((onConflict === "ignore" ? [...feedbacks].reverse() : feedbacks).map((feedback) => [feedback.feedbackId, feedback]));
+	const rows = [...currentFeedback.values()].map(feedbackRecordToRow);
+	const ids = [...currentFeedback.keys()].sort();
+	await withPartitionRetry(client, schema, TABLE_FEEDBACK_EVENTS, rows, partitioning, () => client.tx(async (tx) => {
+		await tx.query(`SELECT pg_advisory_xact_lock(hashtextextended(lock_key, 0))
+         FROM unnest($1::text[]) WITH ORDINALITY AS locks(lock_key, position)
+         ORDER BY position`, [ids.map((id) => JSON.stringify([
+			schema,
+			TABLE_FEEDBACK_EVENTS,
+			id
+		]))]);
+		const existing = await tx.any(`SELECT DISTINCT ON ("feedbackId") "feedbackId", "traceId", "timestamp", "organizationId", "resourceId"
+         FROM ${qualifiedTable(schema, TABLE_FEEDBACK_EVENTS)}
+         WHERE "feedbackId" = ANY($1::text[]) ORDER BY "feedbackId", "cursorId" DESC`, [ids]);
+		const byId = new Map(existing.map((row) => [row.feedbackId, row]));
+		const accepted = rows.filter((row) => {
+			const previous = byId.get(row.feedbackId);
+			if (!previous) return true;
+			if (onConflict === "ignore") return false;
+			if (preserveIdentity && (previous.traceId !== row.traceId || new Date(previous.timestamp).getTime() !== new Date(row.timestamp).getTime() || (previous.organizationId ?? null) !== (row.organizationId ?? null) || (previous.resourceId ?? null) !== (row.resourceId ?? null))) throw new MastraError({
+				id: "OBSERVABILITY_FEEDBACK_IDENTITY_CONFLICT",
+				domain: ErrorDomain.MASTRA_OBSERVABILITY,
+				category: ErrorCategory.USER,
+				text: "Feedback replacement must preserve its trace, timestamp, and tenant"
+			});
+			return true;
+		});
+		const insert = buildInsert(schema, TABLE_FEEDBACK_EVENTS, accepted, FEEDBACK_UPSERT_CLAUSE);
+		if (insert) await tx.query(insert.text, insert.values);
+	}));
 }
 async function getFeedbackById(client, schema, feedbackId) {
 	const row = await client.oneOrNone(`SELECT ${FEEDBACK_SELECT_COLUMNS} FROM ${qualifiedTable(schema, TABLE_FEEDBACK_EVENTS)}
@@ -22003,6 +22139,7 @@ var ObservabilityStoragePostgresVNext = class ObservabilityStoragePostgresVNext
 			"thread-query",
 			"trace-query-tenant-scope",
 			"feedback",
+			"feedback-replacement",
 			"trace-query-context-ids",
 			...contextFieldFeatures,
 			...timingFeatures,
@@ -22031,6 +22168,7 @@ var ObservabilityStoragePostgresVNext = class ObservabilityStoragePostgresVNext
 			"thread-query",
 			"trace-query-tenant-scope",
 			"feedback",
+			"feedback-replacement",
 			"trace-query-context-ids",
 			...contextFieldFeatures,
 			...timingFeatures,
@@ -22130,6 +22268,9 @@ var ObservabilityStoragePostgresVNext = class ObservabilityStoragePostgresVNext
 	async batchCreateScores(args) {
 		await this.#run("BATCH_CREATE_SCORES", () => batchCreateScores(this.#client, this.#schema, args, this.#writePartitioning()), { count: args.scores.length });
 	}
+	async writeFeedbackReplacement(args) {
+		await this.#run("REPLACE_FEEDBACK", () => replaceFeedback(this.#client, this.#schema, args, this.#writePartitioning()));
+	}
 	async createFeedback(args) {
 		await this.#run("CREATE_FEEDBACK", () => createFeedback(this.#client, this.#schema, args, this.#writePartitioning()));
 	}
@@ -22155,7 +22296,7 @@ var ObservabilityStoragePostgresVNext = class ObservabilityStoragePostgresVNext
 		return this.#run("GET_SCORE_BY_ID", () => getScoreById(this.#readClient, this.#schema, scoreId), { scoreId });
 	}
 	async getFeedbackById(feedbackId) {
-		return this.#run("GET_FEEDBACK_BY_ID", () => getFeedbackById(this.#readClient, this.#schema, feedbackId), { feedbackId });
+		return this.#run("GET_FEEDBACK_BY_ID", () => getFeedbackById(this.#client, this.#schema, feedbackId), { feedbackId });
 	}
 	async listFeedback(args) {
 		return this.#run("LIST_FEEDBACK", () => listFeedback(this.#readClient, this.#schema, args));
@@ -25941,15 +26082,16 @@ var WorkflowsPG = class WorkflowsPG extends WorkflowsStorage {
 			this.logger?.warn?.(`Failed to create index ${stateBaseIndexName}:`, error);
 		}
 		if (await this.#db.getColumnType(TABLE_WORKFLOW_SNAPSHOT, "snapshot") !== "jsonb") return;
+		const concurrently = !this.#db.isTableCreatedInThisInit(TABLE_WORKFLOW_SNAPSHOT);
 		const indexName = workflowSnapshotStatusIndexName(this.#schema);
 		try {
-			await this.#db.createIndexFromStatement(indexName, workflowSnapshotStatusIndexSQL(indexName, this.#schema, true));
+			await this.#db.createIndexFromStatement(indexName, workflowSnapshotStatusIndexSQL(indexName, this.#schema, concurrently));
 		} catch (error) {
 			this.logger?.warn?.(`Failed to create index ${indexName}:`, error);
 		}
 		const threadIdIndexName = workflowSnapshotThreadIdIndexName(this.#schema);
 		try {
-			await this.#db.createIndexFromStatement(threadIdIndexName, workflowSnapshotThreadIdIndexSQL(threadIdIndexName, this.#schema, true));
+			await this.#db.createIndexFromStatement(threadIdIndexName, workflowSnapshotThreadIdIndexSQL(threadIdIndexName, this.#schema, concurrently));
 		} catch (error) {
 			this.logger?.warn?.(`Failed to create index ${threadIdIndexName}:`, error);
 		}