Fix Notebook results being out of order (#24102)

* Fix Notebook results being out of order

* deprecated

* Fix tests

src/sql/workbench/contrib/charts/browser/chartView.ts

@@ -234,7 +234,7 @@ export class ChartView extends Disposable implements IPanelView {
 			if (batch) {
 				let summary = batch.resultSetSummaries[this._currentData.resultId];
 				if (summary) {
-					this._queryRunner.getQueryRows(0, Math.min(getChartMaxRowCount(this._configurationService), summary.rowCount), this._currentData.batchId, this._currentData.resultId).then(d => {
+					this._queryRunner.getQueryRowsPaged(0, Math.min(getChartMaxRowCount(this._configurationService), summary.rowCount), this._currentData.batchId, this._currentData.resultId).then(d => {
 						let rows = d.rows;
 						let columns = summary.columnInfo.map(c => c.columnName);
 						this.setData(rows, columns);

src/sql/workbench/contrib/query/browser/keyboardQueryActions.ts

@@ -169,7 +169,7 @@ export class CopyQueryWithResultsKeyboardAction extends Action2 {
 		if (queryRunner && queryRunner.batchSets.length > 0) {
 			for (let i = 0; i < queryRunner.batchSets[0].resultSetSummaries.length; i++) {
 				let resultSummary = queryRunner.batchSets[0].resultSetSummaries[i];
-				let result = await queryRunner.getQueryRows(0, resultSummary.rowCount, resultSummary.batchId, resultSummary.id);
+				let result = await queryRunner.getQueryRowsPaged(0, resultSummary.rowCount, resultSummary.batchId, resultSummary.id);
 				let tableHeaders = resultSummary.columnInfo.map((col, i) => (col.columnName));
 				let htmlTableHeaders = `<thead><tr style="background-color:DarkGray">${resultSummary.columnInfo.map((col, i) => (`<th style="border:1.0pt solid black;padding:3pt;font-size:9pt;font-weight: bold;">${escape(col.columnName)}</th>`)).join('')}</tr></thead>`;
 				let copyString = '\n';

src/sql/workbench/services/insights/browser/insightsDialogController.ts

@@ -161,7 +161,7 @@ export class InsightsDialogController {
 				this._columns = resultset.columnInfo;
 				let rows: ResultSetSubset;
 				try {
-					rows = await this._queryRunner.getQueryRows(0, resultset.rowCount, batch.id, resultset.id);
+					rows = await this._queryRunner.getQueryRowsPaged(0, resultset.rowCount, batch.id, resultset.id);
 				} catch (e) {
 					return Promise.reject(e);
 				}

src/sql/workbench/services/insights/test/browser/insightsDialogController.test.ts

@@ -130,7 +130,7 @@ function getPrimedQueryRunner(data: string[][], columns: string[]): IPrimedQuery
 		];
 	});
 
-	querymock.setup(x => x.getQueryRows(It.isAnyNumber(), It.isAnyNumber(), It.isAnyNumber(), It.isAnyNumber()))
+	querymock.setup(x => x.getQueryRowsPaged(It.isAnyNumber(), It.isAnyNumber(), It.isAnyNumber(), It.isAnyNumber()))
 		.returns(x => Promise.resolve(<ResultSetSubset>{
 			rowCount: data.length,
 			rows: data.map(r => r.map(c => { return { displayValue: c }; }))

src/sql/workbench/services/query/common/queryManagement.ts

@@ -51,7 +51,20 @@ export interface IQueryManagementService {
 	runQueryString(ownerUri: string, queryString: string): Promise<void>;
 	runQueryAndReturn(ownerUri: string, queryString: string): Promise<azdata.SimpleExecuteResult>;
 	parseSyntax(ownerUri: string, query: string): Promise<azdata.SyntaxParseResult>;
-	getQueryRows(rowData: azdata.QueryExecuteSubsetParams, cancellationToken?: CancellationToken, onProgressCallback?: (availableRows: number) => void): Promise<ResultSetSubset>;
+	/**
+	 * Fetches the specified rows - this will fetch all the rows at once.
+	 * @param rowData The rows to fetch
+	 * @deprecated getQueryRowsPaged should be preferred as it is much more performant for large data sets
+	 */
+	getQueryRows(rowData: azdata.QueryExecuteSubsetParams): Promise<ResultSetSubset>;
+	/**
+	 * Fetches the specified rows with paging - getting subsets of the total rows one page at a time and then returning the entire set once
+	 * completed.
+	 * @param rowData The rows to fetch
+	 * @param cancellationToken Optional cancellation token for canceling the operation while fetching rows
+	 * @param onProgressCallback Optional callback that will be called each time a page of rows is fetched
+	 */
+	getQueryRowsPaged(rowData: azdata.QueryExecuteSubsetParams, cancellationToken?: CancellationToken, onProgressCallback?: (availableRows: number) => void): Promise<ResultSetSubset>;
 	disposeQuery(ownerUri: string): Promise<void>;
 	changeConnectionUri(newUri: string, oldUri: string): Promise<void>;
 	saveResults(requestParams: azdata.SaveResultsRequestParams): Promise<azdata.SaveResultRequestResult>;
@@ -272,7 +285,13 @@ export class QueryManagementService implements IQueryManagementService {
 		});
 	}
 
-	public async getQueryRows(rowData: azdata.QueryExecuteSubsetParams, cancellationToken?: CancellationToken, onProgressCallback?: (availableRows: number) => void): Promise<ResultSetSubset> {
+	public async getQueryRows(rowData: azdata.QueryExecuteSubsetParams): Promise<ResultSetSubset> {
+		return this._runAction(rowData.ownerUri, (runner) => {
+			return runner.getQueryRows(rowData).then(r => r.resultSubset);
+		});
+	}
+
+	public async getQueryRowsPaged(rowData: azdata.QueryExecuteSubsetParams, cancellationToken?: CancellationToken, onProgressCallback?: (availableRows: number) => void): Promise<ResultSetSubset> {
 		const pageSize = 500;
 		return this._runAction(rowData.ownerUri, async (runner): Promise<ResultSetSubset> => {
 			const result = [];

src/sql/workbench/services/query/common/queryModelService.ts

@@ -150,7 +150,7 @@ export class QueryModelService implements IQueryModelService {
 	 */
 	public getQueryRows(uri: string, rowStart: number, numberOfRows: number, batchId: number, resultId: number): Promise<ResultSetSubset | undefined> {
 		if (this._queryInfoMap.has(uri)) {
-			return this._getQueryInfo(uri)!.queryRunner!.getQueryRows(rowStart, numberOfRows, batchId, resultId).then(results => {
+			return this._getQueryInfo(uri)!.queryRunner!.getQueryRowsPaged(rowStart, numberOfRows, batchId, resultId).then(results => {
 				return results;
 			});
 		} else {

src/sql/workbench/services/query/common/queryRunner.ts

@@ -348,7 +348,7 @@ export default class QueryRunner extends Disposable {
 			if (hasShowPlan && resultSet.rowCount > 0) {
 				this._isQueryPlan = true;
 
-				this.getQueryRows(0, 1, resultSet.batchId, resultSet.id).then(e => {
+				this.getQueryRowsPaged(0, 1, resultSet.batchId, resultSet.id).then(e => {
 					if (e.rows) {
 						this._planXml.resolve(e.rows[0][0].displayValue);
 					}
@@ -373,7 +373,7 @@ export default class QueryRunner extends Disposable {
 			let hasShowPlan = !!resultSet.columnInfo.find(e => e.columnName === 'Microsoft SQL Server 2005 XML Showplan');
 			if (hasShowPlan) {
 				this._isQueryPlan = true;
-				this.getQueryRows(0, 1, resultSet.batchId, resultSet.id).then(e => {
+				this.getQueryRowsPaged(0, 1, resultSet.batchId, resultSet.id).then(e => {
 
 					if (e.rows) {
 						let planXmlString = e.rows[0][0].displayValue;
@@ -419,6 +419,7 @@ export default class QueryRunner extends Disposable {
 
 	/**
 	 * Get more data rows from the current resultSets from the service layer
+	 * @deprecated getQueryRowsPaged should be used instead as it is much more performant
 	 */
 	public getQueryRows(rowStart: number, numberOfRows: number, batchIndex: number, resultSetIndex: number, cancellationToken?: CancellationToken, onProgressCallback?: (availableRows: number) => void): Promise<ResultSetSubset> {
 		let rowData: QueryExecuteSubsetParams = <QueryExecuteSubsetParams>{
@@ -429,7 +430,22 @@ export default class QueryRunner extends Disposable {
 			batchIndex: batchIndex
 		};
 
-		return this.queryManagementService.getQueryRows(rowData, cancellationToken, onProgressCallback).then(r => r, error => {
+		return this.queryManagementService.getQueryRows(rowData);
+	}
+
+	/**
+	 * Get more data rows from the current resultSets from the service layer with paging, fetching row data in batches until all rows are retrieved.
+	 */
+	public getQueryRowsPaged(rowStart: number, numberOfRows: number, batchIndex: number, resultSetIndex: number, cancellationToken?: CancellationToken, onProgressCallback?: (availableRows: number) => void): Promise<ResultSetSubset> {
+		let rowData: QueryExecuteSubsetParams = <QueryExecuteSubsetParams>{
+			ownerUri: this.uri,
+			resultSetIndex: resultSetIndex,
+			rowsCount: numberOfRows,
+			rowsStartIndex: rowStart,
+			batchIndex: batchIndex
+		};
+
+		return this.queryManagementService.getQueryRowsPaged(rowData, cancellationToken, onProgressCallback).then(r => r, error => {
 			// this._notificationService.notify({
 			// 	severity: Severity.Error,
 			// 	message: nls.localize('query.gettingRowsFailedError', 'Something went wrong getting more rows: {0}', error)
@@ -566,7 +582,7 @@ export class QueryGridDataProvider implements IGridDataProvider {
 	}
 
 	getRowData(rowStart: number, numberOfRows: number, cancellationToken?: CancellationToken, onProgressCallback?: (availableRows: number) => void): Promise<ResultSetSubset> {
-		return this.queryRunner.getQueryRows(rowStart, numberOfRows, this.batchId, this.resultSetId, cancellationToken, onProgressCallback);
+		return this.queryRunner.getQueryRowsPaged(rowStart, numberOfRows, this.batchId, this.resultSetId, cancellationToken, onProgressCallback);
 	}
 
 	copyResults(selection: Slick.Range[], includeHeaders?: boolean, tableView?: IDisposableDataProvider<Slick.SlickData>): Promise<void> {

src/sql/workbench/services/query/test/browser/queryRunner.test.ts

@@ -57,8 +57,8 @@ suite('Query Runner', () => {
 		// get query rows
 		const rowResults: ResultSetSubset = { rowCount: 100, rows: range(100).map(r => range(1).map(c => ({ displayValue: `${r}${c}` }))) };
 		const getRowStub = sinon.stub().returns(Promise.resolve(rowResults));
-		(instantiationService as TestInstantiationService).stub(IQueryManagementService, 'getQueryRows', getRowStub);
+		(instantiationService as TestInstantiationService).stub(IQueryManagementService, 'getQueryRowsPaged', getRowStub);
-		const resultReturn = await runner.getQueryRows(0, 100, 0, 0);
+		const resultReturn = await runner.getQueryRowsPaged(0, 100, 0, 0);
 		assert(getRowStub.calledWithExactly({ ownerUri: uri, batchIndex: 0, resultSetIndex: 0, rowsStartIndex: 0, rowsCount: 100 }, undefined, undefined));
 		assert.deepStrictEqual(resultReturn, rowResults);
 		// batch complete
@@ -121,7 +121,7 @@ suite('Query Runner', () => {
 		assert(runQueryStub.calledWithExactly(uri, undefined, { displayEstimatedQueryPlan: true }));
 		const xmlPlan = 'xml plan';
 		const getRowsStub = sinon.stub().returns(Promise.resolve({ rowCount: 1, rows: [[{ displayValue: xmlPlan }]] } as ResultSetSubset));
-		(instantiationService as TestInstantiationService).stub(IQueryManagementService, 'getQueryRows', getRowsStub);
+		(instantiationService as TestInstantiationService).stub(IQueryManagementService, 'getQueryRowsPaged', getRowsStub);
 		runner.handleBatchStart({ id: 0, executionStart: '' });
 		runner.handleResultSetAvailable({ id: 0, batchId: 0, complete: true, rowCount: 1, columnInfo: [{ columnName: 'Microsoft SQL Server 2005 XML Showplan' }] });
 		const plan = await runner.planXml;
@@ -143,7 +143,7 @@ suite('Query Runner', () => {
 		runner.handleResultSetAvailable({ id: 0, batchId: 0, complete: false, rowCount: 0, columnInfo: [{ columnName: 'Microsoft SQL Server 2005 XML Showplan' }] });
 		const xmlPlan = 'xml plan';
 		const getRowsStub = sinon.stub().returns(Promise.resolve({ rowCount: 1, rows: [[{ displayValue: xmlPlan }]] } as ResultSetSubset));
-		(instantiationService as TestInstantiationService).stub(IQueryManagementService, 'getQueryRows', getRowsStub);
+		(instantiationService as TestInstantiationService).stub(IQueryManagementService, 'getQueryRowsPaged', getRowsStub);
 		runner.handleResultSetUpdated({ id: 0, batchId: 0, complete: true, rowCount: 1, columnInfo: [{ columnName: 'Microsoft SQL Server 2005 XML Showplan' }] });
 		const plan = await runner.planXml;
 		assert(getRowsStub.calledOnce);

src/sql/workbench/services/query/test/common/testQueryManagementService.ts

@@ -51,7 +51,10 @@ export class TestQueryManagementService implements IQueryManagementService {
 	parseSyntax(ownerUri: string, query: string): Promise<azdata.SyntaxParseResult> {
 		throw new Error('Method not implemented.');
 	}
-	getQueryRows(rowData: azdata.QueryExecuteSubsetParams, cancellationToken?: CancellationToken, onProgressCallback?: (availableRows: number) => void): Promise<ResultSetSubset> {
+	getQueryRows(rowData: azdata.QueryExecuteSubsetParams): Promise<ResultSetSubset> {
+		throw new Error('Method not implemented.');
+	}
+	getQueryRowsPaged(rowData: azdata.QueryExecuteSubsetParams, cancellationToken?: CancellationToken, onProgressCallback?: (availableRows: number) => void): Promise<ResultSetSubset> {
 		throw new Error('Method not implemented.');
 	}
 	async disposeQuery(ownerUri: string): Promise<void> {