Support Save As CSV/JSON/Excel/XML from Notebooks (#6627)

Updated existing serialization code so it actually supports serialization. Still needs work to replace the saveAs function when a QueryProvider doesn't support save as, but want to handle in separate PR.
Removed separate MainThread/ExtHostSerializationProvider code as the DataProtocol code is the right place to put this code
Plumbed support through the gridOutputComponent to use the new serialize method in the serialization provider
Refactored the resultSerializer so majority of code can be shared between both implementations (for example file save dialog -> save -> show file on completion)

* Update to latest SQLToolsService release

src/sql/workbench/api/browser/extensionHost.contribution.common.ts

@@ -19,5 +19,4 @@ import './mainThreadNotebookDocumentsAndEditors';
 import './mainThreadObjectExplorer';
 import './mainThreadQueryEditor';
 import './mainThreadResourceProvider';
-import './mainThreadSerializationProvider';
 import './mainThreadTasks';

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

@@ -134,11 +134,11 @@ export class MainThreadDataProtocol implements MainThreadDataProtocolShape {
 				return self._proxy.$disposeQuery(handle, ownerUri);
 			},
 			saveResults(requestParams: azdata.SaveResultsRequestParams): Thenable<azdata.SaveResultRequestResult> {
-				let serializationProvider = self._serializationService.getSerializationFeatureMetadataProvider(requestParams.ownerUri);
-				if (serializationProvider && serializationProvider.enabled) {
+				let saveResultsFeatureInfo = self._serializationService.getSaveResultsFeatureMetadataProvider(requestParams.ownerUri);
+				if (saveResultsFeatureInfo && saveResultsFeatureInfo.enabled) {
 					return self._proxy.$saveResults(handle, requestParams);
 				}
-				else if (serializationProvider && !serializationProvider.enabled) {
+				else if (saveResultsFeatureInfo && !saveResultsFeatureInfo.enabled) {
 					return self._serializationService.disabledSaveAs();
 				}
 				else {
@@ -499,6 +499,20 @@ export class MainThreadDataProtocol implements MainThreadDataProtocolShape {
 		return undefined;
 	}
 
+	public $registerSerializationProvider(providerId: string, handle: number): Promise<any> {
+		const self = this;
+		this._serializationService.registerProvider(providerId, <azdata.SerializationProvider>{
+			startSerialization(requestParams: azdata.SerializeDataStartRequestParams): Thenable<azdata.SerializeDataResult> {
+				return self._proxy.$startSerialization(handle, requestParams);
+			},
+			continueSerialization(requestParams: azdata.SerializeDataContinueRequestParams): Thenable<azdata.SerializeDataResult> {
+				return self._proxy.$continueSerialization(handle, requestParams);
+			},
+		});
+
+		return undefined;
+	}
+
 	// Connection Management handlers
 	public $onConnectionComplete(handle: number, connectionInfoSummary: azdata.ConnectionInfoSummary): void {
 		this._connectionManagementService.onConnectionComplete(handle, connectionInfoSummary);

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

@@ -1,59 +0,0 @@
-/*---------------------------------------------------------------------------------------------
- *  Copyright (c) Microsoft Corporation. All rights reserved.
- *  Licensed under the Source EULA. See License.txt in the project root for license information.
- *--------------------------------------------------------------------------------------------*/
-
-import { IDisposable, dispose } from 'vs/base/common/lifecycle';
-import {
-	SqlExtHostContext, ExtHostSerializationProviderShape,
-	MainThreadSerializationProviderShape, SqlMainContext
-} from 'sql/workbench/api/common/sqlExtHost.protocol';
-import { ISerializationService } from 'sql/platform/serialization/common/serializationService';
-import * as azdata from 'azdata';
-import { IExtHostContext } from 'vs/workbench/api/common/extHost.protocol';
-import { extHostNamedCustomer } from 'vs/workbench/api/common/extHostCustomers';
-
-@extHostNamedCustomer(SqlMainContext.MainThreadSerializationProvider)
-export class MainThreadSerializationProvider implements MainThreadSerializationProviderShape {
-
-	private _proxy: ExtHostSerializationProviderShape;
-
-	private _toDispose: IDisposable[];
-
-	private _registrations: { [handle: number]: IDisposable; } = Object.create(null);
-
-	constructor(
-		extHostContext: IExtHostContext,
-		@ISerializationService private serializationService: ISerializationService
-
-	) {
-		if (extHostContext) {
-			this._proxy = extHostContext.getProxy(SqlExtHostContext.ExtHostSerializationProvider);
-		}
-	}
-
-	public dispose(): void {
-		this._toDispose = dispose(this._toDispose);
-	}
-
-	public $registerSerializationProvider(handle: number): Promise<any> {
-		let self = this;
-
-		this._registrations[handle] = this.serializationService.addEventListener(handle, {
-			onSaveAs(saveFormat: string, savePath: string, results: string, appendToFile: boolean): Thenable<azdata.SaveResultRequestResult> {
-				return self._proxy.$saveAs(saveFormat, savePath, results, appendToFile);
-			}
-		});
-
-		return undefined;
-	}
-
-	public $unregisterSerializationProvider(handle: number): Promise<any> {
-		let registration = this._registrations[handle];
-		if (registration) {
-			registration.dispose();
-			delete this._registrations[handle];
-		}
-		return undefined;
-	}
-}

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

@@ -181,6 +181,12 @@ export class ExtHostDataProtocol extends ExtHostDataProtocolShape {
 		return rt;
 	}
 
+	$registerSerializationProvider(provider: azdata.SerializationProvider): vscode.Disposable {
+		let rt = this.registerProvider(provider, DataProviderType.QueryProvider);
+		this._proxy.$registerSerializationProvider(provider.providerId, provider.handle);
+		return rt;
+	}
+
 	// Capabilities Discovery handlers
 	$getServerCapabilities(handle: number, client: azdata.DataProtocolClientCapabilities): Thenable<azdata.DataProtocolServerCapabilities> {
 		return this._resolveProvider<azdata.CapabilitiesProvider>(handle).getServerCapabilities(client);
@@ -750,4 +756,13 @@ export class ExtHostDataProtocol extends ExtHostDataProtocolShape {
 	public $onJobDataUpdated(handle: Number): void {
 		this._proxy.$onJobDataUpdated(handle);
 	}
+
+	// Serialization methods
+	public $startSerialization(handle: number, requestParams: azdata.SerializeDataStartRequestParams): Thenable<azdata.SerializeDataResult> {
+		return this._resolveProvider<azdata.SerializationProvider>(handle).startSerialization(requestParams);
+	}
+
+	public $continueSerialization(handle: number, requestParams: azdata.SerializeDataContinueRequestParams): Thenable<azdata.SerializeDataResult> {
+		return this._resolveProvider<azdata.SerializationProvider>(handle).continueSerialization(requestParams);
+	}
 }

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

@@ -1,71 +0,0 @@
-/*---------------------------------------------------------------------------------------------
- *  Copyright (c) Microsoft Corporation. All rights reserved.
- *  Licensed under the Source EULA. See License.txt in the project root for license information.
- *--------------------------------------------------------------------------------------------*/
-
-import { IMainContext } from 'vs/workbench/api/common/extHost.protocol';
-import { SqlMainContext, MainThreadSerializationProviderShape, ExtHostSerializationProviderShape } from 'sql/workbench/api/common/sqlExtHost.protocol';
-import * as vscode from 'vscode';
-import * as azdata from 'azdata';
-import { Disposable } from 'vs/workbench/api/common/extHostTypes';
-
-class SerializationAdapter {
-	private _provider: azdata.SerializationProvider;
-
-	constructor(provider: azdata.SerializationProvider) {
-		this._provider = provider;
-	}
-
-	public saveAs(saveFormat: string, savePath: string, results: string, appendToFile: boolean): Thenable<azdata.SaveResultRequestResult> {
-		return this._provider.saveAs(saveFormat, savePath, results, appendToFile);
-	}
-
-}
-
-type Adapter = SerializationAdapter;
-
-export class ExtHostSerializationProvider extends ExtHostSerializationProviderShape {
-
-	private _proxy: MainThreadSerializationProviderShape;
-
-	private static _handlePool: number = 0;
-	private _adapter: { [handle: number]: Adapter } = Object.create(null);
-
-	private _createDisposable(handle: number): Disposable {
-		return new Disposable(() => {
-			delete this._adapter[handle];
-			this._proxy.$unregisterSerializationProvider(handle);
-		});
-	}
-
-	private _nextHandle(): number {
-		return ExtHostSerializationProvider._handlePool++;
-	}
-
-	private _withAdapter<A, R>(handle: number, ctor: { new(...args: any[]): A }, callback: (adapter: A) => Thenable<R>): Thenable<R> {
-		let adapter = this._adapter[handle];
-		if (!(adapter instanceof ctor)) {
-			return Promise.reject(new Error('no adapter found'));
-		}
-		return callback(<any>adapter);
-	}
-
-	constructor(
-		mainContext: IMainContext
-	) {
-		super();
-		this._proxy = mainContext.getProxy(SqlMainContext.MainThreadSerializationProvider);
-	}
-
-	public $registerSerializationProvider(provider: azdata.SerializationProvider): vscode.Disposable {
-		provider.handle = this._nextHandle();
-		this._adapter[provider.handle] = new SerializationAdapter(provider);
-		this._proxy.$registerSerializationProvider(provider.handle);
-		return this._createDisposable(provider.handle);
-	}
-
-	public $saveAs(saveFormat: string, savePath: string, results: string, appendToFile: boolean): Thenable<azdata.SaveResultRequestResult> {
-		return this._withAdapter(0, SerializationAdapter, adapter => adapter.saveAs(saveFormat, savePath, results, appendToFile));
-	}
-
-}

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

@@ -504,6 +504,15 @@ export abstract class ExtHostDataProtocolShape {
 	 */
 	$schemaCompareCancel(handle: number, operationId: string): Thenable<azdata.ResultStatus> { throw ni(); }
 
+	/**
+	 * Serialization start request
+	 */
+	$startSerialization(handle: number, requestParams: azdata.SerializeDataStartRequestParams): Thenable<azdata.SerializeDataResult> { throw ni(); }
+
+	/**
+	 * Serialization continuation request
+	 */
+	$continueSerialization(handle: number, requestParams: azdata.SerializeDataContinueRequestParams): Thenable<azdata.SerializeDataResult> { throw ni(); }
 }
 
 /**
@@ -533,13 +542,6 @@ export abstract class ExtHostCredentialManagementShape {
 	$deleteCredential(credentialId: string): Thenable<boolean> { throw ni(); }
 }
 
-/**
- * Serialization provider extension host class.
- */
-export abstract class ExtHostSerializationProviderShape {
-	$saveAs(saveFormat: string, savePath: string, results: string, appendToFile: boolean): Thenable<azdata.SaveResultRequestResult> { throw ni(); }
-}
-
 export interface MainThreadAccountManagementShape extends IDisposable {
 	$registerAccountProvider(providerMetadata: azdata.AccountProviderMetadata, handle: number): Thenable<any>;
 	$unregisterAccountProvider(handle: number): Thenable<any>;
@@ -575,6 +577,7 @@ export interface MainThreadDataProtocolShape extends IDisposable {
 	$registerAgentServicesProvider(providerId: string, handle: number): Promise<any>;
 	$registerDacFxServicesProvider(providerId: string, handle: number): Promise<any>;
 	$registerSchemaCompareServicesProvider(providerId: string, handle: number): Promise<any>;
+	$registerSerializationProvider(providerId: string, handle: number): Promise<any>;
 	$unregisterProvider(handle: number): Promise<any>;
 	$onConnectionComplete(handle: number, connectionInfoSummary: azdata.ConnectionInfoSummary): void;
 	$onIntelliSenseCacheComplete(handle: number, connectionUri: string): void;
@@ -625,11 +628,6 @@ export interface MainThreadCredentialManagementShape extends IDisposable {
 	$unregisterCredentialProvider(handle: number): Promise<any>;
 }
 
-export interface MainThreadSerializationProviderShape extends IDisposable {
-	$registerSerializationProvider(handle: number): Promise<any>;
-	$unregisterSerializationProvider(handle: number): Promise<any>;
-}
-
 function ni() { return new Error('Not implemented'); }
 
 // --- proxy identifiers
@@ -642,7 +640,6 @@ export const SqlMainContext = {
 	MainThreadDataProtocol: createMainId<MainThreadDataProtocolShape>('MainThreadDataProtocol'),
 	MainThreadObjectExplorer: createMainId<MainThreadObjectExplorerShape>('MainThreadObjectExplorer'),
 	MainThreadBackgroundTaskManagement: createMainId<MainThreadBackgroundTaskManagementShape>('MainThreadBackgroundTaskManagement'),
-	MainThreadSerializationProvider: createMainId<MainThreadSerializationProviderShape>('MainThreadSerializationProvider'),
 	MainThreadResourceProvider: createMainId<MainThreadResourceProviderShape>('MainThreadResourceProvider'),
 	MainThreadModalDialog: createMainId<MainThreadModalDialogShape>('MainThreadModalDialog'),
 	MainThreadTasks: createMainId<MainThreadTasksShape>('MainThreadTasks'),
@@ -662,7 +659,6 @@ export const SqlExtHostContext = {
 	ExtHostCredentialManagement: createExtId<ExtHostCredentialManagementShape>('ExtHostCredentialManagement'),
 	ExtHostDataProtocol: createExtId<ExtHostDataProtocolShape>('ExtHostDataProtocol'),
 	ExtHostObjectExplorer: createExtId<ExtHostObjectExplorerShape>('ExtHostObjectExplorer'),
-	ExtHostSerializationProvider: createExtId<ExtHostSerializationProviderShape>('ExtHostSerializationProvider'),
 	ExtHostResourceProvider: createExtId<ExtHostResourceProviderShape>('ExtHostResourceProvider'),
 	ExtHostModalDialogs: createExtId<ExtHostModalDialogsShape>('ExtHostModalDialogs'),
 	ExtHostTasks: createExtId<ExtHostTasksShape>('ExtHostTasks'),

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

@@ -328,7 +328,8 @@ export enum DataProviderType {
 	DacFxServicesProvider = 'DacFxServicesProvider',
 	SchemaCompareServicesProvider = 'SchemaCompareServicesProvider',
 	ObjectExplorerNodeProvider = 'ObjectExplorerNodeProvider',
-	IconProvider = 'IconProvider'
+	IconProvider = 'IconProvider',
+	SerializationProvider = 'SerializationProvider'
 }
 
 export enum DeclarativeDataType {
@@ -686,4 +687,4 @@ export enum NotebookChangeKind {
 	MetadataUpdated = 1,
 	Save = 2,
 	CellExecuted = 3
-}
+}

src/sql/workbench/api/electron-browser/extensionHost.contribution.ts

@@ -9,7 +9,6 @@ import 'sql/workbench/api/browser/mainThreadCredentialManagement';
 import 'sql/workbench/api/browser/mainThreadDataProtocol';
 import 'sql/workbench/api/browser/mainThreadObjectExplorer';
 import 'sql/workbench/api/browser/mainThreadBackgroundTaskManagement';
-import 'sql/workbench/api/browser/mainThreadSerializationProvider';
 import 'sql/workbench/api/browser/mainThreadResourceProvider';
 import 'sql/workbench/api/browser/mainThreadTasks';
 import 'sql/workbench/api/browser/mainThreadDashboard';

src/sql/workbench/api/node/sqlExtHost.api.impl.ts

@@ -15,7 +15,6 @@ import { SqlExtHostContext } from 'sql/workbench/api/common/sqlExtHost.protocol'
 import { ExtHostAccountManagement } from 'sql/workbench/api/common/extHostAccountManagement';
 import { ExtHostCredentialManagement } from 'sql/workbench/api/common/extHostCredentialManagement';
 import { ExtHostDataProtocol } from 'sql/workbench/api/common/extHostDataProtocol';
-import { ExtHostSerializationProvider } from 'sql/workbench/api/common/extHostSerializationProvider';
 import { ExtHostResourceProvider } from 'sql/workbench/api/common/extHostResourceProvider';
 import * as sqlExtHostTypes from 'sql/workbench/api/common/sqlExtHostTypes';
 import { ExtHostModalDialogs } from 'sql/workbench/api/common/extHostModalDialog';
@@ -71,7 +70,6 @@ export function createApiFactory(
 	const extHostCredentialManagement = rpcProtocol.set(SqlExtHostContext.ExtHostCredentialManagement, new ExtHostCredentialManagement(rpcProtocol));
 	const extHostDataProvider = rpcProtocol.set(SqlExtHostContext.ExtHostDataProtocol, new ExtHostDataProtocol(rpcProtocol, uriTransformer));
 	const extHostObjectExplorer = rpcProtocol.set(SqlExtHostContext.ExtHostObjectExplorer, new ExtHostObjectExplorer(rpcProtocol));
-	const extHostSerializationProvider = rpcProtocol.set(SqlExtHostContext.ExtHostSerializationProvider, new ExtHostSerializationProvider(rpcProtocol));
 	const extHostResourceProvider = rpcProtocol.set(SqlExtHostContext.ExtHostResourceProvider, new ExtHostResourceProvider(rpcProtocol));
 	const extHostModalDialogs = rpcProtocol.set(SqlExtHostContext.ExtHostModalDialogs, new ExtHostModalDialogs(rpcProtocol));
 	const extHostTasks = rpcProtocol.set(SqlExtHostContext.ExtHostTasks, new ExtHostTasks(rpcProtocol, logService));
@@ -188,14 +186,7 @@ export function createApiFactory(
 				}
 			};
 
-			// namespace: serialization
-			const serialization: typeof azdata.serialization = {
-				registerProvider(provider: azdata.SerializationProvider): vscode.Disposable {
-					return extHostSerializationProvider.$registerSerializationProvider(provider);
-				},
-			};
-
-			// namespace: serialization
+			// namespace: resources
 			const resources: typeof azdata.resources = {
 				registerResourceProvider(providerMetadata: azdata.ResourceProviderMetadata, provider: azdata.ResourceProvider): vscode.Disposable {
 					return extHostResourceProvider.$registerResourceProvider(providerMetadata, provider);
@@ -369,6 +360,10 @@ export function createApiFactory(
 				return extHostDataProvider.$registerSchemaCompareServiceProvider(provider);
 			};
 
+			let registerSerializationProvider = (provider: azdata.SerializationProvider): vscode.Disposable => {
+				return extHostDataProvider.$registerSerializationProvider(provider);
+			};
+
 			// namespace: dataprotocol
 			const dataprotocol: typeof azdata.dataprotocol = {
 				registerBackupProvider,
@@ -388,6 +383,7 @@ export function createApiFactory(
 				registerCapabilitiesServiceProvider,
 				registerDacFxServicesProvider,
 				registerSchemaCompareServicesProvider,
+				registerSerializationProvider,
 				onDidChangeLanguageFlavor(listener: (e: azdata.DidChangeLanguageFlavorParams) => any, thisArgs?: any, disposables?: extHostTypes.Disposable[]) {
 					return extHostDataProvider.onDidChangeLanguageFlavor(listener, thisArgs, disposables);
 				},
@@ -516,7 +512,6 @@ export function createApiFactory(
 				credentials,
 				objectexplorer: objectExplorer,
 				resources,
-				serialization,
 				dataprotocol,
 				DataProviderType: sqlExtHostTypes.DataProviderType,
 				DeclarativeDataType: sqlExtHostTypes.DeclarativeDataType,
@@ -630,7 +625,9 @@ export function createApiFactory(
 			// namespace: serialization
 			const serialization: typeof sqlops.serialization = {
 				registerProvider(provider: sqlops.SerializationProvider): vscode.Disposable {
-					return extHostSerializationProvider.$registerSerializationProvider(provider);
+					// No-op this to avoid breaks in existing applications. Tested on Github - no examples,
+					// but I think it's safer to avoid breaking this
+					return undefined;
 				},
 			};
 

src/sql/workbench/parts/notebook/browser/outputs/gridOutput.component.ts

@@ -29,6 +29,9 @@ import { MimeModel } from 'sql/workbench/parts/notebook/common/models/mimemodel'
 import { GridTableState } from 'sql/workbench/parts/query/common/gridPanelState';
 import { GridTableBase } from 'sql/workbench/parts/query/browser/gridPanel';
 import { getErrorMessage } from 'vs/base/common/errors';
+import { ISerializationService, SerializeDataParams } from 'sql/platform/serialization/common/serializationService';
+import { SaveResultAction } from 'sql/workbench/parts/query/browser/actions';
+import { ResultSerializer, SaveResultsResponse } from 'sql/workbench/parts/query/common/resultSerializer';
 
 @Component({
 	selector: GridOutputComponent.SELECTOR,
@@ -110,7 +113,8 @@ class DataResourceTable extends GridTableBase<any> {
 		@IInstantiationService instantiationService: IInstantiationService,
 		@IEditorService editorService: IEditorService,
 		@IUntitledEditorService untitledEditorService: IUntitledEditorService,
-		@IConfigurationService configurationService: IConfigurationService
+		@IConfigurationService configurationService: IConfigurationService,
+		@ISerializationService private _serializationService: ISerializationService
 	) {
 		super(state, createResultSet(source), contextMenuService, instantiationService, editorService, untitledEditorService, configurationService);
 		this._gridDataProvider = this.instantiationService.createInstance(DataResourceDataProvider, source, this.resultSet, documentUri);
@@ -125,7 +129,15 @@ class DataResourceTable extends GridTableBase<any> {
 	}
 
 	protected getContextActions(): IAction[] {
-		return [];
+		if (!this._serializationService.hasProvider()) {
+			return [];
+		}
+		return [
+			this.instantiationService.createInstance(SaveResultAction, SaveResultAction.SAVECSV_ID, SaveResultAction.SAVECSV_LABEL, SaveResultAction.SAVECSV_ICON, SaveFormat.CSV),
+			this.instantiationService.createInstance(SaveResultAction, SaveResultAction.SAVEEXCEL_ID, SaveResultAction.SAVEEXCEL_LABEL, SaveResultAction.SAVEEXCEL_ICON, SaveFormat.EXCEL),
+			this.instantiationService.createInstance(SaveResultAction, SaveResultAction.SAVEJSON_ID, SaveResultAction.SAVEJSON_LABEL, SaveResultAction.SAVEJSON_ICON, SaveFormat.JSON),
+			this.instantiationService.createInstance(SaveResultAction, SaveResultAction.SAVEXML_ID, SaveResultAction.SAVEXML_LABEL, SaveResultAction.SAVEXML_ICON, SaveFormat.XML),
+		];
 	}
 
 	public get maximumSize(): number {
@@ -143,7 +155,9 @@ class DataResourceDataProvider implements IGridDataProvider {
 		@INotificationService private _notificationService: INotificationService,
 		@IClipboardService private _clipboardService: IClipboardService,
 		@IConfigurationService private _configurationService: IConfigurationService,
-		@ITextResourcePropertiesService private _textResourcePropertiesService: ITextResourcePropertiesService
+		@ITextResourcePropertiesService private _textResourcePropertiesService: ITextResourcePropertiesService,
+		@ISerializationService private _serializationService: ISerializationService,
+		@IInstantiationService private _instantiationService: IInstantiationService
 	) {
 		this.transformSource(source);
 	}
@@ -210,14 +224,67 @@ class DataResourceDataProvider implements IGridDataProvider {
 	}
 
 	get canSerialize(): boolean {
-		return false;
+		return this._serializationService.hasProvider();
 	}
 
+
 	serializeResults(format: SaveFormat, selection: Slick.Range[]): Thenable<void> {
-		throw new Error('Method not implemented.');
+		let serializer = this._instantiationService.createInstance(ResultSerializer);
+		return serializer.handleSerialization(this.documentUri, format, (filePath) => this.doSerialize(serializer, filePath, format, selection));
+	}
+
+	private async doSerialize(serializer: ResultSerializer, filePath: string, format: SaveFormat, selection: Slick.Range[]): Promise<SaveResultsResponse> {
+		// TODO implement selection support
+		let columns = this.resultSet.columnInfo;
+		let rowLength = this.rows.length;
+		let minRow = 0;
+		let maxRow = this.rows.length;
+		let singleSelection = selection && selection.length > 0 ? selection[0] : undefined;
+		if (singleSelection && this.isSelected(singleSelection)) {
+			rowLength = singleSelection.toRow - singleSelection.fromRow + 1;
+			minRow = singleSelection.fromRow;
+			maxRow = singleSelection.toRow + 1;
+			columns = columns.slice(singleSelection.fromCell, singleSelection.toCell + 1);
+		}
+		let getRows: ((index: number, rowCount: number) => azdata.DbCellValue[][]) = (index, rowCount) => {
+			// Offset for selections by adding the selection startRow to the index
+			index = index + minRow;
+			if (rowLength === 0 || index < 0 || index >= maxRow) {
+				return [];
+			}
+			let endIndex = index + rowCount;
+			if (endIndex > maxRow) {
+				endIndex = maxRow;
+			}
+			let result = this.rows.slice(index, endIndex).map(row => {
+				if (this.isSelected(singleSelection)) {
+					return row.slice(singleSelection.fromCell, singleSelection.toCell + 1);
+				}
+				return row;
+			});
+			return result;
+		};
+
+		let serializeRequestParams: SerializeDataParams = <SerializeDataParams>Object.assign(serializer.getBasicSaveParameters(format), <Partial<SerializeDataParams>>{
+			saveFormat: format,
+			columns: columns,
+			filePath: filePath,
+			getRowRange: (rowStart, numberOfRows) => getRows(rowStart, numberOfRows),
+			rowCount: rowLength
+		});
+		let result = await this._serializationService.serializeResults(serializeRequestParams);
+		return result;
+	}
+
+	/**
+	 * Check if a range of cells were selected.
+	 */
+	private isSelected(selection: Slick.Range): boolean {
+		return (selection && !((selection.fromCell === selection.toCell) && (selection.fromRow === selection.toRow)));
 	}
 }
 
+
 function createResultSet(source: IDataResource): azdata.ResultSetSummary {
 	let columnInfo: azdata.IDbColumn[] = source.schema.fields.map(field => {
 		let column = new SimpleDbColumn(field.name);

src/sql/workbench/parts/query/browser/actions.ts

@@ -23,6 +23,7 @@ import * as Constants from 'sql/workbench/contrib/extensions/common/constants';
 import { IAdsTelemetryService } from 'sql/platform/telemetry/common/telemetry';
 import * as TelemetryKeys from 'sql/platform/telemetry/common/telemetryKeys';
 import { IEnvironmentService } from 'vs/platform/environment/common/environment';
+import { getErrorMessage } from 'vs/base/common/errors';
 
 export interface IGridActionContext {
 	gridDataProvider: IGridDataProvider;
@@ -79,7 +80,12 @@ export class SaveResultAction extends Action {
 		if (!context.gridDataProvider.canSerialize) {
 			this.notificationService.warn(localize('saveToFileNotSupported', "Save to file is not supported by the backing data source"));
 		}
-		await context.gridDataProvider.serializeResults(this.format, mapForNumberColumn(context.selection));
+		try {
+			await context.gridDataProvider.serializeResults(this.format, mapForNumberColumn(context.selection));
+		} catch (error) {
+			this.notificationService.error(getErrorMessage(error));
+			return false;
+		}
 		return true;
 	}
 }

src/sql/workbench/parts/query/common/resultSerializer.ts

@@ -29,6 +29,12 @@ import { IFileDialogService } from 'vs/platform/dialogs/common/dialogs';
 
 let prevSavePath: string;
 
+
+export interface SaveResultsResponse {
+	succeeded: boolean;
+	messages: string;
+}
+
 interface ICsvConfig {
 	includeHeaders: boolean;
 	delimiter: string;
@@ -48,9 +54,6 @@ interface IXmlConfig {
 export class ResultSerializer {
 	public static tempFileCount: number = 1;
 
-	private _uri: string;
-	private _filePath: string;
-
 	constructor(
 		@IOutputService private _outputService: IOutputService,
 		@IQueryManagementService private _queryManagementService: IQueryManagementService,
@@ -67,12 +70,38 @@ export class ResultSerializer {
 	 */
 	public saveResults(uri: string, saveRequest: ISaveRequest): Thenable<void> {
 		const self = this;
-		this._uri = uri;
-
-		// prompt for filepath
-		return self.promptForFilepath(saveRequest).then(filePath => {
+		return this.promptForFilepath(saveRequest.format, uri).then(filePath => {
 			if (filePath) {
-				return self.sendRequestToService(filePath, saveRequest.batchIndex, saveRequest.resultSetNumber, saveRequest.format, saveRequest.selection ? saveRequest.selection[0] : undefined);
+				if (!path.isAbsolute(filePath)) {
+					filePath = resolveFilePath(uri, filePath, this.rootPath);
+				}
+				let saveResultsParams = this.getParameters(uri, filePath, saveRequest.batchIndex, saveRequest.resultSetNumber, saveRequest.format, saveRequest.selection ? saveRequest.selection[0] : undefined);
+				let sendRequest = () => this.sendSaveRequestToService(saveResultsParams);
+				return self.doSave(filePath, saveRequest.format, sendRequest);
+			}
+			return Promise.resolve(undefined);
+		});
+	}
+
+	private async sendSaveRequestToService(saveResultsParams: SaveResultsRequestParams): Promise<SaveResultsResponse> {
+		let result = await this._queryManagementService.saveResults(saveResultsParams);
+		return {
+			succeeded: !result.messages,
+			messages: result.messages
+		};
+	}
+
+	/**
+	 * Handle save request by getting filename from user and sending request to service
+	 */
+	public handleSerialization(uri: string, format: SaveFormat, sendRequest: ((filePath: string) => Promise<SaveResultsResponse>)): Thenable<void> {
+		const self = this;
+		return this.promptForFilepath(format, uri).then(filePath => {
+			if (filePath) {
+				if (!path.isAbsolute(filePath)) {
+					filePath = resolveFilePath(uri, filePath, this.rootPath);
+				}
+				return self.doSave(filePath, format, () => sendRequest(filePath));
 			}
 			return Promise.resolve(undefined);
 		});
@@ -100,25 +129,26 @@ export class ResultSerializer {
 		this.outputChannel.append(message);
 	}
 
-	private promptForFilepath(saveRequest: ISaveRequest): Thenable<string> {
-		let filepathPlaceHolder = prevSavePath ? path.dirname(prevSavePath) : resolveCurrentDirectory(this._uri, this.rootPath);
+
+	private promptForFilepath(format: SaveFormat, resourceUri: string): Thenable<string> {
+		let filepathPlaceHolder = prevSavePath ? path.dirname(prevSavePath) : resolveCurrentDirectory(resourceUri, this.rootPath);
 		if (filepathPlaceHolder) {
-			filepathPlaceHolder = path.join(filepathPlaceHolder, this.getResultsDefaultFilename(saveRequest));
+			filepathPlaceHolder = path.join(filepathPlaceHolder, this.getResultsDefaultFilename(format));
 		}
 
 		return this.fileDialogService.showSaveDialog({
 			title: nls.localize('resultsSerializer.saveAsFileTitle', "Choose Results File"),
 			defaultUri: filepathPlaceHolder ? URI.file(filepathPlaceHolder) : undefined,
-			filters: this.getResultsFileExtension(saveRequest)
+			filters: this.getResultsFileExtension(format)
 		}).then(filePath => {
 			prevSavePath = filePath.fsPath;
 			return filePath.fsPath;
 		});
 	}
 
-	private getResultsDefaultFilename(saveRequest: ISaveRequest): string {
+	private getResultsDefaultFilename(format: SaveFormat): string {
 		let fileName = 'Results';
-		switch (saveRequest.format) {
+		switch (format) {
 			case SaveFormat.CSV:
 				fileName = fileName + '.csv';
 				break;
@@ -137,11 +167,11 @@ export class ResultSerializer {
 		return fileName;
 	}
 
-	private getResultsFileExtension(saveRequest: ISaveRequest): FileFilter[] {
+	private getResultsFileExtension(format: SaveFormat): FileFilter[] {
 		let fileFilters = new Array<FileFilter>();
 		let fileFilter: { extensions: string[]; name: string } = { extensions: undefined, name: undefined };
 
-		switch (saveRequest.format) {
+		switch (format) {
 			case SaveFormat.CSV:
 				fileFilter.name = nls.localize('resultsSerializer.saveAsFileExtensionCSVTitle', "CSV (Comma delimited)");
 				fileFilter.extensions = ['csv'];
@@ -167,6 +197,22 @@ export class ResultSerializer {
 		return fileFilters;
 	}
 
+	public getBasicSaveParameters(format: string): SaveResultsRequestParams {
+		let saveResultsParams: SaveResultsRequestParams;
+
+		if (format === SaveFormat.CSV) {
+			saveResultsParams = this.getConfigForCsv();
+		} else if (format === SaveFormat.JSON) {
+			saveResultsParams = this.getConfigForJson();
+		} else if (format === SaveFormat.EXCEL) {
+			saveResultsParams = this.getConfigForExcel();
+		} else if (format === SaveFormat.XML) {
+			saveResultsParams = this.getConfigForXml();
+		}
+		return saveResultsParams;
+	}
+
+
 	private getConfigForCsv(): SaveResultsRequestParams {
 		let saveResultsParams = <SaveResultsRequestParams>{ resultFormat: SaveFormat.CSV as string };
 
@@ -231,26 +277,11 @@ export class ResultSerializer {
 		return saveResultsParams;
 	}
 
-	private getParameters(filePath: string, batchIndex: number, resultSetNo: number, format: string, selection: Slick.Range): SaveResultsRequestParams {
-		let saveResultsParams: SaveResultsRequestParams;
-		if (!path.isAbsolute(filePath)) {
-			this._filePath = resolveFilePath(this._uri, filePath, this.rootPath);
-		} else {
-			this._filePath = filePath;
-		}
 
-		if (format === SaveFormat.CSV) {
-			saveResultsParams = this.getConfigForCsv();
-		} else if (format === SaveFormat.JSON) {
-			saveResultsParams = this.getConfigForJson();
-		} else if (format === SaveFormat.EXCEL) {
-			saveResultsParams = this.getConfigForExcel();
-		} else if (format === SaveFormat.XML) {
-			saveResultsParams = this.getConfigForXml();
-		}
-
-		saveResultsParams.filePath = this._filePath;
-		saveResultsParams.ownerUri = this._uri;
+	private getParameters(uri: string, filePath: string, batchIndex: number, resultSetNo: number, format: string, selection: Slick.Range): SaveResultsRequestParams {
+		let saveResultsParams = this.getBasicSaveParameters(format);
+		saveResultsParams.filePath = filePath;
+		saveResultsParams.ownerUri = uri;
 		saveResultsParams.resultSetIndex = resultSetNo;
 		saveResultsParams.batchIndex = batchIndex;
 		if (this.isSelected(selection)) {
@@ -297,13 +328,13 @@ export class ResultSerializer {
 	/**
 	 * Send request to sql tools service to save a result set
 	 */
-	private sendRequestToService(filePath: string, batchIndex: number, resultSetNo: number, format: string, selection: Slick.Range): Thenable<void> {
-		let saveResultsParams = this.getParameters(filePath, batchIndex, resultSetNo, format, selection);
+	private async doSave(filePath: string, format: string, sendRequest: () => Promise<SaveResultsResponse>): Promise<void> {
 
-		this.logToOutputChannel(LocalizedConstants.msgSaveStarted + this._filePath);
+		this.logToOutputChannel(LocalizedConstants.msgSaveStarted + filePath);
 
 		// send message to the sqlserverclient for converting results to the requested format and saving to filepath
-		return this._queryManagementService.saveResults(saveResultsParams).then(result => {
+		try {
+			let result = await sendRequest();
 			if (result.messages) {
 				this._notificationService.notify({
 					severity: Severity.Error,
@@ -311,20 +342,20 @@ export class ResultSerializer {
 				});
 				this.logToOutputChannel(LocalizedConstants.msgSaveFailed + result.messages);
 			} else {
-				this.promptFileSavedNotification(this._filePath);
+				this.promptFileSavedNotification(filePath);
 				this.logToOutputChannel(LocalizedConstants.msgSaveSucceeded + filePath);
-				this.openSavedFile(this._filePath, format);
+				this.openSavedFile(filePath, format);
 			}
 			// TODO telemetry for save results
 			// Telemetry.sendTelemetryEvent('SavedResults', { 'type': format });
 
-		}, error => {
+		} catch (error) {
 			this._notificationService.notify({
 				severity: Severity.Error,
 				message: LocalizedConstants.msgSaveFailed + error
 			});
 			this.logToOutputChannel(LocalizedConstants.msgSaveFailed + error);
-		});
+		}
 	}
 
 	/**