Notebooks: Fast update WYSIWYG support for source update (#12289) (#12399)

* Fast update WYSIWYG support for source update

* Do bracket matching over hardcoding line offsets

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

@@ -105,6 +105,25 @@ export class NotebookTextFileModel {
 				}]);
 			});
 			return true;
+		} else if (contentChange && areRangePropertiesPopulated(cellGuidRange)) {
+			// If no modelContentChanged event, then we're replacing the entire source for that cell
+			let sourceEnd = this.getSourceEndRange(textEditorModel, contentChange.cells[0].cellGuid);
+			if (sourceEnd) {
+				// Need to subtract one because we're going from 1-based to 0-based
+				let startSpaces: string = repeat(' ', cellGuidRange.startColumn - 1);
+				let escapedQuotesAndBackslashes = contentChange.cells[0].source.join('\n').replace(/\\/g, '\\\\').replace(/"/g, '\\"');
+
+				// The text here transforms a string from 'This is a string\n this is another string' to:
+				//     This is a string
+				//     this is another string
+
+				// Note: Adding 1 to startColumn to avoid overwriting first "
+				textEditorModel.textEditorModel.applyEdits([{
+					range: new Range(this._sourceBeginRange.startLineNumber, this._sourceBeginRange.startColumn + 1, sourceEnd.endLineNumber, sourceEnd.endColumn),
+					text: escapedQuotesAndBackslashes.split(/[\r\n]+/gm).join('\\n\",'.concat(this._eol).concat(startSpaces).concat('\"'))
+				}]);
+				return true;
+			}
 		}
 		return false;
 	}
@@ -209,6 +228,45 @@ export class NotebookTextFileModel {
 		}
 	}
 
