From 2d52bc2a49dceeece2198e602daef7a097ee6921 Mon Sep 17 00:00:00 2001
From: Chris LaFreniere <40371649+chlafreniere@users.noreply.github.com>
Date: Wed, 2 Jan 2019 15:20:05 -0800
Subject: [PATCH] Notebooks: Fix Selection/Focus when New Cells Added (#3649)

* Improvemnents to Active Cell

* Fix minor spacing issue

* fix editor focus order

* Fix for add cell above/below

* cleanup logic to have activeCell logic all reside in notebook model
---
 .../notebook/cellViews/code.component.ts      |  1 +
 .../parts/notebook/models/notebookModel.ts    | 10 ++++--
 src/sql/parts/notebook/notebook.component.ts  | 35 +++++++++----------
 3 files changed, 26 insertions(+), 20 deletions(-)

diff --git a/src/sql/parts/notebook/cellViews/code.component.ts b/src/sql/parts/notebook/cellViews/code.component.ts
index 33f9d0937f..3451b8b4c5 100644
--- a/src/sql/parts/notebook/cellViews/code.component.ts
+++ b/src/sql/parts/notebook/cellViews/code.component.ts
@@ -128,6 +128,7 @@ export class CodeComponent extends AngularDisposable implements OnInit, OnChange
 		let uri = this.createUri();
 		this._editorInput = instantiationService.createInstance(UntitledEditorInput, uri, false, this.cellModel.language, '', '');
 		this._editor.setInput(this._editorInput, undefined);
+		this._editor.focus();
 		this._editorInput.resolve().then(model => {
 			this._editorModel = model.textEditorModel;
 			this._modelService.updateModel(this._editorModel, this.cellModel.source);
diff --git a/src/sql/parts/notebook/models/notebookModel.ts b/src/sql/parts/notebook/models/notebookModel.ts
index 14033436d4..959141ca27 100644
--- a/src/sql/parts/notebook/models/notebookModel.ts
+++ b/src/sql/parts/notebook/models/notebookModel.ts
@@ -203,9 +203,9 @@ export class NotebookModel extends Disposable implements INotebookModel {
 		return this._cells.findIndex((cell) => cell.equals(cellModel));
 	}
 
-	public addCell(cellType: CellType, index?: number): void {
+	public addCell(cellType: CellType, index?: number): ICellModel {
 		if (this.inErrorState || !this._cells) {
-			return;
+			return null;
 		}
 		let cell = this.createCell(cellType);
 
@@ -215,12 +215,18 @@ export class NotebookModel extends Disposable implements INotebookModel {
 			this._cells.push(cell);
 			index = undefined;
 		}
+		// Set newly created cell as active cell
+		this._activeCell.active = false;
+		this._activeCell = cell;
+		this._activeCell.active = true;
 
 		this._contentChangedEmitter.fire({
 			changeType: NotebookChangeType.CellsAdded,
 			cells: [cell],
 			cellIndex: index
 		});
+
+		return cell;
 	}
 
 	private createCell(cellType: CellType): ICellModel {
diff --git a/src/sql/parts/notebook/notebook.component.ts b/src/sql/parts/notebook/notebook.component.ts
index b73610379d..4d1e3e0065 100644
--- a/src/sql/parts/notebook/notebook.component.ts
+++ b/src/sql/parts/notebook/notebook.component.ts
@@ -63,14 +63,12 @@ export class NotebookComponent extends AngularDisposable implements OnInit, OnDe
 	private _isInErrorState: boolean = false;
 	private _errorMessage: string;
 	protected _actionBar: Taskbar;
-	private _activeCell: ICellModel;
 	protected isLoading: boolean;
 	private notebookManager: INotebookManager;
 	private _modelReadyDeferred = new Deferred<NotebookModel>();
 	private _modelRegisteredDeferred = new Deferred<NotebookModel>();
 	private profile: IConnectionProfile;
 	private _trustedAction: TrustedAction;
-	private _activeCellId: string;
 
 
 	constructor(
@@ -138,7 +136,7 @@ export class NotebookComponent extends AngularDisposable implements OnInit, OnDe
 	}
 
 	public get activeCellId(): string {
-		return this._activeCellId;
+		return this._model && this._model.activeCell ? this._model.activeCell.id :  '';
 	}
 
 	public get modelRegistered(): Promise<NotebookModel> {
@@ -158,32 +156,29 @@ export class NotebookComponent extends AngularDisposable implements OnInit, OnDe
 		if (event) {
 			event.stopPropagation();
 		}
-		if (cell !== this._activeCell) {
-			if (this._activeCell) {
-				this._activeCell.active = false;
+		if (cell !== this.model.activeCell) {
+			if (this.model.activeCell) {
+				this.model.activeCell.active = false;
 			}
-			this._activeCell = cell;
-			this._activeCell.active = true;
-			this._model.activeCell = this._activeCell;
-			this._activeCellId = cell.id;
+			this._model.activeCell = cell;
+			this._model.activeCell.active = true;
 			this._changeRef.detectChanges();
 		}
 	}
 
 	public unselectActiveCell() {
-		if (this._activeCell) {
-			this._activeCell.active = false;
+		if (this.model.activeCell) {
+			this.model.activeCell.active = false;
 		}
-		this._activeCell = null;
-		this._model.activeCell = null;
-		this._activeCellId = null;
+		this.model.activeCell = null;
 		this._changeRef.detectChanges();
 	}
 
 	// Add cell based on cell type
 	public addCell(cellType: CellType)
 	{
-		this._model.addCell(cellType);
+		let newCell = this._model.addCell(cellType);
+		this.selectCell(newCell);
 	}
 
 	// Updates Notebook model's trust details
@@ -201,12 +196,12 @@ export class NotebookComponent extends AngularDisposable implements OnInit, OnDe
 		switch (event.key) {
 			case 'ArrowDown':
 			case 'ArrowRight':
-				let nextIndex = (this.findCellIndex(this._activeCell) + 1) % this.cells.length;
+				let nextIndex = (this.findCellIndex(this.model.activeCell) + 1) % this.cells.length;
 				this.selectCell(this.cells[nextIndex]);
 				break;
 			case 'ArrowUp':
 			case 'ArrowLeft':
-				let index = this.findCellIndex(this._activeCell);
+				let index = this.findCellIndex(this.model.activeCell);
 				if (index === 0) {
 					index = this.cells.length;
 				}
@@ -252,6 +247,10 @@ export class NotebookComponent extends AngularDisposable implements OnInit, OnDe
 		this._register(model);
 		this._modelRegisteredDeferred.resolve(this._model);
 		model.backgroundStartSession();
+		// Set first cell as default active cell
+		if (this._model && this._model.cells && this._model.cells[0]) {
+			this.selectCell(model.cells[0]);
+		}
 		this._changeRef.detectChanges();
 	}