@mastra/pg 1.9.0-alpha.0 → 1.9.1-alpha.0

package/dist/index.js

@@ -2474,36 +2474,39 @@ var PgDB = class extends MastraBase {
         return getDefaultValue(type);
     }
   }
-  async insert({ tableName, record }) {
-    try {
-      this.addTimestampZColumns(record);
-      const filteredRecord = await this.filterRecordToKnownColumns(tableName, record);
-      const schemaName = getSchemaName(this.schemaName);
-      const columns = Object.keys(filteredRecord).map((col) => parseSqlIdentifier(col, "column name"));
-      if (columns.length === 0) return;
-      const values = this.prepareValuesForInsert(filteredRecord, tableName);
-      const placeholders = values.map((_, i) => `$${i + 1}`).join(", ");
-      const fullTableName = getTableName({ indexName: tableName, schemaName });
-      const columnList = columns.map((c) => `"${c}"`).join(", ");
-      if (tableName === TABLE_SPANS) {
-        const updateColumns = columns.filter((c) => c !== "traceId" && c !== "spanId");
-        if (updateColumns.length > 0) {
-          const updateClause = updateColumns.map((c) => `"${c}" = EXCLUDED."${c}"`).join(", ");
-          await this.client.none(
-            `INSERT INTO ${fullTableName} (${columnList}) VALUES (${placeholders})
+  async executeInsert(client, { tableName, record }) {
+    this.addTimestampZColumns(record);
+    const filteredRecord = await this.filterRecordToKnownColumns(tableName, record);
+    const schemaName = getSchemaName(this.schemaName);
+    const columns = Object.keys(filteredRecord).map((col) => parseSqlIdentifier(col, "column name"));
+    if (columns.length === 0) return;
+    const values = this.prepareValuesForInsert(filteredRecord, tableName);
+    const placeholders = values.map((_, i) => `$${i + 1}`).join(", ");
+    const fullTableName = getTableName({ indexName: tableName, schemaName });
+    const columnList = columns.map((c) => `"${c}"`).join(", ");
+    if (tableName === TABLE_SPANS) {
+      const updateColumns = columns.filter((c) => c !== "traceId" && c !== "spanId");
+      if (updateColumns.length > 0) {
+        const updateClause = updateColumns.map((c) => `"${c}" = EXCLUDED."${c}"`).join(", ");
+        await client.none(
+          `INSERT INTO ${fullTableName} (${columnList}) VALUES (${placeholders})
              ON CONFLICT ("traceId", "spanId") DO UPDATE SET ${updateClause}`,
-            values
-          );
-        } else {
-          await this.client.none(
-            `INSERT INTO ${fullTableName} (${columnList}) VALUES (${placeholders})
-             ON CONFLICT ("traceId", "spanId") DO NOTHING`,
-            values
-          );
-        }
+          values
+        );
       } else {
-        await this.client.none(`INSERT INTO ${fullTableName} (${columnList}) VALUES (${placeholders})`, values);
+        await client.none(
+          `INSERT INTO ${fullTableName} (${columnList}) VALUES (${placeholders})
+             ON CONFLICT ("traceId", "spanId") DO NOTHING`,
+          values
+        );
       }
+    } else {
+      await client.none(`INSERT INTO ${fullTableName} (${columnList}) VALUES (${placeholders})`, values);
+    }
+  }
+  async insert({ tableName, record }) {
+    try {
+      await this.executeInsert(this.client, { tableName, record });
     } catch (error) {
       throw new MastraError(
         {
@@ -2978,13 +2981,12 @@ Note: This migration may take some time for large tables.
   }
   async batchInsert({ tableName, records }) {
     try {
-      await this.client.query("BEGIN");
-      for (const record of records) {
-        await this.insert({ tableName, record });
-      }
-      await this.client.query("COMMIT");
+      await this.client.tx(async (tx) => {
+        for (const record of records) {
+          await this.executeInsert(tx, { tableName, record });
+        }
+      });
     } catch (error) {
-      await this.client.query("ROLLBACK");
       throw new MastraError(
         {
           id: createStorageErrorId("PG", "BATCH_INSERT", "FAILED"),
@@ -3262,30 +3264,7 @@ Note: This migration may take some time for large tables.
     data
   }) {
     try {
-      const filteredData = await this.filterRecordToKnownColumns(tableName, data);
-      if (Object.keys(filteredData).length === 0) return;
-      const setColumns = [];
-      const setValues = [];
-      let paramIndex = 1;
-      Object.entries(filteredData).forEach(([key, value]) => {
-        const parsedKey = parseSqlIdentifier(key, "column name");
-        setColumns.push(`"${parsedKey}" = $${paramIndex++}`);
-        setValues.push(this.prepareValue(value, key, tableName));
-      });
-      const whereConditions = [];
-      const whereValues = [];
-      Object.entries(keys).forEach(([key, value]) => {
-        const parsedKey = parseSqlIdentifier(key, "column name");
-        whereConditions.push(`"${parsedKey}" = $${paramIndex++}`);
-        whereValues.push(this.prepareValue(value, key, tableName));
-      });
-      const tableName_ = getTableName({
-        indexName: tableName,
-        schemaName: getSchemaName(this.schemaName)
-      });
-      const sql = `UPDATE ${tableName_} SET ${setColumns.join(", ")} WHERE ${whereConditions.join(" AND ")}`;
-      const values = [...setValues, ...whereValues];
-      await this.client.none(sql, values);
+      await this.executeUpdate(this.client, { tableName, keys, data });
     } catch (error) {
       throw new MastraError(
         {
@@ -3305,13 +3284,12 @@ Note: This migration may take some time for large tables.
     updates
   }) {
     try {
-      await this.client.query("BEGIN");
-      for (const { keys, data } of updates) {
-        await this.update({ tableName, keys, data });
-      }
-      await this.client.query("COMMIT");
+      await this.client.tx(async (tx) => {
+        for (const { keys, data } of updates) {
+          await this.executeUpdate(tx, { tableName, keys, data });
+        }
+      });
     } catch (error) {
-      await this.client.query("ROLLBACK");
       throw new MastraError(
         {
           id: createStorageErrorId("PG", "BATCH_UPDATE", "FAILED"),
@@ -3370,6 +3348,36 @@ Note: This migration may take some time for large tables.
   async deleteData({ tableName }) {
     return this.clearTable({ tableName });
   }
+  async executeUpdate(client, {
+    tableName,
+    keys,
+    data
+  }) {
+    const filteredData = await this.filterRecordToKnownColumns(tableName, data);
+    if (Object.keys(filteredData).length === 0) return;
+    const setColumns = [];
+    const setValues = [];
+    let paramIndex = 1;
+    Object.entries(filteredData).forEach(([key, value]) => {
+      const parsedKey = parseSqlIdentifier(key, "column name");
+      setColumns.push(`"${parsedKey}" = $${paramIndex++}`);
+      setValues.push(this.prepareValue(value, key, tableName));
+    });
+    const whereConditions = [];
+    const whereValues = [];
+    Object.entries(keys).forEach(([key, value]) => {
+      const parsedKey = parseSqlIdentifier(key, "column name");
+      whereConditions.push(`"${parsedKey}" = $${paramIndex++}`);
+      whereValues.push(this.prepareValue(value, key, tableName));
+    });
+    const tableName_ = getTableName({
+      indexName: tableName,
+      schemaName: getSchemaName(this.schemaName)
+    });
+    const sql = `UPDATE ${tableName_} SET ${setColumns.join(", ")} WHERE ${whereConditions.join(" AND ")}`;
+    const values = [...setValues, ...whereValues];
+    await client.none(sql, values);
+  }
 };
 function getSchemaName2(schema) {
   return schema ? `"${parseSqlIdentifier(schema, "schema name")}"` : void 0;
@@ -5283,6 +5291,16 @@ var ExperimentsPG = class _ExperimentsPG extends ExperimentsStorage {
   async init() {
     await this.#db.createTable({ tableName: TABLE_EXPERIMENTS, schema: EXPERIMENTS_SCHEMA });
     await this.#db.createTable({ tableName: TABLE_EXPERIMENT_RESULTS, schema: EXPERIMENT_RESULTS_SCHEMA });
+    await this.#db.alterTable({
+      tableName: TABLE_EXPERIMENTS,
+      schema: EXPERIMENTS_SCHEMA,
+      ifNotExists: ["agentVersion"]
+    });
+    await this.#db.alterTable({
+      tableName: TABLE_EXPERIMENT_RESULTS,
+      schema: EXPERIMENT_RESULTS_SCHEMA,
+      ifNotExists: ["status", "tags"]
+    });
     await this.createDefaultIndexes();
     await this.createCustomIndexes();
   }