@mastra/pg 1.18.1 → 1.19.0-alpha.1

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
package/dist/index.js CHANGED
@@ -1947,9 +1947,25 @@ var PinnedClientAdapter = class {
1947
1947
  var RoutingDbClient = class {
1948
1948
  #base;
1949
1949
  #pinned = null;
1950
+ #schemaSnapshot = null;
1950
1951
  constructor(base) {
1951
1952
  this.#base = base;
1952
1953
  }
1954
+ /**
1955
+ * Catalog snapshot for the current init window, or `null` outside it.
1956
+ *
1957
+ * It lives here rather than on `PgDB` because every storage domain builds its
1958
+ * own `PgDB` over this one shared client — hanging the snapshot off the
1959
+ * client means one load serves all of them, and its lifetime lines up exactly
1960
+ * with the pinned-init window that `pin()`/`unpin()` already delimit.
1961
+ */
1962
+ get schemaSnapshot() {
1963
+ return this.#schemaSnapshot;
1964
+ }
1965
+ /** Install (or clear, with `null`) the init-window catalog snapshot. */
1966
+ setSchemaSnapshot(snapshot) {
1967
+ this.#schemaSnapshot = snapshot;
1968
+ }
1953
1969
  /** Returns the currently active client (pinned if set, otherwise base). */
1954
1970
  get active() {
1955
1971
  return this.#pinned ?? this.#base;
@@ -2050,6 +2066,67 @@ function isDuplicateSchemaError(error) {
2050
2066
  return /schema .* already exists/i.test(message);
2051
2067
  }
2052
2068
  //#endregion
2069
+ //#region src/storage/db/schema-snapshot.ts
2070
+ /**
2071
+ * Returns the snapshot currently installed on `client`, but only when it was
2072
+ * taken from the schema the caller operates on. A store configured for another
2073
+ * schema (or a client with no snapshot at all) gets `null` and falls back to
2074
+ * probing the live catalog.
2075
+ */
2076
+ function getSchemaSnapshot(client, schemaName) {
2077
+ const snapshot = client?.schemaSnapshot;
2078
+ if (!snapshot) return null;
2079
+ return snapshot.schemaName === (schemaName || "public") ? snapshot : null;
2080
+ }
2081
+ /**
2082
+ * Reads the catalog for `schemaName` in three queries. Must be called on the
2083
+ * pinned init client so the snapshot reflects what that connection will see.
2084
+ */
2085
+ async function loadSchemaSnapshot(client, schemaName) {
2086
+ const schema = schemaName || "public";
2087
+ const [tableRows, columnRows, indexRows] = await Promise.all([
2088
+ client.manyOrNone(`SELECT tablename FROM pg_catalog.pg_tables WHERE schemaname = $1`, [schema]),
2089
+ client.manyOrNone(`SELECT c.relname AS table_name, a.attname AS column_name
2090
+ FROM pg_catalog.pg_class c
2091
+ JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace
2092
+ JOIN pg_catalog.pg_attribute a ON a.attrelid = c.oid
2093
+ WHERE n.nspname = $1
2094
+ AND c.relkind IN ('r', 'p')
2095
+ AND a.attnum > 0
2096
+ AND NOT a.attisdropped`, [schema]),
2097
+ client.manyOrNone(`SELECT c.relname AS indexname, i.indisreplident AS is_replica_identity, i.indisprimary AS is_primary
2098
+ FROM pg_catalog.pg_index i
2099
+ JOIN pg_catalog.pg_class c ON c.oid = i.indexrelid
2100
+ JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace
2101
+ WHERE n.nspname = $1`, [schema])
2102
+ ]);
2103
+ const columns = /* @__PURE__ */ new Map();
2104
+ for (const row of columnRows) {
2105
+ let set = columns.get(row.table_name);
2106
+ if (!set) {
2107
+ set = /* @__PURE__ */ new Set();
2108
+ columns.set(row.table_name, set);
2109
+ }
2110
+ set.add(row.column_name);
2111
+ }
2112
+ const indexes = /* @__PURE__ */ new Set();
2113
+ const replicaIdentityIndexes = /* @__PURE__ */ new Set();
2114
+ const primaryKeyIndexes = /* @__PURE__ */ new Set();
2115
+ for (const row of indexRows) {
2116
+ indexes.add(row.indexname);
2117
+ if (row.is_replica_identity) replicaIdentityIndexes.add(row.indexname);
2118
+ if (row.is_primary) primaryKeyIndexes.add(row.indexname.toLowerCase());
2119
+ }
2120
+ return {
2121
+ schemaName: schema,
2122
+ tables: new Set(tableRows.map((r) => r.tablename)),
2123
+ columns,
2124
+ indexes,
2125
+ replicaIdentityIndexes,
2126
+ primaryKeyIndexes
2127
+ };
2128
+ }
2129
+ //#endregion
2053
2130
  //#region src/storage/db/index.ts
2054
2131
  /**
2055
2132
  * Resolves PgDomainConfig to a database client and schema.
@@ -2229,7 +2306,8 @@ function generateTimestampTriggerSQL(tableName, schemaName) {
2229
2306
  schemaName: quotedSchemaName
2230
2307
  });
2231
2308
  const functionName = `${quotedSchemaName}.trigger_set_timestamps`;
2232
- const triggerName = `"${parseSqlIdentifier(`${tableName}_timestamps`, "trigger name")}"`;
2309
+ const parsedTriggerName = parseSqlIdentifier(`${tableName}_timestamps`, "trigger name");
2310
+ const triggerName = `"${parsedTriggerName}"`;
2233
2311
  return `CREATE OR REPLACE FUNCTION ${functionName}()
2234
2312
  RETURNS TRIGGER AS $$
2235
2313
  BEGIN
@@ -2248,12 +2326,38 @@ BEGIN
2248
2326
  END;
2249
2327
  $$ LANGUAGE plpgsql;
2250
2328
 
2251
- DROP TRIGGER IF EXISTS ${triggerName} ON ${fullTableName};
2329
+ DO $mastra_timestamps_trigger$
2330
+ BEGIN
2331
+ -- Recreating the trigger unconditionally would take an ACCESS EXCLUSIVE
2332
+ -- lock on the table (DROP TRIGGER does, even when nothing changes), and
2333
+ -- init runs on every process start. Skip when the trigger is already
2334
+ -- exactly what the CREATE below would produce.
2335
+ --
2336
+ -- tgtype 23 = ROW (1) | BEFORE (2) | INSERT (4) | UPDATE (16), so a trigger
2337
+ -- whose timing or events differ still falls through and gets rebuilt. The
2338
+ -- behaviour itself lives in the function, which is replaced above on every
2339
+ -- init, so an upgraded function body lands without touching the trigger.
2340
+ IF NOT EXISTS (
2341
+ SELECT 1
2342
+ FROM pg_catalog.pg_trigger tg
2343
+ JOIN pg_catalog.pg_class c ON c.oid = tg.tgrelid
2344
+ JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace
2345
+ WHERE tg.tgname = ${`'${parsedTriggerName}'`}
2346
+ AND c.relname = ${`'${parseSqlIdentifier(tableName, "table name")}'`}
2347
+ AND n.nspname = ${schemaName ? `'${parseSqlIdentifier(schemaName, "schema name")}'` : `'public'`}
2348
+ AND NOT tg.tgisinternal
2349
+ AND tg.tgtype = 23
2350
+ AND tg.tgfoid = '${functionName}()'::regprocedure
2351
+ ) THEN
2352
+ DROP TRIGGER IF EXISTS ${triggerName} ON ${fullTableName};
2252
2353
 
2253
- CREATE TRIGGER ${triggerName}
2254
- BEFORE INSERT OR UPDATE ON ${fullTableName}
2255
- FOR EACH ROW
2256
- EXECUTE FUNCTION ${functionName}();`;
2354
+ CREATE TRIGGER ${triggerName}
2355
+ BEFORE INSERT OR UPDATE ON ${fullTableName}
2356
+ FOR EACH ROW
2357
+ EXECUTE FUNCTION ${functionName}();
2358
+ END IF;
2359
+ END
2360
+ $mastra_timestamps_trigger$;`;
2257
2361
  }
2258
2362
  const schemaSetupRegistry = /* @__PURE__ */ new Map();
2259
2363
  /**
@@ -2280,6 +2384,99 @@ var PgDB = class extends MastraBase {
2280
2384
  this.skipDefaultIndexes = config.skipDefaultIndexes;
2281
2385
  }
2282
2386
  /**
2387
+ * Catalog snapshot for the current init window, or `null` outside it.
2388
+ *
2389
+ * When non-null, the init-path methods below answer existence questions from
2390
+ * it instead of round-tripping to the server, and record the objects they
2391
+ * create so later callers in the same init see them. See
2392
+ * {@link SchemaSnapshot} for why it is scoped to init only.
2393
+ */
2394
+ get schemaSnapshot() {
2395
+ return getSchemaSnapshot(this.client, this.schemaName);
2396
+ }
2397
+ /**
2398
+ * Whether the snapshot proves `generateTableSQL` would be a no-op for this
2399
+ * table — i.e. the CREATE statement can be skipped.
2400
+ *
2401
+ * For most tables that is just "the table exists". `workflow_snapshot` is the
2402
+ * exception: its generated SQL also carries a DO block that back-fills the
2403
+ * `(workflow_name, run_id)` unique constraint and promotes it to the table's
2404
+ * replica identity, so a table created by an older version still needs the
2405
+ * statement to run.
2406
+ */
2407
+ snapshotShowsTableConverged(snapshot, tableName) {
2408
+ if (!snapshot.tables.has(tableName)) return false;
2409
+ if (tableName === TABLE_WORKFLOW_SNAPSHOT) {
2410
+ const constraintName = buildConstraintName({
2411
+ baseName: "mastra_workflow_snapshot_workflow_name_run_id_key",
2412
+ schemaName: this.schemaName ? parseSqlIdentifier(this.schemaName, "schema name") : void 0
2413
+ }).toLowerCase();
2414
+ return snapshot.indexes.has(constraintName) && snapshot.replicaIdentityIndexes.has(constraintName);
2415
+ }
2416
+ return true;
2417
+ }
2418
+ /** Column set for `tableName` in the snapshot, created empty if absent. */
2419
+ snapshotColumns(snapshot, tableName) {
2420
+ let columns = snapshot.columns.get(tableName);
2421
+ if (!columns) {
2422
+ columns = /* @__PURE__ */ new Set();
2423
+ snapshot.columns.set(tableName, columns);
2424
+ }
2425
+ return columns;
2426
+ }
2427
+ /**
2428
+ * Records an out-of-band `ALTER TABLE … RENAME TO` in the init snapshot.
2429
+ *
2430
+ * Init-time migrations that issue raw DDL on `this.client` (instead of going
2431
+ * through createTable/alterTable/createIndex, which maintain the snapshot
2432
+ * themselves) MUST report it through these `note*` methods. A snapshot that
2433
+ * still lists a renamed-away table makes a later createTable() in the same
2434
+ * init skip the rebuild the migration depends on — stranding data. No-op
2435
+ * outside the init window.
2436
+ *
2437
+ * Indexes riding along with a rename keep their names, so the snapshot's
2438
+ * index set stays accurate without changes here.
2439
+ */
2440
+ noteTableRenamed(oldName, newName) {
2441
+ const snapshot = this.schemaSnapshot;
2442
+ if (snapshot) {
2443
+ if (snapshot.tables.delete(oldName)) snapshot.tables.add(newName);
2444
+ const columns = snapshot.columns.get(oldName);
2445
+ if (columns) {
2446
+ snapshot.columns.delete(oldName);
2447
+ snapshot.columns.set(newName, columns);
2448
+ }
2449
+ }
2450
+ this.tableColumnsCache.delete(oldName);
2451
+ this.tableColumnsCache.delete(newName);
2452
+ }
2453
+ /**
2454
+ * Records an out-of-band `DROP TABLE` in the init snapshot. See
2455
+ * {@link noteTableRenamed} for why raw-DDL migrations must call this.
2456
+ *
2457
+ * The dropped table's indexes vanish with it, but the snapshot's flat index
2458
+ * set cannot map names back to tables. Stale entries only make a later
2459
+ * createIndex() skip a recreate until the next init re-reads the catalog —
2460
+ * the same self-healing bound the rest of the snapshot design accepts.
2461
+ */
2462
+ noteTableDropped(tableName) {
2463
+ const snapshot = this.schemaSnapshot;
2464
+ if (snapshot) {
2465
+ snapshot.tables.delete(tableName);
2466
+ snapshot.columns.delete(tableName);
2467
+ }
2468
+ this.tableColumnsCache.delete(tableName);
2469
+ }
2470
+ /**
2471
+ * Records an out-of-band `ALTER TABLE … ADD COLUMN` in the init snapshot.
2472
+ * See {@link noteTableRenamed} for why raw-DDL migrations must call this.
2473
+ */
2474
+ noteColumnAdded(tableName, column) {
2475
+ const snapshot = this.schemaSnapshot;
2476
+ if (snapshot) this.snapshotColumns(snapshot, tableName).add(column);
2477
+ this.tableColumnsCache.delete(tableName);
2478
+ }
2479
+ /**
2283
2480
  * Gets the set of column names that actually exist in the database table.
2284
2481
  * Results are cached; the cache is invalidated when alterTable() adds new columns.
2285
2482
  */
@@ -2306,6 +2503,12 @@ var PgDB = class extends MastraBase {
2306
2503
  }
2307
2504
  async hasColumn(table, column) {
2308
2505
  const schema = this.schemaName || "public";
2506
+ const snapshot = this.schemaSnapshot;
2507
+ if (snapshot) {
2508
+ const columns = snapshot.columns.get(table);
2509
+ if (!columns) return false;
2510
+ return columns.has(column) || columns.has(column.toLowerCase());
2511
+ }
2309
2512
  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)`, [
2310
2513
  schema,
2311
2514
  table,
@@ -2344,6 +2547,14 @@ var PgDB = class extends MastraBase {
2344
2547
  if (!this.schemaName) return;
2345
2548
  let registryEntry = schemaSetupRegistry.get(this.schemaName);
2346
2549
  if (registryEntry?.complete) return;
2550
+ const snapshot = this.schemaSnapshot;
2551
+ if (snapshot && snapshot.tables.size > 0) {
2552
+ schemaSetupRegistry.set(this.schemaName, {
2553
+ promise: null,
2554
+ complete: true
2555
+ });
2556
+ return;
2557
+ }
2347
2558
  const quotedSchemaName = getSchemaName$6(this.schemaName);
2348
2559
  if (!registryEntry?.promise) {
2349
2560
  const schemaNameCapture = this.schemaName;
@@ -2448,16 +2659,28 @@ var PgDB = class extends MastraBase {
2448
2659
  try {
2449
2660
  const timeZColumnNames = Object.entries(schema).filter(([_, def]) => def.type === "timestamp").map(([name]) => name);
2450
2661
  if (this.schemaName) await this.setupSchema();
2451
- const sql = generateTableSQL({
2452
- tableName,
2453
- schema,
2454
- schemaName: this.schemaName,
2455
- compositePrimaryKey
2456
- });
2457
- try {
2458
- await this.client.none(sql);
2459
- } catch (error) {
2460
- if (!isDuplicateRelationError(error)) throw error;
2662
+ const snapshot = this.schemaSnapshot;
2663
+ if (!snapshot || !this.snapshotShowsTableConverged(snapshot, tableName)) {
2664
+ const sql = generateTableSQL({
2665
+ tableName,
2666
+ schema,
2667
+ schemaName: this.schemaName,
2668
+ compositePrimaryKey
2669
+ });
2670
+ try {
2671
+ await this.client.none(sql);
2672
+ } catch (error) {
2673
+ if (!isDuplicateRelationError(error)) throw error;
2674
+ }
2675
+ if (snapshot) {
2676
+ snapshot.tables.add(tableName);
2677
+ const created = this.snapshotColumns(snapshot, tableName);
2678
+ for (const [name, def] of Object.entries(schema)) {
2679
+ const parsedName = parseSqlIdentifier(name, "column name");
2680
+ created.add(parsedName);
2681
+ if (def.type === "timestamp") created.add(`${parsedName}Z`);
2682
+ }
2683
+ }
2461
2684
  }
2462
2685
  await this.alterTable({
2463
2686
  tableName,
@@ -2523,10 +2746,12 @@ MIGRATION REQUIRED: Duplicate spans detected in ${duplicateInfo.tableName}\n====
2523
2746
  const sqlType = mapToSqlType(columnDef.type);
2524
2747
  const alterSql = `ALTER TABLE ${fullTableName} ADD COLUMN IF NOT EXISTS "${parsedColumnName}" ${sqlType} ${columnDef.nullable ? "" : "NOT NULL"} ${!columnDef.nullable ? this.getDefaultValue(columnDef.type) : ""}`.trim();
2525
2748
  await this.client.none(alterSql);
2749
+ this.noteColumnAdded(TABLE_SPANS, columnName);
2526
2750
  this.logger?.debug?.(`Added column '${columnName}' to ${fullTableName}`);
2527
2751
  if (sqlType === "TIMESTAMP") {
2528
2752
  const timestampZSql = `ALTER TABLE ${fullTableName} ADD COLUMN IF NOT EXISTS "${parsedColumnName}Z" TIMESTAMPTZ DEFAULT NOW()`.trim();
2529
2753
  await this.client.none(timestampZSql);
2754
+ this.noteColumnAdded(TABLE_SPANS, `${columnName}Z`);
2530
2755
  this.logger?.debug?.(`Added timezone column '${columnName}Z' to ${fullTableName}`);
2531
2756
  }
2532
2757
  }
@@ -2535,6 +2760,7 @@ MIGRATION REQUIRED: Duplicate spans detected in ${duplicateInfo.tableName}\n====
2535
2760
  if (!await this.hasColumn(TABLE_SPANS, tzColumnName)) {
2536
2761
  const timestampZSql = `ALTER TABLE ${fullTableName} ADD COLUMN IF NOT EXISTS "${parseSqlIdentifier(tzColumnName, "column name")}" TIMESTAMPTZ DEFAULT NOW()`.trim();
2537
2762
  await this.client.none(timestampZSql);
2763
+ this.noteColumnAdded(TABLE_SPANS, tzColumnName);
2538
2764
  this.logger?.debug?.(`Added timezone column '${tzColumnName}' to ${fullTableName}`);
2539
2765
  }
2540
2766
  }
@@ -2657,6 +2883,12 @@ MIGRATION REQUIRED: Duplicate spans detected in ${duplicateInfo.tableName}\n====
2657
2883
  schemaName: (this.schemaName ? parseSqlIdentifier(this.schemaName, "schema name") : "") || void 0
2658
2884
  });
2659
2885
  const schemaFilter = this.schemaName || "public";
2886
+ const snapshot = this.schemaSnapshot;
2887
+ if (snapshot) return snapshot.primaryKeyIndexes.has(constraintName.toLowerCase());
2888
+ return this.spansPrimaryKeyExistsLive(constraintName, schemaFilter);
2889
+ }
2890
+ /** Live-catalog variant of {@link spansPrimaryKeyExists}, bypassing the snapshot. */
2891
+ async spansPrimaryKeyExistsLive(constraintName, schemaFilter) {
2660
2892
  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;
2661
2893
  }
2662
2894
  /**
@@ -2687,10 +2919,12 @@ MIGRATION REQUIRED: Duplicate spans detected in ${duplicateInfo.tableName}\n====
2687
2919
  ADD CONSTRAINT ${constraintName}
2688
2920
  PRIMARY KEY ("traceId", "spanId")
2689
2921
  `);
2922
+ this.schemaSnapshot?.primaryKeyIndexes.add(constraintName.toLowerCase());
2690
2923
  this.logger?.info?.(`Added PRIMARY KEY constraint ${constraintName} to ${fullTableName}`);
2691
2924
  } catch (error) {
2692
2925
  if (isDuplicateRelationError(error)) {
2693
- if (await this.spansPrimaryKeyExists()) {
2926
+ if (await this.spansPrimaryKeyExistsLive(constraintName, schemaFilter)) {
2927
+ this.schemaSnapshot?.primaryKeyIndexes.add(constraintName.toLowerCase());
2694
2928
  this.logger?.debug?.(`PRIMARY KEY constraint ${constraintName} was created by another process`);
2695
2929
  return;
2696
2930
  }
@@ -2772,16 +3006,25 @@ MIGRATION REQUIRED: Duplicate spans detected in ${duplicateInfo.tableName}\n====
2772
3006
  indexName: tableName,
2773
3007
  schemaName: getSchemaName$6(this.schemaName)
2774
3008
  });
3009
+ const snapshot = this.schemaSnapshot;
3010
+ const knownColumns = snapshot ? this.snapshotColumns(snapshot, tableName) : null;
2775
3011
  try {
2776
3012
  for (const columnName of ifNotExists) if (schema[columnName]) {
2777
3013
  const columnDef = schema[columnName];
2778
3014
  const parsedColumnName = parseSqlIdentifier(columnName, "column name");
2779
3015
  const sqlType = mapToSqlType(columnDef.type);
2780
3016
  const alterSql = `ALTER TABLE ${fullTableName} ADD COLUMN IF NOT EXISTS "${parsedColumnName}" ${sqlType} ${columnDef.nullable ? "" : "NOT NULL"} ${!columnDef.nullable ? this.getDefaultValue(columnDef.type) : ""}`.trim();
2781
- await this.client.none(alterSql);
3017
+ if (!knownColumns?.has(parsedColumnName)) {
3018
+ await this.client.none(alterSql);
3019
+ knownColumns?.add(parsedColumnName);
3020
+ }
2782
3021
  if (sqlType === "TIMESTAMP") {
2783
- const timestampZSql = `ALTER TABLE ${fullTableName} ADD COLUMN IF NOT EXISTS "${parsedColumnName}Z" TIMESTAMPTZ DEFAULT NOW()`.trim();
2784
- await this.client.none(timestampZSql);
3022
+ const tzColumnName = `${parsedColumnName}Z`;
3023
+ if (!knownColumns?.has(tzColumnName)) {
3024
+ const timestampZSql = `ALTER TABLE ${fullTableName} ADD COLUMN IF NOT EXISTS "${tzColumnName}" TIMESTAMPTZ DEFAULT NOW()`.trim();
3025
+ await this.client.none(timestampZSql);
3026
+ knownColumns?.add(tzColumnName);
3027
+ }
2785
3028
  }
2786
3029
  this.logger?.debug?.(`Ensured column ${parsedColumnName} exists in table ${fullTableName}`);
2787
3030
  }
@@ -2867,7 +3110,10 @@ MIGRATION REQUIRED: Duplicate spans detected in ${duplicateInfo.tableName}\n====
2867
3110
  indexName: table,
2868
3111
  schemaName: getSchemaName$6(this.schemaName)
2869
3112
  });
2870
- if (await this.client.oneOrNone(`SELECT 1 FROM pg_indexes
3113
+ const snapshot = this.schemaSnapshot;
3114
+ if (snapshot) {
3115
+ if (snapshot.indexes.has(name)) return;
3116
+ } else if (await this.client.oneOrNone(`SELECT 1 FROM pg_indexes
2871
3117
  WHERE indexname = $1
2872
3118
  AND schemaname = $2`, [name, schemaName])) return;
2873
3119
  const uniqueStr = unique ? "UNIQUE " : "";
@@ -2889,6 +3135,7 @@ MIGRATION REQUIRED: Duplicate spans detected in ${duplicateInfo.tableName}\n====
2889
3135
  if (storage && Object.keys(storage).length > 0) withStr = ` WITH (${Object.entries(storage).map(([key, value]) => `${key} = ${value}`).join(", ")})`;
2890
3136
  const sql = `CREATE ${uniqueStr}INDEX ${concurrentStr}${`"${parseSqlIdentifier(name, "index name")}"`} ON ${fullTableName} ${methodStr}(${columnsStr})${withStr}${tablespaceStr}${whereStr}`;
2891
3137
  await this.client.none(sql);
3138
+ snapshot?.indexes.add(name);
2892
3139
  } catch (error) {
2893
3140
  if (error instanceof Error && error.message.includes("CONCURRENTLY")) {
2894
3141
  const retryOptions = {
@@ -2908,15 +3155,34 @@ MIGRATION REQUIRED: Duplicate spans detected in ${duplicateInfo.tableName}\n====
2908
3155
  }, error);
2909
3156
  }
2910
3157
  }
3158
+ /**
3159
+ * Runs a caller-built `CREATE INDEX IF NOT EXISTS` statement, unless the init
3160
+ * snapshot already proves `indexName` exists.
3161
+ *
3162
+ * `createIndex` covers the indexes described by {@link CreateIndexOptions};
3163
+ * this is for the two init paths that hand-write their statement (a partial
3164
+ * or otherwise non-standard index) and would otherwise send a no-op DDL on
3165
+ * every warm init.
3166
+ */
3167
+ async createIndexFromStatement(indexName, sql) {
3168
+ const snapshot = this.schemaSnapshot;
3169
+ if (snapshot?.indexes.has(indexName)) return;
3170
+ await this.client.none(sql);
3171
+ snapshot?.indexes.add(indexName);
3172
+ }
2911
3173
  async dropIndex(indexName) {
2912
3174
  try {
2913
3175
  const schemaName = this.schemaName || "public";
2914
- if (!await this.client.oneOrNone(`SELECT 1 FROM pg_indexes
3176
+ const snapshot = this.schemaSnapshot;
3177
+ if (snapshot) {
3178
+ if (!snapshot.indexes.has(indexName)) return;
3179
+ } else if (!await this.client.oneOrNone(`SELECT 1 FROM pg_indexes
2915
3180
  WHERE indexname = $1
2916
3181
  AND schemaname = $2`, [indexName, schemaName])) return;
2917
3182
  const quotedIndexName = `"${parseSqlIdentifier(indexName, "index name")}"`;
2918
3183
  const sql = `DROP INDEX IF EXISTS ${getSchemaName$6(this.schemaName)}.${quotedIndexName}`;
2919
3184
  await this.client.none(sql);
3185
+ snapshot?.indexes.delete(indexName);
2920
3186
  } catch (error) {
2921
3187
  throw new MastraError({
2922
3188
  id: createStorageErrorId("PG", "INDEX_DROP", "FAILED"),
@@ -3413,7 +3679,9 @@ var AgentsPG = class AgentsPG extends AgentsStorage {
3413
3679
  });
3414
3680
  if (await this.#db.hasColumn(TABLE_AGENTS, "name")) {
3415
3681
  await this.#db.client.none(`ALTER TABLE ${fullTableName} RENAME TO "${TABLE_AGENTS}_legacy"`);
3682
+ this.#db.noteTableRenamed(TABLE_AGENTS, `${TABLE_AGENTS}_legacy`);
3416
3683
  await this.#db.client.none(`DROP TABLE IF EXISTS ${fullVersionsTableName}`);
3684
+ this.#db.noteTableDropped(TABLE_AGENT_VERSIONS);
3417
3685
  }
3418
3686
  if (!await this.#db.hasColumn(`${TABLE_AGENTS}_legacy`, "name")) return;
3419
3687
  const oldAgents = await this.#db.client.manyOrNone(`SELECT * FROM ${legacyTableName}`);
@@ -3470,6 +3738,7 @@ var AgentsPG = class AgentsPG extends AgentsStorage {
3470
3738
  ]);
3471
3739
  }
3472
3740
  await this.#db.client.none(`DROP TABLE IF EXISTS ${legacyTableName}`);
3741
+ this.#db.noteTableDropped(`${TABLE_AGENTS}_legacy`);
3473
3742
  }
3474
3743
  /**
3475
3744
  * Migrates the agent_versions table from the old snapshot-based schema (single `snapshot` JSON column)
@@ -3487,7 +3756,9 @@ var AgentsPG = class AgentsPG extends AgentsStorage {
3487
3756
  schemaName: getSchemaName$5(this.#schema)
3488
3757
  });
3489
3758
  await this.#db.client.none(`DROP TABLE IF EXISTS ${fullVersionsTableName}`);
3759
+ this.#db.noteTableDropped(TABLE_AGENT_VERSIONS);
3490
3760
  await this.#db.client.none(`DROP TABLE IF EXISTS ${legacyTableName}`);
3761
+ this.#db.noteTableDropped(`${TABLE_AGENTS}_legacy`);
3491
3762
  }
3492
3763
  /**
3493
3764
  * Migrates the tools field from string[] format to JSONB format { "tool-key": { "description": "..." } }.
@@ -5003,6 +5274,7 @@ var DatasetsPG = class DatasetsPG extends DatasetsStorage {
5003
5274
  await this.#addColumnIfNotExists(TABLE_DATASET_ITEMS, "organizationId", "TEXT");
5004
5275
  await this.#addColumnIfNotExists(TABLE_DATASET_ITEMS, "projectId", "TEXT");
5005
5276
  await this.#addColumnIfNotExists(TABLE_DATASET_ITEMS, "toolMocks", "JSONB");
5277
+ await this.#addColumnIfNotExists(TABLE_DATASET_ITEMS, "unmockedToolPolicy", "TEXT");
5006
5278
  await this.#addColumnIfNotExists(TABLE_DATASET_ITEMS, "externalId", "TEXT");
5007
5279
  await this.createDefaultIndexes();
5008
5280
  await this.createCustomIndexes();
@@ -5014,6 +5286,7 @@ var DatasetsPG = class DatasetsPG extends DatasetsStorage {
5014
5286
  schemaName: getSchemaName$5(this.#schema)
5015
5287
  });
5016
5288
  await this.#db.client.none(`ALTER TABLE ${fullTableName} ADD COLUMN "${column}" ${sqlType}`);
5289
+ this.#db.noteColumnAdded(table, column);
5017
5290
  }
5018
5291
  }
5019
5292
  getDefaultIndexDefinitions() {
@@ -5124,6 +5397,7 @@ var DatasetsPG = class DatasetsPG extends DatasetsStorage {
5124
5397
  groundTruth: row.groundTruth ? safelyParseJSON(row.groundTruth) : void 0,
5125
5398
  expectedTrajectory: row.expectedTrajectory ? safelyParseJSON(row.expectedTrajectory) : void 0,
5126
5399
  toolMocks: row.toolMocks ? safelyParseJSON(row.toolMocks) : void 0,
5400
+ unmockedToolPolicy: row.unmockedToolPolicy ?? void 0,
5127
5401
  requestContext: row.requestContext ? safelyParseJSON(row.requestContext) : void 0,
5128
5402
  metadata: row.metadata ? safelyParseJSON(row.metadata) : void 0,
5129
5403
  source: row.source ? safelyParseJSON(row.source) : void 0,
@@ -5145,6 +5419,7 @@ var DatasetsPG = class DatasetsPG extends DatasetsStorage {
5145
5419
  groundTruth: row.groundTruth ? safelyParseJSON(row.groundTruth) : void 0,
5146
5420
  expectedTrajectory: row.expectedTrajectory ? safelyParseJSON(row.expectedTrajectory) : void 0,
5147
5421
  toolMocks: row.toolMocks ? safelyParseJSON(row.toolMocks) : void 0,
5422
+ unmockedToolPolicy: row.unmockedToolPolicy ?? void 0,
5148
5423
  requestContext: row.requestContext ? safelyParseJSON(row.requestContext) : void 0,
5149
5424
  metadata: row.metadata ? safelyParseJSON(row.metadata) : void 0,
5150
5425
  source: row.source ? safelyParseJSON(row.source) : void 0,
@@ -5476,7 +5751,7 @@ var DatasetsPG = class DatasetsPG extends DatasetsStorage {
5476
5751
  newVersion = row.version;
5477
5752
  parentOrganizationId = row.organizationId ?? null;
5478
5753
  parentProjectId = row.projectId ?? null;
5479
- 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)`, [
5754
+ 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)`, [
5480
5755
  id,
5481
5756
  args.datasetId,
5482
5757
  newVersion,
@@ -5487,6 +5762,7 @@ var DatasetsPG = class DatasetsPG extends DatasetsStorage {
5487
5762
  jsonbArg(args.groundTruth),
5488
5763
  jsonbArg(args.expectedTrajectory),
5489
5764
  jsonbArg(args.toolMocks),
5765
+ args.unmockedToolPolicy ?? null,
5490
5766
  jsonbArg(args.requestContext),
5491
5767
  jsonbArg(args.metadata),
5492
5768
  jsonbArg(args.source),
@@ -5513,6 +5789,7 @@ var DatasetsPG = class DatasetsPG extends DatasetsStorage {
5513
5789
  groundTruth: args.groundTruth,
5514
5790
  expectedTrajectory: args.expectedTrajectory,
5515
5791
  toolMocks: args.toolMocks,
5792
+ unmockedToolPolicy: args.unmockedToolPolicy,
5516
5793
  requestContext: args.requestContext,
5517
5794
  metadata: args.metadata,
5518
5795
  source: args.source,
@@ -5566,6 +5843,7 @@ var DatasetsPG = class DatasetsPG extends DatasetsStorage {
5566
5843
  const mergedGroundTruth = args.groundTruth !== void 0 ? args.groundTruth : existing.groundTruth;
5567
5844
  const mergedExpectedTrajectory = args.expectedTrajectory !== void 0 ? args.expectedTrajectory : existing.expectedTrajectory;
5568
5845
  const mergedToolMocks = args.toolMocks !== void 0 ? args.toolMocks : existing.toolMocks;
5846
+ const mergedUnmockedToolPolicy = args.unmockedToolPolicy !== void 0 ? args.unmockedToolPolicy : existing.unmockedToolPolicy;
5569
5847
  const mergedRequestContext = args.requestContext !== void 0 ? args.requestContext : existing.requestContext;
5570
5848
  const mergedMetadata = args.metadata !== void 0 ? args.metadata : existing.metadata;
5571
5849
  const mergedSource = args.source !== void 0 ? args.source : existing.source;
@@ -5578,7 +5856,7 @@ var DatasetsPG = class DatasetsPG extends DatasetsStorage {
5578
5856
  parentOrganizationId = row.organizationId ?? null;
5579
5857
  parentProjectId = row.projectId ?? null;
5580
5858
  await t.none(`UPDATE ${itemsTable} SET "validTo" = $1 WHERE "id" = $2 AND "validTo" IS NULL AND "isDeleted" = false`, [newVersion, args.id]);
5581
- 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)`, [
5859
+ 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)`, [
5582
5860
  args.id,
5583
5861
  args.datasetId,
5584
5862
  newVersion,
@@ -5589,6 +5867,7 @@ var DatasetsPG = class DatasetsPG extends DatasetsStorage {
5589
5867
  jsonbArg(mergedGroundTruth),
5590
5868
  jsonbArg(mergedExpectedTrajectory),
5591
5869
  jsonbArg(mergedToolMocks),
5870
+ mergedUnmockedToolPolicy ?? null,
5592
5871
  jsonbArg(mergedRequestContext),
5593
5872
  jsonbArg(mergedMetadata),
5594
5873
  jsonbArg(mergedSource),
@@ -5614,6 +5893,7 @@ var DatasetsPG = class DatasetsPG extends DatasetsStorage {
5614
5893
  groundTruth: mergedGroundTruth,
5615
5894
  expectedTrajectory: mergedExpectedTrajectory,
5616
5895
  toolMocks: mergedToolMocks,
5896
+ unmockedToolPolicy: mergedUnmockedToolPolicy,
5617
5897
  requestContext: mergedRequestContext,
5618
5898
  metadata: mergedMetadata,
5619
5899
  source: mergedSource,
@@ -5662,7 +5942,7 @@ var DatasetsPG = class DatasetsPG extends DatasetsStorage {
5662
5942
  const parentOrganizationId = row.organizationId ?? null;
5663
5943
  const parentProjectId = row.projectId ?? null;
5664
5944
  await t.none(`UPDATE ${itemsTable} SET "validTo" = $1 WHERE "id" = $2 AND "validTo" IS NULL AND "isDeleted" = false`, [newVersion, id]);
5665
- 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)`, [
5945
+ 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)`, [
5666
5946
  id,
5667
5947
  datasetId,
5668
5948
  newVersion,
@@ -5673,6 +5953,7 @@ var DatasetsPG = class DatasetsPG extends DatasetsStorage {
5673
5953
  jsonbArg(existing.groundTruth),
5674
5954
  jsonbArg(existing.expectedTrajectory),
5675
5955
  jsonbArg(existing.toolMocks),
5956
+ existing.unmockedToolPolicy ?? null,
5676
5957
  jsonbArg(existing.requestContext),
5677
5958
  jsonbArg(existing.metadata),
5678
5959
  jsonbArg(existing.source),
@@ -5732,7 +6013,7 @@ var DatasetsPG = class DatasetsPG extends DatasetsStorage {
5732
6013
  const nowIso = now.toISOString();
5733
6014
  await t.none(`UPDATE ${datasetsTable} SET "version" = $2 WHERE "id" = $1`, [input.datasetId, newVersion]);
5734
6015
  for (const { id, item } of plan.inserts) {
5735
- 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)`, [
6016
+ 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)`, [
5736
6017
  id,
5737
6018
  input.datasetId,
5738
6019
  newVersion,
@@ -5743,6 +6024,7 @@ var DatasetsPG = class DatasetsPG extends DatasetsStorage {
5743
6024
  jsonbArg(item.groundTruth),
5744
6025
  jsonbArg(item.expectedTrajectory),
5745
6026
  jsonbArg(item.toolMocks),
6027
+ item.unmockedToolPolicy ?? null,
5746
6028
  jsonbArg(item.requestContext),
5747
6029
  jsonbArg(item.metadata),
5748
6030
  jsonbArg(item.source),
@@ -5762,6 +6044,7 @@ var DatasetsPG = class DatasetsPG extends DatasetsStorage {
5762
6044
  groundTruth: item.groundTruth,
5763
6045
  expectedTrajectory: item.expectedTrajectory,
5764
6046
  toolMocks: item.toolMocks,
6047
+ unmockedToolPolicy: item.unmockedToolPolicy,
5765
6048
  requestContext: item.requestContext,
5766
6049
  metadata: item.metadata,
5767
6050
  source: item.source,
@@ -5824,7 +6107,7 @@ var DatasetsPG = class DatasetsPG extends DatasetsStorage {
5824
6107
  const newVersion = (await t.one(`UPDATE ${datasetsTable} SET "version" = "version" + 1 WHERE "id" = $1 RETURNING "version"`, [input.datasetId])).version;
5825
6108
  for (const item of currentItems) {
5826
6109
  await t.none(`UPDATE ${itemsTable} SET "validTo" = $1 WHERE "id" = $2 AND "validTo" IS NULL AND "isDeleted" = false`, [newVersion, item.id]);
5827
- 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)`, [
6110
+ 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)`, [
5828
6111
  item.id,
5829
6112
  input.datasetId,
5830
6113
  newVersion,
@@ -5835,6 +6118,7 @@ var DatasetsPG = class DatasetsPG extends DatasetsStorage {
5835
6118
  jsonbArg(item.groundTruth),
5836
6119
  jsonbArg(item.expectedTrajectory),
5837
6120
  jsonbArg(item.toolMocks),
6121
+ item.unmockedToolPolicy ?? null,
5838
6122
  jsonbArg(item.requestContext),
5839
6123
  jsonbArg(item.metadata),
5840
6124
  jsonbArg(item.source),
@@ -6129,6 +6413,7 @@ var ExperimentsPG = class ExperimentsPG extends ExperimentsStorage {
6129
6413
  ifNotExists: [
6130
6414
  "status",
6131
6415
  "tags",
6416
+ "comment",
6132
6417
  "toolMockReport",
6133
6418
  "organizationId",
6134
6419
  "projectId"
@@ -6306,6 +6591,7 @@ var ExperimentsPG = class ExperimentsPG extends ExperimentsStorage {
6306
6591
  traceId: row.traceId ?? null,
6307
6592
  status: row.status ?? null,
6308
6593
  tags: row.tags ? safelyParseJSON(row.tags) : null,
6594
+ comment: row.comment ?? null,
6309
6595
  toolMockReport: row.toolMockReport ? safelyParseJSON(row.toolMockReport) : null,
6310
6596
  createdAt: ensureDate(row.createdAtZ || row.createdAt)
6311
6597
  };
@@ -6631,6 +6917,10 @@ var ExperimentsPG = class ExperimentsPG extends ExperimentsStorage {
6631
6917
  setClauses.push(`"tags" = $${paramIndex++}`);
6632
6918
  values.push(JSON.stringify(input.tags));
6633
6919
  }
6920
+ if (input.comment !== void 0) {
6921
+ setClauses.push(`"comment" = $${paramIndex++}`);
6922
+ values.push(input.comment);
6923
+ }
6634
6924
  if (setClauses.length === 0) {
6635
6925
  const existing = await this.getExperimentResultById({ id: input.id });
6636
6926
  if (!existing) throw new MastraError({
@@ -6866,7 +7156,7 @@ var FavoritesPG = class FavoritesPG extends FavoritesStorage {
6866
7156
  indexName: TABLE_FAVORITES,
6867
7157
  schemaName: getSchemaName$5(this.#schema)
6868
7158
  });
6869
- await this.#db.client.none(`CREATE INDEX IF NOT EXISTS idx_favorites_entity ON ${fullTableName} ("entityType", "entityId")`);
7159
+ await this.#db.createIndexFromStatement("idx_favorites_entity", `CREATE INDEX IF NOT EXISTS idx_favorites_entity ON ${fullTableName} ("entityType", "entityId")`);
6870
7160
  }
6871
7161
  async dangerouslyClearAll() {
6872
7162
  const fullTableName = getTableName$5({
@@ -8362,7 +8652,7 @@ var MemoryPG = class MemoryPG extends MemoryStorage {
8362
8652
  indexName: OM_TABLE,
8363
8653
  schemaName: getSchemaName$3(this.#schema)
8364
8654
  });
8365
- await this.#db.client.none(`CREATE INDEX IF NOT EXISTS idx_om_lookup_key ON ${omTableName} ("lookupKey")`);
8655
+ await this.#db.createIndexFromStatement("idx_om_lookup_key", `CREATE INDEX IF NOT EXISTS idx_om_lookup_key ON ${omTableName} ("lookupKey")`);
8366
8656
  }
8367
8657
  await this.createDefaultIndexes();
8368
8658
  await this.createCustomIndexes();
@@ -20083,6 +20373,7 @@ var PostgresStore = class extends MastraCompositeStore {
20083
20373
  pinnedClient = await this.#pool.connect();
20084
20374
  const pinned = new PinnedClientAdapter(this.#pool, pinnedClient);
20085
20375
  this.#db.pin(pinned);
20376
+ this.#db.setSchemaSnapshot(await loadSchemaSnapshot(pinned, this.schema));
20086
20377
  await super.init();
20087
20378
  this.isInitialized = true;
20088
20379
  } catch (error) {
@@ -20094,6 +20385,7 @@ var PostgresStore = class extends MastraCompositeStore {
20094
20385
  category: ErrorCategory.THIRD_PARTY
20095
20386
  }, error);
20096
20387
  } finally {
20388
+ this.#db.setSchemaSnapshot(null);
20097
20389
  if (pinnedClient) {
20098
20390
  this.#db.unpin();
20099
20391
  pinnedClient.release();