Integrate first SQL Notebooks Bits into Master (#3679)

* First crack tsql notebook (no output rendered yet)

* getting messages back

* intellisense working first cell, no connection errors

* sql notebook cell output functioning

* Latest SQL noteobook changes

* Undo change to launch.json

* Plumbing providers through

* Kernels shown from multiple providers, can switch between them. No mementos yet

* Ensure we have a feature flag for SQL notebooks, ensure existing functionality still works

* Fix tslint duplicate imports issue

* Addressing PR comments

* second round of PR feedback to cleanup notebook service manager code

* merge latest from master

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

@@ -38,9 +38,9 @@ export class CodeCellComponent extends CellView implements OnInit, OnChanges {
 
 	ngOnInit() {
 		if (this.cellModel) {
-			this.cellModel.onOutputsChanged(() => {
+			this._register(this.cellModel.onOutputsChanged(() => {
 				this._changeRef.detectChanges();
-			});
+			}));
 		}
 	}
 

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

@@ -33,9 +33,9 @@ export class OutputAreaComponent extends AngularDisposable implements OnInit {
 		this._register(this.themeService.onDidColorThemeChange(this.updateTheme, this));
 		this.updateTheme(this.themeService.getColorTheme());
 		if (this.cellModel) {
-			this.cellModel.onOutputsChanged(() => {
+			this._register(this.cellModel.onOutputsChanged(() => {
 				this._changeRef.detectChanges();
-			});
+			}));
 		}
 	}
 

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

@@ -91,9 +91,9 @@ export class TextCellComponent extends CellView implements OnInit, OnChanges {
 		this.setLoading(false);
 		this._register(this.themeService.onDidColorThemeChange(this.updateTheme, this));
 		this.updateTheme(this.themeService.getColorTheme());
-		this.cellModel.onOutputsChanged(e => {
+		this._register(this.cellModel.onOutputsChanged(e => {
 			this.updatePreview();
-		});
+		}));
 	}
 
 	ngOnChanges(changes: { [propKey: string]: SimpleChange }) {

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

@@ -226,6 +226,9 @@ export class CellModel implements ICellModel {
 			this._outputs.push(this.rewriteOutputUrls(output));
 			this.fireOutputsChanged();
 		}
+		if (!this._future.inProgress) {
+			this._future.dispose();
+		}
 	}
 
 	private rewriteOutputUrls(output: nb.ICellOutput): nb.ICellOutput {
@@ -327,6 +330,7 @@ export class CellModel implements ICellModel {
 		CellModel.LanguageMapping['pyspark3'] = 'python';
 		CellModel.LanguageMapping['python'] = 'python';
 		CellModel.LanguageMapping['scala'] = 'scala';
+		CellModel.LanguageMapping['sql'] = 'sql';
 	}
 
 	private get languageInfo(): nb.ILanguageInfo {

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

@@ -249,9 +249,9 @@ export interface INotebookModel {
 	readonly languageInfo: nb.ILanguageInfo;
 
 	/**
-	 * The notebook service used to call backend APIs
+	 * All notebook managers applicable for a given notebook
 	 */
-	readonly notebookManager: INotebookManager;
+	readonly notebookManagers: INotebookManager[];
 
 	/**
 	 * Event fired on first initialization of the kernel and
@@ -299,6 +299,11 @@ export interface INotebookModel {
 	 */
 	trustedMode: boolean;
 
+	/**
+	 * Current notebook provider id
+	 */
+	providerId: string;
+
 	/**
 	 * Change the current kernel from the Kernel dropdown
 	 * @param displayName kernel name (as displayed in Kernel dropdown)
@@ -411,7 +416,8 @@ export interface INotebookModelOptions {
 	 */
 	factory: IModelFactory;
 
-	notebookManager: INotebookManager;
+	notebookManagers: INotebookManager[];
+	providerId: string;
 
 	notificationService: INotificationService;
 	connectionService: IConnectionManagementService;

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

@@ -16,7 +16,7 @@ import { IClientSession, INotebookModel, IDefaultConnection, INotebookModelOptio
 import { NotebookChangeType, CellTypes, CellType } from 'sql/parts/notebook/models/contracts';
 import { nbversion } from '../notebookConstants';
 import * as notebookUtils from '../notebookUtils';
-import { INotebookManager } from 'sql/services/notebook/notebookService';
+import { INotebookManager, SQL_NOTEBOOK_PROVIDER, DEFAULT_NOTEBOOK_PROVIDER } from 'sql/services/notebook/notebookService';
 import { SparkMagicContexts } from 'sql/parts/notebook/models/sparkMagicContexts';
 import { IConnectionProfile } from 'sql/parts/connection/common/interfaces';
 import { NotebookConnection } from 'sql/parts/notebook/models/notebookConnection';
@@ -45,7 +45,8 @@ export class NotebookModel extends Disposable implements INotebookModel {
 	private _contentChangedEmitter = new Emitter<NotebookContentChange>();
 	private _kernelsChangedEmitter = new Emitter<nb.IKernelSpec>();
 	private _inErrorState: boolean = false;
-	private _clientSession: IClientSession;
+	private _clientSessions: IClientSession[] = [];
+	private _activeClientSession: IClientSession;
 	private _sessionLoadFinished: Promise<void>;
 	private _onClientSessionReady = new Emitter<IClientSession>();
 	private _activeContexts: IDefaultConnection;
@@ -60,20 +61,26 @@ export class NotebookModel extends Disposable implements INotebookModel {
 	private _hadoopConnection: NotebookConnection;
 	private _defaultKernel: nb.IKernelSpec;
 	private _activeCell: ICellModel;
+	private _providerId: string;
 
 	constructor(private notebookOptions: INotebookModelOptions, startSessionImmediately?: boolean, private connectionProfile?: IConnectionProfile) {
 		super();
-		if (!notebookOptions || !notebookOptions.notebookUri || !notebookOptions.notebookManager) {
+		if (!notebookOptions || !notebookOptions.notebookUri || !notebookOptions.notebookManagers) {
 			throw new Error('path or notebook service not defined');
 		}
 		if (startSessionImmediately) {
 			this.backgroundStartSession();
 		}
 		this._trustedMode = false;
+		this._providerId = notebookOptions.providerId;
+	}
+
+	public get notebookManagers(): INotebookManager[] {
+		return this.notebookOptions.notebookManagers.filter(manager => manager.providerId !== DEFAULT_NOTEBOOK_PROVIDER);
 	}
 
 	public get notebookManager(): INotebookManager {
-		return this.notebookOptions.notebookManager;
+		return this.notebookManagers.find(manager => manager.providerId === this._providerId);
 	}
 
 	public get notebookUri() : URI {
@@ -93,7 +100,7 @@ export class NotebookModel extends Disposable implements INotebookModel {
 	}
 
 	public get isSessionReady(): boolean {
-		return !!this._clientSession;
+		return !!this._activeClientSession;
 	}
 
 	/**
@@ -102,11 +109,11 @@ export class NotebookModel extends Disposable implements INotebookModel {
 	 * notebook environment
 	 */
 	public get clientSession(): IClientSession {
-		return this._clientSession;
+		return this._activeClientSession;
 	}
 
 	public get kernelChanged(): Event<nb.IKernelChangedArgs> {
-		return this.clientSession.kernelChanged;
+		return this._activeClientSession.kernelChanged;
 	}
 
 	public get kernelsChanged(): Event<nb.IKernelSpec> {
@@ -130,7 +137,21 @@ export class NotebookModel extends Disposable implements INotebookModel {
 	}
 
 	public get specs(): nb.IAllKernels | undefined {
-		return this.notebookManager.sessionManager.specs;
+		let specs: nb.IAllKernels = {
+			defaultKernel: undefined,
+			kernels: []
+		};
+		this.notebookManagers.forEach(manager => {
+			if (manager.sessionManager && manager.sessionManager.specs && manager.sessionManager.specs.kernels) {
+				manager.sessionManager.specs.kernels.forEach(kernel => {
+					specs.kernels.push(kernel);
+				});
+				if (!specs.defaultKernel) {
+					specs.defaultKernel = manager.sessionManager.specs.defaultKernel;
+				}
+			}
+		});
+		return specs;
 	}
 
 	public get inErrorState(): boolean {
@@ -145,6 +166,10 @@ export class NotebookModel extends Disposable implements INotebookModel {
 		return this._trustedMode;
 	}
 
+	public get providerId(): string {
+		return this._providerId;
+	}
+
 	public set trustedMode(isTrusted: boolean) {
 		this._trustedMode = isTrusted;
 		if (this._cells) {
@@ -178,11 +203,16 @@ export class NotebookModel extends Disposable implements INotebookModel {
 			this._trustedMode = isTrusted;
 			let contents = null;
 			if (this.notebookOptions.notebookUri.scheme !== Schemas.untitled) {
-				contents = await this.notebookManager.contentManager.getNotebookContents(this.notebookOptions.notebookUri);
+				// TODO: separate ContentManager from NotebookManager
+				contents = await this.notebookManagers[0].contentManager.getNotebookContents(this.notebookOptions.notebookUri);
 			}
 			let factory = this.notebookOptions.factory;
 			// if cells already exist, create them with language info (if it is saved)
 			this._cells = undefined;
+			this._defaultLanguageInfo = {
+				name: this._providerId === SQL_NOTEBOOK_PROVIDER ? 'sql' : 'python',
+				version: ''
+			};
 			if (contents) {
 				this._defaultLanguageInfo = this.getDefaultLanguageInfo(contents);
 				this._savedKernelInfo = this.getSavedKernelInfo(contents);
@@ -288,29 +318,36 @@ export class NotebookModel extends Disposable implements INotebookModel {
 	}
 
 	public backgroundStartSession(): void {
-		this._clientSession = this.notebookOptions.factory.createClientSession({
-			notebookUri: this.notebookOptions.notebookUri,
-			notebookManager: this.notebookManager,
-			notificationService: this.notebookOptions.notificationService
-		});
-		let profile = this.connectionProfile as IConnectionProfile;
-
-		if (this.isValidKnoxConnection(profile)) {
-            this._hadoopConnection = new NotebookConnection(this.connectionProfile);
-        } else {
-            this._hadoopConnection = undefined;
-		}
-
-		this._clientSession.initialize(this._hadoopConnection);
-		this._sessionLoadFinished = this._clientSession.ready.then(async () => {
-			if (this._clientSession.isInErrorState) {
-				this.setErrorState(this._clientSession.errorMessage);
-			} else {
-				this._onClientSessionReady.fire(this._clientSession);
-				// Once session is loaded, can use the session manager to retrieve useful info
-				this.loadKernelInfo();
-				await this.loadActiveContexts(undefined);
+		// TODO: only one session should be active at a time, depending on the current provider
+		this.notebookManagers.forEach(manager => {
+			let clientSession = this.notebookOptions.factory.createClientSession({
+				notebookUri: this.notebookOptions.notebookUri,
+				notebookManager: manager,
+				notificationService: this.notebookOptions.notificationService
+			});
+			this._clientSessions.push(clientSession);
+			if (!this._activeClientSession) {
+				this._activeClientSession = clientSession;
 			}
+			let profile = this.connectionProfile as IConnectionProfile;
+
+			if (this.isValidKnoxConnection(profile)) {
+				this._hadoopConnection = new NotebookConnection(this.connectionProfile);
+			} else {
+				this._hadoopConnection = undefined;
+			}
+
+			clientSession.initialize(this._hadoopConnection);
+			this._sessionLoadFinished = clientSession.ready.then(async () => {
+				if (clientSession.isInErrorState) {
+					this.setErrorState(clientSession.errorMessage);
+				} else {
+					this._onClientSessionReady.fire(clientSession);
+					// Once session is loaded, can use the session manager to retrieve useful info
+					this.loadKernelInfo();
+					await this.loadActiveContexts(undefined);
+				}
+			});
 		});
 	}
 
@@ -334,7 +371,8 @@ export class NotebookModel extends Disposable implements INotebookModel {
 	}
 
 	public doChangeKernel(kernelSpec: nb.IKernelSpec): Promise<void> {
-		return this._clientSession.changeKernel(kernelSpec)
+		this.findProviderIdForKernel(kernelSpec);
+		return this._activeClientSession.changeKernel(kernelSpec)
 			.then((kernel) => {
 				kernel.ready.then(() => {
 					if (kernel.info) {
@@ -359,7 +397,7 @@ export class NotebookModel extends Disposable implements INotebookModel {
 			SparkMagicContexts.configureContext(this.notebookOptions);
 			this._hadoopConnection = new NotebookConnection(newConnection);
 			this.refreshConnections(newConnection);
-			this._clientSession.updateConnection(this._hadoopConnection);
+			this._activeClientSession.updateConnection(this._hadoopConnection);
 		} catch (err) {
 			let msg = notebookUtils.getErrorMessage(err);
 			this.notifyError(localize('changeContextFailed', 'Changing context failed: {0}', msg));
@@ -381,20 +419,24 @@ export class NotebookModel extends Disposable implements INotebookModel {
 	}
 
 	private loadKernelInfo(): void {
-		this.clientSession.kernelChanged(async (e) => {
-			await this.loadActiveContexts(e);
+		this._clientSessions.forEach(clientSession => {
+			clientSession.kernelChanged(async (e) => {
+				await this.loadActiveContexts(e);
+			});
 		});
 		try {
 			let sessionManager = this.notebookManager.sessionManager;
 			if (sessionManager) {
 				let defaultKernel = SparkMagicContexts.getDefaultKernel(sessionManager.specs, this.connectionProfile, this._savedKernelInfo, this.notebookOptions.notificationService);
 				this._defaultKernel = defaultKernel;
-				this._clientSession.statusChanged(async (session) => {
-					if (session && session.defaultKernelLoaded === true) {
-						this._kernelsChangedEmitter.fire(defaultKernel);
-					} else if (session && !session.defaultKernelLoaded) {
-						this._kernelsChangedEmitter.fire({ name: notebookConstants.python3, display_name: notebookConstants.python3DisplayName });
-					}
+				this._clientSessions.forEach(clientSession => {
+					clientSession.statusChanged(async (session) => {
+						if (session && session.defaultKernelLoaded === true) {
+							this._kernelsChangedEmitter.fire(defaultKernel);
+						} else if (session && !session.defaultKernelLoaded) {
+							this._kernelsChangedEmitter.fire({ name: notebookConstants.python3, display_name: notebookConstants.python3DisplayName });
+						}
+					});
 				});
 				this.doChangeKernel(defaultKernel);
 			}
@@ -408,9 +450,9 @@ export class NotebookModel extends Disposable implements INotebookModel {
 	// Otherwise, default to python
 	private getDefaultLanguageInfo(notebook: nb.INotebookContents): nb.ILanguageInfo {
 		return notebook!.metadata!.language_info || {
-			name: 'python',
+			name: this._providerId === SQL_NOTEBOOK_PROVIDER ? 'sql' : 'python',
 			version: '',
-			mimetype: 'x-python'
+			mimetype: this._providerId === SQL_NOTEBOOK_PROVIDER ? 'x-sql' : 'x-python'
 		};
 	}
 
@@ -444,9 +486,10 @@ export class NotebookModel extends Disposable implements INotebookModel {
 
 	public async handleClosed(): Promise<void> {
 		try {
-			if (this._clientSession) {
-				await this._clientSession.shutdown();
-				this._clientSession = undefined;
+			if (this._activeClientSession) {
+				await this._activeClientSession.shutdown();
+				this._clientSessions = undefined;
+				this._activeClientSession = undefined;
 			}
 		} catch (err) {
 			this.notifyError(localize('shutdownError', 'An error occurred when closing the notebook: {0}', err));
@@ -482,7 +525,8 @@ export class NotebookModel extends Disposable implements INotebookModel {
 		if (!notebook) {
 			return false;
 		}
-		await this.notebookManager.contentManager.save(this.notebookOptions.notebookUri, notebook);
+		// TODO: refactor ContentManager out from NotebookManager
+		await this.notebookManagers[0].contentManager.save(this.notebookOptions.notebookUri, notebook);
 		this._contentChangedEmitter.fire({
 			changeType: NotebookChangeType.DirtyStateChanged,
 			isDirty: false
@@ -504,6 +548,23 @@ export class NotebookModel extends Disposable implements INotebookModel {
 			}
 		}
 	}
+
+	/**
+	 * Set _providerId and _activeClientSession based on a kernelSpec representing new kernel
+	 * @param kernelSpec KernelSpec for new kernel
+	 */
+	private findProviderIdForKernel(kernelSpec: nb.IKernelSpec): void {
+		for (let i = 0; i < this.notebookManagers.length; i++) {
+			if (this.notebookManagers[i].sessionManager && this.notebookManagers[i].sessionManager.specs && this.notebookManagers[i].sessionManager.specs.kernels) {
+				let index = this.notebookManagers[i].sessionManager.specs.kernels.findIndex(kernel => kernel.name === kernelSpec.name);
+				if (index >= 0) {
+					this._activeClientSession = this._clientSessions[i];
+					this._providerId = this.notebookManagers[i].providerId;
+					break;
+				}
+			}
+		}
+	}
 	/**
 	 * Serialize the model to JSON.
 	 */

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

@@ -34,7 +34,7 @@ import { AngularDisposable } from 'sql/base/common/lifecycle';
 import { CellTypes, CellType } from 'sql/parts/notebook/models/contracts';
 import { ICellModel, IModelFactory, notebookConstants, INotebookModel, NotebookContentChange } from 'sql/parts/notebook/models/modelInterfaces';
 import { IConnectionManagementService, IConnectionDialogService } from 'sql/parts/connection/common/connectionManagement';
-import { INotebookService, INotebookParams, INotebookManager, INotebookEditor, DEFAULT_NOTEBOOK_FILETYPE, DEFAULT_NOTEBOOK_PROVIDER } from 'sql/services/notebook/notebookService';
+import { INotebookService, INotebookParams, INotebookManager, INotebookEditor, DEFAULT_NOTEBOOK_FILETYPE, DEFAULT_NOTEBOOK_PROVIDER, SQL_NOTEBOOK_PROVIDER } from 'sql/services/notebook/notebookService';
 import { IBootstrapParams } from 'sql/services/bootstrap/bootstrapService';
 import { NotebookModel } from 'sql/parts/notebook/models/notebookModel';
 import { ModelFactory } from 'sql/parts/notebook/models/modelFactory';
@@ -64,7 +64,7 @@ export class NotebookComponent extends AngularDisposable implements OnInit, OnDe
 	private _errorMessage: string;
 	protected _actionBar: Taskbar;
 	protected isLoading: boolean;
-	private notebookManager: INotebookManager;
+	private notebookManagers: INotebookManager[] = [];
 	private _modelReadyDeferred = new Deferred<NotebookModel>();
 	private _modelRegisteredDeferred = new Deferred<NotebookModel>();
 	private profile: IConnectionProfile;
@@ -231,13 +231,17 @@ export class NotebookComponent extends AngularDisposable implements OnInit, OnDe
 
 	private async loadModel(): Promise<void> {
 		await this.awaitNonDefaultProvider();
-		this.notebookManager = await this.notebookService.getOrCreateNotebookManager(this._notebookParams.providerId, this._notebookParams.notebookUri);
+		for (let providerId of this._notebookParams.providers) {
+			let notebookManager = await this.notebookService.getOrCreateNotebookManager(providerId, this._notebookParams.notebookUri);
+			this.notebookManagers.push(notebookManager);
+		}
 		let model = new NotebookModel({
 			factory: this.modelFactory,
 			notebookUri: this._notebookParams.notebookUri,
 			connectionService: this.connectionManagementService,
 			notificationService: this.notificationService,
-			notebookManager: this.notebookManager
+			notebookManagers: this.notebookManagers,
+			providerId: notebookUtils.sqlNotebooksEnabled() ? 'sql' : 'jupyter' // this is tricky; really should also depend on the connection profile
 		}, false, this.profile);
 		model.onError((errInfo: INotification) => this.handleModelError(errInfo));
 		await model.requestModelLoad(this._notebookParams.isTrusted);
@@ -258,7 +262,9 @@ export class NotebookComponent extends AngularDisposable implements OnInit, OnDe
 		await this.notebookService.registrationComplete;
 		// Refresh the provider if we had been using default
 		if (DEFAULT_NOTEBOOK_PROVIDER === this._notebookParams.providerId) {
-			this._notebookParams.providerId = notebookUtils.getProviderForFileName(this._notebookParams.notebookUri.fsPath, this.notebookService);
+			let providers= notebookUtils.getProvidersForFileName(this._notebookParams.notebookUri.fsPath, this.notebookService);
+			let tsqlProvider = providers.find(provider => provider === SQL_NOTEBOOK_PROVIDER);
+			this._notebookParams.providerId = tsqlProvider ? SQL_NOTEBOOK_PROVIDER : providers[0];
 		}
 		if (DEFAULT_NOTEBOOK_PROVIDER === this._notebookParams.providerId) {
 			// If it's still the default, warn them they should install an extension

src/sql/parts/notebook/notebookEditor.ts

@@ -88,6 +88,7 @@ export class NotebookEditor extends BaseEditor {
 			notebookUri: input.notebookUri,
 			input: input,
 			providerId: input.providerId ? input.providerId : DEFAULT_NOTEBOOK_PROVIDER,
+			providers: input.providers ? input.providers : [DEFAULT_NOTEBOOK_PROVIDER],
 			isTrusted: input.isTrusted
 		};
 		bootstrapAngular(this.instantiationService,

src/sql/parts/notebook/notebookInput.ts

@@ -24,9 +24,10 @@ export class NotebookInputModel extends EditorModel {
 	private dirty: boolean;
 	private readonly _onDidChangeDirty: Emitter<void> = this._register(new Emitter<void>());
 	private _providerId: string;
-	constructor(public readonly notebookUri: URI, private readonly handle: number, private _isTrusted: boolean = false, private saveHandler?: ModeViewSaveHandler) {
+	constructor(public readonly notebookUri: URI, private readonly handle: number, private _isTrusted: boolean = false, private saveHandler?: ModeViewSaveHandler, provider?: string, private _providers?: string[]) {
 		super();
 		this.dirty = false;
+		this._providerId = provider;
 	}
 
 	public get providerId(): string {
@@ -37,6 +38,14 @@ export class NotebookInputModel extends EditorModel {
 		this._providerId = value;
 	}
 
+	public get providers(): string[] {
+		return this._providers;
+	}
+
+	public set providers(value: string[]) {
+		this._providers = value;
+	}
+
 	get isTrusted(): boolean {
 		return this._isTrusted;
 	}
@@ -99,6 +108,10 @@ export class NotebookInput extends EditorInput {
 		return this._model.providerId;
 	}
 
+	public get providers(): string[] {
+		return this._model.providers;
+	}
+
 	public getTypeId(): string {
 		return NotebookInput.ID;
 	}

src/sql/parts/notebook/notebookUtils.ts

@@ -39,21 +39,26 @@ export async function mkDir(dirPath: string, outputChannel?: IOutputChannel): Pr
 	}
 }
 
-export function getProviderForFileName(fileName: string, notebookService: INotebookService): string {
+export function getProvidersForFileName(fileName: string, notebookService: INotebookService): string[] {
 	let fileExt = path.extname(fileName);
-	let provider: string;
+	let providers: string[];
 	// First try to get provider for actual file type
 	if (fileExt && fileExt.startsWith('.')) {
 		fileExt = fileExt.slice(1,fileExt.length);
-		provider = notebookService.getProviderForFileType(fileExt);
+		providers = notebookService.getProvidersForFileType(fileExt);
 	}
 	// Fallback to provider for default file type (assume this is a global handler)
-	if (!provider) {
-		provider = notebookService.getProviderForFileType(DEFAULT_NOTEBOOK_FILETYPE);
+	if (!providers) {
+		providers = notebookService.getProvidersForFileType(DEFAULT_NOTEBOOK_FILETYPE);
 	}
 	// Finally if all else fails, use the built-in handler
-	if (!provider) {
-		provider = DEFAULT_NOTEBOOK_PROVIDER;
+	if (!providers) {
+		providers = [DEFAULT_NOTEBOOK_PROVIDER];
 	}
-	return provider;
+	return providers;
+}
+
+// Private feature flag to enable Sql Notebook experience
+export function sqlNotebooksEnabled() {
+	return process.env['SQLOPS_SQL_NOTEBOOK'] !== undefined;
 }