Fix #4356 New Notebook from connection doesn't connect (#4364)

* Fix #4356 New Notebook from connection doesn't connect
Fix new notebook error by passing profile instead of ID.
- I could've just sent the ID over, but this fix sets the stage for disconnected connections to work (since we have enough info to properly connect).
- There's a bug in NotebookModel blocking the disconnected connection part working, but Yurong's in progress fixes will unblock this. Hence checking in as-is and working to properly unblock once that's in.

* Support connection profile in commandline service
- Added new context API for things that want to work on commandline and object explorer
- Refactored commandlineservice slightly to be async & have a simpler execution flow (far fewer if/else statements)

* Fix unit tests
- Fixed 2 issues raised by tests (sholdn't do new query if no profile passed, shouldn't error on new query failing)
- Updated unit tests to pass as expected given changes to the APIs.

src/sql/azdata.proposed.d.ts

@@ -3667,19 +3667,30 @@ declare module 'azdata' {
 		export function getProvidersByType<T extends DataProvider>(providerType: DataProviderType): T[];
 	}
 
+
 	/**
 	 * Context object passed as an argument to command callbacks.
-	 * Defines the key properties required to identify a node in the object
-	 * explorer tree and take action against it.
+	 * Defines properties that can be sent for any connected context,
+	 * whether that is the Object Explorer context menu or a command line
+	 * startup argument.
 	 */
