From 9a7353adb969a7d69ae0453521c423967b7cbfe5 Mon Sep 17 00:00:00 2001
From: Shivam Raj <shivam.raj@databricks.com>
Date: Thu, 16 Oct 2025 01:23:49 +0530
Subject: [PATCH] Fix flaky Iterator tests by making metadata fetching
 async-safe

Addresses race condition where concurrent calls to fetchMetadata() would
attempt multiple server requests before the first completed. This caused
failures when GetResultSetMetadata was called after result handler was
consumed.

The fix uses a promise-based locking pattern to ensure only one metadata
fetch occurs, with subsequent concurrent calls waiting for the same promise.

Also adds unit tests to verify async-safety of metadata fetching.

Signed-off-by: Shivam Raj <shivam.raj@databricks.com>
---
 lib/DBSQLOperation.ts             | 31 +++++++++++++++++++++---
 tests/unit/DBSQLOperation.test.ts | 39 +++++++++++++++++++++++++++++++
 2 files changed, 67 insertions(+), 3 deletions(-)

diff --git a/lib/DBSQLOperation.ts b/lib/DBSQLOperation.ts
index 634749bf..fe22995d 100644
--- a/lib/DBSQLOperation.ts
+++ b/lib/DBSQLOperation.ts
@@ -66,6 +66,8 @@ export default class DBSQLOperation implements IOperation {
 
   private metadata?: TGetResultSetMetadataResp;
 
+  private metadataPromise?: Promise<TGetResultSetMetadataResp>;
+
   private state: TOperationState = TOperationState.INITIALIZED_STATE;
 
   // Once operation is finished or fails - cache status response, because subsequent calls
@@ -292,6 +294,12 @@ export default class DBSQLOperation implements IOperation {
       return false;
     }
 
+    // Wait for operation to finish before checking for more rows
+    // This ensures metadata can be fetched successfully
+    if (this.operationHandle.hasResultSet) {
+      await this.waitUntilReady();
+    }
+
     // If we fetched all the data from server - check if there's anything buffered in result handler
     const resultHandler = await this.getResultHandler();
     return resultHandler.hasMore();
@@ -383,16 +391,33 @@ export default class DBSQLOperation implements IOperation {
   }
 
   private async fetchMetadata() {
-    if (!this.metadata) {
+    // If metadata is already cached, return it immediately
+    if (this.metadata) {
+      return this.metadata;
+    }
+
+    // If a fetch is already in progress, wait for it to complete
+    if (this.metadataPromise) {
+      return this.metadataPromise;
+    }
+
+    // Start a new fetch and cache the promise to prevent concurrent fetches
+    this.metadataPromise = (async () => {
       const driver = await this.context.getDriver();
       const metadata = await driver.getResultSetMetadata({
         operationHandle: this.operationHandle,
       });
       Status.assert(metadata.status);
       this.metadata = metadata;
+      return metadata;
+    })();
+
+    try {
+      return await this.metadataPromise;
+    } finally {
+      // Clear the promise once completed (success or failure)
+      this.metadataPromise = undefined;
     }
-
-    return this.metadata;
   }
 
   private async getResultHandler(): Promise<ResultSlicer<any>> {
diff --git a/tests/unit/DBSQLOperation.test.ts b/tests/unit/DBSQLOperation.test.ts
index 94224455..b5f142ba 100644
--- a/tests/unit/DBSQLOperation.test.ts
+++ b/tests/unit/DBSQLOperation.test.ts
@@ -1138,4 +1138,43 @@ describe('DBSQLOperation', () => {
       expect(operation['_data']['hasMoreRowsFlag']).to.be.false;
     });
   });
+
+  describe('metadata fetching (async-safety)', () => {
+    it('should handle concurrent metadata fetch requests without duplicate server calls', async () => {
+      const context = new ClientContextStub();
+      const driver = sinon.spy(context.driver);
+      driver.getOperationStatusResp.operationState = TOperationState.FINISHED_STATE;
+      driver.getOperationStatusResp.hasResultSet = true;
+
+      // Create operation without direct results to force metadata fetching
+      const operation = new DBSQLOperation({ handle: operationHandleStub({ hasResultSet: true }), context });
+
+      // Trigger multiple concurrent metadata fetches
+      const results = await Promise.all([operation.hasMoreRows(), operation.hasMoreRows(), operation.hasMoreRows()]);
+
+      // All should succeed
+      expect(results).to.deep.equal([true, true, true]);
+
+      // But metadata should only be fetched once from server
+      expect(driver.getResultSetMetadata.callCount).to.equal(1);
+    });
+
+    it('should cache metadata after first fetch', async () => {
+      const context = new ClientContextStub();
+      const driver = sinon.spy(context.driver);
+      driver.getOperationStatusResp.operationState = TOperationState.FINISHED_STATE;
+      driver.getOperationStatusResp.hasResultSet = true;
+
+      const operation = new DBSQLOperation({ handle: operationHandleStub({ hasResultSet: true }), context });
+
+      // First call should fetch metadata
+      await operation.hasMoreRows();
+      expect(driver.getResultSetMetadata.callCount).to.equal(1);
+
+      // Subsequent calls should use cached metadata
+      await operation.hasMoreRows();
+      await operation.hasMoreRows();
+      expect(driver.getResultSetMetadata.callCount).to.equal(1);
+    });
+  });
 });