Fix Notebook Viewlet Highlighting (#10868) (#10870)

* Stop double tree data provider registrations

* revert book finding logic, leverage currentBook

extensions/mssql/config.json

@@ -1,6 +1,6 @@
 {
 	"downloadUrl": "https://github.com/Microsoft/sqltoolsservice/releases/download/v{#version#}/microsoft.sqltools.servicelayer-{#fileName#}",
-	"version": "2.0.0-release.68",
+	"version": "2.0.0-release.70",
 	"downloadFileNames": {
 		"Windows_86": "win-x86-netcoreapp2.2.zip",
 		"Windows_64": "win-x64-netcoreapp2.2.zip",

extensions/notebook/src/book/bookModel.ts

@@ -51,7 +51,7 @@ export class BookModel implements azdata.nb.NavigationProvider {
 	}
 
 	public getNotebook(uri: string): BookTreeItem | undefined {
-		return this._allNotebooks.get(uri);
+		return this._allNotebooks.get(this.openAsUntitled ? path.basename(uri) : uri);
 	}
 
 	public async loadTableOfContentFiles(folderPath: string): Promise<void> {
@@ -190,8 +190,8 @@ export class BookModel implements azdata.nb.NavigationProvider {
 						if (this.openAsUntitled) {
 							if (!this._allNotebooks.get(path.basename(pathToNotebook))) {
 								this._allNotebooks.set(path.basename(pathToNotebook), notebook);
-								notebooks.push(notebook);
 							}
+							notebooks.push(notebook);
 						} else {
 							// convert to URI to avoid casing issue with drive letters when getting navigation links
 							let uriToNotebook: vscode.Uri = vscode.Uri.file(pathToNotebook);

extensions/notebook/src/book/bookTreeView.ts

@@ -183,9 +183,6 @@ export class BookTreeViewProvider implements vscode.TreeDataProvider<BookTreeIte
 				this.revealActiveDocumentInViewlet();
 			}
 		});
-		azdata.nb.onDidChangeActiveNotebookEditor(e => {
-			this.revealActiveDocumentInViewlet(e.document.uri, false);
-		});
 	}
 
 	async showPreviewFile(urlToOpen?: string): Promise<void> {

extensions/notebook/src/dialog/configurePython/configurePythonWizard.ts

@@ -63,13 +63,13 @@ export class ConfigurePythonWizard {
 
 		let wizardTitle: string;
 		if (kernelName) {
-			wizardTitle = localize('configurePython.wizardNameWithKernel', 'Configure Python to run {0} kernel', kernelName);
+			wizardTitle = localize('configurePython.wizardNameWithKernel', "Configure Python to run {0} kernel", kernelName);
 		} else {
-			wizardTitle = localize('configurePython.wizardNameWithoutKernel', 'Configure Python to run kernels');
+			wizardTitle = localize('configurePython.wizardNameWithoutKernel', "Configure Python to run kernels");
 		}
-		this._wizard = azdata.window.createWizard(wizardTitle);
-		let page0 = azdata.window.createWizardPage(localize('configurePython.page0Name', 'Configure Python Runtime'));
-		let page1 = azdata.window.createWizardPage(localize('configurePython.page1Name', 'Install Dependencies'));
+		this._wizard = azdata.window.createWizard(wizardTitle, 600);
+		let page0 = azdata.window.createWizardPage(localize('configurePython.page0Name', "Configure Python Runtime"));
+		let page1 = azdata.window.createWizardPage(localize('configurePython.page1Name', "Install Dependencies"));
 
 		page0.registerContent(async (view) => {
 			let configurePathPage = new ConfigurePathPage(this.apiWrapper, this, page0, this.model, view);

extensions/notebook/src/extension.ts

@@ -129,9 +129,15 @@ export async function activate(extensionContext: vscode.ExtensionContext): Promi
 	const providedBookTreeViewProvider = new BookTreeViewProvider(appContext.apiWrapper, [], extensionContext, true, PROVIDED_BOOKS_VIEWID);
 	await providedBookTreeViewProvider.initialized;
 
+	azdata.nb.onDidChangeActiveNotebookEditor(e => {
+		if (e.document.uri.scheme === 'untitled') {
+			providedBookTreeViewProvider.revealActiveDocumentInViewlet(e.document.uri, false);
+		} else {
+			bookTreeViewProvider.revealActiveDocumentInViewlet(e.document.uri, false);
+		}
+
+	});
 
-	extensionContext.subscriptions.push(vscode.window.registerTreeDataProvider(BOOKS_VIEWID, bookTreeViewProvider));
-	extensionContext.subscriptions.push(vscode.window.registerTreeDataProvider(PROVIDED_BOOKS_VIEWID, providedBookTreeViewProvider));
 	return {
 		getJupyterController() {
 			return controller;

package.json

@@ -21,7 +21,7 @@
     "watch-client": "gulp watch-client --max_old_space_size=4095",
     "mocha": "mocha test/unit/node/all.js --delay",
     "precommit": "node build/gulpfile.hygiene.js",
-    "gulp": "gulp --max_old_space_size=8192",
+    "gulp": "gulp --max_old_space_size=4095",
     "electron": "node build/lib/electron",
     "7z": "7z",
     "update-grammars": "node build/npm/update-all-grammars.js",

product.json

@@ -50,6 +50,7 @@
 		"Microsoft.profiler",
 		"Microsoft.schema-compare",
 		"Microsoft.datavirtualization",
+		"Redgate.sql-prompt",
 		"Redgate.sql-search",
 		"IDERA.sqldm-performance-insights",
 		"SentryOne.plan-explorer"

src/sql/base/browser/ui/buttonMenu/buttonMenu.css

@@ -8,6 +8,6 @@
 	padding: 0;
 }
 .button-menu.masked-pseudo-after.dropdown-arrow:after {
-	right: -2px;
+	right: 0;
 	width: 20px;
 }

src/sql/media/icons/common-icons.css

@@ -310,8 +310,10 @@
 
 /* Icons as masked elements for easy theme switching.
 Includes non-masked style declarations. */
-.masked-icon {
-	display: inline-block;
+.masked-icon:before {
+	content: "";
+	display: block;
+	position: absolute;
 	height: 20px;
 	width: 20px;
 	-webkit-mask-position: center;
@@ -378,7 +380,7 @@ Includes non-masked style declarations. */
 .codicon:not(.masked-icon).icon-expand-cells {
 	background-image: url("action-expand.svg");
 }
-.codicon.masked-icon.icon-expand-cells {
+.codicon.masked-icon.icon-expand-cells:before {
 	background-image: none;
 	-webkit-mask-image: url("action-expand.svg");
 	mask-image: url("action-expand.svg");
@@ -387,7 +389,7 @@ Includes non-masked style declarations. */
 .codicon:not(.masked-icon).icon-collapse-cells {
 	background-image: url("action-collapse.svg");
 }
-.codicon.masked-icon.icon-collapse-cells {
+.codicon.masked-icon.icon-collapse-cells:before {
 	-webkit-mask-image: url("action-collapse.svg");
 	mask-image: url("action-collapse.svg");
 }
@@ -395,7 +397,7 @@ Includes non-masked style declarations. */
 .codicon:not(.masked-icon).icon-clear-results {
 	background-image: url("clear.svg");
 }
-.codicon.masked-icon.icon-clear-results {
+.codicon.masked-icon.icon-clear-results:before {
 	-webkit-mask-image: url("clear.svg");
 	mask-image: url("clear.svg");
 }
@@ -403,7 +405,7 @@ Includes non-masked style declarations. */
 .codicon:not(.masked-icon).icon-shield {
 	background-image: url("shield.svg");
 }
-.codicon.masked-icon.icon-shield {
+.codicon.masked-icon.icon-shield:before {
 	-webkit-mask-image: url("shield.svg");
 	mask-image: url("shield.svg");
 }
@@ -411,7 +413,7 @@ Includes non-masked style declarations. */
 .codicon:not(.masked-icon).icon-shield-x {
 	background-image: url("shield-x.svg");
 }
-.codicon.masked-icon.icon-shield-x {
+.codicon.masked-icon.icon-shield-x:before {
 	-webkit-mask-image: url("shield-x.svg");
 	mask-image: url("shield-x.svg");
 }
@@ -419,7 +421,7 @@ Includes non-masked style declarations. */
 .codicon:not(.masked-icon).packages {
 	background-image: url("packages.svg");
 }
-.codicon.masked-icon.packages {
+.codicon.masked-icon.packages:before {
 	background-image: none;
 	-webkit-mask-image: url("packages.svg");
 	mask-image: url("packages.svg");
@@ -504,17 +506,17 @@ Includes non-masked style declarations. */
 	mask-image: url("markdown.svg");
 }
 
-.codicon.masked-icon.new {
+.codicon.masked-icon.new:before {
 	background-image: none;
 	-webkit-mask-image: url('new.svg');
 	mask-image: url('new.svg');
 }
-.codicon.masked-icon.close {
+.codicon.masked-icon.close:before {
 	background-image: none;
 	-webkit-mask-image: url('close-blue.svg');
 	mask-image: url('close-blue.svg');
 }
-.codicon.masked-icon.edit {
+.codicon.masked-icon.edit:before {
 	background-image: none;
 	-webkit-mask-image: url('edit.svg');
 	mask-image: url('edit.svg');
@@ -530,12 +532,12 @@ Includes non-masked style declarations. */
 	-webkit-mask-image: url('down-arrow-blue.svg');
 	mask-image: url('down-arrow-blue.svg');
 }
-.codicon.masked-icon.delete {
+.codicon.masked-icon.delete:before {
 	background-image: none;
 	-webkit-mask-image: url('garbage-can-blue.svg');
 	mask-image: url('garbage-can-blue.svg');
 }
-.codicon.masked-icon.more {
+.codicon.masked-icon.more:before {
 	background-image: none;
 	-webkit-mask-image: url('ellipsis-blue.svg');
 	mask-image: url('ellipsis-blue.svg');
@@ -546,7 +548,6 @@ Includes non-masked style declarations. */
 }
 
 .book.codicon {
-
 	-webkit-mask-image: url("book_image.svg");
 	-webkit-mask-repeat: no-repeat;
 	-webkit-mask-position: 50% 50%;

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

@@ -7,7 +7,7 @@ import 'vs/css!./cellToolbar';
 import * as DOM from 'vs/base/browser/dom';
 import { Component, Inject, ViewChild, ElementRef, Input } from '@angular/core';
 import { localize } from 'vs/nls';
-import { Taskbar } from 'sql/base/browser/ui/taskbar/taskbar';
+import { Taskbar, ITaskbarContent } from 'sql/base/browser/ui/taskbar/taskbar';
 import { IContextMenuService } from 'vs/platform/contextview/browser/contextView';
 import { IInstantiationService } from 'vs/platform/instantiation/common/instantiation';
 import { DeleteCellAction, EditCellAction, CellToggleMoreActions } from 'sql/workbench/contrib/notebook/browser/cellToolbarActions';
@@ -56,10 +56,13 @@ export class CellToolbarComponent {
 		this._actionBar = new Taskbar(taskbar);
 		this._actionBar.context = context;
 
+		let addCellsButton = new AddCellAction('notebook.AddCodeCell', localize('codeCellsPreview', "Add cell"), 'notebook-button masked-pseudo code');
+		addCellsButton.cellType = CellTypes.Code;
+
 		let addCodeCellButton = new AddCellAction('notebook.AddCodeCell', localize('codePreview', "Code cell"), 'notebook-button masked-pseudo code');
 		addCodeCellButton.cellType = CellTypes.Code;
 
-		let addTextCellButton = new AddCellAction('notebook.AddTextCell', localize('textPreview', "Markdown cell"), 'notebook-button masked-pseudo markdown');
+		let addTextCellButton = new AddCellAction('notebook.AddTextCell', localize('textPreview', "Text cell"), 'notebook-button masked-pseudo markdown');
 		addTextCellButton.cellType = CellTypes.Markdown;
 
 		let deleteButton = this.instantiationService.createInstance(DeleteCellAction, 'delete', 'codicon masked-icon delete', localize('delete', "Delete"));
@@ -74,24 +77,27 @@ export class CellToolbarComponent {
 		let buttonDropdownContainer = DOM.$('li.action-item');
 		buttonDropdownContainer.setAttribute('role', 'presentation');
 		let dropdownMenuActionViewItem = new DropdownMenuActionViewItem(
-			addCodeCellButton,
+			addCellsButton,
 			[addCodeCellButton, addTextCellButton],
 			this.contextMenuService,
 			undefined,
 			this._actionBar.actionRunner,
 			undefined,
 			'codicon masked-icon new',
-			localize('addCell', "Cell"),
+			'',
 			undefined
 		);
 		dropdownMenuActionViewItem.render(buttonDropdownContainer);
 		dropdownMenuActionViewItem.setActionContext(context);
 
-		this._actionBar.setContent([
-			{ action: this._editCellAction },
-			{ element: buttonDropdownContainer },
+		let taskbarContent: ITaskbarContent[] = [];
+		if (this.cellModel?.cellType === CellTypes.Markdown) {
+			taskbarContent.push({ action: this._editCellAction });
+		}
+		taskbarContent.push({ element: buttonDropdownContainer },
 			{ action: deleteButton },
-			{ element: moreActionsContainer }
-		]);
+			{ element: moreActionsContainer });
+
+		this._actionBar.setContent(taskbarContent);
 	}
 }

src/sql/workbench/contrib/notebook/browser/cellViews/cellToolbar.css

@@ -39,6 +39,15 @@ cell-toolbar-component .monaco-action-bar .action-label {
 }
 cell-toolbar-component ul.actions-container li a.masked-icon {
 	display: flex;
+}
+cell-toolbar-component ul.actions-container li a.masked-icon,
+cell-toolbar-component ul.actions-container li a.masked-icon:before {
 	height: 24px;
 	width: 29px;
 }
+cell-toolbar-component .codicon.masked-icon,
+.vs cell-toolbar-component .codicon.masked-icon,
+.vs-dark cell-toolbar-component .codicon.masked-icon,
+.hc-black cell-toolbar-component .codicon.masked-icon {
+	background-image: none;
+}

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

@@ -79,6 +79,10 @@ export class CodeComponent extends CellView implements OnInit, OnChanges {
 		this._activeCellId = value;
 	}
 
+	@Input() set hover(value: boolean) {
+		this.cellModel.hover = value;
+	}
+
 	protected _actionBar: Taskbar;
 	private readonly _minimumHeight = 30;
 	private readonly _maximumHeight = 4000;

src/sql/workbench/contrib/notebook/browser/cellViews/code.css

@@ -57,9 +57,6 @@ code-component .editor {
 code-cell-component code-component .monaco-editor .margin-view-overlays .line-numbers {
 	left: 0!important;
 }
-code-cell-component code-component .monaco-scrollable-element.editor-scrollable.vs {
-	left: 40px!important;
-}
 
 code-cell-component .monaco-editor .margin,
 code-cell-component code-component .monaco-editor,

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

@@ -422,7 +422,7 @@ export class NotebookComponent extends AngularDisposable implements OnInit, OnDe
 			let addCodeCellButton = new AddCellAction('notebook.AddCodeCell', localize('codePreview', "Code cell"), 'notebook-button masked-pseudo code');
 			addCodeCellButton.cellType = CellTypes.Code;
 
-			let addTextCellButton = new AddCellAction('notebook.AddTextCell', localize('textPreview', "Markdown cell"), 'notebook-button masked-pseudo markdown');
+			let addTextCellButton = new AddCellAction('notebook.AddTextCell', localize('textPreview', "Text cell"), 'notebook-button masked-pseudo markdown');
 			addTextCellButton.cellType = CellTypes.Markdown;
 
 
@@ -460,8 +460,8 @@ export class NotebookComponent extends AngularDisposable implements OnInit, OnDe
 				{ element: buttonDropdownContainer },
 				{ action: this._runAllCellsAction },
 				{ element: Taskbar.createTaskbarSeparator() },
-				{ element: attachToContainer },
 				{ element: kernelContainer },
+				{ element: attachToContainer },
 				{ element: spacerElement },
 				{ action: collapseCellsAction },
 				{ action: clearResultsButton },

src/sql/workbench/contrib/notebook/browser/notebook.css

@@ -4,7 +4,7 @@
  *--------------------------------------------------------------------------------------------*/
 
 .notebookEditor .scrollable {
-	padding-top: 6px;
+	padding-top: 16px;
 }
 
 .notebookEditor .taskbarSeparator {
@@ -20,7 +20,7 @@
 	border-color: transparent;
 	border-style: solid;
 	border-width: 1px 1px 1px 4px;
-	margin: 16px;
+	margin: 8px;
 	position: relative;
 }
 
@@ -37,21 +37,23 @@
 	min-height: 21px;
 }
 
-.notebookEditor .actions-container .action-item .notebook-button {
+.notebookEditor .editor-toolbar .actions-container .action-item .notebook-button {
 	display: inline-block;
-	width: 100%;
-	padding: 0px;
 	text-align: center;
 	cursor: pointer;
-	padding-left: 18px;
+	padding: 0 18px;
 	background-size: 13px;
-	margin-right: 20px;
 	font-size: 13px;
 	height: 21px;
 }
-.notebookEditor .in-preview .actions-container .action-item .notebook-button {
+.notebookEditor .editor-toolbar .actions-container .action-item .notebook-button.masked-icon {
+	padding: 0;
+	width: 28px;
+}
+
+.notebookEditor .in-preview .actions-container .action-item .notebook-button,
+.notebookEditor .in-preview .actions-container .action-item .notebook-button:before {
 	display: flex;
-	background-size: 16px;
 	height: 100%;
 }
 
@@ -66,7 +68,7 @@
 	.in-preview
 	.actions-container
 	.action-item
-	.notebook-button.masked-icon {
+	.notebook-button.masked-icon:before {
 	margin-right: 0;
 	padding-left: 18px;
 	width: 16px;
@@ -169,12 +171,15 @@
 }
 
 .vs-dark .notebookEditor :not(.in-preview) .notebook-button.icon-clear-results,
-.hc-black
-	.notebookEditor
-	:not(.in-preview)
-	.notebook-button.icon-clear-results {
+.hc-black .notebookEditor :not(.in-preview) .notebook-button.icon-clear-results {
 	background-image: url("./media/dark/clear_results_inverse.svg");
 }
+
+.notebookEditor .in-preview .notebook-button.masked-icon,
+.vs-dark .notebookEditor .in-preview .notebook-button.icon-clear-results,
+.hc-black .notebookEditor .in-preview .notebook-button.icon-clear-results {
+	background-image: none;
+}
 /* non-preview */
 
 .notebookEditor .in-preview .carbon-taskbar.monaco-toolbar .monaco-select-box {

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

@@ -150,7 +150,7 @@ export function registerNotebookThemes(overrideEditorThemeSetting: boolean, conf
 		//Notebook toolbar masked icons
 		const notebookToolbarIconColor = theme.getColor(notebookToolbarIcon);
 		if (notebookToolbarIconColor) {
-			collector.addRule(`.notebookEditor .notebook-button.masked-icon { background-color: ${notebookToolbarIconColor};}`);
+			collector.addRule(`.notebookEditor .notebook-button.masked-icon:before { background-color: ${notebookToolbarIconColor};}`);
 			collector.addRule(`.notebookEditor .notebook-button.masked-pseudo:before { background-color: ${notebookToolbarIconColor};}`);
 		}
 		const notebookToolbarLinesColor = theme.getColor(notebookToolbarLines);
@@ -172,7 +172,7 @@ export function registerNotebookThemes(overrideEditorThemeSetting: boolean, conf
 		if (cellBorderColor) {
 			collector.addRule(`.notebookEditor .notebook-cell.active { border-color: ${cellBorderColor};}`);
 			collector.addRule(`.notebookEditor .notebook-cell.active cell-toolbar-component { border-color: ${cellBorderColor};}`);
-			collector.addRule(`.notebookEditor .notebook-cell.active cell-toolbar-component .codicon { background-color: ${cellBorderColor};}`);
+			collector.addRule(`.notebookEditor .notebook-cell.active cell-toolbar-component .codicon:before { background-color: ${cellBorderColor};}`);
 		}
 		// Cell toolbar background
 		const notebookToolbarSelectBackgroundColor = theme.getColor(notebookToolbarSelectBackground);

src/sql/workbench/contrib/query/browser/queryActions.ts

@@ -242,8 +242,8 @@ export class RunQueryAction extends QueryTaskbarAction {
 		if (this.isConnected(editor)) {
 			// if the selection isn't empty then execute the selection
 			// otherwise, either run the statement or the script depending on parameter
-			let selection = editor.getSelection();
-			if (runCurrentStatement && selection) {
+			let selection = editor.getSelection(false);
+			if (runCurrentStatement && selection && this.isCursorPosition(selection)) {
 				editor.input.runQueryStatement(selection);
 			} else {
 				// get the selection again this time with trimming
@@ -252,6 +252,11 @@ export class RunQueryAction extends QueryTaskbarAction {
 			}
 		}
 	}
+
+	protected isCursorPosition(selection: IRange) {
+		return selection.startLineNumber === selection.endLineNumber
+			&& selection.startColumn === selection.endColumn;
+	}
 }
 
 /**

src/sql/workbench/contrib/query/test/browser/queryActions.test.ts

@@ -560,6 +560,124 @@ suite('SQL QueryAction Tests', () => {
 		// ... The connection should have changed to the provided database
 		assert.equal(listItem.currentDatabaseName, eventParams.connectionProfile.databaseName);
 	});
+
+	test('runCurrent - opens connection dialog when there are no active connections', async () => {
+		// setting up environment
+		let isConnected = false;
+		let predefinedSelection: IRange = { startLineNumber: 1, startColumn: 1, endLineNumber: 1, endColumn: 1 };
+		let calledRunQueryStatementOnInput: boolean = undefined;
+
+		// mocking query editor
+		const contextkeyservice = new MockContextKeyService();
+		let queryEditor = TypeMoq.Mock.ofType(QueryEditor, TypeMoq.MockBehavior.Loose, undefined, new TestThemeService(),
+			new TestStorageService(), contextkeyservice, undefined, new TestFileService(), undefined);
+		queryEditor.setup(x => x.input).returns(() => testQueryInput.object);
+		queryEditor.setup(x => x.getSelection(false)).returns(() => { return predefinedSelection; });
+
+		// Mocking runQueryStatment in unititledQueryEditorInput
+		testQueryInput.setup(x => x.runQueryStatement(TypeMoq.It.isAny())).callback(() => { calledRunQueryStatementOnInput = true; });
+
+		// mocking isConnected in ConnectionManagementService
+		connectionManagementService.setup(x => x.isConnected(TypeMoq.It.isAnyString())).returns(() => isConnected);
+
+		// mocking QueryModelService
+		let queryModelService = TypeMoq.Mock.ofType(QueryModelService, TypeMoq.MockBehavior.Loose);
+		queryModelService.setup(x => x.runQueryStatement(TypeMoq.It.isAny(), TypeMoq.It.isAny()));
+
+		// Calling runCurrent with no open connection
+		let queryAction: RunQueryAction = new RunQueryAction(queryEditor.object, queryModelService.object, connectionManagementService.object);
+		calledRunQueryStatementOnInput = false;
+		await queryAction.runCurrent();
+
+		//connection dialog should open and runQueryStatement should not be called
+		assert.equal(calledRunQueryStatementOnInput, false, 'runCurrent should not call runQueryStatement');
+		testQueryInput.verify(x => x.runQueryStatement(TypeMoq.It.isAny()), TypeMoq.Times.never());
+		connectionManagementService.verify(x => x.showConnectionDialog(TypeMoq.It.isAny()), TypeMoq.Times.once());
+
+
+		// Calling runCurrent with an open connection
+		isConnected = true;
+		await queryAction.runCurrent();
+
+		//connection dialog should not open and runQueryStatement should be called
+		assert.equal(calledRunQueryStatementOnInput, true, 'runCurrent should call runQueryStatement');
+		testQueryInput.verify(x => x.runQueryStatement(TypeMoq.It.isAny()), TypeMoq.Times.once());
+		//show Dialog is not called
+		connectionManagementService.verify(x => x.showConnectionDialog(TypeMoq.It.isAny()), TypeMoq.Times.once());
+
+		// Calling runCurrent with empty Selection
+		isConnected = true;
+		calledRunQueryStatementOnInput = false;
+		await queryAction.runCurrent();
+
+		// Selection is empty
+		queryEditor.setup(x => x.isSelectionEmpty()).returns(() => true);
+
+		//connection dialog should not open and runQueryStatement should not be called
+		assert.equal(calledRunQueryStatementOnInput, false, 'runCurrent should not call runQueryStatemet');
+		connectionManagementService.verify(x => x.showConnectionDialog(TypeMoq.It.isAny()), TypeMoq.Times.once());
+	});
+
+	test('runCurrent- calls appropriate run methods based on different selections', async () => {
+		// setting up environment
+		let calledRunQueryStatementOnInput: boolean = undefined;
+		let predefinedCursorSelection: IRange = { startLineNumber: 1, startColumn: 1, endLineNumber: 1, endColumn: 1 };
+		let predefinedRangeSelection: IRange = { startLineNumber: 1, startColumn: 2, endLineNumber: 3, endColumn: 4 };
+
+		// mocking query editor
+		const contextkeyservice = new MockContextKeyService();
+		let queryEditor = TypeMoq.Mock.ofType(QueryEditor, TypeMoq.MockBehavior.Loose, undefined, new TestThemeService(),
+			new TestStorageService(), contextkeyservice, undefined, new TestFileService(), undefined);
+		queryEditor.setup(x => x.input).returns(() => testQueryInput.object);
+
+		// mocking isConnected in ConnectionManagementService
+		connectionManagementService.setup(x => x.isConnected(TypeMoq.It.isAnyString())).returns(() => true);
+
+		// Mocking runQuery and runQueryStatement in unititledQueryEditorInput
+		testQueryInput.setup(x => x.runQuery(TypeMoq.It.isAny())).callback(() => { calledRunQueryOnInput = true; });
+		testQueryInput.setup(x => x.runQueryStatement(TypeMoq.It.isAny())).callback(() => { calledRunQueryStatementOnInput = true; });
+
+		// mocking QueryModelService
+		let queryModelService = TypeMoq.Mock.ofType(QueryModelService, TypeMoq.MockBehavior.Loose);
+		queryModelService.setup(x => x.runQuery(TypeMoq.It.isAny(), undefined, TypeMoq.It.isAny()));
+		queryModelService.setup(x => x.runQueryStatement(TypeMoq.It.isAny(), TypeMoq.It.isAny()));
+
+		let queryAction: RunQueryAction = new RunQueryAction(queryEditor.object, queryModelService.object, connectionManagementService.object);
+
+		// setting up queryEditor with only a cursor. This case should call runQueryStatement
+		queryEditor.setup(x => x.getSelection(false)).returns(() => { return predefinedCursorSelection; });
+		calledRunQueryOnInput = false;
+		calledRunQueryStatementOnInput = false;
+		await queryAction.runCurrent();
+
+		assert.equal(calledRunQueryStatementOnInput, true, 'runCurrent should call runQueryStatement');
+		assert.equal(calledRunQueryOnInput, false, 'run should not call runQuery');
+
+		// checking if runQuery statement is called with predefinedCursorSelection only
+		testQueryInput.verify(x => x.runQueryStatement(TypeMoq.It.isValue(predefinedCursorSelection)), TypeMoq.Times.once());
+		testQueryInput.verify(x => x.runQueryStatement(TypeMoq.It.isAny()), TypeMoq.Times.once());
+
+		// checking if runQuery is not called at all
+		testQueryInput.verify(x => x.runQuery(TypeMoq.It.isAny()), TypeMoq.Times.never());
+
+		// setting up queryEditor with a selection range. This case should call runQuery
+		queryEditor.setup(x => x.getSelection()).returns(() => { return predefinedRangeSelection; });
+		queryEditor.setup(x => x.getSelection(false)).returns(() => { return predefinedRangeSelection; });
+
+		calledRunQueryOnInput = false;
+		calledRunQueryStatementOnInput = false;
+		await queryAction.runCurrent();
+
+		assert.equal(calledRunQueryStatementOnInput, false, 'runCurrent should not call runQueryStatement');
+		assert.equal(calledRunQueryOnInput, true, 'run should call runQuery');
+
+		// checking if runQuery is called with predefinedRangeSelection only
+		testQueryInput.verify(x => x.runQuery(TypeMoq.It.isValue(predefinedRangeSelection)), TypeMoq.Times.once());
+		testQueryInput.verify(x => x.runQuery(TypeMoq.It.isAny()), TypeMoq.Times.once());
+
+		// checking if runQueryStatement is never called
+		testQueryInput.verify(x => x.runQueryStatement(TypeMoq.It.isAny()), TypeMoq.Times.once());
+	});
 });
 
 class ServiceAccessor {

src/sql/workbench/services/connection/browser/connectionManagementService.ts

@@ -813,7 +813,7 @@ export class ConnectionManagementService extends Disposable implements IConnecti
 		const accounts = await this._accountManagementService.getAccounts();
 		const azureAccounts = accounts.filter(a => a.key.providerId.startsWith('azure'));
 		if (azureAccounts && azureAccounts.length > 0) {
-			let accountName = (connection.authenticationType === Constants.azureMFA) ? connection.azureAccount : connection.userName;
+			let accountName = (connection.authenticationType === Constants.azureMFA || connection.authenticationType === Constants.azureMFAAndUser) ? connection.azureAccount : connection.userName;
 			let account = find(azureAccounts, account => account.key.accountId === accountName);
 			if (account) {
 				this._logService.debug(`Getting security token for Azure account ${account.key.accountId}`);