Save multiline source in notebooks (#6445)

* Save multiline source in notebooks

* Stop demultiline localContentManager unnecessarily

src/sql/azdata.proposed.d.ts

@@ -4990,7 +4990,7 @@ declare module 'azdata' {
 		 * The contents of a requestExecute message sent to the server.
 		 */
 		export interface IExecuteRequest extends IExecuteOptions {
-			code: string;
+			code: string | string[];
 		}
 
 		/**

src/sql/workbench/parts/notebook/electron-browser/cellViews/code.component.ts

@@ -196,7 +196,9 @@ export class CodeComponent extends AngularDisposable implements OnInit, OnChange
 		this._editor.setMinimumHeight(this._minimumHeight);
 		this._editor.setMaximumHeight(this._maximumHeight);
 		let uri = this.cellModel.cellUri;
-		this._editorInput = instantiationService.createInstance(UntitledEditorInput, uri, false, this.cellModel.language, this.cellModel.source, '');
+		let cellModelSource: string;
+		cellModelSource = Array.isArray(this.cellModel.source) ? this.cellModel.source.join('') : this.cellModel.source;
+		this._editorInput = instantiationService.createInstance(UntitledEditorInput, uri, false, this.cellModel.language, cellModelSource, '');
 		await this._editor.setInput(this._editorInput, undefined);
 		this.setFocusAndScroll();
 		let untitledEditorModel: UntitledEditorModel = await this._editorInput.resolve();
@@ -262,7 +264,9 @@ export class CodeComponent extends AngularDisposable implements OnInit, OnChange
 	/// Editor Functions
 	private updateModel() {
 		if (this._editorModel) {
-			this._modelService.updateModel(this._editorModel, this.cellModel.source);
+			let cellModelSource: string;
+			cellModelSource = Array.isArray(this.cellModel.source) ? this.cellModel.source.join('') : this.cellModel.source;
+			this._modelService.updateModel(this._editorModel, cellModelSource);
 		}
 	}
 

src/sql/workbench/parts/notebook/electron-browser/cellViews/textCell.component.ts

@@ -69,7 +69,7 @@ export class TextCellComponent extends CellView implements OnInit, OnChanges {
 		this._model.activeCell = undefined;
 	}
 
-	private _content: string;
+	private _content: string | string[];
 	private _lastTrustedMode: boolean;
 	private isEditMode: boolean;
 	private _sanitizer: ISanitizer;
@@ -178,7 +178,7 @@ export class TextCellComponent extends CellView implements OnInit, OnChanges {
 				this.markdownRenderer.setNotebookURI(this.cellModel.notebookModel.notebookUri);
 				this.markdownResult = this.markdownRenderer.render({
 					isTrusted: true,
-					value: this._content
+					value: Array.isArray(this._content) ? this._content.join('') : this._content
 				});
 				this.markdownResult.element.innerHTML = this.sanitizeContent(this.markdownResult.element.innerHTML);
 				this.setLoading(false);

src/sql/workbench/parts/notebook/node/models/cell.ts

@@ -25,7 +25,7 @@ let modelId = 0;
 
 export class CellModel implements ICellModel {
 	private _cellType: nb.CellType;
-	private _source: string;
+	private _source: string | string[];
 	private _language: string;
 	private _future: FutureInternal;
 	private _outputs: nb.ICellOutput[] = [];
@@ -156,11 +156,12 @@ export class CellModel implements ICellModel {
 		return this._cellType;
 	}
 
-	public get source(): string {
+	public get source(): string | string[] {
 		return this._source;
 	}
 
-	public set source(newSource: string) {
+	public set source(newSource: string | string[]) {
+		newSource = this.getMultilineSource(newSource);
 		if (this._source !== newSource) {
 			this._source = newSource;
 			this.sendChangeToNotebook(NotebookChangeType.CellSourceUpdated);
@@ -533,7 +534,7 @@ export class CellModel implements ICellModel {
 		}
 		this._cellType = cell.cell_type;
 		this.executionCount = cell.execution_count;
-		this._source = Array.isArray(cell.source) ? cell.source.join('') : cell.source;
+		this._source = this.getMultilineSource(cell.source);
 		this._metadata = cell.metadata;
 		this.setLanguageFromContents(cell);
 		if (cell.outputs) {
@@ -594,6 +595,26 @@ export class CellModel implements ICellModel {
 		return endpoint;
 	}
 
+	private getMultilineSource(source: string | string[]): string | string[] {
+		if (typeof source === 'string') {
+			let sourceMultiline = source.split('\n');
+			// If source is one line (i.e. no '\n'), return it immediately
+			if (sourceMultiline.length <= 1) {
+				return source;
+			}
+			// Otherwise, add back all of the newlines here
+			// Note: for Windows machines that require '/r/n',
+			// splitting on '\n' and putting back the '\n' will still
+			// retain the '\r', so that isn't lost in the process
+			// Note: the last line will not include a newline at the end
+			for (let i = 0; i < sourceMultiline.length - 1; i++) {
+				sourceMultiline[i] += '\n';
+			}
+			return sourceMultiline;
+		}
+		return source;
+	}
+
 	// Dispose and set current future to undefined
 	private disposeFuture() {
 		if (this._future) {

src/sql/workbench/parts/notebook/node/models/modelInterfaces.ts

@@ -449,7 +449,7 @@ export interface ICellModel {
 	cellUri: URI;
 	id: string;
 	readonly language: string;
-	source: string;
+	source: string | string[];
 	cellType: CellType;
 	trustedMode: boolean;
 	active: boolean;

src/sql/workbench/parts/notebook/test/node/cell.test.ts

@@ -130,6 +130,121 @@ suite('Cell Model', function (): void {
 		should(cell.language).equal('python');
 	});
 
+	test('Should allow source of type string[] with length 1', async function (): Promise<void> {
+		let cellData: nb.ICellContents = {
+			cell_type: CellTypes.Code,
+			source: ['print(1)'],
+			metadata: { language: 'sql' },
+			execution_count: 1
+		};
+
+		let notebookModel = new NotebookModelStub({
+			name: '',
+			version: '',
+			mimetype: ''
+		});
+		let cell = factory.createCell(cellData, { notebook: notebookModel, isTrusted: false });
+		should(Array.isArray(cell.source)).equal(true);
+		should(cell.source.length).equal(1);
+		should(cell.source[0]).equal('print(1)');
+	});
+
+	test('Should allow source of type string', async function (): Promise<void> {
+		let cellData: nb.ICellContents = {
+			cell_type: CellTypes.Code,
+			source: 'print(1)',
+			metadata: { language: 'sql' },
+			execution_count: 1
+		};
+
+		let notebookModel = new NotebookModelStub({
+			name: '',
+			version: '',
+			mimetype: ''
+		});
+		let cell = factory.createCell(cellData, { notebook: notebookModel, isTrusted: false });
+		should(Array.isArray(cell.source)).equal(false);
+		should(cell.source).equal('print(1)');
+	});
+
+	test('Should allow source of type string with newline and split it', async function (): Promise<void> {
+		let cellData: nb.ICellContents = {
+			cell_type: CellTypes.Code,
+			source: 'print(1)\nprint(2)',
+			metadata: { language: 'sql' },
+			execution_count: 1
+		};
+
+		let notebookModel = new NotebookModelStub({
+			name: '',
+			version: '',
+			mimetype: ''
+		});
+		let cell = factory.createCell(cellData, { notebook: notebookModel, isTrusted: false });
+		should(Array.isArray(cell.source)).equal(true);
+		should(cell.source.length).equal(2);
+		should(cell.source[0]).equal('print(1)\n');
+		should(cell.source[1]).equal('print(2)');
+	});
+
+	test('Should allow source of type string with Windows style newline and split it', async function (): Promise<void> {
+		let cellData: nb.ICellContents = {
+			cell_type: CellTypes.Code,
+			source: 'print(1)\r\nprint(2)',
+			metadata: { language: 'sql' },
+			execution_count: 1
+		};
+
+		let notebookModel = new NotebookModelStub({
+			name: '',
+			version: '',
+			mimetype: ''
+		});
+		let cell = factory.createCell(cellData, { notebook: notebookModel, isTrusted: false });
+		should(Array.isArray(cell.source)).equal(true);
+		should(cell.source.length).equal(2);
+		should(cell.source[0]).equal('print(1)\r\n');
+		should(cell.source[1]).equal('print(2)');
+	});
+
+	test('Should allow source of type string[] with length 2', async function (): Promise<void> {
+		let cellData: nb.ICellContents = {
+			cell_type: CellTypes.Code,
+			source: ['print(1)\n', 'print(2)'],
+			metadata: { language: 'sql' },
+			execution_count: 1
+		};
+
+		let notebookModel = new NotebookModelStub({
+			name: '',
+			version: '',
+			mimetype: ''
+		});
+		let cell = factory.createCell(cellData, { notebook: notebookModel, isTrusted: false });
+		should(Array.isArray(cell.source)).equal(true);
+		should(cell.source.length).equal(2);
+		should(cell.source[0]).equal('print(1)\n');
+		should(cell.source[1]).equal('print(2)');
+	});
+
+	test('Should allow empty string source', async function (): Promise<void> {
+		let cellData: nb.ICellContents = {
+			cell_type: CellTypes.Code,
+			source: '',
+			metadata: { language: 'sql' },
+			execution_count: 1
+		};
+
+		let notebookModel = new NotebookModelStub({
+			name: '',
+			version: '',
+			mimetype: ''
+		});
+		let cell = factory.createCell(cellData, { notebook: notebookModel, isTrusted: false });
+		should(Array.isArray(cell.source)).equal(false);
+		should(cell.source).equal('');
+	});
+
 	suite('Model Future handling', function (): void {
 		let future: TypeMoq.Mock<EmptyFuture>;
 		let cell: ICellModel;

src/sql/workbench/services/notebook/node/localContentManager.ts

@@ -125,7 +125,7 @@ namespace v4 {
 	export function createDefaultCell(cell: nb.ICellContents): nb.ICellContents {
 		return {
 			cell_type: cell.cell_type,
-			source: demultiline(cell.source),
+			source: cell.source,
 			metadata: cell.metadata
 		};
 	}
@@ -133,7 +133,7 @@ namespace v4 {
 	function createCodeCell(cell: nb.ICellContents): nb.ICellContents {
 		return {
 			cell_type: cell.cell_type,
-			source: demultiline(cell.source),
+			source: cell.source,
 			metadata: cell.metadata,
 			execution_count: cell.execution_count,
 			outputs: createOutputs(cell)

src/sql/workbench/services/notebook/sql/sqlSessionManager.ts

@@ -274,7 +274,7 @@ class SqlKernel extends Disposable implements nb.IKernel {
 	}
 
 	private getCodeWithoutCellMagic(content: nb.IExecuteRequest): string {
-		let code = content.code;
+		let code = Array.isArray(content.code) ? content.code.join('') : content.code;
 		let firstLineEnd = code.indexOf(os.EOL);
 		let firstLine = code.substring(0, (firstLineEnd >= 0) ? firstLineEnd : 0).trimLeft();
 		if (firstLine.startsWith('%%')) {