Fix #3420 Analyze in notebook doesn't include text (#3482)

- Add edit API that can be used in the extension
- Separated document and editor classes out since this is the point those get big. I can refactor back in if needed to ease code review
- Based this off text editing APIs but tweaked for the fact this is a cell/array based set of edits

src/sql/workbench/api/node/extHostNotebookDocumentData.ts

@@ -0,0 +1,101 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the Source EULA. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+'use strict';
+
+import * as sqlops from 'sqlops';
+
+import { IDisposable } from 'vs/base/common/lifecycle';
+import URI from 'vs/base/common/uri';
+import { ok } from 'vs/base/common/assert';
+import { Schemas } from 'vs/base/common/network';
+import { TPromise } from 'vs/base/common/winjs.base';
+
+import { MainThreadNotebookDocumentsAndEditorsShape } from 'sql/workbench/api/node/sqlExtHost.protocol';
+import { CellRange } from 'sql/workbench/api/common/sqlExtHostTypes';
+
+
+export class ExtHostNotebookDocumentData implements IDisposable {
+	private _document: sqlops.nb.NotebookDocument;
+	private _cells: sqlops.nb.NotebookCell[];
+	private _isDisposed: boolean = false;
+
+	constructor(private readonly _proxy: MainThreadNotebookDocumentsAndEditorsShape,
+		private readonly _uri: URI,
+		private readonly _providerId: string,
+		private _isDirty: boolean
+	) {
+		// TODO add cell mapping support
+		this._cells = [];
+	}
+
+	dispose(): void {
+		// we don't really dispose documents but let
+		// extensions still read from them. some
+		// operations, live saving, will now error tho
+		ok(!this._isDisposed);
+		this._isDisposed = true;
+		this._isDirty = false;
+	}
+
+
+	get document(): sqlops.nb.NotebookDocument {
+		if (!this._document) {
+			const data = this;
+			this._document = {
+				get uri() { return data._uri; },
+				get fileName() { return data._uri.fsPath; },
+				get isUntitled() { return data._uri.scheme === Schemas.untitled; },
+				get providerId() { return data._providerId; },
+				get isClosed() { return data._isDisposed; },
+				get isDirty() { return data._isDirty; },
+				get cells() { return data._cells; },
+				save() { return data._save(); },
+				validateCellRange(range) { return data._validateRange(range); },
+			};
+		}
+		return Object.freeze(this._document);
+	}
+
+	private _save(): Thenable<boolean> {
+		if (this._isDisposed) {
+			return TPromise.wrapError<boolean>(new Error('Document has been closed'));
+		}
+		return this._proxy.$trySaveDocument(this._uri);
+
+	}
+
+	// ---- range math
+
+	private _validateRange(range: sqlops.nb.CellRange): sqlops.nb.CellRange {
+		if (!(range instanceof CellRange)) {
+			throw new Error('Invalid argument');
+		}
+
+		let start = this._validateIndex(range.start);
+		let end = this._validateIndex(range.end);
+
+		if (start === range.start && end === range.end) {
+			return range;
+		}
+		return new CellRange(start, end);
+	}
+
+	private _validateIndex(index: number): number {
+		if (typeof(index) !== 'number') {
+			throw new Error('Invalid argument');
+		}
+
+		if (index < 0) {
+			index = 0;
+		} else if (this._cells.length > 0 && index > this._cells.length) {
+			// We allow off by 1 as end needs to be outside current length in order to
+			// handle replace scenario. Long term should consider different start vs end validation instead
+			index = this._cells.length;
+		}
+
+		return index;
+	}
+
+}

src/sql/workbench/api/node/extHostNotebookDocumentsAndEditors.ts

@@ -8,111 +8,21 @@ import * as sqlops from 'sqlops';
 import * as vscode from 'vscode';
 
 import { Event, Emitter } from 'vs/base/common/event';
-import { readonly } from 'vs/base/common/errors';
-import { dispose, IDisposable } from 'vs/base/common/lifecycle';
+import { dispose } from 'vs/base/common/lifecycle';
 import URI from 'vs/base/common/uri';
 import { Disposable } from 'vs/workbench/api/node/extHostTypes';
