enable provider/extension side copy to clipboard support (#23363)

* improve copy data experience

* add copy result handler

* refactoring

* updates

* add thirdparty notice for TextCopy nuget package

* add await

* comments

src/sql/workbench/api/browser/mainThreadDataProtocol.ts

@@ -159,6 +159,9 @@ export class MainThreadDataProtocol extends Disposable implements MainThreadData
 					return Promise.resolve(self._serializationService.saveAs(requestParams.resultFormat, requestParams.filePath, undefined, true));
 				}
 			},
+			copyResults(requestParams: azdata.CopyResultsRequestParams): Promise<void> {
+				return Promise.resolve(self._proxy.$copyResults(handle, requestParams));
+			},
 			initializeEdit(ownerUri: string, schemaName: string, objectName: string, objectType: string, rowLimit: number, queryString: string): Promise<void> {
 				return Promise.resolve(self._proxy.$initializeEdit(handle, ownerUri, schemaName, objectName, objectType, rowLimit, queryString));
 			},

src/sql/workbench/api/common/extHostDataProtocol.ts

@@ -414,6 +414,15 @@ export class ExtHostDataProtocol extends ExtHostDataProtocolShape {
 		return this._resolveProvider<azdata.QueryProvider>(handle).saveResults(requestParams);
 	}
 
+	override $copyResults(handle: number, requestParams: azdata.CopyResultsRequestParams): Thenable<void> {
+		const provider = this._resolveProvider<azdata.QueryProvider>(handle);
+		if (provider.copyResults) {
+			return provider.copyResults(requestParams);
+		} else {
+			throw new Error(`copyResults() is not implemented by the provider`);
+		}
+	}
+
 	// Edit Data handlers
 	override $commitEdit(handle: number, ownerUri: string): Thenable<void> {
 		return this._resolveProvider<azdata.QueryProvider>(handle).commitEdit(ownerUri);

src/sql/workbench/api/common/sqlExtHost.protocol.ts

@@ -251,6 +251,11 @@ export abstract class ExtHostDataProtocolShape {
 	 */
 	$saveResults(handle: number, requestParams: azdata.SaveResultsRequestParams): Thenable<azdata.SaveResultRequestResult> { throw ni(); }
 
+	/**
+	 * Copies the selected data to clipboard.
+	 */
+	$copyResults(handle: number, requestParams: azdata.CopyResultsRequestParams): Thenable<void> { throw ni(); }
+
 	/**
 	 * Commits all pending edits in an edit session
 	 */

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

@@ -404,6 +404,11 @@ const queryEditorConfiguration: IConfigurationNode = {
 			'description': localize('queryEditor.results.copyRemoveNewLine', "Configuration options for copying multi-line results from the Results View"),
 			'default': true
 		},
+		'queryEditor.results.preferProvidersCopyHandler': {
+			'type': 'boolean',
+			'description': localize('queryEditor.results.preferProvidersCopyHandler', "Whether the copy result request should be handled by the query provider when it is supported. The default value is true, set this to false to force all copy handling to be done by Azure Data Studio."),
+			'default': true
+		},
 		'queryEditor.results.inMemoryDataProcessingThreshold': {
 			'type': 'number',
 			'default': 5000,

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

@@ -7,7 +7,7 @@ import * as types from 'vs/base/common/types';
 import { SaveFormat } from 'sql/workbench/services/query/common/resultSerializer';
 import { ICellValue, ResultSetSubset } from 'sql/workbench/services/query/common/query';
 import { IDisposableDataProvider } from 'sql/base/common/dataProvider';
-import { INotificationService, Severity } from 'vs/platform/notification/common/notification';
+import { INotificationHandle, INotificationService, Severity } from 'vs/platform/notification/common/notification';
 import * as nls from 'vs/nls';
 import { toAction } from 'vs/base/common/actions';
 import { IClipboardService } from 'vs/platform/clipboard/common/clipboardService';
@@ -79,97 +79,112 @@ function mergeRanges(ranges: Range[]): Range[] {
 	return mergedRanges;
 }
 
-export async function copySelectionToClipboard(clipboardService: IClipboardService, notificationService: INotificationService, provider: IGridDataProvider, selections: Slick.Range[], includeHeaders?: boolean, tableView?: IDisposableDataProvider<Slick.SlickData>): Promise<void> {
-	const batchSize = 100;
-	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 }; }));
+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 }; }));
-
-	const totalRows = rowRanges.map(range => range.end - range.start + 1).reduce((p, c) => p + c);
-
-	let processedRows = 0;
-	const getMessageText = (): string => {
-		return nls.localize('gridDataProvider.loadingRowsInProgress', "Loading the rows to be copied ({0}/{1})...", processedRows, totalRows);
-	};
-
+	const rowCount = rowRanges.map(range => range.end - range.start + 1).reduce((p, c) => p + c);
 	let isCanceled = false;
