From 33c6daaea1ea722bbb141eaa65a6c3ce27bac29c Mon Sep 17 00:00:00 2001
From: Cory Rivera <corivera@microsoft.com>
Date: Wed, 19 Oct 2022 15:37:41 -0700
Subject: [PATCH] Correctly route notebook serialization requests to SQL Tools
 Service. (#20874)

---
 .../serialization/common/serializationService.ts    | 12 +++++++++++-
 .../browser/outputs/gridOutput.component.ts         | 13 +++++++++++++
 .../test/browser/dataResourceDataProvider.test.ts   |  2 +-
 3 files changed, 25 insertions(+), 2 deletions(-)

diff --git a/src/sql/platform/serialization/common/serializationService.ts b/src/sql/platform/serialization/common/serializationService.ts
index f83366a765..4c8b2b5443 100644
--- a/src/sql/platform/serialization/common/serializationService.ts
+++ b/src/sql/platform/serialization/common/serializationService.ts
@@ -17,6 +17,10 @@ const saveAsNotSupported = localize('saveAsNotSupported', "Saving results into d
 const defaultBatchSize = 500;
 
 export interface SerializeDataParams {
+	/**
+	 *  The serializer to use for this request. Typically this is the ID of the connection provider used to run the query.
+	 */
+	serializationProviderId: string;
 	/**
 	 * 'csv', 'json', 'excel', 'xml'
 	 */
@@ -115,7 +119,13 @@ export class SerializationService implements ISerializationService {
 		}
 		try {
 			// Create a new session with the provider and send initial data
-			let provider = this.providers[0].provider;
+			let provider = this.providers.find(provider => provider.providerId === serializationRequest.serializationProviderId)?.provider;
+			if (!provider) {
+				return <azdata.SerializeDataResult>{
+					messages: localize('missingSerializationProviderError', "Could not find a serialization provider with the specified ID '{0}'", serializationRequest.serializationProviderId),
+					succeeded: false
+				};
+			}
 			let index = 0;
 			let startRequestParams = this.createStartRequest(serializationRequest, index);
 			index = index + startRequestParams.rows.length;
diff --git a/src/sql/workbench/contrib/notebook/browser/outputs/gridOutput.component.ts b/src/sql/workbench/contrib/notebook/browser/outputs/gridOutput.component.ts
index 4afcb191bd..ac92943ba8 100644
--- a/src/sql/workbench/contrib/notebook/browser/outputs/gridOutput.component.ts
+++ b/src/sql/workbench/contrib/notebook/browser/outputs/gridOutput.component.ts
@@ -48,6 +48,7 @@ import { getChartMaxRowCount, notifyMaxRowCountExceeded } from 'sql/workbench/co
 import { IAdsTelemetryService } from 'sql/platform/telemetry/common/telemetry';
 import * as TelemetryKeys from 'sql/platform/telemetry/common/telemetryKeys';
 import { IExecutionPlanService } from 'sql/workbench/services/executionPlan/common/interfaces';
+import { mssqlProviderName } from 'sql/platform/connection/common/constants';
 
 @Component({
 	selector: GridOutputComponent.SELECTOR,
@@ -487,8 +488,20 @@ export class DataResourceDataProvider implements IGridDataProvider {
 		// interface than the query execution service's saveResults handlers. Here, we take the
 		// format-specific request params (eg, includeHeaders for CSV) and merge the format-agnostic
 		// request params for the serialization request (eg, saveFormat, filePath).
+		let provider = this.cellModel.notebookModel.context?.providerName;
+		if (!provider) {
+			// If no connection currently exists, then pick the first connection provider for the current kernel.
+			// If there's still no provider, then fallback to the default MSSQL one.
+			let connProviders = this.cellModel.notebookModel.getApplicableConnectionProviderIds(this.cellModel.notebookModel.selectedKernelDisplayName);
+			if (connProviders?.length > 0) {
+				provider = connProviders[0];
+			} else {
+				provider = mssqlProviderName;
+			}
+		}
 		let formatSpecificParams = serializer.getBasicSaveParameters(format);
 		let formatAgnosticParams = <Partial<SerializeDataParams>>{
+			serializationProviderId: provider,
 			saveFormat: format,
 			filePath: filePath.fsPath,
 			columns: columns,
diff --git a/src/sql/workbench/contrib/notebook/test/browser/dataResourceDataProvider.test.ts b/src/sql/workbench/contrib/notebook/test/browser/dataResourceDataProvider.test.ts
index 745ff2763a..8f458d8c25 100644
--- a/src/sql/workbench/contrib/notebook/test/browser/dataResourceDataProvider.test.ts
+++ b/src/sql/workbench/contrib/notebook/test/browser/dataResourceDataProvider.test.ts
@@ -79,7 +79,7 @@ suite('Data Resource Data Provider', function () {
 		fileDialogService = new TestFileDialogService(pathService);
 		notificationService = new TestNotificationService();
 		serializationService = new SerializationService(undefined, undefined); //_connectionService _capabilitiesService
-		serializationService.registerProvider('testProviderId', new TestSerializationProvider());
+		serializationService.registerProvider('MSSQL', new TestSerializationProvider());
 		serializer = new ResultSerializer(
 			undefined, // IQueryManagementService
 			undefined, // IConfigurationService