Wait for BEGIN be send before send RUN, COMMIT or ROLLBACK

bigmontz · bigmontz · commit 0726ae4e82c5 · 2022-09-01T16:54:49.000+02:00
diff --git a/packages/core/src/internal/connection-holder.ts b/packages/core/src/internal/connection-holder.ts
@@ -85,6 +85,7 @@ class ConnectionHolder implements ConnectionHolderInterface {
   private readonly _impersonatedUser?: string
   private readonly _getConnectionAcquistionBookmarks: () => Promise<Bookmarks>
   private readonly _onDatabaseNameResolved?: (databaseName?: string) => void
+  private _connectionError?: Error
 
   /**
    * @constructor
@@ -123,6 +124,7 @@ class ConnectionHolder implements ConnectionHolderInterface {
     this._connectionPromise = Promise.resolve(null)
     this._onDatabaseNameResolved = onDatabaseNameResolved
     this._getConnectionAcquistionBookmarks = getConnectionAcquistionBookmarks ?? (() => Promise.resolve(Bookmarks.empty()))
+    this._connectionError = undefined
   }
 
   mode (): string | undefined {
@@ -161,12 +163,16 @@ class ConnectionHolder implements ConnectionHolderInterface {
   }
 
   private async _createConnectionPromise (connectionProvider: ConnectionProvider): Promise<Connection | null> {
+    this._connectionError = undefined
     return await connectionProvider.acquireConnection({
       accessMode: this._mode,
       database: this._database,
       bookmarks: await this._getBookmarks(),
       impersonatedUser: this._impersonatedUser,
       onDatabaseNameResolved: this._onDatabaseNameResolved
+    }).catch(error => {
+      this._connectionError = error
+      throw error
     })
   }
 
@@ -175,6 +181,13 @@ class ConnectionHolder implements ConnectionHolderInterface {
   }
 
   getConnection (): Promise<Connection | null> {
+    // Failures in the connection promise should be
+    // propagated to the all the .getConnection()
+    // calls util a new connection promise get
+    // set up
+    if (this._connectionError !== undefined) {
+      return Promise.reject(this._connectionError)
+    }
     return this._connectionPromise
   }
 
diff --git a/packages/core/src/transaction.ts b/packages/core/src/transaction.ts
@@ -58,6 +58,8 @@ class Transaction {
   private readonly _lowRecordWatermak: number
   private readonly _highRecordWatermark: number
   private _bookmarks: Bookmarks
+  private readonly _activePromise: Promise<void>
+  private _acceptActive: () => void
 
   /**
    * @constructor
@@ -107,6 +109,9 @@ class Transaction {
     this._lowRecordWatermak = lowRecordWatermark
     this._highRecordWatermark = highRecordWatermark
     this._bookmarks = Bookmarks.empty()
+    this._activePromise = new Promise((resolve, reject) => {
+      this._acceptActive = resolve
+    })
   }
 
   /**
@@ -154,6 +159,10 @@ class Transaction {
         }
         this._onError(error).catch(() => {})
       })
+      // It should make the transaction active anyway
+      // futher errors will be treated by the exiting
+      // observers
+      .finally(() => this._acceptActive())
   }
 
   /**
@@ -178,7 +187,8 @@ class Transaction {
       reactive: this._reactive,
       fetchSize: this._fetchSize,
       highRecordWatermark: this._highRecordWatermark,
-      lowRecordWatermark: this._lowRecordWatermak
+      lowRecordWatermark: this._lowRecordWatermak,
+      preparationJob: this._activePromise
     })
     this._results.push(result)
     return result
@@ -197,7 +207,8 @@ class Transaction {
       onError: this._onError,
       onComplete: (meta: any) => this._onCompleteCallback(meta, this._bookmarks),
       onConnection: this._onConnection,
-      pendingResults: this._results
+      pendingResults: this._results,
+      preparationJob: this._activePromise
     })
     this._state = committed.state
     // clean up
@@ -224,7 +235,8 @@ class Transaction {
       onError: this._onError,
       onComplete: this._onComplete,
       onConnection: this._onConnection,
-      pendingResults: this._results
+      pendingResults: this._results,
+      preparationJob: this._activePromise
     })
     this._state = rolledback.state
     // clean up
@@ -293,6 +305,7 @@ interface StateTransitionParams {
   fetchSize: number
   highRecordWatermark: number
   lowRecordWatermark: number
+  preparationJob?: Promise<any>
 }
 
 const _states = {
@@ -303,7 +316,8 @@ const _states = {
       onError,
       onComplete,
       onConnection,
-      pendingResults
+      pendingResults,
+      preparationJob
     }: StateTransitionParams): any => {
       return {
         result: finishTransaction(
@@ -312,7 +326,8 @@ const _states = {
           onError,
           onComplete,
           onConnection,
-          pendingResults
+          pendingResults,
+          preparationJob
         ),
         state: _states.SUCCEEDED
       }
@@ -322,7 +337,8 @@ const _states = {
       onError,
       onComplete,
       onConnection,
-      pendingResults
+      pendingResults,
+      preparationJob
     }: StateTransitionParams): any => {
       return {
         result: finishTransaction(
@@ -331,7 +347,8 @@ const _states = {
           onError,
           onComplete,
           onConnection,
-          pendingResults
+          pendingResults,
+          preparationJob
         ),
         state: _states.ROLLED_BACK
       }
@@ -347,31 +364,34 @@ const _states = {
         reactive,
         fetchSize,
         highRecordWatermark,
-        lowRecordWatermark
+        lowRecordWatermark,
+        preparationJob
       }: StateTransitionParams
     ): any => {
       // RUN in explicit transaction can't contain bookmarks and transaction configuration
       // No need to include mode and database name as it shall be inclued in begin
-      const observerPromise = connectionHolder
-        .getConnection()
-        .then(conn => {
-          onConnection()
-          if (conn != null) {
-            return conn.protocol().run(query, parameters, {
-              bookmarks: Bookmarks.empty(),
-              txConfig: TxConfig.empty(),
-              beforeError: onError,
-              afterComplete: onComplete,
-              reactive: reactive,
-              fetchSize: fetchSize,
-              highRecordWatermark: highRecordWatermark,
-              lowRecordWatermark: lowRecordWatermark
-            })
-          } else {
-            throw newError('No connection available')
-          }
-        })
-        .catch(error => new FailedObserver({ error, onError }))
+      const requirements = preparationJob ?? Promise.resolve()
+
+      const observerPromise =
+        requirements.then(() => connectionHolder.getConnection())
+          .then(conn => {
+            onConnection()
+            if (conn != null) {
+              return conn.protocol().run(query, parameters, {
+                bookmarks: Bookmarks.empty(),
+                txConfig: TxConfig.empty(),
+                beforeError: onError,
+                afterComplete: onComplete,
+                reactive: reactive,
+                fetchSize: fetchSize,
+                highRecordWatermark: highRecordWatermark,
+                lowRecordWatermark: lowRecordWatermark
+              })
+            } else {
+              throw newError('No connection available')
+            }
+          })
+          .catch(error => new FailedObserver({ error, onError }))
 
       return newCompletedResult(
         observerPromise,
@@ -598,32 +618,36 @@ function finishTransaction (
   onError: (err: Error) => any,
   onComplete: (metadata: any) => any,
   onConnection: () => any,
-  pendingResults: Result[]
+  pendingResults: Result[],
+  preparationJob?: Promise<void>
 ): Result {
-  const observerPromise = connectionHolder
-    .getConnection()
-    .then(connection => {
-      onConnection()
-      pendingResults.forEach(r => r._cancel())
-      return Promise.all(pendingResults.map(result => result.summary())).then(results => {
-        if (connection != null) {
-          if (commit) {
-            return connection.protocol().commitTransaction({
-              beforeError: onError,
-              afterComplete: onComplete
-            })
+  const requirements = preparationJob ?? Promise.resolve()
+
+  const observerPromise =
+    requirements
+      .then(() => connectionHolder.getConnection())
+      .then(connection => {
+        onConnection()
+        pendingResults.forEach(r => r._cancel())
+        return Promise.all(pendingResults.map(result => result.summary())).then(results => {
+          if (connection != null) {
+            if (commit) {
+              return connection.protocol().commitTransaction({
+                beforeError: onError,
+                afterComplete: onComplete
+              })
+            } else {
+              return connection.protocol().rollbackTransaction({
+                beforeError: onError,
+                afterComplete: onComplete
+              })
+            }
           } else {
-            return connection.protocol().rollbackTransaction({
-              beforeError: onError,
-              afterComplete: onComplete
-            })
+            throw newError('No connection available')
           }
-        } else {
-          throw newError('No connection available')
-        }
+        })
       })
-    })
-    .catch(error => new FailedObserver({ error, onError }))
+      .catch(error => new FailedObserver({ error, onError }))
 
   // for commit & rollback we need result that uses real connection holder and notifies it when
   // connection is not needed and can be safely released to the pool
diff --git a/packages/core/test/transaction.test.ts b/packages/core/test/transaction.test.ts
@@ -325,6 +325,8 @@ function testTx<T extends Transaction> (transactionName: string, newTransaction:
           lowRecordWatermark: 300
         })
 
+        tx._begin(async () => Bookmarks.empty(), TxConfig.empty())
+
         await tx.run('RETURN 1')
 
         expect(connection.seenProtocolOptions[0]).toMatchObject({
@@ -343,11 +345,36 @@ function testTx<T extends Transaction> (transactionName: string, newTransaction:
           lowRecordWatermark: 300
         })
 
+        tx._begin(async () => Bookmarks.empty(), TxConfig.empty())
+
         const result = tx.run('RETURN 1')
 
         // @ts-expect-error
         expect(result._watermarks).toEqual({ high: 700, low: 300 })
       })
+
+      it('should wait begin message be send', async () => {
+        const connection = newFakeConnection()
+        const tx = newTransaction({
+          connection
+        })
+
+        const bookmarksPromise: Promise<Bookmarks> = new Promise((resolve) => {
+          setTimeout(() => resolve(Bookmarks.empty()), 1000)
+        })
+
+        tx._begin(async () => await bookmarksPromise, TxConfig.empty())
+
+        const result = tx.run('RETURN 1')
+
+        expect(connection.seenBeginTransaction.length).toEqual(0)
+        expect(connection.seenQueries.length).toEqual(0)
+
+        await result
+
+        expect(connection.seenBeginTransaction.length).toEqual(1)
+        expect(connection.seenQueries.length).toEqual(1)
+      })
     })
 
     describe('.close()', () => {
@@ -356,6 +383,7 @@ function testTx<T extends Transaction> (transactionName: string, newTransaction:
           const connection = newFakeConnection()
           const tx = newTransaction({ connection })
 
+          tx._begin(async () => Bookmarks.empty(), TxConfig.empty())
           await tx.run('RETURN 1')
           await tx.close()
 
@@ -367,6 +395,7 @@ function testTx<T extends Transaction> (transactionName: string, newTransaction:
           const connection = newFakeConnection().withRollbackError(expectedError)
           const tx = newTransaction({ connection })
 
+          tx._begin(async () => Bookmarks.empty(), TxConfig.empty())
           await tx.run('RETURN 1')
 
           try {
@@ -376,6 +405,29 @@ function testTx<T extends Transaction> (transactionName: string, newTransaction:
             expect(error).toEqual(expectedError)
           }
         })
+
+        it('should wait begin message be send', async () => {
+          const connection = newFakeConnection()
+          const tx = newTransaction({
+            connection
+          })
+
+          const bookmarksPromise: Promise<Bookmarks> = new Promise((resolve) => {
+            setTimeout(() => resolve(Bookmarks.empty()), 1000)
+          })
+
+          tx._begin(async () => await bookmarksPromise, TxConfig.empty())
+
+          const result = tx.close()
+
+          expect(connection.seenBeginTransaction.length).toEqual(0)
+          expect(connection.rollbackInvoked).toEqual(0)
+
+          await result
+
+          expect(connection.seenBeginTransaction.length).toEqual(1)
+          expect(connection.rollbackInvoked).toEqual(1)
+        })
       })
 
       describe('when transaction is closed', () => {
@@ -394,6 +446,8 @@ function testTx<T extends Transaction> (transactionName: string, newTransaction:
           const connection = newFakeConnection()
           const tx = newTransaction({ connection })
 
+          tx._begin(async () => Bookmarks.empty(), TxConfig.empty())
+
           await operation(tx, connection)
           const rollbackInvokedAfterOperation = connection.rollbackInvoked
 
