Notebooks: Only have One Connection in Attach To Dropdown (#8582)

* Have only one connection in attach to dropdown

* LGTM fixes

* Test fixes not unnecessarily changing context 2x

* PR Feedback

extensions/integration-tests/src/notebook.test.ts

@@ -9,8 +9,8 @@ import * as azdata from 'azdata';
 import * as vscode from 'vscode';
 import { isTestSetupCompleted } from './testContext';
 import { sqlNotebookContent, writeNotebookToFile, sqlKernelMetadata, getFileName, pySparkNotebookContent, pySparkKernelMetadata, pythonKernelMetadata, sqlNotebookMultipleCellsContent, notebookContentForCellLanguageTest, sqlKernelSpec, pythonKernelSpec, pySparkKernelSpec, CellTypes } from './notebook.util';
-import { getBdcServer, getConfigValue, EnvironmentVariable_PYTHON_PATH } from './testConfig';
+import { getBdcServer, getConfigValue, EnvironmentVariable_PYTHON_PATH, TestServerProfile } from './testConfig';
-import { connectToServer, sleep } from './utils';
+import { connectToServer, sleep, testServerProfileToIConnectionProfile } from './utils';
 import * as fs from 'fs';
 import { stressify } from 'adstest';
 import { isNullOrUndefined } from 'util';
@@ -392,15 +392,20 @@ class NotebookTester {
 	async openNotebook(content: azdata.nb.INotebookContents, kernelMetadata: any, testName: string, connectToDifferentServer?: boolean): Promise<azdata.nb.NotebookEditor> {
 		let notebookConfig = vscode.workspace.getConfiguration('notebook');
 		notebookConfig.update('pythonPath', getConfigValue(EnvironmentVariable_PYTHON_PATH), 1);
+		let server: TestServerProfile;
 		if (!connectToDifferentServer) {
-			let server = await getBdcServer();
+			server = await getBdcServer();
 			assert(server && server.serverName, 'No server could be found in openNotebook');
 			await connectToServer(server, 6000);
 		}
 		let notebookJson = Object.assign({}, content, { metadata: kernelMetadata });
 		let uri = writeNotebookToFile(notebookJson, testName);
 		console.log('Notebook uri ' + uri);
-		let notebook = await azdata.nb.showNotebookDocument(uri);
+		let nbShowOptions: azdata.nb.NotebookShowOptions;
+		if (server) {
+			nbShowOptions = { connectionProfile: testServerProfileToIConnectionProfile(server) };
+		}
+		let notebook = await azdata.nb.showNotebookDocument(uri, nbShowOptions);
 		return notebook;
 	}
 

extensions/integration-tests/src/utils.ts

@@ -271,3 +271,21 @@ export async function assertTableCreationResult(databaseName: string, schema: st
 		assert(result.columnInfo[0].columnName !== 'ErrorMessage', `Checking for table creation threw error ${result.rows[0][0].displayValue}`);
 	}
 }
+
+export function testServerProfileToIConnectionProfile(serverProfile: TestServerProfile): azdata.IConnectionProfile {
+	return {
+		serverName: serverProfile.serverName,
+		databaseName: serverProfile.database,
+		authenticationType: serverProfile.authenticationTypeName,
+		providerName: serverProfile.providerName,
+		connectionName: '',
+		userName: serverProfile.userName,
+		password: serverProfile.password,
+		savePassword: false,
+		groupFullName: undefined,
+		saveProfile: true,
+		id: undefined,
+		groupId: undefined,
+		options: {}
+	};
+}

src/sql/workbench/contrib/notebook/browser/cellViews/code.component.ts

@@ -38,6 +38,7 @@ import { UntitledTextEditorModel } from 'vs/workbench/common/editor/untitledText
 
 export const CODE_SELECTOR: string = 'code-component';
 const MARKDOWN_CLASS = 'markdown';
+const DEFAULT_OR_LOCAL_CONTEXT_ID = '-1';
 
 @Component({
 	selector: CODE_SELECTOR,
@@ -146,8 +147,8 @@ export class CodeComponent extends AngularDisposable implements OnInit, OnChange
 			let connectionService = this.connectionService;
 			if (!shouldConnect && connectionService && connectionService.isConnected(cellUri)) {
 				connectionService.disconnect(cellUri).catch(e => this.logService.error(e));
-			} else if (shouldConnect && this._model.activeConnection && this._model.activeConnection.id !== '-1') {
+			} else if (shouldConnect && this._model.context && this._model.context.id !== DEFAULT_OR_LOCAL_CONTEXT_ID) {
-				connectionService.connect(this._model.activeConnection, cellUri).catch(e => this.logService.error(e));
+				connectionService.connect(this._model.context, cellUri).catch(e => this.logService.error(e));
 			}
 		}
 	}

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

@@ -316,7 +316,7 @@ export class CellModel implements ICellModel {
 				this.sendNotification(notificationService, Severity.Info, localize('runCellCancelled', "Cell execution cancelled"));
 			} else {
 				// TODO update source based on editor component contents
-				if (kernel.requiresConnection && !this.notebookModel.activeConnection) {
+				if (kernel.requiresConnection && !this.notebookModel.context) {
 					let connected = await this.notebookModel.requestConnection();
 					if (!connected) {
 						return false;
@@ -509,11 +509,11 @@ export class CellModel implements ICellModel {
 				let result = output as nb.IDisplayResult;
 				if (result && result.data && result.data['text/html']) {
 					let model = this._options.notebook as NotebookModel;
-					if (model.activeConnection) {
+					if (model.context) {
-						let gatewayEndpointInfo = this.getGatewayEndpoint(model.activeConnection);
+						let gatewayEndpointInfo = this.getGatewayEndpoint(model.context);
 						if (gatewayEndpointInfo) {
 							let hostAndIp = notebookUtils.getHostAndPortFromEndpoint(gatewayEndpointInfo.endpoint);
-							let host = hostAndIp.host ? hostAndIp.host : model.activeConnection.serverName;
+							let host = hostAndIp.host ? hostAndIp.host : model.context.serverName;
 							let port = hostAndIp.port ? ':' + hostAndIp.port : defaultPort;
 							let html = result.data['text/html'];
 							// CTP 3.1 and earlier Spark link

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

@@ -214,11 +214,6 @@ export interface IClientSession extends IDisposable {
 	onKernelChanging(changeHandler: ((kernel: nb.IKernelChangedArgs) => Promise<void>)): void;
 }
 
-export interface IDefaultConnection {
-	defaultConnection: ConnectionProfile;
-	otherConnections: ConnectionProfile[];
-}
-
 /**
  * A kernel preference.
  */
@@ -323,10 +318,10 @@ export interface INotebookModel {
 	readonly specs: nb.IAllKernels | undefined;
 
 	/**
-	 * The specs for available contexts, or undefined if these have
+	 * The specs for available context, or undefined if this has
 	 * not been loaded yet
 	 */
-	readonly contexts: IDefaultConnection | undefined;
+	readonly context: ConnectionProfile | undefined;
 
 	/**
 	 * Event fired on first initialization of the cells and

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

@@ -3,113 +3,44 @@
  *  Licensed under the Source EULA. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
-import { nb } from 'azdata';
 
 import { localize } from 'vs/nls';
-import { IDefaultConnection } from 'sql/workbench/contrib/notebook/browser/models/modelInterfaces';
-import { IConnectionManagementService } from 'sql/platform/connection/common/connectionManagement';
 import { ConnectionProfile } from 'sql/platform/connection/common/connectionProfile';
-import { IConnectionProfile } from 'sql/platform/connection/common/interfaces';
 import { mssqlProviderName } from 'sql/platform/connection/common/constants';
-import { find } from 'vs/base/common/arrays';
+
 
 export class NotebookContexts {
 
-	public static get DefaultContext(): IDefaultConnection {
+	public static get DefaultContext(): ConnectionProfile {
-		let defaultConnection: ConnectionProfile = <any>{
+		return <any>{
 			providerName: mssqlProviderName,
 			id: '-1',
 			serverName: localize('selectConnection', "Select Connection")
 		};
-
-		return {
-			// default context if no other contexts are applicable
-			defaultConnection: defaultConnection,
-			otherConnections: [defaultConnection]
-		};
 	}
 
-	public static get LocalContext(): IDefaultConnection {
+	public static get LocalContext(): ConnectionProfile {
-		let localConnection: ConnectionProfile = <any>{
+		return <any>{
 			providerName: mssqlProviderName,
 			id: '-1',
 			serverName: localize('localhost', "localhost")
 		};
-
-		return {
-			// default context if no other contexts are applicable
-			defaultConnection: localConnection,
-			otherConnections: [localConnection]
-		};
 	}
 
 	/**
-	 * Get all of the applicable contexts for a given kernel
+	 * Get the applicable context for a given kernel
-	 * @param connectionService connection management service
+	 * @param context current connection profile
 	 * @param connProviderIds array of connection provider ids applicable for a kernel
-	 * @param kernelChangedArgs kernel changed args (both old and new kernel info)
-	 * @param profile current connection profile
 	 */
-	public static getContextsForKernel(connectionService: IConnectionManagementService, connProviderIds: string[], kernelChangedArgs?: nb.IKernelChangedArgs, profile?: IConnectionProfile): IDefaultConnection {
+	public static getContextForKernel(context: ConnectionProfile, connProviderIds: string[]): ConnectionProfile {
-		let connections: IDefaultConnection = this.DefaultContext;
-		if (!profile) {
-			if (!kernelChangedArgs || !kernelChangedArgs.newValue ||
-				(kernelChangedArgs.oldValue && kernelChangedArgs.newValue.id === kernelChangedArgs.oldValue.id)) {
-				// nothing to do, kernels are the same or new kernel is undefined
-				return connections;
-			}
-		}
-		if (kernelChangedArgs && kernelChangedArgs.newValue && kernelChangedArgs.newValue.name && connProviderIds.length < 1) {
-			return connections;
-		} else {
-			connections = this.getActiveContexts(connectionService, connProviderIds, profile);
-		}
-		return connections;
-	}
-
-	/**
-	 * Get all active contexts and sort them
-	 * @param connectionService connection service
-	 * @param connProviderIds array of applicable connection providers to filter connections
-	 * @param profile connection profile passed when launching notebook
-	 */
-	public static getActiveContexts(connectionService: IConnectionManagementService, connProviderIds: string[], profile?: IConnectionProfile): IDefaultConnection {
-		let defaultConnection: ConnectionProfile = NotebookContexts.DefaultContext.defaultConnection;
-		let activeConnections: ConnectionProfile[] = connectionService.getActiveConnections();
-		if (activeConnections && activeConnections.length > 0) {
-			activeConnections = activeConnections.filter(conn => conn.id !== '-1');
-		}
 		// If no connection provider ids exist for a given kernel, the attach to should show localhost
 		if (connProviderIds.length === 0) {
 			return NotebookContexts.LocalContext;
 		}
-		// If no active connections exist, show "Select connection" as the default value
+		if (context && context.providerName && connProviderIds.filter(p => p === context.providerName).length > 0) {
-		if (activeConnections.length === 0) {
+			return context;
-			return NotebookContexts.DefaultContext;
 		}
-		// Filter active connections by their provider ids to match kernel's supported connection providers
+		return NotebookContexts.DefaultContext;
-		else if (activeConnections.length > 0) {
-			let connections = activeConnections.filter(connection => {
-				return connProviderIds.some(x => x === connection.providerName);
-			});
-			if (connections && connections.length > 0) {
-				defaultConnection = connections[0];
-				if (profile && profile.options) {
-					let matchingConn = find(connections, connection => connection.serverName === profile.serverName);
-					if (matchingConn) {
-						defaultConnection = matchingConn;
-					}
-				}
-			} else if (connections.length === 0) {
-				return NotebookContexts.DefaultContext;
-			}
-			activeConnections = [];
-			connections.forEach(connection => activeConnections.push(connection));
-		}
-
-		return {
-			otherConnections: activeConnections,
-			defaultConnection: defaultConnection
-		};
 	}
+
 }

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

@@ -9,7 +9,7 @@ import { localize } from 'vs/nls';
 import { Event, Emitter } from 'vs/base/common/event';
 import { Disposable } from 'vs/base/common/lifecycle';
 
-import { IClientSession, INotebookModel, IDefaultConnection, INotebookModelOptions, ICellModel, NotebookContentChange, notebookConstants } from 'sql/workbench/contrib/notebook/browser/models/modelInterfaces';
+import { IClientSession, INotebookModel, INotebookModelOptions, ICellModel, NotebookContentChange, notebookConstants } from 'sql/workbench/contrib/notebook/browser/models/modelInterfaces';
 import { NotebookChangeType, CellType, CellTypes } from 'sql/workbench/contrib/notebook/common/models/contracts';
 import { nbversion } from 'sql/workbench/contrib/notebook/common/models/notebookConstants';
 import * as notebookUtils from 'sql/workbench/contrib/notebook/browser/models/notebookUtils';
@@ -56,7 +56,6 @@ export class NotebookModel extends Disposable implements INotebookModel {
 	private _sessionLoadFinished: Promise<void>;
 	private _onClientSessionReady = new Emitter<IClientSession>();
 	private _onProviderIdChanged = new Emitter<string>();
-	private _activeContexts: IDefaultConnection;
 	private _trustedMode: boolean;
 	private _onActiveCellChanged = new Emitter<ICellModel>();
 
@@ -69,7 +68,6 @@ export class NotebookModel extends Disposable implements INotebookModel {
 	private readonly _nbformat: number = nbversion.MAJOR_VERSION;
 	private readonly _nbformatMinor: number = nbversion.MINOR_VERSION;
 	private _activeConnection: ConnectionProfile;
-	private _otherConnections: ConnectionProfile[] = [];
 	private _activeCell: ICellModel;
 	private _providerId: string;
 	private _defaultKernel: nb.IKernelSpec;
@@ -190,8 +188,8 @@ export class NotebookModel extends Disposable implements INotebookModel {
 		return this._cells;
 	}
 
-	public get contexts(): IDefaultConnection {
+	public get context(): ConnectionProfile {
-		return this._activeContexts;
+		return this._activeConnection;
 	}
 
 	public get specs(): nb.IAllKernels | undefined {
@@ -244,10 +242,6 @@ export class NotebookModel extends Disposable implements INotebookModel {
 		});
 	}
 
-	public get activeConnection(): IConnectionProfile {
-		return this._activeConnection;
-	}
-
 	/**
 	 * Indicates the server has finished loading. It may have failed to load in
 	 * which case the view will be in an error state.
@@ -457,14 +451,8 @@ export class NotebookModel extends Disposable implements INotebookModel {
 			}
 			let profile = new ConnectionProfile(this._notebookOptions.capabilitiesService, this.connectionProfile);
 
-			// TODO: this code needs to be fixed since it is called before the this._savedKernelInfo is set.
-			// This means it always fails, and we end up using the default connection instead. If you right-click
-			// and run "New Notebook" on a disconnected server this means you get the wrong connection (global active)
-			// instead of the one you chose, or it'll fail to connect in general
 			if (this.isValidConnection(profile)) {
 				this._activeConnection = profile;
-			} else {
-				this._activeConnection = undefined;
 			}
 
 			clientSession.onKernelChanging(async (e) => {
@@ -687,19 +675,13 @@ export class NotebookModel extends Disposable implements INotebookModel {
 
 	public async changeContext(title: string, newConnection?: ConnectionProfile, hideErrorMessage?: boolean): Promise<void> {
 		try {
-			if (!newConnection) {
+			if ((!newConnection) && this._activeConnection && (this._activeConnection.title === title)) {
-				newConnection = find(this._activeContexts.otherConnections, (connection) => connection.title === title);
+				newConnection = this._activeConnection;
-			}
-			if ((!newConnection) && (this._activeContexts.defaultConnection.title === title)) {
-				newConnection = this._activeContexts.defaultConnection;
 			}
 
 			if (newConnection) {
-				if (this._activeConnection && this._activeConnection.id !== newConnection.id) {
-					this._otherConnections.push(this._activeConnection);
-				}
 				this._activeConnection = newConnection;
-				this.refreshConnections(newConnection);
+				this.setActiveConnectionIfDifferent(newConnection);
 				this._activeClientSession.updateConnection(newConnection.toIConnectionProfile()).then(
 					result => {
 						//Remove 'Select connection' from 'Attach to' drop-down since its a valid connection
@@ -723,17 +705,12 @@ export class NotebookModel extends Disposable implements INotebookModel {
 		}
 	}
 
-	private refreshConnections(newConnection: ConnectionProfile) {
+	private setActiveConnectionIfDifferent(newConnection: ConnectionProfile) {
 		if (this.isValidConnection(newConnection) &&
 			this._activeConnection.id !== '-1' &&
-			this._activeConnection.id !== this._activeContexts.defaultConnection.id) {
+			this._activeConnection.id !== newConnection.id) {
-			// Put the defaultConnection to the head of otherConnections
+			// Change the active connection to newConnection
-			if (this.isValidConnection(this._activeContexts.defaultConnection)) {
+			this._activeConnection = newConnection;
-				this._activeContexts.otherConnections = this._activeContexts.otherConnections.filter(conn => conn.id !== this._activeContexts.defaultConnection.id);
-				this._activeContexts.otherConnections.unshift(this._activeContexts.defaultConnection);
-			}
-			// Change the defaultConnection to newConnection
-			this._activeContexts.defaultConnection = newConnection;
 		}
 	}
 
@@ -803,17 +780,9 @@ export class NotebookModel extends Disposable implements INotebookModel {
 
 	public async handleClosed(): Promise<void> {
 		try {
-			if (this.notebookOptions && this.notebookOptions.connectionService) {
+			if (this.notebookOptions && this.notebookOptions.connectionService && this._activeConnection) {
-				if (this._otherConnections) {
+				await this.disconnectNotebookConnection(this._activeConnection);
-					notebookUtils.asyncForEach(this._otherConnections, async (conn) => {
+				this._activeConnection = undefined;
-						await this.disconnectNotebookConnection(conn);
-					});
-					this._otherConnections = [];
-				}
-				if (this._activeConnection) {
-					await this.disconnectNotebookConnection(this._activeConnection);
-					this._activeConnection = undefined;
-				}
 			}
 			await this.shutdownActiveSession();
 		} catch (err) {
@@ -837,11 +806,11 @@ export class NotebookModel extends Disposable implements INotebookModel {
 	private async loadActiveContexts(kernelChangedArgs: nb.IKernelChangedArgs): Promise<void> {
 		if (kernelChangedArgs && kernelChangedArgs.newValue && kernelChangedArgs.newValue.name) {
 			let kernelDisplayName = this.getDisplayNameFromSpecName(kernelChangedArgs.newValue);
-			this._activeContexts = NotebookContexts.getContextsForKernel(this._notebookOptions.connectionService, this.getApplicableConnectionProviderIds(kernelDisplayName), kernelChangedArgs, this.connectionProfile);
+			let context = NotebookContexts.getContextForKernel(this._activeConnection, this.getApplicableConnectionProviderIds(kernelDisplayName));
-			this._contextsChangedEmitter.fire();
+			if (context !== undefined && context.serverName !== undefined && context.title !== undefined) {
-			if (this.contexts.defaultConnection !== undefined && this.contexts.defaultConnection.serverName !== undefined && this.contexts.defaultConnection.title !== undefined) {
+				await this.changeContext(context.title, context);
-				await this.changeContext(this.contexts.defaultConnection.title, this.contexts.defaultConnection);
 			}
+			this._contextsChangedEmitter.fire();
 		}
 	}
 
@@ -938,7 +907,7 @@ export class NotebookModel extends Disposable implements INotebookModel {
 		}
 	}
 
-	// Disconnect any connections that were added through the "Add new connection" functionality in the Attach To dropdown
+	// Disconnect any connections that were added through the "Change connection" functionality in the Attach To dropdown
 	private async disconnectAttachToConnections(): Promise<void> {
 		notebookUtils.asyncForEach(this._connectionUrisToDispose, async conn => {
 			await this.notebookOptions.connectionService.disconnect(conn).catch(e => this.logService.error(e));

src/sql/workbench/contrib/notebook/browser/notebook.component.ts

@@ -406,7 +406,7 @@ export class NotebookComponent extends AngularDisposable implements OnInit, OnDe
 
 		let attachToContainer = document.createElement('div');
 		let attachToDropdown = new AttachToDropdown(attachToContainer, this.contextViewService, this.modelReady,
-			this.connectionManagementService, this.connectionDialogService, this.notificationService, this.capabilitiesService, this.logService);
+			this.connectionManagementService, this.connectionDialogService, this.notificationService, this.capabilitiesService);
 		attachToDropdown.render(attachToContainer);
 		attachSelectBoxStyler(attachToDropdown, this.themeService);
 

src/sql/workbench/contrib/notebook/browser/notebookActions.ts

@@ -17,7 +17,6 @@ import { ConnectionProfile } from 'sql/platform/connection/common/connectionProf
 import { noKernel } from 'sql/workbench/services/notebook/browser/sessionManager';
 import { IConnectionDialogService } from 'sql/workbench/services/connection/common/connectionDialogService';
 import { NotebookModel } from 'sql/workbench/contrib/notebook/browser/models/notebookModel';
-import { generateUri } from 'sql/platform/connection/common/utils';
 import { IKeybindingService } from 'vs/platform/keybinding/common/keybinding';
 import { ILogService } from 'vs/platform/log/common/log';
 import { ICommandService } from 'vs/platform/commands/common/commands';
@@ -34,7 +33,7 @@ const msgChanging = localize('changing', "Changing kernel...");
 const kernelLabel: string = localize('Kernel', "Kernel: ");
 const attachToLabel: string = localize('AttachTo', "Attach To: ");
 const msgLoadingContexts = localize('loadingContexts', "Loading contexts...");
-const msgAddNewConnection = localize('addNewConnection', "Add New Connection");
+const msgChangeConnection = localize('changeConnection', "Change Connection");
 const msgSelectConnection = localize('selectConnection', "Select Connection");
 const msgLocalHost = localize('localhost', "localhost");
 const HIDE_ICON_CLASS = ' hideIcon';
@@ -369,21 +368,22 @@ export class AttachToDropdown extends SelectBox {
 		@IConnectionDialogService private _connectionDialogService: IConnectionDialogService,
 		@INotificationService private _notificationService: INotificationService,
 		@ICapabilitiesService private _capabilitiesService: ICapabilitiesService,
-		@ILogService private readonly logService: ILogService
 	) {
 		super([msgLoadingContexts], msgLoadingContexts, contextViewProvider, container, { labelText: attachToLabel, labelOnTop: false, ariaLabel: attachToLabel } as ISelectBoxOptionsWithLabel);
 		if (modelReady) {
 			modelReady
 				.then(model => {
 					this.updateModel(model);
-					this.updateAttachToDropdown(model);
+					this._register(model.onValidConnectionSelected(validConnection => {
+						this.handleContextsChanged(!validConnection);
+					}));
 				})
 				.catch(err => {
 					// No-op for now
 				});
 		}
 		this.onDidSelect(e => {
-			this.doChangeContext(this.getSelectedConnection(e.selected));
+			this.doChangeContext();
 		});
 	}
 
@@ -408,25 +408,6 @@ export class AttachToDropdown extends SelectBox {
 		}
 	}
 
-	private updateAttachToDropdown(model: INotebookModel): void {
-		if (this.model.connectionProfile && this.model.connectionProfile.serverName) {
-			let connectionUri = generateUri(this.model.connectionProfile, 'notebook');
-			this.model.notebookOptions.connectionService.connect(this.model.connectionProfile, connectionUri).then(result => {
-				if (result.connected) {
-					let connectionProfile = new ConnectionProfile(this._capabilitiesService, result.connectionProfile);
-					this.model.addAttachToConnectionsToBeDisposed(connectionUri);
-					this.doChangeContext(connectionProfile);
-				} else {
-					this.openConnectionDialog(true);
-				}
-			}).catch(err =>
-				this.logService.error(err));
-		}
-		this._register(model.onValidConnectionSelected(validConnection => {
-			this.handleContextsChanged(!validConnection);
-		}));
-	}
-
 	private getKernelDisplayName(): string {
 		let kernelDisplayName: string;
 		if (this.model.clientSession && this.model.clientSession.kernel && this.model.clientSession.kernel.name) {
@@ -446,81 +427,17 @@ export class AttachToDropdown extends SelectBox {
 			this.setOptions([msgLocalHost]);
 		}
 		else {
-			let connections = this.getConnections(model);
+			let connections: string[] = model.context && model.context.title ? [model.context.title] : [msgSelectConnection];
-			this.enable();
+			if (!find(connections, x => x === msgChangeConnection)) {
-			if (showSelectConnection) {
+				connections.push(msgChangeConnection);
-				this.loadWithSelectConnection(connections);
-			}
-			else {
-				if (connections.length === 1 && connections[0] === msgAddNewConnection) {
-					connections.unshift(msgSelectConnection);
-				}
-				else {
-					if (!find(connections, x => x === msgAddNewConnection)) {
-						connections.push(msgAddNewConnection);
-					}
-				}
-				this.setOptions(connections, 0);
-			}
-		}
-	}
-
-	private loadWithSelectConnection(connections: string[]): void {
-		if (connections && connections.length > 0) {
-			if (!find(connections, x => x === msgSelectConnection)) {
-				connections.unshift(msgSelectConnection);
-			}
-
-			if (!find(connections, x => x === msgAddNewConnection)) {
-				connections.push(msgAddNewConnection);
 			}
 			this.setOptions(connections, 0);
-		}
+			this.enable();
-	}
-
-	//Get connections from context
-	public getConnections(model: INotebookModel): string[] {
-		let otherConnections: ConnectionProfile[] = [];
-		model.contexts.otherConnections.forEach((conn) => { otherConnections.push(conn); });
-		// If current connection connects to master, select the option in the dropdown that doesn't specify a database
-		if (!model.contexts.defaultConnection.databaseName) {
-			this.selectWithOptionName(model.contexts.defaultConnection.serverName);
-		} else {
-			if (model.contexts.defaultConnection) {
-				this.selectWithOptionName(model.contexts.defaultConnection.title ? model.contexts.defaultConnection.title : model.contexts.defaultConnection.serverName);
-			} else {
-				this.select(0);
-			}
-		}
-		otherConnections = this.setConnectionsList(model.contexts.defaultConnection, model.contexts.otherConnections);
-		let connections = otherConnections.map((context) => context.title ? context.title : context.serverName);
-		return connections;
-	}
-
-	private setConnectionsList(defaultConnection: ConnectionProfile, otherConnections: ConnectionProfile[]) {
-		if (defaultConnection.serverName !== msgSelectConnection) {
-			otherConnections = otherConnections.filter(conn => conn.id !== defaultConnection.id);
-			otherConnections.unshift(defaultConnection);
-			if (otherConnections.length > 1) {
-				otherConnections = otherConnections.filter(val => val.serverName !== msgSelectConnection);
-			}
-		}
-		return otherConnections;
-	}
-
-	public getSelectedConnection(selection: string): ConnectionProfile {
-		// Find all connections with the the same server as the selected option
-		let connections = this.model.contexts.otherConnections.filter((c) => selection === c.title);
-		// If only one connection exists with the same server name, use that one
-		if (connections.length === 1) {
-			return connections[0];
-		} else {
-			return find(this.model.contexts.otherConnections, (c) => selection === c.title);
 		}
 	}
 
 	public doChangeContext(connection?: ConnectionProfile, hideErrorMessage?: boolean): void {
-		if (this.value === msgAddNewConnection) {
+		if (this.value === msgChangeConnection || this.value === msgSelectConnection) {
 			this.openConnectionDialog();
 		} else {
 			this.model.changeContext(this.value, connection, hideErrorMessage).then(ok => undefined, err => this._notificationService.error(getErrorMessage(err)));

src/sql/workbench/contrib/notebook/test/browser/notebookContexts.test.ts

@@ -5,14 +5,12 @@
 
 import * as TypeMoq from 'typemoq';
 import * as assert from 'assert';
-import { nb } from 'azdata';
 import { NotebookContexts } from 'sql/workbench/contrib/notebook/browser/models/notebookContexts';
 import { IConnectionManagementService } from 'sql/platform/connection/common/connectionManagement';
 import { TestConnectionManagementService } from 'sql/platform/connection/test/common/testConnectionManagementService';
 import { ConnectionProfile } from 'sql/platform/connection/common/connectionProfile';
 import { TestCapabilitiesService } from 'sql/platform/capabilities/test/common/testCapabilitiesService';
 import { mssqlProviderName } from 'sql/platform/connection/common/constants';
-import { IDefaultConnection } from 'sql/workbench/contrib/notebook/browser/models/modelInterfaces';
 
 suite('Notebook Contexts', function (): void {
 	const defaultContext = NotebookContexts.DefaultContext;
@@ -40,90 +38,28 @@ suite('Notebook Contexts', function (): void {
 		const connService: TypeMoq.Mock<IConnectionManagementService>
 			= TypeMoq.Mock.ofType<IConnectionManagementService>(TestConnectionManagementService, TypeMoq.MockBehavior.Strict);
 
-		// No kernel or profile info provided
+		// No profile info provided
-		let conns = NotebookContexts.getContextsForKernel(connService.object, [mssqlProviderName]);
+		let conns = NotebookContexts.getContextForKernel(undefined, [mssqlProviderName]);
-		assert.deepStrictEqual(conns, defaultContext);
+		assert.deepStrictEqual(conns, defaultContext, 'Contexts not the same when no profile info passed in');
 
-		// No Profile, Kernels are the same
+		// Profile, but no provider IDs
-		let kernelChangeArgs = <nb.IKernelChangedArgs>{
-			oldValue: <nb.IKernel>{
-				id: '1',
-				name: 'TestKernel'
-			},
-			newValue: <nb.IKernel>{
-				id: '1',
-				name: 'TestKernel'
-			}
-		};
-		conns = NotebookContexts.getContextsForKernel(connService.object, [mssqlProviderName], kernelChangeArgs);
-		assert.deepStrictEqual(conns, defaultContext);
-
-		// Kernel Info and Profile, but no provider IDs
 		let testConn = createTestConnProfile();
-		conns = NotebookContexts.getContextsForKernel(connService.object, [], kernelChangeArgs, testConn);
+		conns = NotebookContexts.getContextForKernel(testConn, []);
-		assert.deepStrictEqual(conns, defaultContext);
+		assert.deepStrictEqual(conns, localContext, 'Contexts not the same when no provider ids passed in');
 
 		// Normal use case
 		connService.setup(c => c.getActiveConnections()).returns(() => [testConn]);
-		conns = NotebookContexts.getContextsForKernel(connService.object, [mssqlProviderName], kernelChangeArgs, testConn);
+		conns = NotebookContexts.getContextForKernel(testConn, [mssqlProviderName]);
-		assert.deepStrictEqual(conns, <IDefaultConnection>{
+		assert.deepStrictEqual(conns, testConn, 'Contexts not the same when testing mssql provider connections');
-			otherConnections: [testConn],
-			defaultConnection: testConn
-		});
-	});
 
-	test('Get Active Contexts', async function (): Promise<void> {
+		// Multiple provider IDs including mssql
-		const connService: TypeMoq.Mock<IConnectionManagementService>
-			= TypeMoq.Mock.ofType<IConnectionManagementService>(TestConnectionManagementService, TypeMoq.MockBehavior.Strict);
-
-		let testConn = createTestConnProfile();
-
-		// No provider IDs
 		connService.setup(c => c.getActiveConnections()).returns(() => [testConn]);
-		let conns = NotebookContexts.getActiveContexts(connService.object, [], testConn);
+		conns = NotebookContexts.getContextForKernel(testConn, [mssqlProviderName, 'fakeProvider']);
-		assert.deepStrictEqual(conns, localContext);
+		assert.deepStrictEqual(conns, testConn, 'Contexts not the same when multiple providers passed in');
 
-		// No connections
+		// Connection provider IDs do not match
-		connService.setup(c => c.getActiveConnections()).returns(() => []);
-		conns = NotebookContexts.getActiveContexts(connService.object, [mssqlProviderName], testConn);
-		assert.deepStrictEqual(conns, defaultContext);
-
-		// No valid connection IDs
-		testConn.id = '-1';
 		connService.setup(c => c.getActiveConnections()).returns(() => [testConn]);
-		conns = NotebookContexts.getActiveContexts(connService.object, [mssqlProviderName], testConn);
+		conns = NotebookContexts.getContextForKernel(testConn, ['fakeProvider']);
-		assert.deepStrictEqual(conns, defaultContext);
+		assert.deepStrictEqual(conns, defaultContext, 'Contexts not the same when provider ids do not match');
-
-		// No matching provider IDs
-		testConn.id = 'testId';
-		connService.setup(c => c.getActiveConnections()).returns(() => [testConn]);
-		conns = NotebookContexts.getActiveContexts(connService.object, ['notARealProvider'], testConn);
-		assert.deepStrictEqual(conns, defaultContext);
-
-		// Normal behavior, valid connection present
-		connService.setup(c => c.getActiveConnections()).returns(() => [testConn]);
-		conns = NotebookContexts.getActiveContexts(connService.object, [mssqlProviderName], testConn);
-		assert.deepStrictEqual(conns, <IDefaultConnection>{
-			otherConnections: [testConn],
-			defaultConnection: testConn
-		});
-
-		// Multiple active connections
-		let newTestConn = createTestConnProfile();
-		newTestConn.serverName = 'otherTestServerName';
-		connService.setup(c => c.getActiveConnections()).returns(() => [newTestConn, testConn]);
-		conns = NotebookContexts.getActiveContexts(connService.object, [mssqlProviderName], testConn);
-		assert.deepStrictEqual(conns, <IDefaultConnection>{
-			otherConnections: [newTestConn, testConn],
-			defaultConnection: testConn
-		});
-
-		// Multiple connections, no profile provided
-		connService.setup(c => c.getActiveConnections()).returns(() => [newTestConn, testConn]);
-		conns = NotebookContexts.getActiveContexts(connService.object, [mssqlProviderName], undefined);
-		assert.deepStrictEqual(conns, <IDefaultConnection>{
-			otherConnections: [newTestConn, testConn],
-			defaultConnection: newTestConn
-		});
 	});
 });

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

@@ -6,13 +6,14 @@
 import { nb, IConnectionProfile } from 'azdata';
 
 import { Event, Emitter } from 'vs/base/common/event';
-import { INotebookModel, ICellModel, IClientSession, IDefaultConnection, NotebookContentChange, IKernelPreference } from 'sql/workbench/contrib/notebook/browser/models/modelInterfaces';
+import { INotebookModel, ICellModel, IClientSession, NotebookContentChange, IKernelPreference } from 'sql/workbench/contrib/notebook/browser/models/modelInterfaces';
 import { NotebookChangeType, CellType } from 'sql/workbench/contrib/notebook/common/models/contracts';
 import { INotebookManager, INotebookService, INotebookEditor, ILanguageMagic, INotebookProvider, INavigationProvider } from 'sql/workbench/services/notebook/browser/notebookService';
 import { ISingleNotebookEditOperation } from 'sql/workbench/api/common/sqlExtHostTypes';
 import { IStandardKernelWithProvider } from 'sql/workbench/contrib/notebook/browser/models/notebookUtils';
 import { URI } from 'vs/workbench/workbench.web.api';
 import { RenderMimeRegistry } from 'sql/workbench/contrib/notebook/browser/outputs/registry';
+import { ConnectionProfile } from 'sql/platform/connection/common/connectionProfile';
 
 export class NotebookModelStub implements INotebookModel {
 	constructor(private _languageInfo?: nb.ILanguageInfo) {
@@ -63,7 +64,7 @@ export class NotebookModelStub implements INotebookModel {
 	get specs(): nb.IAllKernels {
 		throw new Error('method not implemented.');
 	}
-	get contexts(): IDefaultConnection {
+	get context(): ConnectionProfile {
 		throw new Error('method not implemented.');
 	}
 	get providerId(): string {