-
 	const notificationHandle = notificationService.notify({
-		message: getMessageText(),
+		message: nls.localize('gridDataProvider.copying', "Copying..."),
 		severity: Severity.Info,
 		progress: {
 			infinite: true
 		},
 		actions: {
-			primary: [
+			primary: isCancelable ? [
 				toAction({
 					id: 'cancelCopyResults',
 					label: nls.localize('gridDataProvider.cancelCopyResults', "Cancel"),
 					run: () => {
 						isCanceled = true;
+						onCanceled!();
 						notificationHandle.close();
 					}
-				})]
+				})] : []
 		}
 	});
-
-	let resultString = '';
-	if (includeHeaders) {
-		const headers: string[] = [];
-		columnRanges.forEach(range => {
-			headers.push(...provider.getColumnHeaders(<Slick.Range>{
-				fromCell: range.start,
-				toCell: range.end
-			}));
-		});
-		resultString = Array.from(headers.values()).join(valueSeparator).concat(eol);
-	}
-
-	const batchResult: string[] = [];
-	for (const range of rowRanges) {
-		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));
-			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)
+	try {
+		await copyHandler(notificationHandle, rowCount);
+		if (!isCanceled) {
+			notificationHandle.progress.done();
+			notificationHandle.updateActions({
+				primary: [
+					toAction({
+						id: 'closeCopyResultsNotification',
+						label: nls.localize('gridDataProvider.closeNotification', "Close"),
+						run: () => { notificationHandle.close(); }
+					})]
+			});
+			notificationHandle.updateMessage(nls.localize('gridDataProvider.copyResultsCompleted', "Selected data has been copied to the clipboard. Row count: {0}.", rowCount));
 		}
 	}
-	if (!isCanceled) {
-		resultString += batchResult.join(eol);
-		notificationHandle.progress.done();
-		notificationHandle.updateActions({
-			primary: [
-				toAction({
-					id: 'closeCopyResultsNotification',
-					label: nls.localize('gridDataProvider.closeNotification', "Close"),
-					run: () => { notificationHandle.close(); }
-				})]
-		});
-		await clipboardService.writeText(resultString);
-		notificationHandle.updateMessage(nls.localize('gridDataProvider.copyResultsCompleted', "Selected data has been copied to the clipboard. Row count: {0}.", totalRows));
+	catch (err) {
+		notificationHandle.close();
+		throw err;
 	}
 }
 