-	export interface ObjectExplorerContext {
 
+	export interface ConnectedContext {
 		/**
 		 * The connection information for the selected object.
 		 * Note that the connection is not guaranteed to be in a connected
 		 * state on click.
 		 */
 		connectionProfile: IConnectionProfile;
+	}
+
+	/**
+	 * Context object passed as an argument to command callbacks.
+	 * Defines the key properties required to identify a node in the object
+	 * explorer tree and take action against it.
+	 */
+	export interface ObjectExplorerContext extends ConnectedContext {
+
 		/**
 		 * Defines whether this is a Connection-level object.
 		 * If not, the object is expected to be a child object underneath
@@ -4016,9 +4027,9 @@ declare module 'azdata' {
 			providerId?: string;
 
 			/**
-			 * Optional ID indicating the initial connection to use for this editor
+			 * Optional profile indicating the initial connection to use for this editor
 			 */
-			connectionId?: string;
+			connectionProfile?: IConnectionProfile;
 
 			/**
 			 * Default kernel for notebook

src/sql/parts/common/customInputConverter.ts

@@ -59,7 +59,6 @@ export function convertEditorInput(input: EditorInput, options: IQueryEditorOpti
 		if (uri) {
 			return withService<INotebookService, NotebookInput>(instantiationService, INotebookService, notebookService => {
 				let fileName: string = 'untitled';
-				let providerIds: string[] = [DEFAULT_NOTEBOOK_PROVIDER];
 				if (input) {
 					fileName = input.getName();
 				}

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

@@ -408,6 +408,10 @@ 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 {

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

@@ -104,23 +104,18 @@ export class NotebookComponent extends AngularDisposable implements OnInit, OnDe
 
 	private updateProfile(): void {
 		this.profile = this.notebookParams ? this.notebookParams.profile : undefined;
-		let profile: IConnectionProfile;
 		if (!this.profile) {
-			// Use connectionProfile passed in first
-			if (this._notebookParams.connectionProfileId !== undefined && this._notebookParams.connectionProfileId) {
-				profile = this.connectionManagementService.getConnectionProfileById(this._notebookParams.connectionProfileId);
-			} else {
-				// Second use global connection if possible
-				profile = TaskUtilities.getCurrentGlobalConnection(this.objectExplorerService, this.connectionManagementService, this.editorService);
-			}
+			// Second use global connection if possible
+			let profile: IConnectionProfile = TaskUtilities.getCurrentGlobalConnection(this.objectExplorerService, this.connectionManagementService, this.editorService);
+
 			// TODO use generic method to match kernel with valid connection that's compatible. For now, we only have 1
 			if (profile && profile.providerName) {
 				this.profile = profile;
 			} else {
 				// if not, try 1st active connection that matches our filter
-				let profiles = this.connectionManagementService.getActiveConnections();
-				if (profiles && profiles.length > 0) {
-					this.profile = profiles[0];
+				let activeProfiles = this.connectionManagementService.getActiveConnections();
+				if (activeProfiles && activeProfiles.length > 0) {
+					this.profile = activeProfiles[0];
 				}
 			}
 		}

src/sql/parts/notebook/notebookEditor.ts

@@ -95,7 +95,7 @@ export class NotebookEditor extends BaseEditor {
 			providerId: input.providerId ? input.providerId : DEFAULT_NOTEBOOK_PROVIDER,
 			providers: input.providers ? input.providers : [DEFAULT_NOTEBOOK_PROVIDER],
 			isTrusted: input.isTrusted,
-			connectionProfileId: input.connectionProfileId
+			profile: input.connectionProfile
 		};
 		bootstrapAngular(this.instantiationService,
 			NotebookModule,

src/sql/parts/notebook/notebookInput.ts

@@ -27,6 +27,7 @@ import { IUntitledEditorService } from 'vs/workbench/services/untitled/common/un
 import { notebookModeId } from 'sql/common/constants';
 import { ITextFileService, ISaveOptions } from 'vs/workbench/services/textfile/common/textfiles';
 import { LocalContentManager } from 'sql/workbench/services/notebook/node/localContentManager';
+import { IConnectionProfile } from 'sql/platform/connection/common/interfaces';
 
 export type ModeViewSaveHandler = (handle: number) => Thenable<boolean>;
 
@@ -128,7 +129,7 @@ export class NotebookInput extends EditorInput {
 	private _providerId: string;
 	private _providers: string[];
 	private _standardKernels: IStandardKernelWithProvider[];
-	private _connectionProfileId: string;
+	private _connectionProfile: IConnectionProfile;
 	private _defaultKernel: azdata.nb.IKernelSpec;
 	private _isTrusted: boolean = false;
 	public hasBootstrapped = false;
@@ -191,12 +192,12 @@ export class NotebookInput extends EditorInput {
 		this._isTrusted = value;
 	}
 
-	public set connectionProfileId(value: string) {
-		this._connectionProfileId = value;
+	public set connectionProfile(value: IConnectionProfile) {
+		this._connectionProfile = value;
 	}
 
-	public get connectionProfileId(): string {
-		return this._connectionProfileId;
+	public get connectionProfile(): IConnectionProfile {
+		return this._connectionProfile;
 	}
 
 	public get standardKernels(): IStandardKernelWithProvider[] {

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

@@ -167,7 +167,7 @@ export class ExtHostNotebookDocumentsAndEditors implements ExtHostNotebookDocume
 			options.preview = showOptions.preview;
 			options.position = showOptions.viewColumn;
 			options.providerId = showOptions.providerId;
-			options.connectionId = showOptions.connectionId;
+			options.connectionProfile = showOptions.connectionProfile;
 			options.defaultKernel = showOptions.defaultKernel;
 		}
 		let id = await this._proxy.$tryShowNotebookDocument(uri, options);

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

@@ -16,6 +16,7 @@ import { IInstantiationService } from 'vs/platform/instantiation/common/instanti
 import { ITextEditorOptions } from 'vs/platform/editor/common/editor';
 import { viewColumnToEditorGroup } from 'vs/workbench/api/shared/editor';
 import { Schemas } from 'vs/base/common/network';
+import { ConnectionProfile } from 'sql/platform/connection/common/connectionProfile';
 
 import {
 	SqlMainContext, MainThreadNotebookDocumentsAndEditorsShape, SqlExtHostContext, ExtHostNotebookDocumentsAndEditorsShape,
@@ -28,6 +29,7 @@ import { ISingleNotebookEditOperation } from 'sql/workbench/api/common/sqlExtHos
 import { disposed } from 'vs/base/common/errors';
 import { ICellModel, NotebookContentChange, INotebookModel } from 'sql/parts/notebook/models/modelInterfaces';
 import { NotebookChangeType, CellTypes } from 'sql/parts/notebook/models/contracts';
+import { ICapabilitiesService } from 'sql/platform/capabilities/common/capabilitiesService';
 
 class MainThreadNotebookEditor extends Disposable {
 	private _contentChangedEmitter = new Emitter<NotebookContentChange>();
@@ -291,7 +293,7 @@ export class MainThreadNotebookDocumentsAndEditors extends Disposable implements
 		@IInstantiationService private _instantiationService: IInstantiationService,
 		@IEditorService private _editorService: IEditorService,
 		@IEditorGroupsService private _editorGroupService: IEditorGroupsService,
-		@INotebookService private readonly _notebookService: INotebookService
+		@ICapabilitiesService private _capabilitiesService: ICapabilitiesService
 	) {
 		super();
 		if (extHostContext) {
@@ -363,7 +365,7 @@ export class MainThreadNotebookDocumentsAndEditors extends Disposable implements
 		let input = this._instantiationService.createInstance(NotebookInput, uri.fsPath, uri);
 		input.isTrusted = trusted;
 		input.defaultKernel = options.defaultKernel;
-		input.connectionProfileId = options.connectionId;
+		input.connectionProfile = new ConnectionProfile(this._capabilitiesService, options.connectionProfile);
 
 		let editor = await this._editorService.openEditor(input, editorOptions, viewColumnToEditorGroup(this._editorGroupService, options.position));
 		if (!editor) {

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

@@ -850,7 +850,7 @@ export interface INotebookShowOptions {
 	preserveFocus?: boolean;
 	preview?: boolean;
 	providerId?: string;
-	connectionId?: string;
+	connectionProfile?: azdata.IConnectionProfile;
 	defaultKernel?: azdata.nb.IKernelSpec;
 }
 

src/sql/workbench/services/commandLine/common/commandLineService.ts

@@ -3,6 +3,7 @@
  *  Licensed under the Source EULA. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 'use strict';
+import * as azdata from 'azdata';
 import { ConnectionProfile } from 'sql/platform/connection/common/connectionProfile';
 import { ICommandLineProcessing } from 'sql/workbench/services/commandLine/common/commandLine';
 import { IConnectionManagementService } from 'sql/platform/connection/common/connectionManagement';
@@ -19,8 +20,10 @@ import { IWorkspaceConfigurationService } from 'vs/workbench/services/configurat
 import { ICommandService } from 'vs/platform/commands/common/commands';
 import { warn } from 'sql/base/common/log';
 import { ipcRenderer as ipc} from 'electron';
+import { IConnectionProfile } from 'sql/platform/connection/common/interfaces';
 
 export class CommandLineService implements ICommandLineProcessing {
+	public _serviceBrand: any;
 
 	constructor(
 		@ICapabilitiesService private _capabilitiesService: ICapabilitiesService,
@@ -57,14 +60,13 @@ export class CommandLineService implements ICommandLineProcessing {
 		}
 	}
 
-	public _serviceBrand: any;
 	// We base our logic on the combination of (server, command) values.
 	// (serverName, commandName) => Connect object explorer and execute the command, passing the connection profile to the command. Do not load query editor.
 	// (null, commandName) => Launch the command with a null connection. If the command implementation needs a connection, it will need to create it.
 	// (serverName, null) => Connect object explorer and open a new query editor
 	// (null, null) => Prompt for a connection unless there are registered servers
-	public processCommandLine(args: ParsedArgs): Promise<void> {
-		let profile = undefined;
+	public async processCommandLine(args: ParsedArgs): Promise<void> {
+		let profile: IConnectionProfile = undefined;
 		let commandName = undefined;
 		if (args) {
 			if (this._commandService) {
@@ -72,66 +74,58 @@ export class CommandLineService implements ICommandLineProcessing {
 			}
 
 			if (args.server) {
-				profile = new ConnectionProfile(this._capabilitiesService, null);
-				// We want connection store to use any matching password it finds
-				profile.savePassword = true;
-				profile.providerName = Constants.mssqlProviderName;
-				profile.serverName = args.server;
-				profile.databaseName = args.database ? args.database : '';
-				profile.userName = args.user ? args.user : '';
-				profile.authenticationType = args.integrated ? 'Integrated' : 'SqlLogin';
-				profile.connectionName = '';
-				profile.setOptionValue('applicationName', Constants.applicationName);
-				profile.setOptionValue('databaseDisplayName', profile.databaseName);
-				profile.setOptionValue('groupId', profile.groupId);
+				profile = this.readProfileFromArgs(args);
 			}
 	    }
-		let self = this;
-		return new Promise<void>((resolve, reject) => {
-			let showConnectDialogOnStartup: boolean = self._configurationService.getValue('workbench.showConnectDialogOnStartup');
-			if (showConnectDialogOnStartup && !commandName && !profile && !self._connectionManagementService.hasRegisteredServers()) {
-				// prompt the user for a new connection on startup if no profiles are registered
-				self._connectionManagementService.showConnectionDialog()
-					.then(() => {
-						resolve();
-					},
-						error => {
-							reject(error);
-						});
-			} else if (profile) {
-				if (!commandName) {
-					self._connectionManagementService.connectIfNotConnected(profile, 'connection', true)
-						.then(() => {
-							TaskUtilities.newQuery(profile,
-								self._connectionManagementService,
-								self._queryEditorService,
-								self._objectExplorerService,
-								self._editorService)
-								.then(() => {
-									resolve();
-								}, error => {
-									// ignore query editor failing to open.
-									// the tests don't mock this out
-									warn('unable to open query editor ' + error);
-									resolve();
-								});
-						}, error => {
-							reject(error);
-						});
-				} else {
-					self._connectionManagementService.connectIfNotConnected(profile, 'connection', true)
-						.then(() => {
-							self._commandService.executeCommand(commandName, profile.id).then(() => resolve(), error => reject(error));
-						}, error => {
-							reject(error);
-						});
-				}
-			} else if (commandName) {
-				self._commandService.executeCommand(commandName).then(() => resolve(), error => reject(error));
+		let showConnectDialogOnStartup: boolean = this._configurationService.getValue('workbench.showConnectDialogOnStartup');
+		if (showConnectDialogOnStartup && !commandName && !profile && !this._connectionManagementService.hasRegisteredServers()) {
+			// prompt the user for a new connection on startup if no profiles are registered
+			await this._connectionManagementService.showConnectionDialog();
+			return;
+		}
+		let connectedContext: azdata.ConnectedContext = undefined;
+		if (profile) {
+			try {
+				await this._connectionManagementService.connectIfNotConnected(profile, 'connection', true);
+				// Before sending to extensions, we should a) serialize to IConnectionProfile or things will fail,
+				// and b) use the latest version of the profile from the service so most fields are filled in.
+				let updatedProfile = this._connectionManagementService.getConnectionProfileById(profile.id);
+				connectedContext = { connectionProfile: new ConnectionProfile(this._capabilitiesService, updatedProfile).toIConnectionProfile() };
+			} catch (err) {
+				warn('Failed to connect due to error' + err.message);
 			}
-			else {
-				resolve();
+		}
+		if (commandName) {
+			await this._commandService.executeCommand(commandName, connectedContext);
+		} else if (profile) {
+			// Default to showing new query
+			try {
+				await TaskUtilities.newQuery(profile,
+					this._connectionManagementService,
+					this._queryEditorService,
+					this._objectExplorerService,
+					this._editorService);
+			} catch (error) {
+				warn('unable to open query editor ' + error);
+				// Note: we are intentionally swallowing this error.
+				// In part this is to accommodate unit testing where we don't want to set up the query stack
 			}
-		});
+		}
+	}
+
+	private readProfileFromArgs(args: ParsedArgs) {
+		let profile = new ConnectionProfile(this._capabilitiesService, null);
+		// We want connection store to use any matching password it finds
+		profile.savePassword = true;
+		profile.providerName = Constants.mssqlProviderName;
+		profile.serverName = args.server;
+		profile.databaseName = args.database ? args.database : '';
+		profile.userName = args.user ? args.user : '';
+		profile.authenticationType = args.integrated ? 'Integrated' : 'SqlLogin';
+		profile.connectionName = '';
+		profile.setOptionValue('applicationName', Constants.applicationName);
+		profile.setOptionValue('databaseDisplayName', profile.databaseName);
+		profile.setOptionValue('groupId', profile.groupId);
+		return profile;
 	}
 }

src/sql/workbench/services/notebook/common/notebookService.ts

@@ -95,7 +95,6 @@ export interface INotebookParams extends IBootstrapParams {
 	isTrusted: boolean;
 	profile?: IConnectionProfile;
 	modelFactory?: ModelFactory;
-	connectionProfileId?: string;
 }
 
 export interface INotebookEditor {