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/platform/serialization/common/serializationService.ts

@@ -8,25 +8,56 @@ import { createDecorator } from 'vs/platform/instantiation/common/instantiation'
 import { IConnectionManagementService } from 'sql/platform/connection/common/connectionManagement';
 import { ICapabilitiesService } from 'sql/platform/capabilities/common/capabilitiesService';
 import * as azdata from 'azdata';
+import { localize } from 'vs/nls';
+import { getErrorMessage } from 'vs/base/common/errors';
 
 export const SERVICE_ID = 'serializationService';
 
-export interface SerializationProviderEvents {
-	onSaveAs(saveFormat: string, savePath: string, results: string, appendToFile: boolean): Thenable<azdata.SaveResultRequestResult>;
-}
-
 export const ISerializationService = createDecorator<ISerializationService>(SERVICE_ID);
+const saveAsNotSupported = localize('saveAsNotSupported', "Saving results into different format disabled for this data provider.");
+const defaultBatchSize = 500;
+
+export interface SerializeDataParams {
+	/**
+	 * 'csv', 'json', 'excel', 'xml'
+	 */
+	saveFormat: string;
+	filePath: string;
+	/**
+	 * Gets an array of rows to be sent for serialization
+	 * @param rowStart Index in the array to start copying rows from
+	 * @param numberOfRows Total number of rows to copy. If 0 or undefined, will copy all
+	 */
+	getRowRange(rowStart: number, numberOfRows?: number): azdata.DbCellValue[][];
+	rowCount: number;
+	columns: azdata.IDbColumn[];
+	includeHeaders?: boolean;
+	delimiter?: string;
+	lineSeperator?: string;
+	textIdentifier?: string;
+	encoding?: string;
+	formatted?: boolean;
+}
 
 export interface ISerializationService {
 	_serviceBrand: any;
 
+	registerProvider(providerId: string, provider: azdata.SerializationProvider): void;
+
+	hasProvider(): boolean;
+
 	saveAs(saveFormat: string, savePath: string, results: string, appendToFile: boolean): Thenable<azdata.SaveResultRequestResult>;
 
 	disabledSaveAs(): Thenable<azdata.SaveResultRequestResult>;
 
-	addEventListener(handle: number, events: SerializationProviderEvents): IDisposable;
+	serializeResults(request: SerializeDataParams): Promise<azdata.SerializeDataResult>;
 
-	getSerializationFeatureMetadataProvider(ownerUri: string): azdata.FeatureMetadataProvider;
+	getSaveResultsFeatureMetadataProvider(ownerUri: string): azdata.FeatureMetadataProvider;
+}
+
+function getBatchSize(totalRows: number, currentIndex: number): number {
+	let rowsAvailable = totalRows - currentIndex;
+	return (defaultBatchSize < rowsAvailable) ? defaultBatchSize : rowsAvailable;
 }
 
 export class SerializationService implements ISerializationService {
@@ -35,9 +66,7 @@ export class SerializationService implements ISerializationService {
 
 	private disposables: IDisposable[] = [];
 
-	private _serverEvents: { [handle: number]: SerializationProviderEvents; } = Object.create(null);
-
-	private _lastHandle: number;
+	private providers: { providerId: string, provider: azdata.SerializationProvider }[] = [];
 
 	constructor(
 		@IConnectionManagementService private _connectionService: IConnectionManagementService,
@@ -45,31 +74,26 @@ export class SerializationService implements ISerializationService {
 	) {
 	}
 
-	public addEventListener(handle: number, events: SerializationProviderEvents): IDisposable {
-		this._lastHandle = handle;
-
-		this._serverEvents[handle] = events;
-
-		return {
-			dispose: () => {
-			}
-		};
+	public registerProvider(providerId: string, provider: azdata.SerializationProvider): void {
+		this.providers.push({ providerId: providerId, provider: provider });
 	}
 
-	public saveAs(saveFormat: string, savePath: string, results: string, appendToFile: boolean): Thenable<azdata.SaveResultRequestResult> {
-		if (this._serverEvents === undefined || this._serverEvents[this._lastHandle] === undefined) {
-			return this.disabledSaveAs();
-		}
+	public hasProvider(): boolean {
+		return this.providers.length > 0;
+	}
 
-		return this._serverEvents[this._lastHandle].onSaveAs(saveFormat, savePath, results, appendToFile);
+
+	public saveAs(saveFormat: string, savePath: string, results: string, appendToFile: boolean): Thenable<azdata.SaveResultRequestResult> {
+		// ideally, should read data from source & route to the serialization provider, but not implemented yet
+		return Promise.reject(new Error(saveAsNotSupported));
 	}
 
 	public disabledSaveAs(): Thenable<azdata.SaveResultRequestResult> {
-		return Promise.resolve({ messages: 'Saving results into different format disabled for this data provider.' });
+		return Promise.resolve({ messages: saveAsNotSupported });
 
 	}
 
-	public getSerializationFeatureMetadataProvider(ownerUri: string): azdata.FeatureMetadataProvider {
+	public getSaveResultsFeatureMetadataProvider(ownerUri: string): azdata.FeatureMetadataProvider {
 		let providerId: string = this._connectionService.getProviderIdFromUri(ownerUri);
 		let providerCapabilities = this._capabilitiesService.getLegacyCapabilities(providerId);
 
@@ -80,6 +104,100 @@ export class SerializationService implements ISerializationService {
 		return undefined;
 	}
 
+	public async serializeResults(serializationRequest: SerializeDataParams): Promise<azdata.SerializeDataResult> {
+		// Validate inputs
+		if (!serializationRequest) {
+			// Throwing here as this should only be a development time error
+			throw new Error('request data for serialization is missing');
+		}
+		if (!this.hasProvider()) {
+			return <azdata.SerializeDataResult>{
+				messages: localize('noSerializationProvider', "Cannot serialize data as no provider has been registered"),
+				succeeded: false
+			};
+		}
+		try {
+			// Create a new session with the provider and send initial data
+			let provider = this.providers[0].provider;
+			let index = 0;
+			let startRequestParams = this.createStartRequest(serializationRequest, index);
+			index = index + startRequestParams.rows.length;
+
+			let startResult = await provider.startSerialization(startRequestParams);
+
+			if (!startResult) {
+				return <azdata.SerializeDataResult>{
+					messages: localize('unknownSerializationError', "Serialization failed with an unknown error"),
+					succeeded: false
+				};
+			}
+			if (!startResult.succeeded) {
+				return startResult;
+			}
+			// Continue to send additional data
+			while (index < serializationRequest.rowCount) {
+				let continueRequestParams = this.createContinueRequest(serializationRequest, index);
+				index += continueRequestParams.rows.length;
+				let continueResult = await provider.continueSerialization(continueRequestParams);
+				if (!continueResult.succeeded) {
+					return continueResult;
+				}
+			}
+
+			// Complete the request
+			return <azdata.SerializeDataResult>{
+				messages: undefined,
+				succeeded: true
+			};
+		} catch (error) {
+			return <azdata.SerializeDataResult>{
+				messages: getErrorMessage(error),
+				succeeded: false
+			};
+		}
+	}
+
+	private createStartRequest(serializationRequest: SerializeDataParams, index: number): azdata.SerializeDataStartRequestParams {
+		let batchSize = getBatchSize(serializationRequest.rowCount, index);
+		let rows = serializationRequest.getRowRange(index, batchSize);
+		let columns: azdata.SimpleColumnInfo[] = serializationRequest.columns.map(c => {
+			// For now treat all as strings. In the future, would like to use the
+			// type info for correct data type mapping
+			let simpleCol: azdata.SimpleColumnInfo = {
+				name: c.columnName,
+				dataTypeName: 'NVarChar'
+			};
+			return simpleCol;
+		});
+		let isLastBatch = index + rows.length >= serializationRequest.rowCount;
+		let startSerializeRequest: azdata.SerializeDataStartRequestParams = {
+			saveFormat: serializationRequest.saveFormat,
+			filePath: serializationRequest.filePath,
+			columns: columns,
+			rows: rows,
+			isLastBatch: isLastBatch,
+			delimiter: serializationRequest.delimiter,
+			encoding: serializationRequest.encoding,
+			formatted: serializationRequest.formatted,
+			includeHeaders: serializationRequest.includeHeaders,
+			lineSeperator: serializationRequest.lineSeperator,
+			textIdentifier: serializationRequest.textIdentifier,
+		};
+		return startSerializeRequest;
+	}
+
+	private createContinueRequest(serializationRequest: SerializeDataParams, index: number): azdata.SerializeDataContinueRequestParams {
+		let numberOfRows = getBatchSize(serializationRequest.rowCount, index);
+		let rows = serializationRequest.getRowRange(index, numberOfRows);
+		let isLastBatch = index + rows.length >= serializationRequest.rowCount;
+		let continueSerializeRequest: azdata.SerializeDataContinueRequestParams = {
+			filePath: serializationRequest.filePath,
+			rows: rows,
+			isLastBatch: isLastBatch
+		};
+		return continueSerializeRequest;
+	}
+
 	public dispose(): void {
 		this.disposables = dispose(this.disposables);
 	}