Skip to content

Commit 9a7353a

Browse files
committed
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 <[email protected]>
1 parent 3f2eec1 commit 9a7353a

File tree

2 files changed

+67
-3
lines changed

2 files changed

+67
-3
lines changed

lib/DBSQLOperation.ts

Lines changed: 28 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -66,6 +66,8 @@ export default class DBSQLOperation implements IOperation {
6666

6767
private metadata?: TGetResultSetMetadataResp;
6868

69+
private metadataPromise?: Promise<TGetResultSetMetadataResp>;
70+
6971
private state: TOperationState = TOperationState.INITIALIZED_STATE;
7072

7173
// Once operation is finished or fails - cache status response, because subsequent calls
@@ -292,6 +294,12 @@ export default class DBSQLOperation implements IOperation {
292294
return false;
293295
}
294296

297+
// Wait for operation to finish before checking for more rows
298+
// This ensures metadata can be fetched successfully
299+
if (this.operationHandle.hasResultSet) {
300+
await this.waitUntilReady();
301+
}
302+
295303
// If we fetched all the data from server - check if there's anything buffered in result handler
296304
const resultHandler = await this.getResultHandler();
297305
return resultHandler.hasMore();
@@ -383,16 +391,33 @@ export default class DBSQLOperation implements IOperation {
383391
}
384392

385393
private async fetchMetadata() {
386-
if (!this.metadata) {
394+
// If metadata is already cached, return it immediately
395+
if (this.metadata) {
396+
return this.metadata;
397+
}
398+
399+
// If a fetch is already in progress, wait for it to complete
400+
if (this.metadataPromise) {
401+
return this.metadataPromise;
402+
}
403+
404+
// Start a new fetch and cache the promise to prevent concurrent fetches
405+
this.metadataPromise = (async () => {
387406
const driver = await this.context.getDriver();
388407
const metadata = await driver.getResultSetMetadata({
389408
operationHandle: this.operationHandle,
390409
});
391410
Status.assert(metadata.status);
392411
this.metadata = metadata;
412+
return metadata;
413+
})();
414+
415+
try {
416+
return await this.metadataPromise;
417+
} finally {
418+
// Clear the promise once completed (success or failure)
419+
this.metadataPromise = undefined;
393420
}
394-
395-
return this.metadata;
396421
}
397422

398423
private async getResultHandler(): Promise<ResultSlicer<any>> {

tests/unit/DBSQLOperation.test.ts

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1138,4 +1138,43 @@ describe('DBSQLOperation', () => {
11381138
expect(operation['_data']['hasMoreRowsFlag']).to.be.false;
11391139
});
11401140
});
1141+
1142+
describe('metadata fetching (async-safety)', () => {
1143+
it('should handle concurrent metadata fetch requests without duplicate server calls', async () => {
1144+
const context = new ClientContextStub();
1145+
const driver = sinon.spy(context.driver);
1146+
driver.getOperationStatusResp.operationState = TOperationState.FINISHED_STATE;
1147+
driver.getOperationStatusResp.hasResultSet = true;
1148+
1149+
// Create operation without direct results to force metadata fetching
1150+
const operation = new DBSQLOperation({ handle: operationHandleStub({ hasResultSet: true }), context });
1151+
1152+
// Trigger multiple concurrent metadata fetches
1153+
const results = await Promise.all([operation.hasMoreRows(), operation.hasMoreRows(), operation.hasMoreRows()]);
1154+
1155+
// All should succeed
1156+
expect(results).to.deep.equal([true, true, true]);
1157+
1158+
// But metadata should only be fetched once from server
1159+
expect(driver.getResultSetMetadata.callCount).to.equal(1);
1160+
});
1161+
1162+
it('should cache metadata after first fetch', async () => {
1163+
const context = new ClientContextStub();
1164+
const driver = sinon.spy(context.driver);
1165+
driver.getOperationStatusResp.operationState = TOperationState.FINISHED_STATE;
1166+
driver.getOperationStatusResp.hasResultSet = true;
1167+
1168+
const operation = new DBSQLOperation({ handle: operationHandleStub({ hasResultSet: true }), context });
1169+
1170+
// First call should fetch metadata
1171+
await operation.hasMoreRows();
1172+
expect(driver.getResultSetMetadata.callCount).to.equal(1);
1173+
1174+
// Subsequent calls should use cached metadata
1175+
await operation.hasMoreRows();
1176+
await operation.hasMoreRows();
1177+
expect(driver.getResultSetMetadata.callCount).to.equal(1);
1178+
});
1179+
});
11411180
});

0 commit comments

Comments
 (0)