+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 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 }; }));
+
+		let processedRows = 0;
+		const getMessageText = (): string => {
+			return nls.localize('gridDataProvider.loadingRowsInProgress', "Loading the rows to be copied ({0}/{1})...", processedRows, rowCount);
+		};
+		let resultString = '';
+		if (includeHeaders) {
+			const headers: string[] = [];
+			columnRanges.forEach(range => {
+				headers.push(...provider.getColumnHeaders(<Slick.Range>{
+					fromCell: range.start,
+					toCell: range.end
+				}));
+			});
+			resultString = Array.from(headers.values()).join(valueSeparator).concat(eol);
+		}
+
+		const batchResult: string[] = [];
+		for (const range of rowRanges) {
+			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));
+				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)
+			}
+		}
+		if (!isCanceled) {
+			resultString += batchResult.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) => {

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

@@ -55,6 +55,7 @@ export interface IQueryManagementService {
 	changeConnectionUri(newUri: string, oldUri: string): Promise<void>;
 	saveResults(requestParams: azdata.SaveResultsRequestParams): Promise<azdata.SaveResultRequestResult>;
 	setQueryExecutionOptions(uri: string, options: azdata.QueryExecutionOptions): Promise<void>;
+	copyResults(params: azdata.CopyResultsRequestParams): Promise<void>;
 
 	// Callbacks
 	onQueryComplete(result: azdata.QueryExecuteCompleteNotificationResult): void;
@@ -93,6 +94,7 @@ export interface IQueryRequestHandler {
 	disposeQuery(ownerUri: string): Promise<void>;
 	connectionUriChanged(newUri: string, oldUri: string): Promise<void>;
 	saveResults(requestParams: azdata.SaveResultsRequestParams): Promise<azdata.SaveResultRequestResult>;
+	copyResults(requestParams: azdata.CopyResultsRequestParams): Promise<void>;
 	setQueryExecutionOptions(ownerUri: string, options: azdata.QueryExecutionOptions): Promise<void>;
 
 	// Edit Data actions
@@ -311,6 +313,12 @@ export class QueryManagementService implements IQueryManagementService {
 		});
 	}
 
+	public copyResults(requestParams: azdata.CopyResultsRequestParams): Promise<void> {
+		return this._runAction(requestParams.ownerUri, (runner) => {
+			return runner.copyResults(requestParams);
+		});
+	}
+
 	public onQueryComplete(result: azdata.QueryExecuteCompleteNotificationResult): void {
 		this._notify(result.ownerUri, (runner: QueryRunner) => {
 			runner.handleQueryComplete(result.batchSummaries.map(s => ({ ...s, range: selectionDataToRange(s.selection) })));

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

@@ -21,7 +21,7 @@ import { IConfigurationService } from 'vs/platform/configuration/common/configur
 import { URI } from 'vs/base/common/uri';
 import * as perf from 'vs/base/common/performance';
 import { mssqlProviderName } from 'sql/platform/connection/common/constants';
-import { IGridDataProvider, copySelectionToClipboard, getTableHeaderString } from 'sql/workbench/services/query/common/gridDataProvider';
+import { IGridDataProvider, copySelectionToClipboard, executeCopyWithNotification, getTableHeaderString } from 'sql/workbench/services/query/common/gridDataProvider';
 import { getErrorMessage } from 'vs/base/common/errors';
 import { ILogService } from 'vs/platform/log/common/log';
 import { IRange, Range } from 'vs/editor/common/core/range';
@@ -29,6 +29,7 @@ import { BatchSummary, IQueryMessage, ResultSetSummary, QueryExecuteSubsetParams
 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';
 
 /*
 * Query Runner class which handles running a query, reports the results to the content manager,
@@ -460,17 +461,30 @@ export default class QueryRunner extends Disposable {
 
 	/**
 	 * Sends a copy request
-	 * @param selection The selection range to copy
+	 * @param selections The selection range to copy
 	 * @param batchId The batch id of the result to copy from
 	 * @param resultId The result id of the result to copy from
+	 * @param removeNewLines Whether to remove line breaks from values.
 	 * @param includeHeaders [Optional]: Should column headers be included in the copy selection
 	 */
-	async copyResults(selection: Slick.Range[], batchId: number, resultId: number, includeHeaders?: boolean): Promise<void> {
-		let provider = this.getGridDataProvider(batchId, resultId);
-		return provider.copyResults(selection, includeHeaders);
+	async copyResults(selections: Slick.Range[], batchId: number, resultId: number, removeNewLines: boolean, includeHeaders?: boolean): Promise<void> {
+		await this.queryManagementService.copyResults({
+			ownerUri: this.uri,
+			batchIndex: batchId,
+			resultSetIndex: resultId,
+			removeNewLines: removeNewLines,
+			includeHeaders: includeHeaders,
+			selections: selections.map(selection => {
+				return {
+					fromRow: selection.fromRow,
+					toRow: selection.toRow,
+					fromColumn: selection.fromCell,
+					toColumn: selection.toCell
+				};
+			})
+		});
 	}
 
-
 	public getColumnHeaders(batchId: number, resultId: number, range: Slick.Range): string[] | undefined {
 		let headers: string[] | undefined = undefined;
 		let batchSummary: BatchSummary = this._batchSets[batchId];
@@ -547,7 +561,8 @@ export class QueryGridDataProvider implements IGridDataProvider {
 		@INotificationService private _notificationService: INotificationService,
 		@IClipboardService private _clipboardService: IClipboardService,
 		@IConfigurationService private _configurationService: IConfigurationService,
-		@ITextResourcePropertiesService private _textResourcePropertiesService: ITextResourcePropertiesService
+		@ITextResourcePropertiesService private _textResourcePropertiesService: ITextResourcePropertiesService,
+		@ICapabilitiesService private _capabilitiesService: ICapabilitiesService
 	) {
 	}
 
@@ -559,14 +574,27 @@ export class QueryGridDataProvider implements IGridDataProvider {
 		return this.copyResultsAsync(selection, includeHeaders, tableView);
 	}
 
-	private async copyResultsAsync(selection: Slick.Range[], includeHeaders?: boolean, tableView?: IDisposableDataProvider<Slick.SlickData>): Promise<void> {
+	private async copyResultsAsync(selections: Slick.Range[], includeHeaders?: boolean, tableView?: IDisposableDataProvider<Slick.SlickData>): Promise<void> {
 		try {
-			await copySelectionToClipboard(this._clipboardService, this._notificationService, this, selection, includeHeaders, tableView);
+			const providerId = this.queryRunner.getProviderId();
+			const providerSupportCopyResults = this._capabilitiesService.getCapabilities(providerId).connection.supportCopyResultsToClipboard;
+			const preferProvidersCopyHandler = this._configurationService.getValue<IQueryEditorConfiguration>('queryEditor').results.preferProvidersCopyHandler;
+			if (preferProvidersCopyHandler && providerSupportCopyResults && (tableView === undefined || !tableView.isDataInMemory)) {
+				await this.handleCopyRequestByProvider(selections, includeHeaders);
+			} else {
+				await copySelectionToClipboard(this._clipboardService, this._notificationService, this, selections, includeHeaders, tableView);
+			}
 		} catch (error) {
 			this._notificationService.error(nls.localize('copyFailed', "Copy failed with error: {0}", getErrorMessage(error)));
 		}
 	}
 
+	private async handleCopyRequestByProvider(selections: Slick.Range[], includeHeaders?: boolean): Promise<void> {
+		executeCopyWithNotification(this._notificationService, selections, false, async () => {
+			await this.queryRunner.copyResults(selections, this.batchId, this.resultSetId, this.shouldRemoveNewLines(), this.shouldIncludeHeaders(includeHeaders));
+		});
+	}
+
 	async copyHeaders(selection: Slick.Range[]): Promise<void> {
 		try {
 			const results = getTableHeaderString(this, selection);

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

@@ -62,6 +62,9 @@ export class TestQueryManagementService implements IQueryManagementService {
 	saveResults(requestParams: azdata.SaveResultsRequestParams): Promise<azdata.SaveResultRequestResult> {
 		throw new Error('Method not implemented.');
 	}
+	copyResults(params: azdata.CopyResultsRequestParams): Promise<void> {
+		throw new Error('Method not implemented.');
+	}
 	setQueryExecutionOptions(uri: string, options: azdata.QueryExecutionOptions): Promise<void> {
 		throw new Error('Method not implemented.');
 	}