Prevent executing mutations twice

Author: oguimbal

src/interfaces-private.ts

@@ -131,11 +131,17 @@ export interface _Transaction {
     fullCommit(): _Transaction;
     rollback(): _Transaction;
     delete(identity: symbol): void;
+    /** Set data persisted in this transaction */
     set<T>(identity: symbol, data: T): T;
+    /** Get data persisted in this transaction */
     get<T>(identity: symbol): T;
     getMap<T extends ImMap<any, any>>(identity: symbol): T;
     getSet<T>(identity: symbol): ImSet<T>;
-    affectedRows: number;
+    /** Get transient data, which will only exist within the scope of the current statement */
+    setTransient<T>(identity: symbol, data: T): T;
+    /** Set transient data, which will only exist within the scope of the current statement */
+    getTransient<T>(identity: symbol): T;
+    clearTransientData(): void;
 }
 
 export interface Stats {

src/mutations/mutation-base.ts

@@ -9,13 +9,16 @@ import { buildValue } from '../expression-builder';
 
 type MutationStatement = InsertStatement | UpdateStatement | DeleteStatement;
 
+
 export abstract class MutationDataSourceBase<T> extends DataSourceBase<T> {
+    public static readonly affectedRows = Symbol('affectedRows');
 
     /** Perform the mutation, and returns the affected elements */
     protected abstract performMutation(t: _Transaction): T[];
 
     private returningRows?: ArrayFilter;
     private returning?: _ISelection;
+    private mutationResult = Symbol('mutationResult');
 
     get columns() {
         return this.returning?.columns ?? [];
@@ -32,9 +35,19 @@ export abstract class MutationDataSourceBase<T> extends DataSourceBase<T> {
     }
 
     *enumerate(t: _Transaction): Iterable<any> {
-        const affected = this.performMutation(t);
+        // check if this mutation has already been executed in the statement being executed
+        // and get the result from cache to avoid re-excuting it
+        // see unit test "can use delete result multiple times in select"
+        let affected = t.getTransient<T[]>(this.mutationResult);
+        if (!affected) {
+            // execute mutation if nescessary
+            affected = this.performMutation(t);
+            t.setTransient(this.mutationResult, affected);
+        }
+
 
-        t.affectedRows = affected.length;
+        // set the result count
+        t.setTransient(MutationDataSourceBase.affectedRows, affected.length);
         if (!this.returning) {
             return;
         }

src/schema.ts

@@ -20,6 +20,7 @@ import { OverloadResolver } from './overload-resolver';
 import { Deletion } from './mutations/deletion';
 import { Update } from './mutations/update';
 import { Insert } from './mutations/insert';
+import { MutationDataSourceBase } from './mutations/mutation-base';
 
 export class DbSchema implements _ISchema, ISchema {
 
@@ -137,6 +138,7 @@ export class DbSchema implements _ISchema, ISchema {
             const { checked: p, check } = this.db.options.noAstCoverageCheck
                 ? { checked: _p, check: null }
                 : watchUse(_p);
+            t.clearTransientData();
 
             switch (p.type) {
                 case 'start transaction':
@@ -536,7 +538,7 @@ but the resulting statement cannot be executed → Probably not a pg-mem error.`
             : cleanResults([...last.enumerate(t)]);
         return {
             rows,
-            rowCount: typeof last === 'number' ? last : rows.length,
+            rowCount: t.getTransient(MutationDataSourceBase.affectedRows) ?? rows.length,
             command: p.type.toUpperCase(),
             fields: [],
             location: this.locOf(p),

src/transaction.ts

@@ -4,7 +4,7 @@ import { NotSupported, QueryError } from './interfaces';
 
 export class Transaction implements _Transaction {
     private origData: ImMap<symbol, any>;
-    affectedRows: number = 0;
+    private transientData: any = {};
 
     static root() {
         return new Transaction(null, ImMap());
@@ -78,4 +78,17 @@ export class Transaction implements _Transaction {
         return this.parent ?? this;
     }
 
+    setTransient<T>(identity: symbol, data: T): T {
+        this.transientData[identity] = data as any;
+        return data;
+    }
+
+    /** Set transient data, which will only exist within the scope of the current statement */
+    getTransient<T>(identity: symbol): T {
+        return this.transientData[identity] as T;
+    }
+
+    clearTransientData(): void {
+        this.transientData = {};
+    }
 }