@mastra/pg 1.18.1 → 1.19.0-alpha.1

package/dist/index.cjs

@@ -1971,9 +1971,25 @@ var PinnedClientAdapter = class {
 var RoutingDbClient = class {
 	#base;
 	#pinned = null;
+	#schemaSnapshot = null;
 	constructor(base) {
 		this.#base = base;
 	}
+	/**
+	* Catalog snapshot for the current init window, or `null` outside it.
+	*
+	* It lives here rather than on `PgDB` because every storage domain builds its
+	* own `PgDB` over this one shared client — hanging the snapshot off the
+	* client means one load serves all of them, and its lifetime lines up exactly
+	* with the pinned-init window that `pin()`/`unpin()` already delimit.
+	*/
+	get schemaSnapshot() {
+		return this.#schemaSnapshot;
+	}
+	/** Install (or clear, with `null`) the init-window catalog snapshot. */
+	setSchemaSnapshot(snapshot) {
+		this.#schemaSnapshot = snapshot;
+	}
 	/** Returns the currently active client (pinned if set, otherwise base). */
 	get active() {
 		return this.#pinned ?? this.#base;
@@ -2074,6 +2090,67 @@ function isDuplicateSchemaError(error) {
 	return /schema .* already exists/i.test(message);
 }
 //#endregion
+//#region src/storage/db/schema-snapshot.ts
+/**
+* Returns the snapshot currently installed on `client`, but only when it was
+* taken from the schema the caller operates on. A store configured for another
+* schema (or a client with no snapshot at all) gets `null` and falls back to
+* probing the live catalog.
+*/
+function getSchemaSnapshot(client, schemaName) {
+	const snapshot = client?.schemaSnapshot;
+	if (!snapshot) return null;
+	return snapshot.schemaName === (schemaName || "public") ? snapshot : null;
+}
+/**
+* Reads the catalog for `schemaName` in three queries. Must be called on the
+* pinned init client so the snapshot reflects what that connection will see.
+*/
+async function loadSchemaSnapshot(client, schemaName) {
+	const schema = schemaName || "public";
+	const [tableRows, columnRows, indexRows] = await Promise.all([
+		client.manyOrNone(`SELECT tablename FROM pg_catalog.pg_tables WHERE schemaname = $1`, [schema]),
+		client.manyOrNone(`SELECT c.relname AS table_name, a.attname AS column_name
+         FROM pg_catalog.pg_class c
+         JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace
+         JOIN pg_catalog.pg_attribute a ON a.attrelid = c.oid
+        WHERE n.nspname = $1
+          AND c.relkind IN ('r', 'p')
+          AND a.attnum > 0
+          AND NOT a.attisdropped`, [schema]),
+		client.manyOrNone(`SELECT c.relname AS indexname, i.indisreplident AS is_replica_identity, i.indisprimary AS is_primary
+         FROM pg_catalog.pg_index i
+         JOIN pg_catalog.pg_class c ON c.oid = i.indexrelid
+         JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace
+        WHERE n.nspname = $1`, [schema])
+	]);
+	const columns = /* @__PURE__ */ new Map();
+	for (const row of columnRows) {
+		let set = columns.get(row.table_name);
+		if (!set) {
+			set = /* @__PURE__ */ new Set();
+			columns.set(row.table_name, set);
+		}
+		set.add(row.column_name);
+	}
+	const indexes = /* @__PURE__ */ new Set();
+	const replicaIdentityIndexes = /* @__PURE__ */ new Set();
+	const primaryKeyIndexes = /* @__PURE__ */ new Set();
+	for (const row of indexRows) {
+		indexes.add(row.indexname);
+		if (row.is_replica_identity) replicaIdentityIndexes.add(row.indexname);
+		if (row.is_primary) primaryKeyIndexes.add(row.indexname.toLowerCase());
+	}
+	return {
+		schemaName: schema,
+		tables: new Set(tableRows.map((r) => r.tablename)),
+		columns,
+		indexes,
+		replicaIdentityIndexes,
+		primaryKeyIndexes
+	};
+}
+//#endregion
 //#region src/storage/db/index.ts
 /**
 * Resolves PgDomainConfig to a database client and schema.
@@ -2253,7 +2330,8 @@ function generateTimestampTriggerSQL(tableName, schemaName) {
 		schemaName: quotedSchemaName
 	});
 	const functionName = `${quotedSchemaName}.trigger_set_timestamps`;
-	const triggerName = `"${(0, _mastra_core_utils.parseSqlIdentifier)(`${tableName}_timestamps`, "trigger name")}"`;
+	const parsedTriggerName = (0, _mastra_core_utils.parseSqlIdentifier)(`${tableName}_timestamps`, "trigger name");
+	const triggerName = `"${parsedTriggerName}"`;
 	return `CREATE OR REPLACE FUNCTION ${functionName}()
 RETURNS TRIGGER AS $$
 BEGIN
@@ -2272,12 +2350,38 @@ BEGIN
 END;
 $$ LANGUAGE plpgsql;
 
-DROP TRIGGER IF EXISTS ${triggerName} ON ${fullTableName};
+DO $mastra_timestamps_trigger$
+BEGIN
+    -- Recreating the trigger unconditionally would take an ACCESS EXCLUSIVE
+    -- lock on the table (DROP TRIGGER does, even when nothing changes), and
+    -- init runs on every process start. Skip when the trigger is already
+    -- exactly what the CREATE below would produce.
+    --
+    -- tgtype 23 = ROW (1) | BEFORE (2) | INSERT (4) | UPDATE (16), so a trigger
+    -- whose timing or events differ still falls through and gets rebuilt. The
+    -- behaviour itself lives in the function, which is replaced above on every
+    -- init, so an upgraded function body lands without touching the trigger.
+    IF NOT EXISTS (
+        SELECT 1
+        FROM pg_catalog.pg_trigger tg
+        JOIN pg_catalog.pg_class c ON c.oid = tg.tgrelid
+        JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace
+        WHERE tg.tgname = ${`'${parsedTriggerName}'`}
+          AND c.relname = ${`'${(0, _mastra_core_utils.parseSqlIdentifier)(tableName, "table name")}'`}
+          AND n.nspname = ${schemaName ? `'${(0, _mastra_core_utils.parseSqlIdentifier)(schemaName, "schema name")}'` : `'public'`}
+          AND NOT tg.tgisinternal
+          AND tg.tgtype = 23
+          AND tg.tgfoid = '${functionName}()'::regprocedure
+    ) THEN
+        DROP TRIGGER IF EXISTS ${triggerName} ON ${fullTableName};
 
-CREATE TRIGGER ${triggerName}
-    BEFORE INSERT OR UPDATE ON ${fullTableName}
-    FOR EACH ROW
-    EXECUTE FUNCTION ${functionName}();`;
+        CREATE TRIGGER ${triggerName}
+            BEFORE INSERT OR UPDATE ON ${fullTableName}
+            FOR EACH ROW
+            EXECUTE FUNCTION ${functionName}();
+    END IF;
+END
+$mastra_timestamps_trigger$;`;
 }
 const schemaSetupRegistry = /* @__PURE__ */ new Map();
 /**
@@ -2304,6 +2408,99 @@ var PgDB = class extends _mastra_core_base.MastraBase {
 		this.skipDefaultIndexes = config.skipDefaultIndexes;
 	}
 	/**
+	* Catalog snapshot for the current init window, or `null` outside it.
+	*
+	* When non-null, the init-path methods below answer existence questions from
+	* it instead of round-tripping to the server, and record the objects they
+	* create so later callers in the same init see them. See
+	* {@link SchemaSnapshot} for why it is scoped to init only.
+	*/
+	get schemaSnapshot() {
+		return getSchemaSnapshot(this.client, this.schemaName);
+	}
+	/**
+	* Whether the snapshot proves `generateTableSQL` would be a no-op for this
+	* table — i.e. the CREATE statement can be skipped.
+	*
+	* For most tables that is just "the table exists". `workflow_snapshot` is the
+	* exception: its generated SQL also carries a DO block that back-fills the
+	* `(workflow_name, run_id)` unique constraint and promotes it to the table's
+	* replica identity, so a table created by an older version still needs the
+	* statement to run.
+	*/
+	snapshotShowsTableConverged(snapshot, tableName) {
+		if (!snapshot.tables.has(tableName)) return false;
+		if (tableName === _mastra_core_storage.TABLE_WORKFLOW_SNAPSHOT) {
+			const constraintName = buildConstraintName({
+				baseName: "mastra_workflow_snapshot_workflow_name_run_id_key",
+				schemaName: this.schemaName ? (0, _mastra_core_utils.parseSqlIdentifier)(this.schemaName, "schema name") : void 0
+			}).toLowerCase();
+			return snapshot.indexes.has(constraintName) && snapshot.replicaIdentityIndexes.has(constraintName);
+		}
+		return true;
+	}
+	/** Column set for `tableName` in the snapshot, created empty if absent. */
+	snapshotColumns(snapshot, tableName) {
+		let columns = snapshot.columns.get(tableName);
+		if (!columns) {
+			columns = /* @__PURE__ */ new Set();
+			snapshot.columns.set(tableName, columns);
+		}
+		return columns;
+	}
+	/**
+	* Records an out-of-band `ALTER TABLE … RENAME TO` in the init snapshot.
+	*
+	* Init-time migrations that issue raw DDL on `this.client` (instead of going
+	* through createTable/alterTable/createIndex, which maintain the snapshot
+	* themselves) MUST report it through these `note*` methods. A snapshot that
+	* still lists a renamed-away table makes a later createTable() in the same
+	* init skip the rebuild the migration depends on — stranding data. No-op
+	* outside the init window.
+	*
+	* Indexes riding along with a rename keep their names, so the snapshot's
+	* index set stays accurate without changes here.
+	*/
+	noteTableRenamed(oldName, newName) {
+		const snapshot = this.schemaSnapshot;
+		if (snapshot) {
+			if (snapshot.tables.delete(oldName)) snapshot.tables.add(newName);
+			const columns = snapshot.columns.get(oldName);
+			if (columns) {
+				snapshot.columns.delete(oldName);
+				snapshot.columns.set(newName, columns);
+			}
+		}
+		this.tableColumnsCache.delete(oldName);
+		this.tableColumnsCache.delete(newName);
+	}
+	/**
+	* Records an out-of-band `DROP TABLE` in the init snapshot. See
+	* {@link noteTableRenamed} for why raw-DDL migrations must call this.
+	*
+	* The dropped table's indexes vanish with it, but the snapshot's flat index
+	* set cannot map names back to tables. Stale entries only make a later
+	* createIndex() skip a recreate until the next init re-reads the catalog —
+	* the same self-healing bound the rest of the snapshot design accepts.
+	*/
+	noteTableDropped(tableName) {
+		const snapshot = this.schemaSnapshot;
+		if (snapshot) {
+			snapshot.tables.delete(tableName);
+			snapshot.columns.delete(tableName);
+		}
+		this.tableColumnsCache.delete(tableName);
+	}
+	/**
+	* Records an out-of-band `ALTER TABLE … ADD COLUMN` in the init snapshot.
+	* See {@link noteTableRenamed} for why raw-DDL migrations must call this.
+	*/
+	noteColumnAdded(tableName, column) {
+		const snapshot = this.schemaSnapshot;
+		if (snapshot) this.snapshotColumns(snapshot, tableName).add(column);
+		this.tableColumnsCache.delete(tableName);
+	}
+	/**
 	* Gets the set of column names that actually exist in the database table.
 	* Results are cached; the cache is invalidated when alterTable() adds new columns.
 	*/
@@ -2330,6 +2527,12 @@ var PgDB = class extends _mastra_core_base.MastraBase {
 	}
 	async hasColumn(table, column) {
 		const schema = this.schemaName || "public";
+		const snapshot = this.schemaSnapshot;
+		if (snapshot) {
+			const columns = snapshot.columns.get(table);
+			if (!columns) return false;
+			return columns.has(column) || columns.has(column.toLowerCase());
+		}
 		return !!await this.client.oneOrNone(`SELECT 1 FROM information_schema.columns WHERE table_schema = $1 AND table_name = $2 AND (column_name = $3 OR column_name = $4)`, [
 			schema,
 			table,
@@ -2368,6 +2571,14 @@ var PgDB = class extends _mastra_core_base.MastraBase {
 		if (!this.schemaName) return;
 		let registryEntry = schemaSetupRegistry.get(this.schemaName);
 		if (registryEntry?.complete) return;
+		const snapshot = this.schemaSnapshot;
+		if (snapshot && snapshot.tables.size > 0) {
+			schemaSetupRegistry.set(this.schemaName, {
+				promise: null,
+				complete: true
+			});
+			return;
+		}
 		const quotedSchemaName = getSchemaName$6(this.schemaName);
 		if (!registryEntry?.promise) {
 			const schemaNameCapture = this.schemaName;
@@ -2472,16 +2683,28 @@ var PgDB = class extends _mastra_core_base.MastraBase {
 		try {
 			const timeZColumnNames = Object.entries(schema).filter(([_, def]) => def.type === "timestamp").map(([name]) => name);
 			if (this.schemaName) await this.setupSchema();
-			const sql = generateTableSQL({
-				tableName,
-				schema,
-				schemaName: this.schemaName,
-				compositePrimaryKey
-			});
-			try {
-				await this.client.none(sql);
-			} catch (error) {
-				if (!isDuplicateRelationError(error)) throw error;
+			const snapshot = this.schemaSnapshot;
+			if (!snapshot || !this.snapshotShowsTableConverged(snapshot, tableName)) {
+				const sql = generateTableSQL({
+					tableName,
+					schema,
+					schemaName: this.schemaName,
+					compositePrimaryKey
+				});
+				try {
+					await this.client.none(sql);
+				} catch (error) {
+					if (!isDuplicateRelationError(error)) throw error;
+				}
+				if (snapshot) {
+					snapshot.tables.add(tableName);
+					const created = this.snapshotColumns(snapshot, tableName);
+					for (const [name, def] of Object.entries(schema)) {
+						const parsedName = (0, _mastra_core_utils.parseSqlIdentifier)(name, "column name");
+						created.add(parsedName);
+						if (def.type === "timestamp") created.add(`${parsedName}Z`);
+					}
+				}
 			}
 			await this.alterTable({
 				tableName,
@@ -2547,10 +2770,12 @@ MIGRATION REQUIRED: Duplicate spans detected in ${duplicateInfo.tableName}\n====
 				const sqlType = mapToSqlType(columnDef.type);
 				const alterSql = `ALTER TABLE ${fullTableName} ADD COLUMN IF NOT EXISTS "${parsedColumnName}" ${sqlType} ${columnDef.nullable ? "" : "NOT NULL"} ${!columnDef.nullable ? this.getDefaultValue(columnDef.type) : ""}`.trim();
 				await this.client.none(alterSql);
+				this.noteColumnAdded(_mastra_core_storage.TABLE_SPANS, columnName);
 				this.logger?.debug?.(`Added column '${columnName}' to ${fullTableName}`);
 				if (sqlType === "TIMESTAMP") {
 					const timestampZSql = `ALTER TABLE ${fullTableName} ADD COLUMN IF NOT EXISTS "${parsedColumnName}Z" TIMESTAMPTZ DEFAULT NOW()`.trim();
 					await this.client.none(timestampZSql);
+					this.noteColumnAdded(_mastra_core_storage.TABLE_SPANS, `${columnName}Z`);
 					this.logger?.debug?.(`Added timezone column '${columnName}Z' to ${fullTableName}`);
 				}
 			}
@@ -2559,6 +2784,7 @@ MIGRATION REQUIRED: Duplicate spans detected in ${duplicateInfo.tableName}\n====
 				if (!await this.hasColumn(_mastra_core_storage.TABLE_SPANS, tzColumnName)) {
 					const timestampZSql = `ALTER TABLE ${fullTableName} ADD COLUMN IF NOT EXISTS "${(0, _mastra_core_utils.parseSqlIdentifier)(tzColumnName, "column name")}" TIMESTAMPTZ DEFAULT NOW()`.trim();
 					await this.client.none(timestampZSql);
+					this.noteColumnAdded(_mastra_core_storage.TABLE_SPANS, tzColumnName);
 					this.logger?.debug?.(`Added timezone column '${tzColumnName}' to ${fullTableName}`);
 				}
 			}
@@ -2681,6 +2907,12 @@ MIGRATION REQUIRED: Duplicate spans detected in ${duplicateInfo.tableName}\n====
 			schemaName: (this.schemaName ? (0, _mastra_core_utils.parseSqlIdentifier)(this.schemaName, "schema name") : "") || void 0
 		});
 		const schemaFilter = this.schemaName || "public";
+		const snapshot = this.schemaSnapshot;
+		if (snapshot) return snapshot.primaryKeyIndexes.has(constraintName.toLowerCase());
+		return this.spansPrimaryKeyExistsLive(constraintName, schemaFilter);
+	}
+	/** Live-catalog variant of {@link spansPrimaryKeyExists}, bypassing the snapshot. */
+	async spansPrimaryKeyExistsLive(constraintName, schemaFilter) {
 		return (await this.client.oneOrNone(`SELECT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = lower($1) AND connamespace = (SELECT oid FROM pg_namespace WHERE nspname = $2)) as exists`, [constraintName, schemaFilter]))?.exists ?? false;
 	}
 	/**
@@ -2711,10 +2943,12 @@ MIGRATION REQUIRED: Duplicate spans detected in ${duplicateInfo.tableName}\n====
         ADD CONSTRAINT ${constraintName}
         PRIMARY KEY ("traceId", "spanId")
       `);
+			this.schemaSnapshot?.primaryKeyIndexes.add(constraintName.toLowerCase());
 			this.logger?.info?.(`Added PRIMARY KEY constraint ${constraintName} to ${fullTableName}`);
 		} catch (error) {
 			if (isDuplicateRelationError(error)) {
-				if (await this.spansPrimaryKeyExists()) {
+				if (await this.spansPrimaryKeyExistsLive(constraintName, schemaFilter)) {
+					this.schemaSnapshot?.primaryKeyIndexes.add(constraintName.toLowerCase());
 					this.logger?.debug?.(`PRIMARY KEY constraint ${constraintName} was created by another process`);
 					return;
 				}
@@ -2796,16 +3030,25 @@ MIGRATION REQUIRED: Duplicate spans detected in ${duplicateInfo.tableName}\n====
 			indexName: tableName,
 			schemaName: getSchemaName$6(this.schemaName)
 		});
+		const snapshot = this.schemaSnapshot;
+		const knownColumns = snapshot ? this.snapshotColumns(snapshot, tableName) : null;
 		try {
 			for (const columnName of ifNotExists) if (schema[columnName]) {
 				const columnDef = schema[columnName];
 				const parsedColumnName = (0, _mastra_core_utils.parseSqlIdentifier)(columnName, "column name");
 				const sqlType = mapToSqlType(columnDef.type);
 				const alterSql = `ALTER TABLE ${fullTableName} ADD COLUMN IF NOT EXISTS "${parsedColumnName}" ${sqlType} ${columnDef.nullable ? "" : "NOT NULL"} ${!columnDef.nullable ? this.getDefaultValue(columnDef.type) : ""}`.trim();
-				await this.client.none(alterSql);
+				if (!knownColumns?.has(parsedColumnName)) {
+					await this.client.none(alterSql);
+					knownColumns?.add(parsedColumnName);
+				}
 				if (sqlType === "TIMESTAMP") {
-					const timestampZSql = `ALTER TABLE ${fullTableName} ADD COLUMN IF NOT EXISTS "${parsedColumnName}Z" TIMESTAMPTZ DEFAULT NOW()`.trim();
-					await this.client.none(timestampZSql);
+					const tzColumnName = `${parsedColumnName}Z`;
+					if (!knownColumns?.has(tzColumnName)) {
+						const timestampZSql = `ALTER TABLE ${fullTableName} ADD COLUMN IF NOT EXISTS "${tzColumnName}" TIMESTAMPTZ DEFAULT NOW()`.trim();
+						await this.client.none(timestampZSql);
+						knownColumns?.add(tzColumnName);
+					}
 				}
 				this.logger?.debug?.(`Ensured column ${parsedColumnName} exists in table ${fullTableName}`);
 			}
@@ -2891,7 +3134,10 @@ MIGRATION REQUIRED: Duplicate spans detected in ${duplicateInfo.tableName}\n====
 				indexName: table,
 				schemaName: getSchemaName$6(this.schemaName)
 			});
-			if (await this.client.oneOrNone(`SELECT 1 FROM pg_indexes
+			const snapshot = this.schemaSnapshot;
+			if (snapshot) {
+				if (snapshot.indexes.has(name)) return;
+			} else if (await this.client.oneOrNone(`SELECT 1 FROM pg_indexes
          WHERE indexname = $1
          AND schemaname = $2`, [name, schemaName])) return;
 			const uniqueStr = unique ? "UNIQUE " : "";
@@ -2913,6 +3159,7 @@ MIGRATION REQUIRED: Duplicate spans detected in ${duplicateInfo.tableName}\n====
 			if (storage && Object.keys(storage).length > 0) withStr = ` WITH (${Object.entries(storage).map(([key, value]) => `${key} = ${value}`).join(", ")})`;
 			const sql = `CREATE ${uniqueStr}INDEX ${concurrentStr}${`"${(0, _mastra_core_utils.parseSqlIdentifier)(name, "index name")}"`} ON ${fullTableName} ${methodStr}(${columnsStr})${withStr}${tablespaceStr}${whereStr}`;
 			await this.client.none(sql);
+			snapshot?.indexes.add(name);
 		} catch (error) {
 			if (error instanceof Error && error.message.includes("CONCURRENTLY")) {
 				const retryOptions = {
@@ -2932,15 +3179,34 @@ MIGRATION REQUIRED: Duplicate spans detected in ${duplicateInfo.tableName}\n====
 			}, error);
 		}
 	}
+	/**
+	* Runs a caller-built `CREATE INDEX IF NOT EXISTS` statement, unless the init
+	* snapshot already proves `indexName` exists.
+	*
+	* `createIndex` covers the indexes described by {@link CreateIndexOptions};
+	* this is for the two init paths that hand-write their statement (a partial
+	* or otherwise non-standard index) and would otherwise send a no-op DDL on
+	* every warm init.
+	*/
+	async createIndexFromStatement(indexName, sql) {
+		const snapshot = this.schemaSnapshot;
+		if (snapshot?.indexes.has(indexName)) return;
+		await this.client.none(sql);
+		snapshot?.indexes.add(indexName);
+	}
 	async dropIndex(indexName) {
 		try {
 			const schemaName = this.schemaName || "public";
-			if (!await this.client.oneOrNone(`SELECT 1 FROM pg_indexes
+			const snapshot = this.schemaSnapshot;
+			if (snapshot) {
+				if (!snapshot.indexes.has(indexName)) return;
+			} else if (!await this.client.oneOrNone(`SELECT 1 FROM pg_indexes
          WHERE indexname = $1
          AND schemaname = $2`, [indexName, schemaName])) return;
 			const quotedIndexName = `"${(0, _mastra_core_utils.parseSqlIdentifier)(indexName, "index name")}"`;
 			const sql = `DROP INDEX IF EXISTS ${getSchemaName$6(this.schemaName)}.${quotedIndexName}`;
 			await this.client.none(sql);
+			snapshot?.indexes.delete(indexName);
 		} catch (error) {
 			throw new _mastra_core_error.MastraError({
 				id: (0, _mastra_core_storage.createStorageErrorId)("PG", "INDEX_DROP", "FAILED"),
@@ -3437,7 +3703,9 @@ var AgentsPG = class AgentsPG extends _mastra_core_storage.AgentsStorage {
 		});
 		if (await this.#db.hasColumn(_mastra_core_storage.TABLE_AGENTS, "name")) {
 			await this.#db.client.none(`ALTER TABLE ${fullTableName} RENAME TO "${_mastra_core_storage.TABLE_AGENTS}_legacy"`);
+			this.#db.noteTableRenamed(_mastra_core_storage.TABLE_AGENTS, `${_mastra_core_storage.TABLE_AGENTS}_legacy`);
 			await this.#db.client.none(`DROP TABLE IF EXISTS ${fullVersionsTableName}`);
+			this.#db.noteTableDropped(_mastra_core_storage.TABLE_AGENT_VERSIONS);
 		}
 		if (!await this.#db.hasColumn(`${_mastra_core_storage.TABLE_AGENTS}_legacy`, "name")) return;
 		const oldAgents = await this.#db.client.manyOrNone(`SELECT * FROM ${legacyTableName}`);
@@ -3494,6 +3762,7 @@ var AgentsPG = class AgentsPG extends _mastra_core_storage.AgentsStorage {
 			]);
 		}
 		await this.#db.client.none(`DROP TABLE IF EXISTS ${legacyTableName}`);
+		this.#db.noteTableDropped(`${_mastra_core_storage.TABLE_AGENTS}_legacy`);
 	}
 	/**
 	* Migrates the agent_versions table from the old snapshot-based schema (single `snapshot` JSON column)
@@ -3511,7 +3780,9 @@ var AgentsPG = class AgentsPG extends _mastra_core_storage.AgentsStorage {
 			schemaName: getSchemaName$5(this.#schema)
 		});
 		await this.#db.client.none(`DROP TABLE IF EXISTS ${fullVersionsTableName}`);
+		this.#db.noteTableDropped(_mastra_core_storage.TABLE_AGENT_VERSIONS);
 		await this.#db.client.none(`DROP TABLE IF EXISTS ${legacyTableName}`);
+		this.#db.noteTableDropped(`${_mastra_core_storage.TABLE_AGENTS}_legacy`);
 	}
 	/**
 	* Migrates the tools field from string[] format to JSONB format { "tool-key": { "description": "..." } }.
@@ -5027,6 +5298,7 @@ var DatasetsPG = class DatasetsPG extends _mastra_core_storage.DatasetsStorage {
 		await this.#addColumnIfNotExists(_mastra_core_storage.TABLE_DATASET_ITEMS, "organizationId", "TEXT");
 		await this.#addColumnIfNotExists(_mastra_core_storage.TABLE_DATASET_ITEMS, "projectId", "TEXT");
 		await this.#addColumnIfNotExists(_mastra_core_storage.TABLE_DATASET_ITEMS, "toolMocks", "JSONB");
+		await this.#addColumnIfNotExists(_mastra_core_storage.TABLE_DATASET_ITEMS, "unmockedToolPolicy", "TEXT");
 		await this.#addColumnIfNotExists(_mastra_core_storage.TABLE_DATASET_ITEMS, "externalId", "TEXT");
 		await this.createDefaultIndexes();
 		await this.createCustomIndexes();
@@ -5038,6 +5310,7 @@ var DatasetsPG = class DatasetsPG extends _mastra_core_storage.DatasetsStorage {
 				schemaName: getSchemaName$5(this.#schema)
 			});
 			await this.#db.client.none(`ALTER TABLE ${fullTableName} ADD COLUMN "${column}" ${sqlType}`);
+			this.#db.noteColumnAdded(table, column);
 		}
 	}
 	getDefaultIndexDefinitions() {
@@ -5148,6 +5421,7 @@ var DatasetsPG = class DatasetsPG extends _mastra_core_storage.DatasetsStorage {
 			groundTruth: row.groundTruth ? (0, _mastra_core_storage.safelyParseJSON)(row.groundTruth) : void 0,
 			expectedTrajectory: row.expectedTrajectory ? (0, _mastra_core_storage.safelyParseJSON)(row.expectedTrajectory) : void 0,
 			toolMocks: row.toolMocks ? (0, _mastra_core_storage.safelyParseJSON)(row.toolMocks) : void 0,
+			unmockedToolPolicy: row.unmockedToolPolicy ?? void 0,
 			requestContext: row.requestContext ? (0, _mastra_core_storage.safelyParseJSON)(row.requestContext) : void 0,
 			metadata: row.metadata ? (0, _mastra_core_storage.safelyParseJSON)(row.metadata) : void 0,
 			source: row.source ? (0, _mastra_core_storage.safelyParseJSON)(row.source) : void 0,
@@ -5169,6 +5443,7 @@ var DatasetsPG = class DatasetsPG extends _mastra_core_storage.DatasetsStorage {
 			groundTruth: row.groundTruth ? (0, _mastra_core_storage.safelyParseJSON)(row.groundTruth) : void 0,
 			expectedTrajectory: row.expectedTrajectory ? (0, _mastra_core_storage.safelyParseJSON)(row.expectedTrajectory) : void 0,
 			toolMocks: row.toolMocks ? (0, _mastra_core_storage.safelyParseJSON)(row.toolMocks) : void 0,
+			unmockedToolPolicy: row.unmockedToolPolicy ?? void 0,
 			requestContext: row.requestContext ? (0, _mastra_core_storage.safelyParseJSON)(row.requestContext) : void 0,
 			metadata: row.metadata ? (0, _mastra_core_storage.safelyParseJSON)(row.metadata) : void 0,
 			source: row.source ? (0, _mastra_core_storage.safelyParseJSON)(row.source) : void 0,
@@ -5500,7 +5775,7 @@ var DatasetsPG = class DatasetsPG extends _mastra_core_storage.DatasetsStorage {
 				newVersion = row.version;
 				parentOrganizationId = row.organizationId ?? null;
 				parentProjectId = row.projectId ?? null;
-				await t.none(`INSERT INTO ${itemsTable} ("id","datasetId","datasetVersion","externalId","organizationId","projectId","validTo","isDeleted","input","groundTruth","expectedTrajectory","toolMocks","requestContext","metadata","source","createdAt","createdAtZ","updatedAt","updatedAtZ") VALUES ($1,$2,$3,$4,$5,$6,NULL,false,$7,$8,$9,$10,$11,$12,$13,$14,$15,$16,$17)`, [
+				await t.none(`INSERT INTO ${itemsTable} ("id","datasetId","datasetVersion","externalId","organizationId","projectId","validTo","isDeleted","input","groundTruth","expectedTrajectory","toolMocks","unmockedToolPolicy","requestContext","metadata","source","createdAt","createdAtZ","updatedAt","updatedAtZ") VALUES ($1,$2,$3,$4,$5,$6,NULL,false,$7,$8,$9,$10,$11,$12,$13,$14,$15,$16,$17,$18)`, [
 					id,
 					args.datasetId,
 					newVersion,
@@ -5511,6 +5786,7 @@ var DatasetsPG = class DatasetsPG extends _mastra_core_storage.DatasetsStorage {
 					jsonbArg(args.groundTruth),
 					jsonbArg(args.expectedTrajectory),
 					jsonbArg(args.toolMocks),
+					args.unmockedToolPolicy ?? null,
 					jsonbArg(args.requestContext),
 					jsonbArg(args.metadata),
 					jsonbArg(args.source),
@@ -5537,6 +5813,7 @@ var DatasetsPG = class DatasetsPG extends _mastra_core_storage.DatasetsStorage {
 				groundTruth: args.groundTruth,
 				expectedTrajectory: args.expectedTrajectory,
 				toolMocks: args.toolMocks,
+				unmockedToolPolicy: args.unmockedToolPolicy,
 				requestContext: args.requestContext,
 				metadata: args.metadata,
 				source: args.source,
@@ -5590,6 +5867,7 @@ var DatasetsPG = class DatasetsPG extends _mastra_core_storage.DatasetsStorage {
 			const mergedGroundTruth = args.groundTruth !== void 0 ? args.groundTruth : existing.groundTruth;
 			const mergedExpectedTrajectory = args.expectedTrajectory !== void 0 ? args.expectedTrajectory : existing.expectedTrajectory;
 			const mergedToolMocks = args.toolMocks !== void 0 ? args.toolMocks : existing.toolMocks;
+			const mergedUnmockedToolPolicy = args.unmockedToolPolicy !== void 0 ? args.unmockedToolPolicy : existing.unmockedToolPolicy;
 			const mergedRequestContext = args.requestContext !== void 0 ? args.requestContext : existing.requestContext;
 			const mergedMetadata = args.metadata !== void 0 ? args.metadata : existing.metadata;
 			const mergedSource = args.source !== void 0 ? args.source : existing.source;
@@ -5602,7 +5880,7 @@ var DatasetsPG = class DatasetsPG extends _mastra_core_storage.DatasetsStorage {
 				parentOrganizationId = row.organizationId ?? null;
 				parentProjectId = row.projectId ?? null;
 				await t.none(`UPDATE ${itemsTable} SET "validTo" = $1 WHERE "id" = $2 AND "validTo" IS NULL AND "isDeleted" = false`, [newVersion, args.id]);
-				await t.none(`INSERT INTO ${itemsTable} ("id","datasetId","datasetVersion","externalId","organizationId","projectId","validTo","isDeleted","input","groundTruth","expectedTrajectory","toolMocks","requestContext","metadata","source","createdAt","createdAtZ","updatedAt","updatedAtZ") VALUES ($1,$2,$3,$4,$5,$6,NULL,false,$7,$8,$9,$10,$11,$12,$13,$14,$15,$16,$17)`, [
+				await t.none(`INSERT INTO ${itemsTable} ("id","datasetId","datasetVersion","externalId","organizationId","projectId","validTo","isDeleted","input","groundTruth","expectedTrajectory","toolMocks","unmockedToolPolicy","requestContext","metadata","source","createdAt","createdAtZ","updatedAt","updatedAtZ") VALUES ($1,$2,$3,$4,$5,$6,NULL,false,$7,$8,$9,$10,$11,$12,$13,$14,$15,$16,$17,$18)`, [
 					args.id,
 					args.datasetId,
 					newVersion,
@@ -5613,6 +5891,7 @@ var DatasetsPG = class DatasetsPG extends _mastra_core_storage.DatasetsStorage {
 					jsonbArg(mergedGroundTruth),
 					jsonbArg(mergedExpectedTrajectory),
 					jsonbArg(mergedToolMocks),
+					mergedUnmockedToolPolicy ?? null,
 					jsonbArg(mergedRequestContext),
 					jsonbArg(mergedMetadata),
 					jsonbArg(mergedSource),
@@ -5638,6 +5917,7 @@ var DatasetsPG = class DatasetsPG extends _mastra_core_storage.DatasetsStorage {
 				groundTruth: mergedGroundTruth,
 				expectedTrajectory: mergedExpectedTrajectory,
 				toolMocks: mergedToolMocks,
+				unmockedToolPolicy: mergedUnmockedToolPolicy,
 				requestContext: mergedRequestContext,
 				metadata: mergedMetadata,
 				source: mergedSource,
@@ -5686,7 +5966,7 @@ var DatasetsPG = class DatasetsPG extends _mastra_core_storage.DatasetsStorage {
 				const parentOrganizationId = row.organizationId ?? null;
 				const parentProjectId = row.projectId ?? null;
 				await t.none(`UPDATE ${itemsTable} SET "validTo" = $1 WHERE "id" = $2 AND "validTo" IS NULL AND "isDeleted" = false`, [newVersion, id]);
-				await t.none(`INSERT INTO ${itemsTable} ("id","datasetId","datasetVersion","externalId","organizationId","projectId","validTo","isDeleted","input","groundTruth","expectedTrajectory","toolMocks","requestContext","metadata","source","createdAt","createdAtZ","updatedAt","updatedAtZ") VALUES ($1,$2,$3,$4,$5,$6,NULL,true,$7,$8,$9,$10,$11,$12,$13,$14,$15,$16,$17)`, [
+				await t.none(`INSERT INTO ${itemsTable} ("id","datasetId","datasetVersion","externalId","organizationId","projectId","validTo","isDeleted","input","groundTruth","expectedTrajectory","toolMocks","unmockedToolPolicy","requestContext","metadata","source","createdAt","createdAtZ","updatedAt","updatedAtZ") VALUES ($1,$2,$3,$4,$5,$6,NULL,true,$7,$8,$9,$10,$11,$12,$13,$14,$15,$16,$17,$18)`, [
 					id,
 					datasetId,
 					newVersion,
@@ -5697,6 +5977,7 @@ var DatasetsPG = class DatasetsPG extends _mastra_core_storage.DatasetsStorage {
 					jsonbArg(existing.groundTruth),
 					jsonbArg(existing.expectedTrajectory),
 					jsonbArg(existing.toolMocks),
+					existing.unmockedToolPolicy ?? null,
 					jsonbArg(existing.requestContext),
 					jsonbArg(existing.metadata),
 					jsonbArg(existing.source),
@@ -5756,7 +6037,7 @@ var DatasetsPG = class DatasetsPG extends _mastra_core_storage.DatasetsStorage {
 					const nowIso = now.toISOString();
 					await t.none(`UPDATE ${datasetsTable} SET "version" = $2 WHERE "id" = $1`, [input.datasetId, newVersion]);
 					for (const { id, item } of plan.inserts) {
-						await t.none(`INSERT INTO ${itemsTable} ("id","datasetId","datasetVersion","externalId","organizationId","projectId","validTo","isDeleted","input","groundTruth","expectedTrajectory","toolMocks","requestContext","metadata","source","createdAt","createdAtZ","updatedAt","updatedAtZ") VALUES ($1,$2,$3,$4,$5,$6,NULL,false,$7,$8,$9,$10,$11,$12,$13,$14,$15,$16,$17)`, [
+						await t.none(`INSERT INTO ${itemsTable} ("id","datasetId","datasetVersion","externalId","organizationId","projectId","validTo","isDeleted","input","groundTruth","expectedTrajectory","toolMocks","unmockedToolPolicy","requestContext","metadata","source","createdAt","createdAtZ","updatedAt","updatedAtZ") VALUES ($1,$2,$3,$4,$5,$6,NULL,false,$7,$8,$9,$10,$11,$12,$13,$14,$15,$16,$17,$18)`, [
 							id,
 							input.datasetId,
 							newVersion,
@@ -5767,6 +6048,7 @@ var DatasetsPG = class DatasetsPG extends _mastra_core_storage.DatasetsStorage {
 							jsonbArg(item.groundTruth),
 							jsonbArg(item.expectedTrajectory),
 							jsonbArg(item.toolMocks),
+							item.unmockedToolPolicy ?? null,
 							jsonbArg(item.requestContext),
 							jsonbArg(item.metadata),
 							jsonbArg(item.source),
@@ -5786,6 +6068,7 @@ var DatasetsPG = class DatasetsPG extends _mastra_core_storage.DatasetsStorage {
 							groundTruth: item.groundTruth,
 							expectedTrajectory: item.expectedTrajectory,
 							toolMocks: item.toolMocks,
+							unmockedToolPolicy: item.unmockedToolPolicy,
 							requestContext: item.requestContext,
 							metadata: item.metadata,
 							source: item.source,
@@ -5848,7 +6131,7 @@ var DatasetsPG = class DatasetsPG extends _mastra_core_storage.DatasetsStorage {
 				const newVersion = (await t.one(`UPDATE ${datasetsTable} SET "version" = "version" + 1 WHERE "id" = $1 RETURNING "version"`, [input.datasetId])).version;
 				for (const item of currentItems) {
 					await t.none(`UPDATE ${itemsTable} SET "validTo" = $1 WHERE "id" = $2 AND "validTo" IS NULL AND "isDeleted" = false`, [newVersion, item.id]);
-					await t.none(`INSERT INTO ${itemsTable} ("id","datasetId","datasetVersion","externalId","organizationId","projectId","validTo","isDeleted","input","groundTruth","expectedTrajectory","toolMocks","requestContext","metadata","source","createdAt","createdAtZ","updatedAt","updatedAtZ") VALUES ($1,$2,$3,$4,$5,$6,NULL,true,$7,$8,$9,$10,$11,$12,$13,$14,$15,$16,$17)`, [
+					await t.none(`INSERT INTO ${itemsTable} ("id","datasetId","datasetVersion","externalId","organizationId","projectId","validTo","isDeleted","input","groundTruth","expectedTrajectory","toolMocks","unmockedToolPolicy","requestContext","metadata","source","createdAt","createdAtZ","updatedAt","updatedAtZ") VALUES ($1,$2,$3,$4,$5,$6,NULL,true,$7,$8,$9,$10,$11,$12,$13,$14,$15,$16,$17,$18)`, [
 						item.id,
 						input.datasetId,
 						newVersion,
@@ -5859,6 +6142,7 @@ var DatasetsPG = class DatasetsPG extends _mastra_core_storage.DatasetsStorage {
 						jsonbArg(item.groundTruth),
 						jsonbArg(item.expectedTrajectory),
 						jsonbArg(item.toolMocks),
+						item.unmockedToolPolicy ?? null,
 						jsonbArg(item.requestContext),
 						jsonbArg(item.metadata),
 						jsonbArg(item.source),
@@ -6153,6 +6437,7 @@ var ExperimentsPG = class ExperimentsPG extends _mastra_core_storage.Experiments
 			ifNotExists: [
 				"status",
 				"tags",
+				"comment",
 				"toolMockReport",
 				"organizationId",
 				"projectId"
@@ -6330,6 +6615,7 @@ var ExperimentsPG = class ExperimentsPG extends _mastra_core_storage.Experiments
 			traceId: row.traceId ?? null,
 			status: row.status ?? null,
 			tags: row.tags ? (0, _mastra_core_storage.safelyParseJSON)(row.tags) : null,
+			comment: row.comment ?? null,
 			toolMockReport: row.toolMockReport ? (0, _mastra_core_storage.safelyParseJSON)(row.toolMockReport) : null,
 			createdAt: (0, _mastra_core_storage.ensureDate)(row.createdAtZ || row.createdAt)
 		};
@@ -6655,6 +6941,10 @@ var ExperimentsPG = class ExperimentsPG extends _mastra_core_storage.Experiments
 				setClauses.push(`"tags" = $${paramIndex++}`);
 				values.push(JSON.stringify(input.tags));
 			}
+			if (input.comment !== void 0) {
+				setClauses.push(`"comment" = $${paramIndex++}`);
+				values.push(input.comment);
+			}
 			if (setClauses.length === 0) {
 				const existing = await this.getExperimentResultById({ id: input.id });
 				if (!existing) throw new _mastra_core_error.MastraError({
@@ -6890,7 +7180,7 @@ var FavoritesPG = class FavoritesPG extends _mastra_core_storage.FavoritesStorag
 			indexName: _mastra_core_storage.TABLE_FAVORITES,
 			schemaName: getSchemaName$5(this.#schema)
 		});
-		await this.#db.client.none(`CREATE INDEX IF NOT EXISTS idx_favorites_entity ON ${fullTableName} ("entityType", "entityId")`);
+		await this.#db.createIndexFromStatement("idx_favorites_entity", `CREATE INDEX IF NOT EXISTS idx_favorites_entity ON ${fullTableName} ("entityType", "entityId")`);
 	}
 	async dangerouslyClearAll() {
 		const fullTableName = getTableName$5({
@@ -8386,7 +8676,7 @@ var MemoryPG = class MemoryPG extends _mastra_core_storage.MemoryStorage {
 				indexName: OM_TABLE,
 				schemaName: getSchemaName$3(this.#schema)
 			});
-			await this.#db.client.none(`CREATE INDEX IF NOT EXISTS idx_om_lookup_key ON ${omTableName} ("lookupKey")`);
+			await this.#db.createIndexFromStatement("idx_om_lookup_key", `CREATE INDEX IF NOT EXISTS idx_om_lookup_key ON ${omTableName} ("lookupKey")`);
 		}
 		await this.createDefaultIndexes();
 		await this.createCustomIndexes();
@@ -20107,6 +20397,7 @@ var PostgresStore = class extends _mastra_core_storage.MastraCompositeStore {
 			pinnedClient = await this.#pool.connect();
 			const pinned = new PinnedClientAdapter(this.#pool, pinnedClient);
 			this.#db.pin(pinned);
+			this.#db.setSchemaSnapshot(await loadSchemaSnapshot(pinned, this.schema));
 			await super.init();
 			this.isInitialized = true;
 		} catch (error) {
@@ -20118,6 +20409,7 @@ var PostgresStore = class extends _mastra_core_storage.MastraCompositeStore {
 				category: _mastra_core_error.ErrorCategory.THIRD_PARTY
 			}, error);
 		} finally {
+			this.#db.setSchemaSnapshot(null);
 			if (pinnedClient) {
 				this.#db.unpin();
 				pinnedClient.release();