fix: report a close or cancel during a wait as OperationStateError

CHANGELOG.md

@@ -23,6 +23,7 @@
 - Fix `DECIMAL` parameters losing their fractional digits on the Thrift backend: `"123.45"` was bound as `decimal(3,0)` and came back as `123`. Both backends now send `DECIMAL(p,s)` derived from the value (also for the `DEC` / `NUMERIC` synonyms), stringify `Int64` values beyond 2^53 exactly instead of as `Infinity`, accept exponent forms such as `1e-7` (rewritten as a plain numeral), and throw a `ParameterError` for non-numeric values or values that need more than 38 digits. (#474)
 - `customHeaders` are now sent on every Thrift request (session, statement and metadata calls), as the kernel backend already did, instead of only on telemetry and feature-flag requests. On these requests, caller values for driver-managed headers (`Authorization`, `User-Agent`, `Content-Type`, `Content-Length`, `Transfer-Encoding`, `Connection`) are ignored on both backends, and the Thrift backend rejects header names or values that cannot be sent at connect time. (#491)
 - Error telemetry now reports only the error class name and the server SQLSTATE. It no longer includes the error message or stack trace, which could contain table and column names from the query. On the Thrift backend, `statement.complete` telemetry now also reports the result delivery format (inline Arrow, inline columnar, or CloudFetch) and whether the server LZ4-compressed the result. (#508)
+- Closing or cancelling an operation (directly, or through `session.close()` / `client.close()`) while `fetchAll()`, `fetchChunk()` or `finished()` is waiting on it now rejects that wait with `OperationStateError` (`CLOSED` or `CANCELED`). Previously the request in flight could fail with the server's answer for the released statement, such as a Thrift `THTTPException` for HTTP 409 (`A concurrent GetOperationStatus request updated the command. Please retry.`).
 
 ## 3.0.0
 

lib/DBSQLOperation.ts

@@ -57,6 +57,9 @@ export default class DBSQLOperation implements IOperation {
 
   private cancelled: boolean = false;
 
+  // The latest close or cancel request (see `clientStopError`).
+  private stopRequest?: Promise<Status>;
+
   public readonly _data?: unknown;
 
   private metadata?: ResultMetadata;
@@ -159,11 +162,16 @@ export default class DBSQLOperation implements IOperation {
 
     const defaultMaxRows = this.context.getConfig().fetchChunkDefaultMaxRows;
     const limit = options?.maxRows ?? defaultMaxRows;
-    const result = await this.backend.fetchChunk({
-      limit,
-      disableBuffering: options?.disableBuffering,
-      isClosed: () => this.closed || this.cancelled,
-    });
+    let result: Array<object>;
+    try {
+      result = await this.backend.fetchChunk({
+        limit,
+        disableBuffering: options?.disableBuffering,
+        isClosed: () => this.closed || this.cancelled,
+      });
+    } catch (err) {
+      throw (err instanceof OperationStateError ? undefined : await this.clientStopError()) ?? err;
+    }
     await this.failIfClosed();
 
     this.context.getLogger().log(LogLevel.debug, `Fetched chunk of size: ${limit} from operation with id: ${this.id}`);
@@ -207,7 +215,8 @@ export default class DBSQLOperation implements IOperation {
     if (this.closed || this.cancelled) {
       return Status.success();
     }
-    const result = await this.backend.cancel();
+    this.stopRequest = this.backend.cancel();
+    const result = await this.stopRequest;
     this.cancelled = true;
     this.onClose?.();
     return result;
@@ -226,9 +235,10 @@ export default class DBSQLOperation implements IOperation {
     if (this.closed || this.cancelled) {
       return Status.success();
     }
+    this.stopRequest = this.backend.close();
     let result: Status;
     try {
-      result = await this.backend.close();
+      result = await this.stopRequest;
     } catch (error: any) {
       // Until the statement finished, closing it is what stops it server-side,
       // so the failure is the caller's to handle (and the close can be retried).
@@ -383,6 +393,26 @@ export default class DBSQLOperation implements IOperation {
     }
   }
 
+  /**
+   * The error for a backend call that failed because this operation was closed
+   * or cancelled while the call was in flight. Such a call fails with whatever
+   * the server answers for the released statement (Thrift answers an in-flight
+   * GetOperationStatus with HTTP 409 ABORTED), so report the client's action,
+   * but only once that close or cancel has succeeded.
+   */
+  private async clientStopError(): Promise<OperationStateError | undefined> {
+    // `closeInternal` / `cancelInternal` awaited the same request first, so
+    // their flags are already updated when this resumes.
+    await this.stopRequest?.catch(() => undefined);
+    if (this.closed) {
+      return new OperationStateError(OperationStateErrorCode.Closed);
+    }
+    if (this.cancelled) {
+      return new OperationStateError(OperationStateErrorCode.Canceled);
+    }
+    return undefined;
+  }
+
   private async waitUntilReadyThroughBackend(options?: WaitUntilReadyOptions) {
     // The backend-facing `waitUntilReady` takes a neutral
     // `IOperationBackendWaitOptions` whose `callback` receives an
@@ -412,8 +442,9 @@ export default class DBSQLOperation implements IOperation {
         } else if (err.errorCode === OperationStateErrorCode.Closed) {
           this.closed = true;
         }
+        throw err;
       }
-      throw err;
+      throw (await this.clientStopError()) ?? err;
     }
   }
 

tests/unit/DBSQLOperation.test.ts

@@ -608,6 +608,164 @@ describe('DBSQLOperation', () => {
     });
   });
 
+  describe('close or cancel while a request is in flight', () => {
+    // How the HTTP layer reports the 409 the server answers an in-flight
+    // GetOperationStatus with once the operation is closed.
+    const abortedRequest = () =>
+      new Error(
+        'Received a response with a bad HTTP status code: 409: ABORTED: A concurrent GetOperationStatus request updated the command. Please retry.',
+      );
+
+    const flush = () => new Promise((resolve) => setTimeout(resolve, 0));
+
+    function pending<T>(stub: sinon.SinonStub) {
+      let failRequest: (error: Error) => void = () => {};
+      stub.callsFake(
+        () =>
+          new Promise<T>((_resolve, reject) => {
+            failRequest = reject;
+          }),
+      );
+      return (error: Error) => failRequest(error);
+    }
+
+    function runningOperation() {
+      const context = new ClientContextStub();
+      context.driver.getOperationStatusResp.operationState = TOperationState.RUNNING_STATE;
+      const getOperationStatus = sinon.stub(context.driver, 'getOperationStatus');
+      const failPoll = pending(getOperationStatus);
+      const operation = createOperationForTest({ handle: operationHandleStub({ hasResultSet: true }), context });
+      return { context, operation, getOperationStatus, failPoll };
+    }
+
+    async function expectStateError(promise: Promise<unknown>, errorCode: OperationStateErrorCode) {
+      try {
+        await promise;
+        expect.fail('It should throw an OperationStateError');
+      } catch (e) {
+        if (e instanceof AssertionError) {
+          throw e;
+        }
+        expect(e).to.be.instanceOf(OperationStateError);
+        expect((e as OperationStateError).errorCode).to.equal(errorCode);
+      }
+    }
+
+    it('should report a status poll failed by close() as closed', async () => {
+      const { operation, getOperationStatus, failPoll } = runningOperation();
+
+      const fetching = operation.fetchAll();
+      await flush();
+      expect(getOperationStatus.calledOnce).to.be.true;
+
+      const closing = operation.close();
+      failPoll(abortedRequest());
+
+      await closing;
+      await expectStateError(fetching, OperationStateErrorCode.Closed);
+    });
+
+    it('should report a status poll failed by cancel() as canceled', async () => {
+      const { operation, getOperationStatus, failPoll } = runningOperation();
+
+      const waiting = operation.finished();
+      await flush();
+      expect(getOperationStatus.calledOnce).to.be.true;
+
+      const cancelling = operation.cancel();
+      failPoll(abortedRequest());
+
+      await cancelling;
+      await expectStateError(waiting, OperationStateErrorCode.Canceled);
+    });
+
+    it('should keep the error of a status poll that failed on its own', async () => {
+      const { operation, failPoll } = runningOperation();
+
+      const fetching = operation.fetchAll();
+      await flush();
+      const error = abortedRequest();
+      failPoll(error);
+
+      try {
+        await fetching;
+        expect.fail('It should throw the poll error');
+      } catch (e) {
+        if (e instanceof AssertionError) {
+          throw e;
+        }
+        expect(e).to.equal(error);
+      }
+    });
+
+    it('should keep the error of a status poll when the close request failed', async () => {
+      const { context, operation, failPoll } = runningOperation();
+      context.driver.closeOperationResp.status.statusCode = TStatusCode.ERROR_STATUS;
+
+      const fetching = operation.fetchAll();
+      await flush();
+      await expectFailure(() => operation.close());
+      const error = abortedRequest();
+      failPoll(error);
+
+      try {
+        await fetching;
+        expect.fail('It should throw the poll error');
+      } catch (e) {
+        if (e instanceof AssertionError) {
+          throw e;
+        }
+        expect(e).to.equal(error);
+      }
+    });
+
+    it('should keep the error of a status poll that failed while a close request that then failed was pending', async () => {
+      const { context, operation, failPoll } = runningOperation();
+      const failClose = pending(sinon.stub(context.driver, 'closeOperation'));
+
+      const fetching = operation.fetchAll();
+      await flush();
+      const closing = operation.close();
+      const error = abortedRequest();
+      failPoll(error);
+      await flush();
+      failClose(new Error('close failed'));
+
+      await expectFailure(() => closing);
+      expect(operation['closed']).to.be.false;
+      try {
+        await fetching;
+        expect.fail('It should throw the poll error');
+      } catch (e) {
+        if (e instanceof AssertionError) {
+          throw e;
+        }
+        expect(e).to.equal(error);
+      }
+    });
+
+    it('should report a result fetch failed by close() as closed', async () => {
+      const context = new ClientContextStub();
+      context.driver.getOperationStatusResp.operationState = TOperationState.FINISHED_STATE;
+      context.driver.getOperationStatusResp.hasResultSet = true;
+      const fetchResults = sinon.stub(context.driver, 'fetchResults');
+      const failFetch = pending(fetchResults);
+      const operation = createOperationForTest({ handle: operationHandleStub({ hasResultSet: true }), context });
+
+      const fetching = operation.fetchChunk({ disableBuffering: true });
+      for (let i = 0; i < 20 && !fetchResults.called; i += 1) {
+        await flush(); // eslint-disable-line no-await-in-loop
+      }
+      expect(fetchResults.calledOnce).to.be.true;
+
+      const closing = operation.close();
+      failFetch(abortedRequest());
+
+      await closing;
+      await expectStateError(fetching, OperationStateErrorCode.Closed);
+    });
+  });
+
   describe('Symbol.asyncDispose', () => {
     it('should close the operation', async () => {
       const context = new ClientContextStub();