+	private getSourceEndRange(textEditorModel: ITextEditorModel, cellGuid: string): IRange | undefined {
+		if (!cellGuid) {
+			return undefined;
+		}
+		let cellGuidMatches = findOrSetCellGuidMatch(textEditorModel, cellGuid);
+		if (cellGuidMatches?.length > 0) {
+			if (!this._sourceBeginRange) {
+				this.updateSourceBeginRange(textEditorModel, cellGuid);
+			}
+			// Source begin range tracks where the first " in exists.
+			// The line before that will always include '"source": ['
+			let sourceBeforeLineNumber = this._sourceBeginRange?.startLineNumber - 1;
+			if (sourceBeforeLineNumber) {
+				// The 2nd to last column (ie before newline) is guaranteed to be [
+				let sourceBeforeColumn = textEditorModel.textEditorModel.getLineMaxColumn(sourceBeforeLineNumber);
+				if (sourceBeforeColumn) {
+					// Match the end of the source array
+					let sourceEnd = textEditorModel.textEditorModel.matchBracket({ column: sourceBeforeColumn - 1, lineNumber: sourceBeforeLineNumber });
+					if (sourceEnd?.length === 2) {
+						// Last quote in the source array will end the line before the source array
+						// e.g.
+						// "source": [
+						//	   "SELECT 12" <-- Looking for this " position
+						// ],
+						let lineForSourceEnd = sourceEnd[1].endLineNumber - 1;
+						let lastCharacterPosition = textEditorModel.textEditorModel.getLineLength(lineForSourceEnd);
+						return {
+							startColumn: lastCharacterPosition,
+							startLineNumber: lineForSourceEnd,
+							endLineNumber: lineForSourceEnd,
+							endColumn: lastCharacterPosition
+						};
+					}
+				}
+			}
+		}
+		return undefined;
+	}
+
 	// Find the beginning of a cell's outputs in the text editor model
 	private updateOutputBeginRange(textEditorModel: ITextEditorModel, cellGuid: string): void {
 		if (!cellGuid) {

src/sql/workbench/contrib/notebook/test/electron-browser/notebookEditorModel.test.ts

@@ -400,6 +400,87 @@ suite('Notebook Editor Model', function (): void {
 		assert(!notebookEditorModel.lastEditFullReplacement);
 	});
 
+	test('should not replace entire text model but replace entire source when no modelContentChangedEvent passed in', async function (): Promise<void> {
+		await createNewNotebookModel();
+		let notebookEditorModel = await createTextEditorModel(this);
+		notebookEditorModel.replaceEntireTextEditorModel(notebookModel, undefined);
+
+		let newCell = notebookModel.addCell(CellTypes.Code);
+
+		let contentChange: NotebookContentChange = {
+			changeType: NotebookChangeType.CellsModified,
+			cells: [newCell],
+			cellIndex: 0
+		};
+		notebookEditorModel.updateModel(contentChange, NotebookChangeType.CellsModified);
+		assert(notebookEditorModel.lastEditFullReplacement);
+
+		assert.equal(notebookEditorModel.editorModel.textEditorModel.getLineContent(14), '            "outputs": [');
+
+		newCell.source = 'This is a test';
+
+		contentChange = {
+			changeType: NotebookChangeType.CellSourceUpdated,
+			cells: [newCell],
+			cellIndex: 0,
+			modelContentChangedEvent: undefined
+		};
+
+		notebookEditorModel.updateModel(contentChange, NotebookChangeType.CellSourceUpdated);
+
+		assert(!notebookEditorModel.lastEditFullReplacement, 'should not do a full replacement for a source update');
+
+		assert.equal(notebookEditorModel.editorModel.textEditorModel.getLineContent(8), '            "source": [');
+		assert.equal(notebookEditorModel.editorModel.textEditorModel.getLineContent(9), '                "This is a test"');
+		assert.equal(notebookEditorModel.editorModel.textEditorModel.getLineContent(10), '            ],');
+		assert.equal(notebookEditorModel.editorModel.textEditorModel.getLineContent(12), '                "azdata_cell_guid": "' + newCell.cellGuid + '"');
+		assert.equal(notebookEditorModel.editorModel.textEditorModel.getLineContent(14), '            "outputs": [');
+		assert.equal(notebookEditorModel.editorModel.textEditorModel.getLineContent(25), '            "execution_count": null');
+		assert.equal(notebookEditorModel.editorModel.textEditorModel.getLineContent(26), '        }');
+
+	});
+
+	test('should not replace entire text model but replace entire source when no modelContentChangedEvent passed in multiline change', async function (): Promise<void> {
+		await createNewNotebookModel();
+		let notebookEditorModel = await createTextEditorModel(this);
+		notebookEditorModel.replaceEntireTextEditorModel(notebookModel, undefined);
+
+		let newCell = notebookModel.addCell(CellTypes.Code);
+
+		let contentChange: NotebookContentChange = {
+			changeType: NotebookChangeType.CellsModified,
+			cells: [newCell],
+			cellIndex: 0
+		};
+		notebookEditorModel.updateModel(contentChange, NotebookChangeType.CellsModified);
+		assert(notebookEditorModel.lastEditFullReplacement);
+
+		assert.equal(notebookEditorModel.editorModel.textEditorModel.getLineContent(14), '            "outputs": [');
+
+		newCell.source = 'This is a test' + os.EOL + 'Line 2 test' + os.EOL + 'Line 3 test';
+
+		contentChange = {
+			changeType: NotebookChangeType.CellSourceUpdated,
+			cells: [newCell],
+			cellIndex: 0,
+			modelContentChangedEvent: undefined
+		};
+
+		notebookEditorModel.updateModel(contentChange, NotebookChangeType.CellSourceUpdated);
+
+		assert(!notebookEditorModel.lastEditFullReplacement, 'should not do a full replacement for a source update');
+
+		assert.equal(notebookEditorModel.editorModel.textEditorModel.getLineContent(8), '            "source": [');
+		assert.equal(notebookEditorModel.editorModel.textEditorModel.getLineContent(9), '                "This is a test\\n",');
+		assert.equal(notebookEditorModel.editorModel.textEditorModel.getLineContent(10), '                "Line 2 test\\n",');
+		assert.equal(notebookEditorModel.editorModel.textEditorModel.getLineContent(11), '                "Line 3 test"');
+		assert.equal(notebookEditorModel.editorModel.textEditorModel.getLineContent(12), '            ],');
+		assert.equal(notebookEditorModel.editorModel.textEditorModel.getLineContent(14), '                "azdata_cell_guid": "' + newCell.cellGuid + '"');
+		assert.equal(notebookEditorModel.editorModel.textEditorModel.getLineContent(16), '            "outputs": [');
+		assert.equal(notebookEditorModel.editorModel.textEditorModel.getLineContent(27), '            "execution_count": null');
+		assert.equal(notebookEditorModel.editorModel.textEditorModel.getLineContent(28), '        }');
+	});
+
 	test('should not replace entire text model for single line source change then delete', async function (): Promise<void> {
 		await createNewNotebookModel();
 		let notebookEditorModel = await createTextEditorModel(this);