From dd5adad772c00de20c1410904009ffd64dc5f256 Mon Sep 17 00:00:00 2001
From: Lucy Zhang <luczhan@microsoft.com>
Date: Mon, 29 Mar 2021 07:14:49 -0700
Subject: [PATCH] Notebooks: fix save as csv/excel/json/xml (#14882)

* add getEncoding method to notebooks

* fix row 501 missing in serialization

* fix row index

* pr comment
---
 .../platform/serialization/common/serializationService.ts  | 4 ++++
 .../contrib/notebook/browser/models/fileNotebookInput.ts   | 4 ++++
 .../notebook/browser/models/untitledNotebookInput.ts       | 4 ++++
 src/sql/workbench/contrib/query/browser/actions.ts         | 7 +++----
 4 files changed, 15 insertions(+), 4 deletions(-)

diff --git a/src/sql/platform/serialization/common/serializationService.ts b/src/sql/platform/serialization/common/serializationService.ts
index 58b8a474c3..f83366a765 100644
--- a/src/sql/platform/serialization/common/serializationService.ts
+++ b/src/sql/platform/serialization/common/serializationService.ts
@@ -119,6 +119,10 @@ export class SerializationService implements ISerializationService {
 			let index = 0;
 			let startRequestParams = this.createStartRequest(serializationRequest, index);
 			index = index + startRequestParams.rows.length;
+			// Adjust row index based on whether or not header row is included
+			if (serializationRequest.includeHeaders) {
+				index--;
+			}
 
 			let startResult = await provider.startSerialization(startRequestParams);
 
diff --git a/src/sql/workbench/contrib/notebook/browser/models/fileNotebookInput.ts b/src/sql/workbench/contrib/notebook/browser/models/fileNotebookInput.ts
index 5aeec05f89..44bc0ac092 100644
--- a/src/sql/workbench/contrib/notebook/browser/models/fileNotebookInput.ts
+++ b/src/sql/workbench/contrib/notebook/browser/models/fileNotebookInput.ts
@@ -45,4 +45,8 @@ export class FileNotebookInput extends NotebookInput {
 	public getTypeId(): string {
 		return FileNotebookInput.ID;
 	}
+
+	public getEncoding(): string | undefined {
+		return this.textInput.getEncoding();
+	}
 }
diff --git a/src/sql/workbench/contrib/notebook/browser/models/untitledNotebookInput.ts b/src/sql/workbench/contrib/notebook/browser/models/untitledNotebookInput.ts
index 3e421191fc..127d41f0c6 100644
--- a/src/sql/workbench/contrib/notebook/browser/models/untitledNotebookInput.ts
+++ b/src/sql/workbench/contrib/notebook/browser/models/untitledNotebookInput.ts
@@ -42,4 +42,8 @@ export class UntitledNotebookInput extends NotebookInput {
 	public getTypeId(): string {
 		return UntitledNotebookInput.ID;
 	}
+
+	public getEncoding(): string | undefined {
+		return this.textInput.getEncoding();
+	}
 }
diff --git a/src/sql/workbench/contrib/query/browser/actions.ts b/src/sql/workbench/contrib/query/browser/actions.ts
index 8c6452da0c..01e61e8d58 100644
--- a/src/sql/workbench/contrib/query/browser/actions.ts
+++ b/src/sql/workbench/contrib/query/browser/actions.ts
@@ -8,7 +8,6 @@ import { localize } from 'vs/nls';
 import { IEditorService } from 'vs/workbench/services/editor/common/editorService';
 import { Table } from 'sql/base/browser/ui/table/table';
 import { QueryEditor } from './queryEditor';
-import { UntitledQueryEditorInput } from 'sql/workbench/common/editor/query/untitledQueryEditorInput';
 import { CellSelectionModel } from 'sql/base/browser/ui/table/plugins/cellSelectionModel.plugin';
 import { IGridDataProvider } from 'sql/workbench/services/query/common/gridDataProvider';
 import { INotificationService, Severity, NeverShowAgainScope } from 'vs/platform/notification/common/notification';
@@ -20,6 +19,7 @@ import * as TelemetryKeys from 'sql/platform/telemetry/common/telemetryKeys';
 import { getErrorMessage } from 'vs/base/common/errors';
 import { SaveFormat } from 'sql/workbench/services/query/common/resultSerializer';
 import { IExtensionRecommendationsService } from 'vs/workbench/services/extensionRecommendations/common/extensionRecommendations';
+import { IEncodingSupport } from 'vs/workbench/common/editor';
 
 export interface IGridActionContext {
 	gridDataProvider: IGridDataProvider;
@@ -70,9 +70,8 @@ export class SaveResultAction extends Action {
 
 	public async run(context: IGridActionContext): Promise<boolean> {
 
-		const activeEditor = this.editorService.activeEditorPane as QueryEditor;
-		let input = activeEditor.input as UntitledQueryEditorInput;
-		if (input.getEncoding() !== 'utf8') {
+		const activeEditor = this.editorService.activeEditorPane as unknown as IEncodingSupport;
+		if (typeof activeEditor.getEncoding === 'function' && activeEditor.getEncoding() !== 'utf8') {
 			this.notificationService.notify({
 				severity: Severity.Info,
 				message: localize('jsonEncoding', "Results encoding will not be saved when exporting to JSON, remember to save with desired encoding once file is created."),