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/parts/notebook/models/modelInterfaces.ts

@@ -19,6 +19,7 @@ import { INotebookManager } from 'sql/services/notebook/notebookService';
 import { IConnectionProfile } from 'sql/parts/connection/common/interfaces';
 import { NotebookConnection } from 'sql/parts/notebook/models/notebookConnection';
 import { IConnectionManagementService } from 'sql/parts/connection/common/connectionManagement';
+import { ISingleNotebookEditOperation } from 'sql/workbench/api/common/sqlExtHostTypes';
 
 export interface IClientSessionOptions {
 	notebookUri: URI;
@@ -328,6 +329,14 @@ export interface INotebookModel {
 	 * Notifies the notebook of a change in the cell
 	 */
 	onCellChange(cell: ICellModel, change: NotebookChangeType): void;
+
+
+	/**
+	 * Push edit operations, basically editing the model. This is the preferred way of
+	 * editing the model. Long-term, this will ensure edit operations can be added to the undo stack
+	 * @param edits The edit operations to perform
+	 */
+	pushEditOperations(edits: ISingleNotebookEditOperation[]): void;
 }
 
 export interface ICellModelOptions {

src/sql/parts/notebook/models/notebookModel.ts

@@ -22,6 +22,7 @@ import { IConnectionProfile } from 'sql/parts/connection/common/interfaces';
 import { NotebookConnection } from 'sql/parts/notebook/models/notebookConnection';
 import { INotification, Severity } from 'vs/platform/notification/common/notification';
 import { Schemas } from 'vs/base/common/network';
+import { ISingleNotebookEditOperation } from 'sql/workbench/api/common/sqlExtHostTypes';
 
 /*
 * Used to control whether a message in a dialog/wizard is displayed as an error,
@@ -263,6 +264,25 @@ export class NotebookModel extends Disposable implements INotebookModel {
 		}
 	}
 
+	pushEditOperations(edits: ISingleNotebookEditOperation[]): void {
+		if (this.inErrorState || !this._cells) {
+			return;
+		}
+
+		for (let edit of edits) {
+			let newCells: ICellModel[] = [];
+			if (edit.cell) {
+				// TODO: should we validate and complete required missing parameters?
+				let contents: nb.ICellContents = edit.cell as nb.ICellContents;
+				newCells.push(this.notebookOptions.factory.createCell(contents, { notebook: this, isTrusted: this._trustedMode }));
+			}
+			this._cells.splice(edit.range.start, edit.range.end - edit.range.start, ...newCells);
+			this._contentChangedEmitter.fire({
+				changeType: NotebookChangeType.CellsAdded
+			});
+		}
+	}
+
 	public get activeCell(): ICellModel {
 		return this._activeCell;
 	}

src/sql/parts/notebook/notebook.component.ts

@@ -40,6 +40,7 @@ import { fillInActions, LabeledMenuItemActionItem } from 'vs/platform/actions/br
 import { IObjectExplorerService } from 'sql/parts/objectExplorer/common/objectExplorerService';
 import * as TaskUtilities from 'sql/workbench/common/taskUtilities';
 import { IEditorService } from 'vs/workbench/services/editor/common/editorService';
+import { ISingleNotebookEditOperation } from 'sql/workbench/api/common/sqlExtHostTypes';
 
 export const NOTEBOOK_SELECTOR: string = 'notebook-component';
 
@@ -379,4 +380,12 @@ export class NotebookComponent extends AngularDisposable implements OnInit, OnDe
 	isDirty(): boolean {
 		return this.notebookParams.input.isDirty();
 	}
+
+	executeEdits(edits: ISingleNotebookEditOperation[]): boolean {
+		if (!edits || edits.length === 0) {
+			return false;
+		}
+		this._model.pushEditOperations(edits);
+		return true;
+	}
 }

src/sql/services/notebook/notebookService.ts

@@ -15,6 +15,7 @@ import { RenderMimeRegistry } from 'sql/parts/notebook/outputs/registry';
 import { ModelFactory } from 'sql/parts/notebook/models/modelFactory';
 import { IConnectionProfile } from 'sql/parts/connection/common/interfaces';
 import { NotebookInput } from 'sql/parts/notebook/notebookInput';
+import { ISingleNotebookEditOperation } from 'sql/workbench/api/common/sqlExtHostTypes';
 
 export const SERVICE_ID = 'notebookService';
 export const INotebookService = createDecorator<INotebookService>(SERVICE_ID);
@@ -87,4 +88,5 @@ export interface INotebookEditor {
 	isActive(): boolean;
 	isVisible(): boolean;
 	save(): Promise<boolean>;
+	executeEdits(edits: ISingleNotebookEditOperation[]): boolean;
 }

src/sql/sqlops.proposed.d.ts

@@ -1482,6 +1482,42 @@ declare module 'sqlops' {
 			 * will return false.
 			 */
 			save(): Thenable<boolean>;
+
+			/**
+			 * Ensure a cell range is completely contained in this document.
+			 *
+			 * @param range A cell range.
+			 * @return The given range or a new, adjusted range.
+			 */
+			validateCellRange(range: CellRange): CellRange;
+		}
+
+		/**
+		 * A cell range represents an ordered pair of two positions in a list of cells.
+		 * It is guaranteed that [start](#CellRange.start).isBeforeOrEqual([end](#CellRange.end))
+		 *
+		 * CellRange objects are __immutable__.
+		 */
+		export class CellRange {
+
+			/**
+			 * The start index. It is before or equal to [end](#CellRange.end).
+			 */
+			readonly start: number;
+
+			/**
+			 * The end index. It is after or equal to [start](#CellRange.start).
+			 */
+			readonly end: number;
+
+			/**
+			 * Create a new range from two positions. If `start` is not
+			 * before or equal to `end`, the values will be swapped.
+			 *
+			 * @param start A number.
+			 * @param end A number.
+			 */
+			constructor(start: number, end: number);
 		}
 
 		export interface NotebookEditor {
@@ -1495,6 +1531,19 @@ declare module 'sqlops' {
 			 * column is larger than three.
 			 */
 			viewColumn?: vscode.ViewColumn;
+
+			/**
+			 * Perform an edit on the document associated with this notebook editor.
+			 *
+			 * The given callback-function is invoked with an [edit-builder](#NotebookEditorEdit) which must
+			 * be used to make edits. Note that the edit-builder is only valid while the
+			 * callback executes.
+			 *
+			 * @param callback A function which can create edits using an [edit-builder](#NotebookEditorEdit).
+			 * @param options The undo/redo behavior around this edit. By default, undo stops will be created before and after this edit.
+			 * @return A promise that resolves with a value indicating if the edits could be applied.
+			 */
+			edit(callback: (editBuilder: NotebookEditorEdit) => void, options?: { undoStopBefore: boolean; undoStopAfter: boolean; }): Thenable<boolean>;
 		}
 
 		export interface NotebookCell {
@@ -1552,6 +1601,38 @@ declare module 'sqlops' {
 			kind?: vscode.TextEditorSelectionChangeKind;
 		}
 
+		/**
+		 * A complex edit that will be applied in one transaction on a TextEditor.
+		 * This holds a description of the edits and if the edits are valid (i.e. no overlapping regions, document was not changed in the meantime, etc.)
+		 * they can be applied on a [document](#TextDocument) associated with a [text editor](#TextEditor).
+		 *
+		 */
+		export interface NotebookEditorEdit {
+			/**
+			 * Replace a cell range with a new cell.
+			 *
+			 * @param location The range this operation should remove.
+			 * @param value The new cell this operation should insert after removing `location`.
+			 */
+			replace(location: number | CellRange, value: ICellContents): void;
+
+			/**
+			 * Insert a cell (optionally) at a specific index. Any index outside of the length of the cells
+			 * will result in the cell being added at the end.
+			 *
+			 * @param index The position where the new text should be inserted.
+			 * @param value The new text this operation should insert.
+			 */
+			insertCell(value: ICellContents, index?: number): void;
+
+			/**
+			 * Delete a certain cell.
+			 *
+			 * @param index The index of the cell to remove.
+			 */
+			deleteCell(index: number): void;
+		}
+
 		/**
 		 * Register a notebook provider. The supported file types handled by this
 		 * provider are defined in the `package.json:

src/sql/workbench/api/common/sqlExtHostTypes.ts

@@ -4,6 +4,7 @@
  *--------------------------------------------------------------------------------------------*/
 'use strict';
 
+import { nb } from 'sqlops';
 import { TreeItem } from 'vs/workbench/api/node/extHostTypes';
 
 // SQL added extension host types
@@ -457,4 +458,44 @@ export enum FutureMessageType {
 export interface INotebookFutureDone {
 	succeeded: boolean;
 	rejectReason: string;
-}
+}
+
+export interface ICellRange {
+	readonly start: number;
+	readonly end: number;
+}
+
+export class CellRange {
+
+	protected _start: number;
+	protected _end: number;
+
+	get start(): number {
+		return this._start;
+	}
+
+	get end(): number {
+		return this._end;
+	}
+
+	constructor(start: number, end: number) {
+		if (typeof(start) !== 'number' || typeof(start) !== 'number' || start < 0 || end < 0) {
+			throw new Error('Invalid arguments');
+		}
+
+		// Logic taken from range handling.
+		if (start <= end) {
+			this._start = start;
+			this._end = end;
+		} else {
+			this._start = end;
+			this._end = start;
+		}
+	}
+}
+
+export interface ISingleNotebookEditOperation {
+	range: ICellRange;
+	cell: Partial<nb.ICellContents>;
+	forceMoveMarkers: boolean;
+}

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 } from 'vs/base/common/lifecycle';
-import { dispose, IDisposable } 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,
+	SqlMainContext, INotebookDocumentsAndEditorsDelta, ExtHostNotebookDocumentsAndEditorsShape,
-	ExtHostNotebookDocumentsAndEditorsShape, MainThreadNotebookDocumentsAndEditorsShape, INotebookShowOptions
+	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>;
 }

src/sqltest/parts/notebook/common.ts

@@ -11,6 +11,7 @@ import { Event, Emitter } from 'vs/base/common/event';
 import { INotebookModel, ICellModel, IClientSession, IDefaultConnection } from 'sql/parts/notebook/models/modelInterfaces';
 import { NotebookChangeType, CellType } from 'sql/parts/notebook/models/contracts';
 import { INotebookManager } from 'sql/services/notebook/notebookService';
+import { ISingleNotebookEditOperation } from 'sql/workbench/api/common/sqlExtHostTypes';
 
 export class NotebookModelStub implements INotebookModel {
     constructor(private _languageInfo?: nb.ILanguageInfo) {
@@ -67,6 +68,9 @@ export class NotebookModelStub implements INotebookModel {
     saveModel(): Promise<boolean> {
         throw new Error('Method not implemented.');
     }
+    pushEditOperations(edits: ISingleNotebookEditOperation[]): void {
+        throw new Error('Method not implemented.');
+    }
 }
 
 export class NotebookManagerStub implements INotebookManager {