query result selection summary improvement (both perf and usability) (#23378)

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

@@ -33,11 +33,11 @@ import { ActionBar, ActionsOrientation } from 'vs/base/browser/ui/actionbar/acti
 import { isInDOM, Dimension } from 'vs/base/browser/dom';
 import { IInstantiationService } from 'vs/platform/instantiation/common/instantiation';
 import { IEditorService } from 'vs/workbench/services/editor/common/editorService';
-import { IAction, Separator } from 'vs/base/common/actions';
+import { IAction, Separator, toAction } from 'vs/base/common/actions';
 import { ILogService } from 'vs/platform/log/common/log';
 import { localize } from 'vs/nls';
 import { IGridDataProvider } from 'sql/workbench/services/query/common/gridDataProvider';
-import { CancellationToken } from 'vs/base/common/cancellation';
+import { CancellationToken, CancellationTokenSource } from 'vs/base/common/cancellation';
 import { GridPanelState, GridTableState } from 'sql/workbench/common/editor/query/gridTableState';
 import { IUntitledTextEditorService } from 'vs/workbench/services/untitled/common/untitledTextEditorService';
 import { SaveFormat } from 'sql/workbench/services/query/common/resultSerializer';
@@ -48,7 +48,7 @@ import { Orientation } from 'vs/base/browser/ui/splitview/splitview';
 import { IQueryModelService } from 'sql/workbench/services/query/common/queryModel';
 import { FilterButtonWidth, HeaderFilter } from 'sql/base/browser/ui/table/plugins/headerFilter.plugin';
 import { HybridDataProvider } from 'sql/base/browser/ui/table/hybridDataProvider';
-import { INotificationService } from 'vs/platform/notification/common/notification';
+import { INotificationHandle, INotificationService, Severity } from 'vs/platform/notification/common/notification';
 import { alert, status } from 'vs/base/browser/ui/aria/aria';
 import { IExecutionPlanService } from 'sql/workbench/services/executionPlan/common/interfaces';
 import { ExecutionPlanInput } from 'sql/workbench/contrib/executionPlan/browser/executionPlanInput';
@@ -58,6 +58,8 @@ import { IAccessibilityService } from 'vs/platform/accessibility/common/accessib
 import { IQuickInputService } from 'vs/platform/quickinput/common/quickInput';
 import { queryEditorNullBackground } from 'sql/platform/theme/common/colorRegistry';
 import { IComponentContextService } from 'sql/workbench/services/componentContext/browser/componentContextService';
+import { GridRange } from 'sql/base/common/gridRange';
+import { onUnexpectedError } from 'vs/base/common/errors';
 
 const ROW_HEIGHT = 29;
 const HEADER_HEIGHT = 26;
@@ -373,6 +375,7 @@ export abstract class GridTableBase<T> extends Disposable implements IView, IQue
 	private filterPlugin: HeaderFilter<T>;
 	private isDisposed: boolean = false;
 	private gridConfig: IResultGridConfiguration;
+	private selectionChangeHandlerTokenSource: CancellationTokenSource | undefined;
 
 	private columns: Slick.Column<T>[];
 
@@ -664,7 +667,7 @@ export abstract class GridTableBase<T> extends Disposable implements IView, IQue
 			if (this.state) {
 				this.state.selection = this.selectionModel.getSelectedRanges();
 			}
-			await this.notifyTableSelectionChanged();
+			await this.handleTableSelectionChange();
 		});
 
 		this.table.grid.onScroll.subscribe((e, data) => {
@@ -763,7 +766,7 @@ export abstract class GridTableBase<T> extends Disposable implements IView, IQue
 		this._state = val;
 	}
 
-	private async getRowData(start: number, length: number): Promise<ICellValue[][]> {
+	private async getRowData(start: number, length: number, cancellationToken?: CancellationToken, onProgressCallback?: (availableRows: number) => void): Promise<ICellValue[][]> {
 		let subset;
 		if (this.dataProvider.isDataInMemory) {
 			// handle the scenario when the data is sorted/filtered,
@@ -771,23 +774,102 @@ export abstract class GridTableBase<T> extends Disposable implements IView, IQue
 			const data = await this.dataProvider.getRangeAsync(start, length);
 			subset = data.map(item => Object.keys(item).map(key => item[key]));
 		} else {
-			subset = (await this.gridDataProvider.getRowData(start, length)).rows;
+			subset = (await this.gridDataProvider.getRowData(start, length, cancellationToken, onProgressCallback)).rows;
 		}
 		return subset;
 	}
 
-	private async notifyTableSelectionChanged() {
-		const selectedCells = [];
-		for (const range of this.state.selection) {
-			const subset = await this.getRowData(range.fromRow, range.toRow - range.fromRow + 1);
-			subset.forEach(row => {
-				// start with range.fromCell -1 because we have row number column which is not available in the actual data
-				for (let i = range.fromCell - 1; i < range.toCell; i++) {
-					selectedCells.push(row[i]);
-				}
-			});
+	private async handleTableSelectionChange(): Promise<void> {
+		if (this.selectionChangeHandlerTokenSource) {
+			this.selectionChangeHandlerTokenSource.cancel();
+		}
+		this.selectionChangeHandlerTokenSource = new CancellationTokenSource();
+		await this.notifyTableSelectionChanged(this.selectionChangeHandlerTokenSource);
+	}
+
+	private async notifyTableSelectionChanged(cancellationTokenSource: CancellationTokenSource): Promise<void> {
+		const gridRanges = GridRange.fromSlickRanges(this.state.selection ?? []);
+		const rowRanges = GridRange.getUniqueRows(gridRanges);
+		const columnRanges = GridRange.getUniqueColumns(gridRanges);
+		const rowCount = rowRanges.map(range => range.end - range.start + 1).reduce((p, c) => p + c);
+		const runAction = async (proceed: boolean) => {
+			const selectedCells = [];
+			if (proceed && !cancellationTokenSource.token.isCancellationRequested) {
+				let notificationHandle: INotificationHandle = undefined;
+				const timeout = setTimeout(() => {
+					notificationHandle = this.notificationService.notify({
+						message: localize('resultsGrid.loadingData', "Loading selected rows for calculation..."),
+						severity: Severity.Info,
+						progress: {
+							infinite: true
+						},
+						actions: {
+							primary: [
+								toAction({
+									id: 'cancelLoadingCells',
+									label: localize('resultsGrid.cancel', "Cancel"),
+									run: () => {
+										cancellationTokenSource.cancel();
+										notificationHandle.close();
+									}
+								})]
+						}
+					});
+				}, 1000);
+				this.queryModelService.notifyCellSelectionChanged([]);
+				let rowsInProcessedRanges = 0;
+				for (const range of rowRanges) {
+					if (cancellationTokenSource.token.isCancellationRequested) {
+						break;
+					}
+					const rows = await this.getRowData(range.start, range.end - range.start + 1, cancellationTokenSource.token, (availableRows: number) => {
+						notificationHandle?.updateMessage(localize('resultsGrid.loadingDataWithProgress', "Loading selected rows for calculation ({0}/{1})...", rowsInProcessedRanges + availableRows, rowCount));
+					});
+					rows.forEach((row, rowIndex) => {
+						columnRanges.forEach(cr => {
+							for (let i = cr.start; i <= cr.end; i++) {
+								if (this.state.selection.some(selection => selection.contains(rowIndex + range.start, i))) {
+									// need to reduce the column index by 1 because we have row number column which is not available in the actual data
+									selectedCells.push(row[i - 1]);
+								}
+							}
+						});
+					});
+					rowsInProcessedRanges += range.end - range.start + 1;
+				}
+				clearTimeout(timeout);
+				notificationHandle?.close();
+			}
+			cancellationTokenSource.dispose();
+			if (!cancellationTokenSource.token.isCancellationRequested) {
+				this.queryModelService.notifyCellSelectionChanged(selectedCells);
+			}
+		};
+		const showPromptConfigValue = this.configurationService.getValue<IQueryEditorConfiguration>('queryEditor').results.promptForLargeRowSelection;
+		if (this.options.inMemoryDataCountThreshold && rowCount > this.options.inMemoryDataCountThreshold && showPromptConfigValue) {
+			this.notificationService.prompt(Severity.Warning, localize('resultsGrid.largeRowSelectionPrompt.', 'You have selected {0} rows, it might take a while to load the data and calculate the summary, do you want to continue?', rowCount), [
+				{
+					label: localize('resultsGrid.confirmLargeRowSelection', "Yes"),
+					run: async () => {
+						await runAction(true);
+					}
+				}, {
+					label: localize('resultsGrid.cancelLargeRowSelection', "Cancel"),
+					run: async () => {
+						await runAction(false);
+					}
+				}, {
+					label: localize('resultsGrid.donotShowLargeRowSelectionPromptAgain', "Don't show again"),
+					run: async () => {
+						this.configurationService.updateValue('queryEditor.results.promptForLargeRowSelection', false).catch(e => onUnexpectedError(e));
+						await runAction(true);
+					},
+					isSecondary: true
+				}
+			]);
+		} else {
+			await runAction(true);
 		}
-		this.queryModelService.notifyCellSelectionChanged(selectedCells);
 	}
 
 	private async onTableClick(event: ITableMouseEvent) {

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

@@ -424,6 +424,11 @@ const queryEditorConfiguration: IConfigurationNode = {
 			'description': localize('queryEditor.results.showActionBar', "Whether to show the action bar in the query results view"),
 			'default': true
 		},
+		'queryEditor.results.promptForLargeRowSelection': {
+			'type': 'boolean',
+			'default': true,
+			'description': localize('queryEditor.results.promptForLargeRowSelection', "When cells are selected in the results grid, ADS will calculate the summary for them, This setting controls whether to show the a confirmation when the number of rows selected is larger than the value specified in the 'inMemoryDataProcessingThreshold' setting. The default value is true.")
+		},
 		'queryEditor.messages.showBatchTime': {
 			'type': 'boolean',
 			'description': localize('queryEditor.messages.showBatchTime', "Should execution time be shown for individual batches"),

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

@@ -316,10 +316,12 @@ export class QueryResultSelectionSummaryStatusBarContribution extends Disposable
 		const nullCount = selectedCells.filter(cell => cell.isNull).length;
 		let summaryText, tooltipText;
 		if (numericValues.length >= 2) {
-			const sum = numericValues.reduce((previous, current, idx, array) => previous + current);
+			const sum = numericValues.reduce((previous, current) => previous + current);
+			const min = numericValues.reduce((previous, current) => Math.min(previous, current));
+			const max = numericValues.reduce((previous, current) => Math.max(previous, current));
 			summaryText = localize('status.query.summaryText', "Average: {0}  Count: {1}  Sum: {2}", Number((sum / numericValues.length).toFixed(3)), selectedCells.length, sum);
 			tooltipText = localize('status.query.summaryTooltip', "Average: {0}  Count: {1}  Distinct Count: {2}  Max: {3}  Min: {4}  Null Count: {5}  Sum: {6}",
-				Number((sum / numericValues.length).toFixed(3)), selectedCells.length, distinctValues.size, Math.max(...numericValues), Math.min(...numericValues), nullCount, sum);
+				Number((sum / numericValues.length).toFixed(3)), selectedCells.length, distinctValues.size, max, min, nullCount, sum);
 		} else {
 			summaryText = summaryText = localize('status.query.summaryTextNonNumeric', "Count: {0}  Distinct Count: {1}  Null Count: {2}", selectedCells.length, distinctValues.size, nullCount);
 		}

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

@@ -11,6 +11,8 @@ import { INotificationHandle, INotificationService, Severity } from 'vs/platform
 import * as nls from 'vs/nls';
 import { toAction } from 'vs/base/common/actions';
 import { IClipboardService } from 'vs/platform/clipboard/common/clipboardService';
+import { CancellationToken, CancellationTokenSource } from 'vs/base/common/cancellation';
+import { GridRange } from 'sql/base/common/gridRange';
 
 export interface IGridDataProvider {
 
@@ -19,7 +21,7 @@ export interface IGridDataProvider {
 	 * @param rowStart 0-indexed start row to retrieve data from
 	 * @param numberOfRows total number of rows of data to retrieve
 	 */
-	getRowData(rowStart: number, numberOfRows: number): Thenable<ResultSetSubset>;
+	getRowData(rowStart: number, numberOfRows: number, cancellationToken?: CancellationToken, onProgressCallback?: (availableRows: number) => void): Thenable<ResultSetSubset>;
 
 	/**
 	 * Sends a copy request to copy data to the clipboard
@@ -51,38 +53,9 @@ export interface IGridDataProvider {
 	serializeResults(format: SaveFormat, selection: Slick.Range[]): Thenable<void>;
 }
 
-interface Range {
-	start: number;
-	end: number;
-}
-
-/**
- * Merge the ranges and get the sorted ranges.
- */
-function mergeRanges(ranges: Range[]): Range[] {
-	const mergedRanges: Range[] = [];
-	const orderedRanges = ranges.sort((s1, s2) => { return s1.start - s2.start; });
-	orderedRanges.forEach(range => {
-		let merged = false;
-		for (let i = 0; i < mergedRanges.length; i++) {
-			const mergedRange = mergedRanges[i];
-			if (range.start <= mergedRange.end) {
-				mergedRange.end = Math.max(range.end, mergedRange.end);
-				merged = true;
-				break;
-			}
-		}
-		if (!merged) {
-			mergedRanges.push(range);
-		}
-	});
-	return mergedRanges;
-}
-
-export async function executeCopyWithNotification(notificationService: INotificationService, selections: Slick.Range[], isCancelable: boolean, copyHandler: (notification: INotificationHandle, rowCount: number) => Promise<void>, onCanceled?: () => void): Promise<void> {
-	const rowRanges: Range[] = mergeRanges(selections.map(selection => { return { start: selection.fromRow, end: selection.toRow }; }));
+export async function executeCopyWithNotification(notificationService: INotificationService, selections: Slick.Range[], copyHandler: (notification: INotificationHandle, rowCount: number) => Promise<void>, cancellationTokenSource?: CancellationTokenSource): Promise<void> {
+	const rowRanges = GridRange.getUniqueRows(GridRange.fromSlickRanges(selections));
 	const rowCount = rowRanges.map(range => range.end - range.start + 1).reduce((p, c) => p + c);
-	let isCanceled = false;
 	const notificationHandle = notificationService.notify({
 		message: nls.localize('gridDataProvider.copying', "Copying..."),
 		severity: Severity.Info,
@@ -90,13 +63,12 @@ export async function executeCopyWithNotification(notificationService: INotifica
 			infinite: true
 		},
 		actions: {
-			primary: isCancelable ? [
+			primary: cancellationTokenSource ? [
 				toAction({
 					id: 'cancelCopyResults',
 					label: nls.localize('gridDataProvider.cancelCopyResults', "Cancel"),
 					run: () => {
-						isCanceled = true;
-						onCanceled!();
+						cancellationTokenSource.cancel();
 						notificationHandle.close();
 					}
 				})] : []
@@ -104,7 +76,7 @@ export async function executeCopyWithNotification(notificationService: INotifica
 	});
 	try {
 		await copyHandler(notificationHandle, rowCount);
-		if (!isCanceled) {
+		if (cancellationTokenSource === undefined || !cancellationTokenSource.token.isCancellationRequested) {
 			notificationHandle.progress.done();
 			notificationHandle.updateActions({
 				primary: [
@@ -124,16 +96,16 @@ export async function executeCopyWithNotification(notificationService: INotifica
 }
 
 export async function copySelectionToClipboard(clipboardService: IClipboardService, notificationService: INotificationService, provider: IGridDataProvider, selections: Slick.Range[], includeHeaders?: boolean, tableView?: IDisposableDataProvider<Slick.SlickData>): Promise<void> {
-	let isCanceled = false;
-	await executeCopyWithNotification(notificationService, selections, true, async (notificationHandle, rowCount) => {
-		const batchSize = 100;
+	const cancellationTokenSource = new CancellationTokenSource()
+	await executeCopyWithNotification(notificationService, selections, async (notificationHandle, rowCount) => {
 		const eol = provider.getEolString();
 		const valueSeparator = '\t';
 		const shouldRemoveNewLines = provider.shouldRemoveNewLines();
 
-		// Merge the selections to get the columns and rows.
-		const columnRanges: Range[] = mergeRanges(selections.map(selection => { return { start: selection.fromCell, end: selection.toCell }; }));
-		const rowRanges: Range[] = mergeRanges(selections.map(selection => { return { start: selection.fromRow, end: selection.toRow }; }));
+		// Merge the selections to get the unique columns and unique rows.
+		const gridRanges = GridRange.fromSlickRanges(selections);
+		const columnRanges = GridRange.getUniqueColumns(gridRanges);
+		const rowRanges = GridRange.getUniqueRows(gridRanges);
 
 		let processedRows = 0;
 		const getMessageText = (): string => {
@@ -151,57 +123,41 @@ export async function copySelectionToClipboard(clipboardService: IClipboardServi
 			resultString = Array.from(headers.values()).join(valueSeparator).concat(eol);
 		}
 
-		const batchResult: string[] = [];
+		const rowValues: string[] = [];
 		for (const range of rowRanges) {
+			let rows: ICellValue[][];
+			let processedRowsSnapshot = processedRows;
+			const rangeLength = range.end - range.start + 1;
 			if (tableView && tableView.isDataInMemory) {
-				const rangeLength = range.end - range.start + 1;
 				// If the data is sorted/filtered in memory, we need to get the data that is currently being displayed
 				const tableData = await tableView.getRangeAsync(range.start, rangeLength);
-				const rowSet = tableData.map(item => Object.keys(item).map(key => item[key]));
-				batchResult.push(getStringValueForRowSet(rowSet, columnRanges, selections, range.start, eol, valueSeparator, shouldRemoveNewLines));
+				rows = tableData.map(item => Object.keys(item).map(key => item[key]));
 				processedRows += rangeLength;
 				notificationHandle.updateMessage(getMessageText());
 			} else {
-				let start = range.start;
-				do {
-					const end = Math.min(start + batchSize - 1, range.end);
-					const batchLength = end - start + 1
-					const rowSet = (await provider.getRowData(start, batchLength)).rows;
-					batchResult.push(getStringValueForRowSet(rowSet, columnRanges, selections, range.start, eol, valueSeparator, shouldRemoveNewLines));
-					start = end + 1;
-					processedRows = processedRows + batchLength;
-					if (!isCanceled) {
-						notificationHandle.updateMessage(getMessageText());
-					}
-				} while (start < range.end && !isCanceled)
+				rows = (await provider.getRowData(range.start, rangeLength, cancellationTokenSource.token, (fetchedRows) => {
+					processedRows = processedRowsSnapshot + fetchedRows;
+					notificationHandle.updateMessage(getMessageText());
+				})).rows;
 			}
+			rows.forEach((values, index) => {
+				const rowIndex = index + range.start;
+				const columnValues = [];
+				columnRanges.forEach(cr => {
+					for (let i = cr.start; i <= cr.end; i++) {
+						if (selections.some(selection => selection.contains(rowIndex, i))) {
+							columnValues.push(shouldRemoveNewLines ? removeNewLines(values[i].displayValue) : values[i].displayValue);
+						}
+					}
+				});
+				rowValues.push(columnValues.join(valueSeparator));
+			});
 		}
-		if (!isCanceled) {
-			resultString += batchResult.join(eol);
+		if (!cancellationTokenSource.token.isCancellationRequested) {
+			resultString += rowValues.join(eol);
 			await clipboardService.writeText(resultString);
 		}
-	}, () => {
-		isCanceled = true;
-	});
-}
-
-function getStringValueForRowSet(rows: ICellValue[][], columnRanges: Range[], selections: Slick.Range[], rowSetStartIndex: number, eol: string, valueSeparator: string, shouldRemoveNewLines: boolean): string {
-	let rowStrings: string[] = [];
-	rows.forEach((values, index) => {
-		const rowIndex = index + rowSetStartIndex;
-		const rowValues = [];
-		columnRanges.forEach(cr => {
-			for (let i = cr.start; i <= cr.end; i++) {
-				if (selections.some(selection => selection.contains(rowIndex, i))) {
-					rowValues.push(shouldRemoveNewLines ? removeNewLines(values[i].displayValue) : values[i].displayValue);
-				} else {
-					rowValues.push('');
-				}
-			}
-		});
-		rowStrings.push(rowValues.join(valueSeparator));
-	});
-	return rowStrings.join(eol);
+	}, cancellationTokenSource);
 }
 
 export function getTableHeaderString(provider: IGridDataProvider, selection: Slick.Range[]): string {

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

@@ -18,6 +18,7 @@ import { ResultSetSubset } from 'sql/workbench/services/query/common/query';
 import { isUndefined } from 'vs/base/common/types';
 import { ILogService } from 'vs/platform/log/common/log';
 import * as nls from 'vs/nls';
+import { CancellationToken } from 'vs/base/common/cancellation';
 
 export const SERVICE_ID = 'queryManagementService';
 
@@ -50,7 +51,7 @@ 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): Promise<ResultSetSubset>;
+	getQueryRows(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>;
@@ -271,9 +272,38 @@ export class QueryManagementService implements IQueryManagementService {
 		});
 	}
 
-	public async getQueryRows(rowData: azdata.QueryExecuteSubsetParams): Promise<ResultSetSubset> {
-		return this._runAction(rowData.ownerUri, (runner) => {
-			return runner.getQueryRows(rowData).then(r => r.resultSubset);
+	public async getQueryRows(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 = [];
+			let start = rowData.rowsStartIndex;
+			this._logService.trace(`Getting ${rowData.rowsCount} rows starting from index: ${rowData.rowsStartIndex}.`);
+			do {
+				const rowCount = Math.min(pageSize, rowData.rowsStartIndex + rowData.rowsCount - start);
+				this._logService.trace(`Paged Fetch - Getting ${rowCount} rows starting from index: ${start}.`);
+				const rowSet = await runner.getQueryRows({
+					ownerUri: rowData.ownerUri,
+					batchIndex: rowData.batchIndex,
+					resultSetIndex: rowData.resultSetIndex,
+					rowsCount: rowCount,
+					rowsStartIndex: start
+				});
+				this._logService.trace(`Paged Fetch - Received ${rowSet.resultSubset.rows} rows starting from index: ${start}.`);
+				result.push(...rowSet.resultSubset.rows);
+				start += rowCount;
+				if (onProgressCallback) {
+					onProgressCallback(start - rowData.rowsStartIndex);
+				}
+			} while (start < rowData.rowsStartIndex + rowData.rowsCount && (cancellationToken === undefined || !cancellationToken.isCancellationRequested));
+			if (cancellationToken?.isCancellationRequested) {
+				this._logService.trace(`Stop getting more rows since cancellation has been requested.`);
+			} else {
+				this._logService.trace(`Successfully fetched ${result.length} rows. Expected Rows: ${rowData.rowsCount}.`);
+			}
+			return {
+				rows: result,
+				rowCount: result.length
+			};
 		});
 	}
 

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

@@ -30,6 +30,7 @@ import { IQueryEditorConfiguration } from 'sql/platform/query/common/query';
 import { IDisposableDataProvider } from 'sql/base/common/dataProvider';
 import { ITextResourcePropertiesService } from 'vs/editor/common/services/textResourceConfiguration';
 import { ICapabilitiesService } from 'sql/platform/capabilities/common/capabilitiesService';
+import { CancellationToken } from 'vs/base/common/cancellation';
 
 /*
 * Query Runner class which handles running a query, reports the results to the content manager,
@@ -419,7 +420,7 @@ export default class QueryRunner extends Disposable {
 	/**
 	 * Get more data rows from the current resultSets from the service layer
 	 */
-	public getQueryRows(rowStart: number, numberOfRows: number, batchIndex: number, resultSetIndex: number): Promise<ResultSetSubset> {
+	public getQueryRows(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,
@@ -428,7 +429,7 @@ export default class QueryRunner extends Disposable {
 			batchIndex: batchIndex
 		};
 
-		return this.queryManagementService.getQueryRows(rowData).then(r => r, error => {
+		return this.queryManagementService.getQueryRows(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,8 +567,8 @@ export class QueryGridDataProvider implements IGridDataProvider {
 	) {
 	}
 
-	getRowData(rowStart: number, numberOfRows: number): Promise<ResultSetSubset> {
-		return this.queryRunner.getQueryRows(rowStart, numberOfRows, this.batchId, this.resultSetId);
+	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);
 	}
 
 	copyResults(selection: Slick.Range[], includeHeaders?: boolean, tableView?: IDisposableDataProvider<Slick.SlickData>): Promise<void> {
@@ -590,7 +591,7 @@ export class QueryGridDataProvider implements IGridDataProvider {
 	}
 
 	private async handleCopyRequestByProvider(selections: Slick.Range[], includeHeaders?: boolean): Promise<void> {
-		executeCopyWithNotification(this._notificationService, selections, false, async () => {
+		executeCopyWithNotification(this._notificationService, selections, async () => {
 			await this.queryRunner.copyResults(selections, this.batchId, this.resultSetId, this.shouldRemoveNewLines(), this.shouldIncludeHeaders(includeHeaders));
 		});
 	}

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

@@ -59,7 +59,7 @@ suite('Query Runner', () => {
 		const getRowStub = sinon.stub().returns(Promise.resolve(rowResults));
 		(instantiationService as TestInstantiationService).stub(IQueryManagementService, 'getQueryRows', getRowStub);
 		const resultReturn = await runner.getQueryRows(0, 100, 0, 0);
-		assert(getRowStub.calledWithExactly({ ownerUri: uri, batchIndex: 0, resultSetIndex: 0, rowsStartIndex: 0, rowsCount: 100 }));
+		assert(getRowStub.calledWithExactly({ ownerUri: uri, batchIndex: 0, resultSetIndex: 0, rowsStartIndex: 0, rowsCount: 100 }, undefined, undefined));
 		assert.deepStrictEqual(resultReturn, rowResults);
 		// batch complete
 		const batchComplete: CompleteBatchSummary = { ...batch, executionEnd: 'endstring', executionElapsed: 'elapsedstring' };

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

@@ -10,6 +10,7 @@ import QueryRunner from 'sql/workbench/services/query/common/queryRunner';
 import * as azdata from 'azdata';
 import { IRange } from 'vs/editor/common/core/range';
 import { ResultSetSubset } from 'sql/workbench/services/query/common/query';
+import { CancellationToken } from 'vs/base/common/cancellation';
 
 export class TestQueryManagementService implements IQueryManagementService {
 	_serviceBrand: undefined;
@@ -50,7 +51,7 @@ export class TestQueryManagementService implements IQueryManagementService {
 	parseSyntax(ownerUri: string, query: string): Promise<azdata.SyntaxParseResult> {
 		throw new Error('Method not implemented.');
 	}
-	getQueryRows(rowData: azdata.QueryExecuteSubsetParams): Promise<ResultSetSubset> {
+	getQueryRows(rowData: azdata.QueryExecuteSubsetParams, cancellationToken?: CancellationToken, onProgressCallback?: (availableRows: number) => void): Promise<ResultSetSubset> {
 		throw new Error('Method not implemented.');
 	}
 	async disposeQuery(ownerUri: string): Promise<void> {