-import { Schemas } from 'vs/base/common/network';
-import { TPromise } from 'vs/base/common/winjs.base';
 import * as typeConverters from 'vs/workbench/api/node/extHostTypeConverters';
 import { IMainContext } from 'vs/workbench/api/node/extHost.protocol';
 import { ok } from 'vs/base/common/assert';
 
 import {
-	MainThreadNotebookShape, SqlMainContext, INotebookDocumentsAndEditorsDelta,
-	ExtHostNotebookDocumentsAndEditorsShape, MainThreadNotebookDocumentsAndEditorsShape, INotebookShowOptions
+	SqlMainContext, INotebookDocumentsAndEditorsDelta, ExtHostNotebookDocumentsAndEditorsShape,
+	MainThreadNotebookDocumentsAndEditorsShape, INotebookShowOptions
 } from 'sql/workbench/api/node/sqlExtHost.protocol';
+import { ExtHostNotebookDocumentData } from 'sql/workbench/api/node/extHostNotebookDocumentData';
+import { ExtHostNotebookEditor } from 'sql/workbench/api/node/extHostNotebookEditor';
 
 
-export class ExtHostNotebookDocumentData implements IDisposable {
-	private _document: sqlops.nb.NotebookDocument;
-	private _cells: sqlops.nb.NotebookCell[];
-	private _isDisposed: boolean = false;
-
-	constructor(private readonly _proxy: MainThreadNotebookDocumentsAndEditorsShape,
-		private readonly _uri: URI,
-		private readonly _providerId: string,
-		private _isDirty: boolean
-	) {
-		// TODO add cell mapping support
-		this._cells = [];
-	}
-
-	dispose(): void {
-		// we don't really dispose documents but let
-		// extensions still read from them. some
-		// operations, live saving, will now error tho
-		ok(!this._isDisposed);
-		this._isDisposed = true;
-		this._isDirty = false;
-	}
-
-
-	get document(): sqlops.nb.NotebookDocument {
-		if (!this._document) {
-			const data = this;
-			this._document = {
-				get uri() { return data._uri; },
-				get fileName() { return data._uri.fsPath; },
-				get isUntitled() { return data._uri.scheme === Schemas.untitled; },
-				get providerId() { return data._providerId; },
-				get isClosed() { return data._isDisposed; },
-				get isDirty() { return data._isDirty; },
-				get cells() { return data._cells; },
-				save() { return data._save(); },
-			};
-		}
-		return Object.freeze(this._document);
-	}
-
-	private _save(): Thenable<boolean> {
-		if (this._isDisposed) {
-			return TPromise.wrapError<boolean>(new Error('Document has been closed'));
-		}
-		return this._proxy.$trySaveDocument(this._uri);
-
-	}
-}
-
-export class ExtHostNotebookEditor implements sqlops.nb.NotebookEditor, IDisposable {
-	private _disposed: boolean = false;
-
-	constructor(
-		private _proxy: MainThreadNotebookShape,
-		private _id: string,
-		private readonly _documentData: ExtHostNotebookDocumentData,
-		private _viewColumn: vscode.ViewColumn
-	) {
-
-	}
-
-	dispose() {
-		ok(!this._disposed);
-		this._disposed = true;
-	}
-
-	get document(): sqlops.nb.NotebookDocument {
-		return this._documentData.document;
-	}
-
-	set document(value) {
-		throw readonly('document');
-	}
-
-	get viewColumn(): vscode.ViewColumn {
-		return this._viewColumn;
-	}
-
-	set viewColumn(value) {
-		throw readonly('viewColumn');
-	}
-
-
-	get id(): string {
-		return this._id;
-	}
-}
-
 export class ExtHostNotebookDocumentsAndEditors implements ExtHostNotebookDocumentsAndEditorsShape {
 
 	private _disposables: Disposable[] = [];
@@ -194,7 +104,7 @@ export class ExtHostNotebookDocumentsAndEditors implements ExtHostNotebookDocume
 
 				const documentData = this._documents.get(resource.toString());
 				const editor = new ExtHostNotebookEditor(
-					this._mainContext.getProxy(SqlMainContext.MainThreadNotebook),
+					this._mainContext.getProxy(SqlMainContext.MainThreadNotebookDocumentsAndEditors),
 					data.id,
 					documentData,
 					typeConverters.ViewColumn.to(data.editorPosition)

src/sql/workbench/api/node/extHostNotebookEditor.ts

@@ -0,0 +1,210 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the Source EULA. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+'use strict';
+
+import * as sqlops from 'sqlops';
+import * as vscode from 'vscode';
+
+import { ok } from 'vs/base/common/assert';
+import { IDisposable } from 'vs/base/common/lifecycle';
+import { readonly } from 'vs/base/common/errors';
+import { TPromise } from 'vs/base/common/winjs.base';
+
+import { MainThreadNotebookDocumentsAndEditorsShape } from 'sql/workbench/api/node/sqlExtHost.protocol';
+import { ExtHostNotebookDocumentData } from 'sql/workbench/api/node/extHostNotebookDocumentData';
+import { CellRange, ISingleNotebookEditOperation, ICellRange } from 'sql/workbench/api/common/sqlExtHostTypes';
+
+export interface INotebookEditOperation {
+	range: sqlops.nb.CellRange;
+	cell: Partial<sqlops.nb.ICellContents>;
+	forceMoveMarkers: boolean;
+}
+
+export interface INotebookEditData {
+	documentVersionId: number;
+	edits: INotebookEditOperation[];
+	undoStopBefore: boolean;
+	undoStopAfter: boolean;
+}
+
+function toICellRange(range: sqlops.nb.CellRange): ICellRange {
+	return {
+		start: range.start,
+		end: range.end
+	};
+}
+
+export class NotebookEditorEdit {
+
+	private readonly _document: sqlops.nb.NotebookDocument;
+	private readonly _documentVersionId: number;
+	private _collectedEdits: INotebookEditOperation[];
+	private readonly _undoStopBefore: boolean;
+	private readonly _undoStopAfter: boolean;
+
+	constructor(document: sqlops.nb.NotebookDocument, options: { undoStopBefore: boolean; undoStopAfter: boolean; }) {
+		this._document = document;
+		// TODO add version handling
+		this._documentVersionId = 0;
+		// this._documentVersionId = document.version;
+		this._collectedEdits = [];
+		this._undoStopBefore = options ? options.undoStopBefore : true;
+		this._undoStopAfter = options ? options.undoStopAfter : false;
+	}
+
+	finalize(): INotebookEditData {
+		return {
+			documentVersionId: this._documentVersionId,
+			edits: this._collectedEdits,
+			undoStopBefore: this._undoStopBefore,
+			undoStopAfter: this._undoStopAfter
+		};
+	}
+
+	replace(location: number | CellRange, value: Partial<sqlops.nb.ICellContents>): void {
+		let range: CellRange = this.getAsRange(location);
+		this._pushEdit(range, value, false);
+	}
+
+	private getAsRange(location: number | CellRange): CellRange {
+		let range: CellRange = null;
+		if (typeof (location) === 'number') {
+			range = new CellRange(location, location+1);
+		}
+		else if (location instanceof CellRange) {
+			range = location;
+		}
+		else {
+			throw new Error('Unrecognized location');
+		}
+		return range;
+	}
+
+	insertCell(value:  Partial<sqlops.nb.ICellContents>, location?: number): void {
+		if (location === null || location === undefined) {
+			// If not specified, assume adding to end of list
+			location = this._document.cells.length - 1;
+		}
+		this._pushEdit(new CellRange(location, location), value, true);
+	}
+
+	deleteCell(index: number): void {
+		let range: CellRange = null;
+
+		if (typeof(index) === 'number') {
+			// Currently only allowing single-cell deletion.
+			// Do this by saying the range extends over 1 cell so on the main thread
+			// we can delete that cell, then handle insertions
+			range = new CellRange(index, index+1);
+		} else {
+			throw new Error('Unrecognized index');
+		}
+
+		this._pushEdit(range, null, true);
+	}
+
+	private _pushEdit(range: sqlops.nb.CellRange, cell:  Partial<sqlops.nb.ICellContents>, forceMoveMarkers: boolean): void {
+		let validRange = this._document.validateCellRange(range);
+		this._collectedEdits.push({
+			range: validRange,
+			cell: cell,
+			forceMoveMarkers: forceMoveMarkers
+		});
+	}
+}
+
+export class ExtHostNotebookEditor implements sqlops.nb.NotebookEditor, IDisposable {
+	private _disposed: boolean = false;
+
+	constructor(
+		private _proxy: MainThreadNotebookDocumentsAndEditorsShape,
+		private _id: string,
+		private readonly _documentData: ExtHostNotebookDocumentData,
+		private _viewColumn: vscode.ViewColumn
+	) {
+
+	}
+
+	dispose() {
+		ok(!this._disposed);
+		this._disposed = true;
+	}
+
+	get document(): sqlops.nb.NotebookDocument {
+		return this._documentData.document;
+	}
+
+	set document(value) {
+		throw readonly('document');
+	}
+
+	get viewColumn(): vscode.ViewColumn {
+		return this._viewColumn;
+	}
+
+	set viewColumn(value) {
+		throw readonly('viewColumn');
+	}
+
+	get id(): string {
+		return this._id;
+	}
+
+	edit(callback: (editBuilder: sqlops.nb.NotebookEditorEdit) => void, options?: { undoStopBefore: boolean; undoStopAfter: boolean; }): Thenable<boolean> {
+		if (this._disposed) {
+			return TPromise.wrapError<boolean>(new Error('NotebookEditor#edit not possible on closed editors'));
+		}
+		let edit = new NotebookEditorEdit(this._documentData.document, options);
+		callback(edit);
+		return this._applyEdit(edit);
+	}
+
+	private _applyEdit(editBuilder: NotebookEditorEdit): TPromise<boolean> {
+		let editData = editBuilder.finalize();
+
+		// return when there is nothing to do
+		if (editData.edits.length === 0) {
+			return TPromise.wrap(true);
+		}
+
+		// check that the edits are not overlapping (i.e. illegal)
+		let editRanges = editData.edits.map(edit => edit.range);
+
+		// sort ascending (by end and then by start)
+		editRanges.sort((a, b) => {
+			if (a.end === b.end) {
+				return a.start - b.start;
+			}
+			return a.end - b.end;
+		});
+
+		// check that no edits are overlapping
+		for (let i = 0, count = editRanges.length - 1; i < count; i++) {
+			const rangeEnd = editRanges[i].end;
+			const nextRangeStart = editRanges[i + 1].start;
+
+			if (nextRangeStart < rangeEnd) {
+				// overlapping ranges
+				return TPromise.wrapError<boolean>(
+					new Error('Overlapping ranges are not allowed!')
+				);
+			}
+		}
+
+		// prepare data for serialization
+		let edits: ISingleNotebookEditOperation[] = editData.edits.map((edit) => {
+			return {
+				range: toICellRange(edit.range),
+				cell: edit.cell,
+				forceMoveMarkers: edit.forceMoveMarkers
+			};
+		});
+
+		return this._proxy.$tryApplyEdits(this._id, editData.documentVersionId, edits, {
+			undoStopBefore: editData.undoStopBefore,
+			undoStopAfter: editData.undoStopAfter
+		});
+	}
+}

src/sql/workbench/api/node/mainThreadNotebookDocumentsAndEditors.ts

@@ -7,9 +7,8 @@
 import * as sqlops from 'sqlops';
 import { extHostNamedCustomer } from 'vs/workbench/api/electron-browser/extHostCustomers';
 import { Disposable } from 'vs/base/common/lifecycle';
-import { Registry } from 'vs/platform/registry/common/platform';
 import URI, { UriComponents } from 'vs/base/common/uri';
-import { IExtHostContext } from 'vs/workbench/api/node/extHost.protocol';
+import { IExtHostContext, IUndoStopOptions } from 'vs/workbench/api/node/extHost.protocol';
 import { IEditorService } from 'vs/workbench/services/editor/common/editorService';
 import { IEditorGroupsService } from 'vs/workbench/services/group/common/editorGroupsService';
 import { IInstantiationService } from 'vs/platform/instantiation/common/instantiation';
@@ -21,10 +20,11 @@ import {
 	INotebookDocumentsAndEditorsDelta, INotebookEditorAddData, INotebookShowOptions, INotebookModelAddedData
 } from 'sql/workbench/api/node/sqlExtHost.protocol';
 import { NotebookInputModel, NotebookInput } from 'sql/parts/notebook/notebookInput';
-import { INotebookService, INotebookEditor, DEFAULT_NOTEBOOK_FILETYPE, DEFAULT_NOTEBOOK_PROVIDER } from 'sql/services/notebook/notebookService';
+import { INotebookService, INotebookEditor } from 'sql/services/notebook/notebookService';
 import { TPromise } from 'vs/base/common/winjs.base';
-import { INotebookProviderRegistry, Extensions } from 'sql/services/notebook/notebookRegistry';
 import { getProviderForFileName } from 'sql/parts/notebook/notebookUtils';
+import { ISingleNotebookEditOperation } from 'sql/workbench/api/common/sqlExtHostTypes';
+import { disposed } from 'vs/base/common/errors';
 
 class MainThreadNotebookEditor extends Disposable {
 
@@ -58,6 +58,31 @@ class MainThreadNotebookEditor extends Disposable {
 		}
 		return input === this.editor.notebookParams.input;
 	}
+
+	public applyEdits(versionIdCheck: number, edits: ISingleNotebookEditOperation[], opts: IUndoStopOptions): boolean {
+		// TODO Handle version tracking
+		// if (this._model.getVersionId() !== versionIdCheck) {
+		// 	// throw new Error('Model has changed in the meantime!');
+		// 	// model changed in the meantime
+		// 	return false;
+		// }
+
+		if (!this.editor) {
+			// console.warn('applyEdits on invisible editor');
+			return false;
+		}
+
+		// TODO handle undo tracking
+		// if (opts.undoStopBefore) {
+		// 	this._codeEditor.pushUndoStop();
+		// }
+
+		this.editor.executeEdits(edits);
+		// if (opts.undoStopAfter) {
+		// 	this._codeEditor.pushUndoStop();
+		// }
+		return true;
+	}
 }
 
 function wait(timeMs: number): Promise<void> {
@@ -218,6 +243,7 @@ class MainThreadNotebookDocumentAndEditorStateComputer extends Disposable {
 
 @extHostNamedCustomer(SqlMainContext.MainThreadNotebookDocumentsAndEditors)
 export class MainThreadNotebookDocumentsAndEditors extends Disposable implements MainThreadNotebookDocumentsAndEditorsShape {
+
 	private _proxy: ExtHostNotebookDocumentsAndEditorsShape;
 	private _notebookEditors = new Map<string, MainThreadNotebookEditor>();
 
@@ -250,6 +276,14 @@ export class MainThreadNotebookDocumentsAndEditors extends Disposable implements
 	$tryShowNotebookDocument(resource: UriComponents, options: INotebookShowOptions): TPromise<string> {
 		return TPromise.wrap(this.doOpenEditor(resource, options));
 	}
+
+	$tryApplyEdits(id: string, modelVersionId: number, edits: ISingleNotebookEditOperation[], opts: IUndoStopOptions): TPromise<boolean> {
+		let editor = this.getEditor(id);
+		if (!editor) {
+			return TPromise.wrapError<boolean>(disposed(`TextEditor(${id})`));
+		}
+		return TPromise.as(editor.applyEdits(modelVersionId, edits, opts));
+	}
 	//#endregion
 
 	private async doOpenEditor(resource: UriComponents, options: INotebookShowOptions): Promise<string> {

src/sql/workbench/api/node/sqlExtHost.api.impl.ts

@@ -442,7 +442,8 @@ export function createApiFactory(
 				},
 				registerNotebookProvider(provider: sqlops.nb.NotebookProvider): vscode.Disposable {
 					return extHostNotebook.registerNotebookProvider(provider);
-				}
+				},
+				CellRange: sqlExtHostTypes.CellRange
 			};
 
 			return {

src/sql/workbench/api/node/sqlExtHost.protocol.ts

@@ -21,9 +21,10 @@ import { ITreeComponentItem } from 'sql/workbench/common/views';
 import { ITaskHandlerDescription } from 'sql/platform/tasks/common/tasks';
 import {
 	IItemConfig, ModelComponentTypes, IComponentShape, IModelViewDialogDetails, IModelViewTabDetails, IModelViewButtonDetails,
-	IModelViewWizardDetails, IModelViewWizardPageDetails, INotebookManagerDetails, INotebookSessionDetails, INotebookKernelDetails, INotebookFutureDetails, FutureMessageType, INotebookFutureDone
+	IModelViewWizardDetails, IModelViewWizardPageDetails, INotebookManagerDetails, INotebookSessionDetails, INotebookKernelDetails, INotebookFutureDetails, FutureMessageType, INotebookFutureDone, ISingleNotebookEditOperation
 } from 'sql/workbench/api/common/sqlExtHostTypes';
 import { EditorViewColumn } from 'vs/workbench/api/shared/editor';
+import { IUndoStopOptions } from 'vs/workbench/api/node/extHost.protocol';
 
 export abstract class ExtHostAccountManagementShape {
 	$autoOAuthCancelled(handle: number): Thenable<void> { throw ni(); }
@@ -819,4 +820,5 @@ export interface ExtHostNotebookDocumentsAndEditorsShape {
 export interface MainThreadNotebookDocumentsAndEditorsShape extends IDisposable {
 	$trySaveDocument(uri: UriComponents): Thenable<boolean>;
 	$tryShowNotebookDocument(resource: UriComponents, options: INotebookShowOptions): TPromise<string>;
+	$tryApplyEdits(id: string, modelVersionId: number, edits: ISingleNotebookEditOperation[], opts: IUndoStopOptions): TPromise<boolean>;
 }