Notebooks: Ensure quotes and backslashes are escaped properly in text editor model (#7497)

* Ensure quotes and backslashes are escaped properly

* PR comment

* PR comments

* Reliably fix

src/sql/workbench/parts/notebook/browser/models/notebookTextFileModel.ts

@@ -39,11 +39,60 @@ export class NotebookTextFileModel {
 		// convert the range to leverage offsets in the json
 		if (contentChange && contentChange.modelContentChangedEvent && areRangePropertiesPopulated(cellGuidRange)) {
 			contentChange.modelContentChangedEvent.changes.forEach(change => {
+				// When writing to JSON we need to escape double quotes and backslashes
+				let textEscapedQuotesAndBackslashes = change.text.replace(/\\/g, '\\\\').replace(/"/g, '\\"');
+
+				let startLineNumber = change.range.startLineNumber + cellGuidRange.startLineNumber - 1;
+				let endLineNumber = change.range.endLineNumber + cellGuidRange.startLineNumber - 1;
+				let startLineText = textEditorModel.textEditorModel.getLineContent(startLineNumber);
+				let endLineText = textEditorModel.textEditorModel.getLineContent(endLineNumber);
+
+				/* This gets the text on the start and end lines of the range where we'll be inserting text. We need to convert the escaped strings to unescaped strings.
+					Example:
+
+					Previous state
+					EDITOR:
+					""""
+
+					TEXTEDITORMODEL:
+					'        "\"\"\"\""'
+
+					Now, user wants to insert text after the 4 double quotes, like so:
+					EDITOR:
+					""""sample text
+
+					TEXTEDITORMODEL (result):
+					'        "\"\"\"\"sample text"'
+
+					Notice that we don't have a 1:1 mapping for characters from the editor to the text editor model, because the double quotes need to be escaped
+					(the same is true for backslashes).
+
+					Therefore, we need to determine (at both the start and end lines) the "real" start and end columns in the text editor model by counting escaped characters.
+
+					We do this by doing the following:
+					- Start with (escaped) text in the text editor model
+					- Unescape this text
+					- Get a substring of that text from the column in JSON until the change's startColumn (starting from the first " in the text editor model)
+					- Escape this substring
+					- Leverage the substring's length to calculate the "real" start/end columns
+				 */
+				let unescapedStartSubstring = startLineText.replace(/\\"/g, '"').replace(/\\\\/g, '\\').substr(cellGuidRange.startColumn, change.range.startColumn - 1);
+				let unescapedEndSubstring = endLineText.replace(/\\"/g, '"').replace(/\\\\/g, '\\').substr(cellGuidRange.startColumn, change.range.endColumn - 1);
+
+				// now we have the portion before the text to be inserted for both the start and end lines
+				// so next step is to escape " and \
+
+				let escapedStartSubstring = unescapedStartSubstring.replace(/\\/g, '\\\\').replace(/"/g, '\\"');
+				let escapedEndSubstring = unescapedEndSubstring.replace(/\\/g, '\\\\').replace(/"/g, '\\"');
+
+				let computedStartColumn = escapedStartSubstring.length + cellGuidRange.startColumn + 1;
+				let computedEndColumn = escapedEndSubstring.length + cellGuidRange.startColumn + 1;
+
 				let convertedRange: IRange = {
-					startLineNumber: change.range.startLineNumber + cellGuidRange.startLineNumber - 1,
+					startLineNumber: startLineNumber,
-					endLineNumber: change.range.endLineNumber + cellGuidRange.startLineNumber - 1,
+					endLineNumber: endLineNumber,
-					startColumn: change.range.startColumn + cellGuidRange.startColumn,
+					startColumn: computedStartColumn,
-					endColumn: change.range.endColumn + cellGuidRange.startColumn
+					endColumn: computedEndColumn
 				};
 				// Need to subtract one because we're going from 1-based to 0-based
 				let startSpaces: string = ' '.repeat(cellGuidRange.startColumn - 1);
@@ -52,14 +101,13 @@ export class NotebookTextFileModel {
 				//     this is another string
 				textEditorModel.textEditorModel.applyEdits([{
 					range: new Range(convertedRange.startLineNumber, convertedRange.startColumn, convertedRange.endLineNumber, convertedRange.endColumn),
-					text: change.text.split(this._eol).join('\\n\",'.concat(this._eol).concat(startSpaces).concat('\"'))
+					text: textEscapedQuotesAndBackslashes.split(/[\r\n]+/gm).join('\\n\",'.concat(this._eol).concat(startSpaces).concat('\"'))
 				}]);
 			});
-		} else {
-			return false;
-		}
 			return true;
 		}
+		return false;
+	}
 
 	public transformAndApplyEditForOutputUpdate(contentChange: NotebookContentChange, textEditorModel: TextFileEditorModel | UntitledEditorModel): boolean {
 		if (Array.isArray(contentChange.cells[0].outputs) && contentChange.cells[0].outputs.length > 0) {
@@ -117,9 +165,7 @@ export class NotebookTextFileModel {
 		if (!textEditorModel || !contentChange || !contentChange.cells || !contentChange.cells[0] || !contentChange.cells[0].cellGuid) {
 			return false;
 		}
-		if (!this.getOutputNodeByGuid(textEditorModel, contentChange.cells[0].cellGuid)) {
 		this.updateOutputBeginRange(textEditorModel, contentChange.cells[0].cellGuid);
-		}
 		let outputEndRange = this.getEndOfOutputs(textEditorModel, contentChange.cells[0].cellGuid);
 		let outputStartRange = this.getOutputNodeByGuid(textEditorModel, contentChange.cells[0].cellGuid);
 		if (outputStartRange && outputEndRange) {
@@ -253,7 +299,7 @@ export class NotebookTextFileModel {
 }
 
 function areRangePropertiesPopulated(range: Range) {
-	return range && range.startLineNumber && range.startColumn && range.endLineNumber && range.endColumn;
+	return range && range.startLineNumber !== 0 && range.startColumn !== 0 && range.endLineNumber !== 0 && range.endColumn !== 0;
 }
 
 function findOrSetCellGuidMatch(textEditorModel: TextFileEditorModel | UntitledEditorModel, cellGuid: string): FindMatch[] {

src/sql/workbench/parts/notebook/test/browser/notebookEditorModel.test.ts

@@ -12,7 +12,7 @@ import { ConnectionManagementService } from 'sql/platform/connection/browser/con
 import { CellModel } from 'sql/workbench/parts/notebook/browser/models/cell';
 import { CellTypes, NotebookChangeType } from 'sql/workbench/parts/notebook/common/models/contracts';
 import { ModelFactory } from 'sql/workbench/parts/notebook/browser/models/modelFactory';
-import { INotebookModelOptions, NotebookContentChange } from 'sql/workbench/parts/notebook/browser/models/modelInterfaces';
+import { INotebookModelOptions, NotebookContentChange, ICellModel } from 'sql/workbench/parts/notebook/browser/models/modelInterfaces';
 import { NotebookEditorModel } from 'sql/workbench/parts/notebook/browser/models/notebookInput';
 import { NotebookModel } from 'sql/workbench/parts/notebook/browser/models/notebookModel';
 import { NotebookService } from 'sql/workbench/services/notebook/browser/notebookServiceImpl';
@@ -691,6 +691,170 @@ suite('Notebook Editor Model', function (): void {
 		should(notebookEditorModel.lastEditFullReplacement).equal(false);
 	});
 
+	test('should not replace entire text model for content change with double quotes', async function (): Promise<void> {
+		await createNewNotebookModel();
+		let notebookEditorModel = await createTextEditorModel(this);
+		notebookEditorModel.replaceEntireTextEditorModel(notebookModel, undefined);
+
+		let newCell = notebookModel.addCell(CellTypes.Code);
+		setupTextEditorModelWithEmptyOutputs(notebookEditorModel, newCell);
+
+		addTextToBeginningOfTextEditorModel(notebookEditorModel, newCell, '"This text is in quotes"');
+		should(notebookEditorModel.editorModel.textEditorModel.getLineContent(9)).equal('                "\\"This text is in quotes\\""');
+
+		ensureStaticContentInOneLineCellIsCorrect(notebookEditorModel, newCell);
+		should(notebookEditorModel.lastEditFullReplacement).equal(false);
+	});
+
+	test('should not replace entire text model for content change with many double quotes', async function (): Promise<void> {
+		await createNewNotebookModel();
+		let notebookEditorModel = await createTextEditorModel(this);
+		notebookEditorModel.replaceEntireTextEditorModel(notebookModel, undefined);
+
+		let newCell = notebookModel.addCell(CellTypes.Code);
+		setupTextEditorModelWithEmptyOutputs(notebookEditorModel, newCell);
+
+		addTextToBeginningOfTextEditorModel(notebookEditorModel, newCell, '""""""""""');
+		should(notebookEditorModel.editorModel.textEditorModel.getLineContent(9)).equal('                "\\"\\"\\"\\"\\"\\"\\"\\"\\"\\""');
+
+		ensureStaticContentInOneLineCellIsCorrect(notebookEditorModel, newCell);
+		should(notebookEditorModel.lastEditFullReplacement).equal(false);
+	});
+
+	test('should not replace entire text model for content change with many backslashes', async function (): Promise<void> {
+		await createNewNotebookModel();
+		let notebookEditorModel = await createTextEditorModel(this);
+		notebookEditorModel.replaceEntireTextEditorModel(notebookModel, undefined);
+
+		let newCell = notebookModel.addCell(CellTypes.Code);
+		setupTextEditorModelWithEmptyOutputs(notebookEditorModel, newCell);
+
+		addTextToBeginningOfTextEditorModel(notebookEditorModel, newCell, '\\\\\\\\\\');
+		should(notebookEditorModel.editorModel.textEditorModel.getLineContent(9)).equal('                "\\\\\\\\\\\\\\\\\\\\\"');
+
+		ensureStaticContentInOneLineCellIsCorrect(notebookEditorModel, newCell);
+		should(notebookEditorModel.lastEditFullReplacement).equal(false);
+	});
+
+	test('should not replace entire text model for content change with many backslashes and double quotes', async function (): Promise<void> {
+		await createNewNotebookModel();
+		let notebookEditorModel = await createTextEditorModel(this);
+		notebookEditorModel.replaceEntireTextEditorModel(notebookModel, undefined);
+
+		let newCell = notebookModel.addCell(CellTypes.Code);
+		setupTextEditorModelWithEmptyOutputs(notebookEditorModel, newCell);
+
+		addTextToBeginningOfTextEditorModel(notebookEditorModel, newCell, '\"\"\"\"\"\"\"\"\"\"');
+		should(notebookEditorModel.editorModel.textEditorModel.getLineContent(9)).equal('                "\\\"\\\"\\\"\\\"\\\"\\\"\\\"\\\"\\\"\\\""');
+
+		ensureStaticContentInOneLineCellIsCorrect(notebookEditorModel, newCell);
+		should(notebookEditorModel.lastEditFullReplacement).equal(false);
+	});
+
+	test('should not replace entire text model for content change with no special characters', async function (): Promise<void> {
+		await createNewNotebookModel();
+		let notebookEditorModel = await createTextEditorModel(this);
+		notebookEditorModel.replaceEntireTextEditorModel(notebookModel, undefined);
+
+		let newCell = notebookModel.addCell(CellTypes.Code);
+		setupTextEditorModelWithEmptyOutputs(notebookEditorModel, newCell);
+
+		addTextToBeginningOfTextEditorModel(notebookEditorModel, newCell, 'this is a long line in a cell test. Everything should serialize correctly! # Comments here: adding more tests is fun?');
+		should(notebookEditorModel.editorModel.textEditorModel.getLineContent(9)).equal('                "this is a long line in a cell test. Everything should serialize correctly! # Comments here: adding more tests is fun?"');
+
+		ensureStaticContentInOneLineCellIsCorrect(notebookEditorModel, newCell);
+		should(notebookEditorModel.lastEditFullReplacement).equal(false);
+	});
+
+	test('should not replace entire text model for content change with variety of characters', async function (): Promise<void> {
+		await createNewNotebookModel();
+		let notebookEditorModel = await createTextEditorModel(this);
+		notebookEditorModel.replaceEntireTextEditorModel(notebookModel, undefined);
+
+		let newCell = notebookModel.addCell(CellTypes.Code);
+		setupTextEditorModelWithEmptyOutputs(notebookEditorModel, newCell);
+
+		addTextToBeginningOfTextEditorModel(notebookEditorModel, newCell, '`~1!2@3#4$5%6^7&8*9(0)-_=+[{]}\\|;:",<.>/?\'');
+		should(notebookEditorModel.editorModel.textEditorModel.getLineContent(9)).equal('                "`~1!2@3#4$5%6^7&8*9(0)-_=+[{]}\\\\|;:\\",<.>/?\'"');
+
+		ensureStaticContentInOneLineCellIsCorrect(notebookEditorModel, newCell);
+		should(notebookEditorModel.lastEditFullReplacement).equal(false);
+	});
+
+	test('should not replace entire text model for content change with single quotes', async function (): Promise<void> {
+		await createNewNotebookModel();
+		let notebookEditorModel = await createTextEditorModel(this);
+		notebookEditorModel.replaceEntireTextEditorModel(notebookModel, undefined);
+
+		let newCell = notebookModel.addCell(CellTypes.Code);
+		setupTextEditorModelWithEmptyOutputs(notebookEditorModel, newCell);
+
+		addTextToBeginningOfTextEditorModel(notebookEditorModel, newCell, '\'\'\'\'');
+		should(notebookEditorModel.editorModel.textEditorModel.getLineContent(9)).equal('                "\'\'\'\'"');
+
+		ensureStaticContentInOneLineCellIsCorrect(notebookEditorModel, newCell);
+		should(notebookEditorModel.lastEditFullReplacement).equal(false);
+	});
+
+	test('should not replace entire text model for content change with empty content', async function (): Promise<void> {
+		await createNewNotebookModel();
+		let notebookEditorModel = await createTextEditorModel(this);
+		notebookEditorModel.replaceEntireTextEditorModel(notebookModel, undefined);
+
+		let newCell = notebookModel.addCell(CellTypes.Code);
+		setupTextEditorModelWithEmptyOutputs(notebookEditorModel, newCell);
+
+		addTextToBeginningOfTextEditorModel(notebookEditorModel, newCell, '');
+		should(notebookEditorModel.editorModel.textEditorModel.getLineContent(9)).equal('                ""');
+
+		ensureStaticContentInOneLineCellIsCorrect(notebookEditorModel, newCell);
+		should(notebookEditorModel.lastEditFullReplacement).equal(false);
+	});
+
+	test('should not replace entire text model for content change with multiline content', async function (): Promise<void> {
+		await createNewNotebookModel();
+		let notebookEditorModel = await createTextEditorModel(this);
+		notebookEditorModel.replaceEntireTextEditorModel(notebookModel, undefined);
+
+		let newCell = notebookModel.addCell(CellTypes.Code);
+		setupTextEditorModelWithEmptyOutputs(notebookEditorModel, newCell);
+
+		addTextToBeginningOfTextEditorModel(notebookEditorModel, newCell, '"test"' + os.EOL + 'test""');
+		should(notebookEditorModel.editorModel.textEditorModel.getLineContent(9)).equal('                "\\"test\\"\\n",');
+
+		should(notebookEditorModel.editorModel.textEditorModel.getLineContent(8)).equal('            "source": [');
+		should(notebookEditorModel.editorModel.textEditorModel.getLineContent(10)).equal('                "test\\"\\""');
+		should(notebookEditorModel.editorModel.textEditorModel.getLineContent(11)).equal('            ],');
+		should(notebookEditorModel.editorModel.textEditorModel.getLineContent(13)).equal('                "azdata_cell_guid": "' + newCell.cellGuid + '"');
+		should(notebookEditorModel.editorModel.textEditorModel.getLineContent(15)).equal('            "outputs": [],');
+		should(notebookEditorModel.editorModel.textEditorModel.getLineContent(16)).equal('            "execution_count": 0');
+		should(notebookEditorModel.editorModel.textEditorModel.getLineContent(17)).equal('        }');
+
+		should(notebookEditorModel.lastEditFullReplacement).equal(false);
+	});
+
+	test('should not replace entire text model for content change with multiline content different escaped characters', async function (): Promise<void> {
+		await createNewNotebookModel();
+		let notebookEditorModel = await createTextEditorModel(this);
+		notebookEditorModel.replaceEntireTextEditorModel(notebookModel, undefined);
+
+		let newCell = notebookModel.addCell(CellTypes.Code);
+		setupTextEditorModelWithEmptyOutputs(notebookEditorModel, newCell);
+
+		addTextToBeginningOfTextEditorModel(notebookEditorModel, newCell, '"""""test"' + os.EOL + '"""""""test\\""');
+		should(notebookEditorModel.editorModel.textEditorModel.getLineContent(9)).equal('                "\\"\\"\\"\\"\\"test\\"\\n",');
+
+		should(notebookEditorModel.editorModel.textEditorModel.getLineContent(8)).equal('            "source": [');
+		should(notebookEditorModel.editorModel.textEditorModel.getLineContent(10)).equal('                "\\"\\"\\"\\"\\"\\"\\"test\\\\\\"\\""');
+		should(notebookEditorModel.editorModel.textEditorModel.getLineContent(11)).equal('            ],');
+		should(notebookEditorModel.editorModel.textEditorModel.getLineContent(13)).equal('                "azdata_cell_guid": "' + newCell.cellGuid + '"');
+		should(notebookEditorModel.editorModel.textEditorModel.getLineContent(15)).equal('            "outputs": [],');
+		should(notebookEditorModel.editorModel.textEditorModel.getLineContent(16)).equal('            "execution_count": 0');
+		should(notebookEditorModel.editorModel.textEditorModel.getLineContent(17)).equal('        }');
+
+		should(notebookEditorModel.lastEditFullReplacement).equal(false);
+	});
+
 	async function createNewNotebookModel() {
 		let options: INotebookModelOptions = Object.assign({}, defaultModelOptions, <Partial<INotebookModelOptions>><unknown>{
 			factory: mockModelFactory.object
@@ -705,4 +869,47 @@ suite('Notebook Editor Model', function (): void {
 		await textFileEditorModel.load();
 		return new NotebookEditorModel(defaultUri, textFileEditorModel, mockNotebookService.object, accessor.textFileService, testResourcePropertiesService);
 	}
+
+	function setupTextEditorModelWithEmptyOutputs(notebookEditorModel: NotebookEditorModel, newCell: ICellModel) {
+		// clear outputs
+		newCell[<any>'_outputs'] = [];
+
+		let contentChange: NotebookContentChange = {
+			changeType: NotebookChangeType.CellsModified,
+			cells: [newCell],
+			cellIndex: 0
+		};
+
+		notebookEditorModel.updateModel(contentChange, NotebookChangeType.CellsModified);
+		should(notebookEditorModel.lastEditFullReplacement).equal(true);
+
+		should(notebookEditorModel.editorModel.textEditorModel.getLineContent(14)).equal('            "outputs": [],');
+	}
+
+	function addTextToBeginningOfTextEditorModel(notebookEditorModel: NotebookEditorModel, newCell: ICellModel, textToAdd: string) {
+		let contentChange: NotebookContentChange = {
+			changeType: NotebookChangeType.CellSourceUpdated,
+			cells: [newCell],
+			cellIndex: 0,
+			modelContentChangedEvent: {
+				changes: [{ range: new Range(1, 1, 1, 1), rangeLength: 0, rangeOffset: 0, text: textToAdd }],
+				eol: '\n',
+				isFlush: false,
+				isRedoing: false,
+				isUndoing: false,
+				versionId: 2
+			}
+		};
+
+		notebookEditorModel.updateModel(contentChange, NotebookChangeType.CellSourceUpdated);
+	}
+
+	function ensureStaticContentInOneLineCellIsCorrect(notebookEditorModel: NotebookEditorModel, newCell: ICellModel) {
+		should(notebookEditorModel.editorModel.textEditorModel.getLineContent(8)).equal('            "source": [');
+		should(notebookEditorModel.editorModel.textEditorModel.getLineContent(10)).equal('            ],');
+		should(notebookEditorModel.editorModel.textEditorModel.getLineContent(12)).equal('                "azdata_cell_guid": "' + newCell.cellGuid + '"');
+		should(notebookEditorModel.editorModel.textEditorModel.getLineContent(14)).equal('            "outputs": [],');
+		should(notebookEditorModel.editorModel.textEditorModel.getLineContent(15)).equal('            "execution_count": 0');
+		should(notebookEditorModel.editorModel.textEditorModel.getLineContent(16)).equal('        }');
+	}
 });