only show new query when query is supported (#17346)

* only show new query when query is supported

* test case fix and pr comments

* rename context key

samples/sqlservices/package.json

@@ -101,7 +101,7 @@
     "connectionProvider": {
       "providerId": "TESTPROVIDER",
       "languageMode": "sql",
-      "notebookKernelAlias": "Test Provider",
+      "isQueryProvider": false,
       "displayName": "Test Provider",
       "iconPath": [
         {

src/sql/platform/capabilities/common/capabilitiesService.ts

@@ -27,6 +27,7 @@ export interface ConnectionProviderProperties {
 	notebookKernelAlias?: string;
 	azureResource?: string;
 	connectionOptions: azdata.ConnectionOption[];
+	isQueryProvider?: boolean;
 }
 
 export interface ProviderFeatures {

src/sql/platform/capabilities/common/capabilitiesServiceImpl.ts

@@ -38,7 +38,8 @@ export class CapabilitiesService extends Disposable implements ICapabilitiesServ
 			};
 			this._providers.set(id, provider);
 		}
-
+		// By default isQueryProvider is true.
+		provider.connection.isQueryProvider = provider.connection.isQueryProvider !== false;
 		this._onCapabilitiesRegistered.fire({ id, features: provider });
 	}
 

src/sql/workbench/contrib/connection/common/connectionProviderExtension.ts

@@ -30,6 +30,10 @@ const ConnectionProviderContrib: IJSONSchema = {
 			type: 'string',
 			description: localize('schema.notebookKernelAlias', "Notebook Kernel Alias for the provider")
 		},
+		isQueryProvider: {
+			type: 'boolean',
+			description: localize('schema.isQueryProvider', "Whether the provider is also a query provider. The default value is true.")
+		},
 		iconPath: {
 			description: localize('schema.iconPath', "Icon path for the server type"),
 			oneOf: [

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

@@ -134,7 +134,7 @@ MenuRegistry.appendMenuItem(MenuId.ObjectExplorerItemContext, {
 		id: OE_NEW_NOTEBOOK_COMMAND_ID,
 		title: localize('newQuery', "New Notebook")
 	},
-	when: ContextKeyExpr.or(ContextKeyExpr.and(TreeNodeContextKey.Status.notEqualsTo('Unavailable'), TreeNodeContextKey.NodeType.isEqualTo('Server')), ContextKeyExpr.and(TreeNodeContextKey.Status.notEqualsTo('Unavailable'), TreeNodeContextKey.NodeType.isEqualTo('Database')))
+	when: ContextKeyExpr.and(TreeNodeContextKey.Status.notEqualsTo('Unavailable'), TreeNodeContextKey.IsQueryProvider.isEqualTo(true), ContextKeyExpr.or(TreeNodeContextKey.NodeType.isEqualTo('Server'), TreeNodeContextKey.NodeType.isEqualTo('Database')))
 });
 
 const ExplorerNotebookActionID = 'explorer.notebook';

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

@@ -78,7 +78,7 @@ MenuRegistry.appendMenuItem(MenuId.ObjectExplorerItemContext, {
 		id: OE_NEW_QUERY_ACTION_ID,
 		title: localize('newQuery', "New Query")
 	},
-	when: ContextKeyExpr.or(ContextKeyExpr.and(TreeNodeContextKey.Status.notEqualsTo('Unavailable'), TreeNodeContextKey.NodeType.isEqualTo('Server')), ContextKeyExpr.and(TreeNodeContextKey.Status.notEqualsTo('Unavailable'), TreeNodeContextKey.NodeType.isEqualTo('Database')))
+	when: ContextKeyExpr.and(TreeNodeContextKey.Status.notEqualsTo('Unavailable'), TreeNodeContextKey.IsQueryProvider.isEqualTo(true), ContextKeyExpr.or(TreeNodeContextKey.NodeType.isEqualTo('Server'), TreeNodeContextKey.NodeType.isEqualTo('Database')))
 });
 
 // New Query
@@ -89,7 +89,7 @@ MenuRegistry.appendMenuItem(MenuId.DataExplorerContext, {
 		id: DE_NEW_QUERY_COMMAND_ID,
 		title: localize('newQuery', "New Query")
 	},
-	when: MssqlNodeContext.IsDatabaseOrServer
+	when: ContextKeyExpr.and(MssqlNodeContext.IsDatabaseOrServer, MssqlNodeContext.IsQueryProvider)
 });
 
 const ExplorerNewQueryActionID = 'explorer.query';

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

@@ -7,7 +7,7 @@ import 'vs/css!./media/connectionDialog';
 import { Button } from 'sql/base/browser/ui/button/button';
 import { SelectBox } from 'sql/base/browser/ui/selectBox/selectBox';
 import { HideReason, Modal } from 'sql/workbench/browser/modal/modal';
-import { IConnectionManagementService, INewConnectionParams } from 'sql/platform/connection/common/connectionManagement';
+import { ConnectionType, IConnectionManagementService, INewConnectionParams } from 'sql/platform/connection/common/connectionManagement';
 import * as DialogHelper from 'sql/workbench/browser/modal/dialogHelper';
 import { TreeCreationUtils } from 'sql/workbench/services/objectExplorer/browser/treeCreationUtils';
 import { TabbedPanel, PanelTabIdentifier } from 'sql/base/browser/ui/panel/panel';
@@ -43,6 +43,7 @@ import { AsyncServerTree } from 'sql/workbench/services/objectExplorer/browser/a
 import { IConfigurationService } from 'vs/platform/configuration/common/configuration';
 import { ConnectionBrowseTab } from 'sql/workbench/services/connection/browser/connectionBrowseTab';
 import { ElementSizeObserver } from 'vs/editor/browser/config/elementSizeObserver';
+import { ICapabilitiesService } from 'sql/platform/capabilities/common/capabilitiesService';
 
 export interface OnShowUIResponse {
 	selectedProviderDisplayName: string;
@@ -122,7 +123,8 @@ export class ConnectionDialogWidget extends Modal {
 		@IClipboardService clipboardService: IClipboardService,
 		@ILogService logService: ILogService,
 		@ITextResourcePropertiesService textResourcePropertiesService: ITextResourcePropertiesService,
-		@IConfigurationService private _configurationService: IConfigurationService
+		@IConfigurationService private _configurationService: IConfigurationService,
+		@ICapabilitiesService private _capabilitiesService: ICapabilitiesService
 	) {
 		super(
 			localize('connection', "Connection"),
@@ -155,8 +157,20 @@ export class ConnectionDialogWidget extends Modal {
 
 	public refresh(): void {
 		let filteredProviderMap = this.providerNameToDisplayNameMap;
-		if (this._newConnectionParams && this._newConnectionParams.providers) {
+		if (this._newConnectionParams && (this._newConnectionParams.providers || this._newConnectionParams.connectionType === ConnectionType.editor)) {
-			const validProviderMap = entries(this.providerNameToDisplayNameMap).filter(x => this.includeProvider(x[0], this._newConnectionParams));
+			const validProviderMap = entries(this.providerNameToDisplayNameMap).filter(x => {
+				const providerName = x[0];
+				const capabilities = this._capabilitiesService.getCapabilities(providerName).connection;
+				// If the connection is for an editor (e.g. query editor or notebook), the provider must be a query provider
+				if (this._newConnectionParams.connectionType === ConnectionType.editor && !capabilities.isQueryProvider) {
+					return false;
+				}
+				// If the provider list is provided, the provider must be in the list.
+				if (this._newConnectionParams.providers && !this._newConnectionParams.providers.some(p => p === providerName)) {
+					return false;
+				}
+				return true;
+			});
 			if (validProviderMap && validProviderMap.length > 0) {
 				let map: { [providerDisplayName: string]: string } = {};
 				validProviderMap.forEach(v => {
@@ -171,10 +185,6 @@ export class ConnectionDialogWidget extends Modal {
 		this._providerTypeSelectBox.setOptions(Object.keys(uniqueProvidersMap).map(k => uniqueProvidersMap[k]));
 	}
 
-	private includeProvider(providerName: string, params?: INewConnectionParams): Boolean {
-		return params === undefined || params.providers === undefined || params.providers.some(x => x === providerName);
-	}
-
 	protected renderBody(container: HTMLElement): void {
 		this._body = DOM.append(container, DOM.$('.connection-dialog'));
 

src/sql/workbench/services/connection/test/browser/connectionDialogService.test.ts

@@ -52,6 +52,7 @@ import { TestTreeView } from 'sql/workbench/services/connection/test/browser/tes
 import { TestConfigurationService } from 'sql/platform/connection/test/common/testConfigurationService';
 import { ConnectionTreeService, IConnectionTreeService } from 'sql/workbench/services/connection/common/connectionTreeService';
 import { ConnectionBrowserView } from 'sql/workbench/services/connection/browser/connectionBrowseTab';
+import { ICapabilitiesService } from 'sql/platform/capabilities/common/capabilitiesService';
 
 suite('ConnectionDialogService tests', () => {
 	const testTreeViewId = 'testTreeView';
@@ -102,6 +103,7 @@ suite('ConnectionDialogService tests', () => {
 		testInstantiationService.stub(ILayoutService, new TestLayoutService());
 		testInstantiationService.stub(IAdsTelemetryService, new NullAdsTelemetryService());
 		testInstantiationService.stub(IConnectionTreeService, new ConnectionTreeService());
+		testInstantiationService.stub(ICapabilitiesService, new TestCapabilitiesService());
 		connectionDialogService = new ConnectionDialogService(testInstantiationService, capabilitiesService, errorMessageService.object,
 			new TestConfigurationService(), new BrowserClipboardService(), NullCommandService, new NullLogService());
 		(connectionDialogService as any)._connectionManagementService = mockConnectionManagementService.object;
@@ -113,7 +115,7 @@ suite('ConnectionDialogService tests', () => {
 		mockConnectionManagementService.setup(x => x.getConnectionGroups(TypeMoq.It.isAny())).returns(() => {
 			return [new ConnectionProfileGroup('test_group', undefined, 'test_group')];
 		});
-		testConnectionDialog = new TestConnectionDialogWidget(providerDisplayNames, providerNameToDisplayMap['MSSQL'], providerNameToDisplayMap, testInstantiationService, mockConnectionManagementService.object, undefined, undefined, viewDescriptorService, new TestThemeService(), new TestLayoutService(), new NullAdsTelemetryService(), new MockContextKeyService(), undefined, new NullLogService(), new TestTextResourcePropertiesService(new TestConfigurationService), new TestConfigurationService());
+		testConnectionDialog = new TestConnectionDialogWidget(providerDisplayNames, providerNameToDisplayMap['MSSQL'], providerNameToDisplayMap, testInstantiationService, mockConnectionManagementService.object, undefined, undefined, viewDescriptorService, new TestThemeService(), new TestLayoutService(), new NullAdsTelemetryService(), new MockContextKeyService(), undefined, new NullLogService(), new TestTextResourcePropertiesService(new TestConfigurationService), new TestConfigurationService(), new TestCapabilitiesService());
 		testConnectionDialog.render();
 		testConnectionDialog['renderBody'](DOM.createStyleSheet());
 		(connectionDialogService as any)._connectionDialog = testConnectionDialog;

src/sql/workbench/services/connection/test/browser/connectionDialogWidget.test.ts

@@ -65,7 +65,7 @@ suite('ConnectionDialogWidget tests', () => {
 			new TestCapabilitiesService());
 		let providerDisplayNames = ['Mock SQL Server'];
 		let providerNameToDisplayMap = { 'MSSQL': 'Mock SQL Server' };
-		connectionDialogWidget = new TestConnectionDialogWidget(providerDisplayNames, providerNameToDisplayMap['MSSQL'], providerNameToDisplayMap, cmInstantiationService, mockConnectionManagementService.object, undefined, undefined, viewDescriptorService, new TestThemeService(), new TestLayoutService(), new NullAdsTelemetryService(), new MockContextKeyService(), undefined, new NullLogService(), new TestTextResourcePropertiesService(new TestConfigurationService()), new TestConfigurationService());
+		connectionDialogWidget = new TestConnectionDialogWidget(providerDisplayNames, providerNameToDisplayMap['MSSQL'], providerNameToDisplayMap, cmInstantiationService, mockConnectionManagementService.object, undefined, undefined, viewDescriptorService, new TestThemeService(), new TestLayoutService(), new NullAdsTelemetryService(), new MockContextKeyService(), undefined, new NullLogService(), new TestTextResourcePropertiesService(new TestConfigurationService()), new TestConfigurationService(), new TestCapabilitiesService());
 		element = DOM.createStyleSheet();
 		connectionDialogWidget.render();
 		connectionDialogWidget['renderBody'](element);

src/sql/workbench/services/connection/test/browser/testConnectionDialogWidget.ts

@@ -16,6 +16,7 @@ import { ILogService } from 'vs/platform/log/common/log';
 import { ITextResourcePropertiesService } from 'vs/editor/common/services/textResourceConfigurationService';
 import { IViewDescriptorService } from 'vs/workbench/common/views';
 import { IConfigurationService } from 'vs/platform/configuration/common/configuration';
+import { ICapabilitiesService } from 'sql/platform/capabilities/common/capabilitiesService';
 
 export class TestConnectionDialogWidget extends ConnectionDialogWidget {
 	constructor(
@@ -34,8 +35,9 @@ export class TestConnectionDialogWidget extends ConnectionDialogWidget {
 		@IClipboardService clipboardService: IClipboardService,
 		@ILogService logService: ILogService,
 		@ITextResourcePropertiesService textResourcePropertiesService: ITextResourcePropertiesService,
-		@IConfigurationService configurationService: IConfigurationService
+		@IConfigurationService configurationService: IConfigurationService,
+		@ICapabilitiesService capabilitiesService: ICapabilitiesService
 	) {
-		super(providerDisplayNameOptions, selectedProviderType, providerNameToDisplayNameMap, _instantiationService, _connectionManagementService, _contextMenuService, _contextViewService, themeService, layoutService, telemetryService, contextKeyService, clipboardService, logService, textResourcePropertiesService, configurationService);
+		super(providerDisplayNameOptions, selectedProviderType, providerNameToDisplayNameMap, _instantiationService, _connectionManagementService, _contextMenuService, _contextViewService, themeService, layoutService, telemetryService, contextKeyService, clipboardService, logService, textResourcePropertiesService, configurationService, capabilitiesService);
 	}
 }

src/sql/workbench/services/objectExplorer/browser/mssqlNodeContext.ts

@@ -48,6 +48,7 @@ export class MssqlNodeContext extends Disposable {
 	static CanScriptAsCreateOrDelete = new RawContextKey<boolean>('canScriptAsCreateOeDelete', false);
 	static CanScriptAsExecute = new RawContextKey<boolean>('canScriptAsExecute', false);
 	static CanScriptAsAlter = new RawContextKey<boolean>('canScriptAsAlter', false);
+	static IsQueryProvider = new RawContextKey<boolean>('isQueryProvider', false);
 
 	private nodeProviderKey!: IContextKey<string>;
 	private isCloudKey!: IContextKey<boolean>;
@@ -62,6 +63,8 @@ export class MssqlNodeContext extends Disposable {
 	private canScriptAsCreateOrDeleteKey!: IContextKey<boolean>;
 	private canScriptAsExecuteKey!: IContextKey<boolean>;
 	private canScriptAsAlterKey!: IContextKey<boolean>;
+	private isQueryProviderKey!: IContextKey<boolean>;
+
 
 	constructor(
 		private nodeContextValue: INodeContextValue,
@@ -88,6 +91,7 @@ export class MssqlNodeContext extends Disposable {
 					this.setScriptingContextKeys();
 					this.nodeTypeKey.set(node.contextValue);
 				}
+				this.setQueryEnabledKey();
 			}
 			if (node.label) {
 				this.nodeLabelKey.set(node.label.label);
@@ -108,6 +112,7 @@ export class MssqlNodeContext extends Disposable {
 		this.canScriptAsAlterKey = MssqlNodeContext.CanScriptAsAlter.bindTo(this.contextKeyService);
 		this.nodeProviderKey = MssqlNodeContext.NodeProvider.bindTo(this.contextKeyService);
 		this.canOpenInAzurePortal = MssqlNodeContext.CanOpenInAzurePortal.bindTo(this.contextKeyService);
+		this.isQueryProviderKey = MssqlNodeContext.IsQueryProvider.bindTo(this.contextKeyService);
 	}
 
 	/**
@@ -203,4 +208,13 @@ export class MssqlNodeContext extends Disposable {
 			this.canScriptAsSelectKey.set(true);
 		}
 	}
+
+	/**
+	 * Set whether the current node's provider is also a query provider.
+	 */
+	private setQueryEnabledKey(): void {
+		const provider = this.nodeContextValue?.node?.payload?.providerName || this.nodeContextValue.node.childProvider;
+		const capabilities = provider ? this.capabilitiesService.getCapabilities(provider) : undefined;
+		this.isQueryProviderKey.set(capabilities?.connection.isQueryProvider);
+	}
 }

src/sql/workbench/services/objectExplorer/browser/serverTreeActionProvider.ts

@@ -25,6 +25,7 @@ import { IQueryManagementService } from 'sql/workbench/services/query/common/que
 import { ServerInfoContextKey } from 'sql/workbench/services/connection/common/serverInfoContextKey';
 import { fillInActions } from 'vs/platform/actions/browser/menuEntryActionViewItem';
 import { AsyncServerTree, ServerTreeElement } from 'sql/workbench/services/objectExplorer/browser/asyncServerTree';
+import { ICapabilitiesService } from 'sql/platform/capabilities/common/capabilitiesService';
 
 /**
  *  Provides actions for the server tree elements
@@ -36,7 +37,8 @@ export class ServerTreeActionProvider {
 		@IConnectionManagementService private _connectionManagementService: IConnectionManagementService,
 		@IQueryManagementService private _queryManagementService: IQueryManagementService,
 		@IMenuService private menuService: IMenuService,
-		@IContextKeyService private _contextKeyService: IContextKeyService
+		@IContextKeyService private _contextKeyService: IContextKeyService,
+		@ICapabilitiesService private _capabilitiesService: ICapabilitiesService
 	) {
 	}
 
@@ -68,6 +70,7 @@ export class ServerTreeActionProvider {
 	 */
 	private getConnectionActions(tree: AsyncServerTree | ITree, profile: ConnectionProfile): IAction[] {
 		let node = new TreeNode(NodeType.Server, '', false, '', '', '', undefined, undefined, undefined, undefined);
+		node.connection = profile;
 		return this.getAllActions({
 			tree: tree,
 			profile: profile,
@@ -146,7 +149,7 @@ export class ServerTreeActionProvider {
 				serverInfoContextKey.set(serverInfo);
 			}
 		}
-		let treeNodeContextKey = new TreeNodeContextKey(scopedContextService);
+		let treeNodeContextKey = new TreeNodeContextKey(scopedContextService, this._capabilitiesService);
 		if (context.treeNode) {
 			treeNodeContextKey.set(context.treeNode);
 		}

src/sql/workbench/services/objectExplorer/common/treeNodeContextKey.ts

@@ -5,6 +5,7 @@
 
 import { RawContextKey, IContextKeyService, IContextKey } from 'vs/platform/contextkey/common/contextkey';
 import { TreeNode } from 'sql/workbench/services/objectExplorer/common/treeNode';
+import { ICapabilitiesService } from 'sql/platform/capabilities/common/capabilitiesService';
 
 export class TreeNodeContextKey implements IContextKey<TreeNode> {
 
@@ -14,6 +15,7 @@ export class TreeNodeContextKey implements IContextKey<TreeNode> {
 	static TreeNode = new RawContextKey<TreeNode>('treeNode', undefined);
 	static NodeLabel = new RawContextKey<string>('nodeLabel', undefined);
 	static NodePath = new RawContextKey<string>('nodePath', undefined);
+	static IsQueryProvider = new RawContextKey<boolean>('isQueryProvider', false);
 
 	private _nodeTypeKey: IContextKey<string>;
 	private _subTypeKey: IContextKey<string>;
@@ -21,9 +23,11 @@ export class TreeNodeContextKey implements IContextKey<TreeNode> {
 	private _treeNodeKey: IContextKey<TreeNode>;
 	private _nodeLabelKey: IContextKey<string>;
 	private _nodePathKey: IContextKey<string>;
+	private _isQueryProvider: IContextKey<boolean>;
 
 	constructor(
-		@IContextKeyService contextKeyService: IContextKeyService
+		@IContextKeyService contextKeyService: IContextKeyService,
+		@ICapabilitiesService private _capabilitiesService: ICapabilitiesService
 	) {
 		this._nodeTypeKey = TreeNodeContextKey.NodeType.bindTo(contextKeyService);
 		this._subTypeKey = TreeNodeContextKey.SubType.bindTo(contextKeyService);
@@ -31,6 +35,7 @@ export class TreeNodeContextKey implements IContextKey<TreeNode> {
 		this._treeNodeKey = TreeNodeContextKey.TreeNode.bindTo(contextKeyService);
 		this._nodeLabelKey = TreeNodeContextKey.NodeLabel.bindTo(contextKeyService);
 		this._nodePathKey = TreeNodeContextKey.NodePath.bindTo(contextKeyService);
+		this._isQueryProvider = TreeNodeContextKey.IsQueryProvider.bindTo(contextKeyService);
 	}
 
 	set(value: TreeNode) {
@@ -42,6 +47,9 @@ export class TreeNodeContextKey implements IContextKey<TreeNode> {
 		}
 		this._nodeLabelKey.set(value && value.label);
 		this._nodePathKey.set(value && value.nodePath);
+		const connectionProfile = value.getConnectionProfile();
+		const capabilities = connectionProfile ? this._capabilitiesService.getCapabilities(connectionProfile.providerName) : undefined;
+		this._isQueryProvider.set(capabilities?.connection.isQueryProvider);
 	}
 
 	reset(): void {
@@ -51,6 +59,7 @@ export class TreeNodeContextKey implements IContextKey<TreeNode> {
 		this._treeNodeKey.reset();
 		this._nodeLabelKey.reset();
 		this._nodePathKey.reset();
+		this._isQueryProvider.reset();
 	}
 
 	public get(): TreeNode | undefined {