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

package/dist/index.js

@@ -7302,7 +7302,7 @@ var FavoritesPG = class _FavoritesPG extends FavoritesStorage {
     }
     const fullFavoritesTable = getTableName2({ indexName: TABLE_FAVORITES, schemaName: getSchemaName2(schemaName) });
     statements.push(
-      `CREATE INDEX IF NOT EXISTS idx_favorites_entity ON ${fullFavoritesTable} ("entityType", "entityId")`
+      `CREATE INDEX IF NOT EXISTS idx_favorites_entity ON ${fullFavoritesTable} ("entityType", "entityId");`
     );
     return statements;
   }
@@ -20231,10 +20231,12 @@ function hashAdvisoryLockKey(key) {
   return [digest.readInt32BE(0), digest.readInt32BE(4)];
 }
 var PgFactoryStorageOps = class {
-  #pool;
+  #queryable;
+  #transactionClient;
   #schemas;
-  constructor(pool, schemas) {
-    this.#pool = pool;
+  constructor(queryable, schemas, transactionClient) {
+    this.#queryable = queryable;
+    this.#transactionClient = transactionClient;
     this.#schemas = schemas;
   }
   #schema(collection) {
@@ -20382,11 +20384,11 @@ var PgFactoryStorageOps = class {
     return result.rows.map((raw) => this.#deserializeRow(schema, raw));
   }
   async findOne(collection, where) {
-    const rows = await this.#select(this.#pool, collection, where, { limit: 1 });
+    const rows = await this.#select(this.#queryable, collection, where, { limit: 1 });
     return rows[0] ?? null;
   }
   async findMany(collection, where, opts) {
-    return this.#select(this.#pool, collection, where, opts);
+    return this.#select(this.#queryable, collection, where, opts);
   }
   async insertOne(collection, row) {
     const schema = this.#schema(collection);
@@ -20400,7 +20402,7 @@ var PgFactoryStorageOps = class {
     const sql = `INSERT INTO "${schema.name}" (${columns.map((c) => `"${c}"`).join(", ")}) VALUES (${columns.map((_, i) => `$${i + 1}`).join(", ")}) RETURNING *`;
     const args = columns.map((column) => this.#serialize(this.#column(schema, column), values[column]));
     try {
-      const result = await this.#pool.query(sql, args);
+      const result = await this.#queryable.query(sql, args);
       return this.#deserializeRow(schema, result.rows[0]);
     } catch (error) {
       if (isUniqueViolation(error)) throw new UniqueViolationError(collection, { cause: error });
@@ -20438,7 +20440,7 @@ var PgFactoryStorageOps = class {
     throw lastError ?? new Error(`PgFactoryStorage: upsert into '${collection}' did not converge`);
   }
   async updateMany(collection, where, set) {
-    return this.#updateMany(this.#pool, collection, where, set);
+    return this.#updateMany(this.#queryable, collection, where, set);
   }
   async #updateMany(queryable, collection, where, set) {
     const schema = this.#schema(collection);
@@ -20453,30 +20455,30 @@ var PgFactoryStorageOps = class {
   async deleteMany(collection, where) {
     const schema = this.#schema(collection);
     const filter = this.#buildWhere(schema, where);
-    const result = await this.#pool.query(`DELETE FROM "${schema.name}" WHERE ${filter.sql}`, filter.args);
+    const result = await this.#queryable.query(`DELETE FROM "${schema.name}" WHERE ${filter.sql}`, filter.args);
     return result.rowCount ?? 0;
   }
   async updateAtomic(collection, where, fn) {
-    const schema = this.#schema(collection);
-    const pk = primaryKeyOf(schema);
-    const client = await this.#pool.connect();
+    const run = async (client2) => {
+      const schema = this.#schema(collection);
+      const pk = primaryKeyOf(schema);
+      const rows = await this.#select(client2, collection, where, { limit: 1 }, true);
+      const row = rows[0];
+      if (!row) return null;
+      const patch = await fn(row);
+      if (patch === null) return row;
+      const pkWhere = { [pk]: row[pk] };
+      await this.#updateMany(client2, collection, pkWhere, patch);
+      const updated = await this.#select(client2, collection, pkWhere, { limit: 1 });
+      return updated[0] ?? null;
+    };
+    if (this.#transactionClient) return run(this.#transactionClient);
+    const pool = this.#queryable;
+    const client = await pool.connect();
     try {
       await client.query("BEGIN");
       try {
-        const rows = await this.#select(client, collection, where, { limit: 1 }, true);
-        const row = rows[0];
-        if (!row) {
-          await client.query("COMMIT");
-          return null;
-        }
-        const patch = await fn(row);
-        let result = row;
-        if (patch !== null) {
-          const pkWhere = { [pk]: row[pk] };
-          await this.#updateMany(client, collection, pkWhere, patch);
-          const updated = await this.#select(client, collection, pkWhere, { limit: 1 });
-          result = updated[0] ?? null;
-        }
+        const result = await run(client);
         await client.query("COMMIT");
         return result;
       } catch (error) {
@@ -20522,6 +20524,22 @@ var PgFactoryStorage = class extends FactoryStorage {
   async initStorage() {
     await this.#pool.query("SELECT 1");
   }
+  async withTransaction(fn) {
+    const client = await this.#pool.connect();
+    try {
+      await client.query("BEGIN");
+      try {
+        const result = await fn(new PgFactoryStorageOps(client, this.#schemas, client));
+        await client.query("COMMIT");
+        return result;
+      } catch (error) {
+        await client.query("ROLLBACK");
+        throw error;
+      }
+    } finally {
+      client.release();
+    }
+  }
   async ensureCollections(schemas) {
     for (const schema of schemas) {
       await this.#ensureCollection(schema);