enable 'New Notebook' entry points for PG SQL (#5194)

* enable 'New Notebook' entry points for PG SQL

* fix comment

* permanent fix for context menu entry point

extensions/notebook/package.json

@@ -201,7 +201,7 @@
 			"objectExplorer/item/context": [
 				{
 					"command": "notebook.command.new",
-					"when": "connectionProvider == MSSQL && nodeType == Server",
+					"when": "isQueryProvider && nodeType == Server",
 					"group": "1root@1"
 				},
 				{

src/sql/platform/connection/common/utils.ts

@@ -136,5 +136,8 @@ export function findProfileInGroup(og: IConnectionProfile, groups: ConnectionPro
 }
 
 export function isMaster(profile: IConnectionProfile): boolean {
-	return profile.providerName.toLowerCase() === 'mssql' && profile.databaseName.toLowerCase() === 'master';
+	// TODO: the connection profile should have a property to indicate whether the connection is a server connection or database connection
+	// created issue to track the problem: https://github.com/Microsoft/azuredatastudio/issues/5193.
+	return (profile.providerName.toLowerCase() === 'mssql' && profile.databaseName.toLowerCase() === 'master')
+		|| (profile.providerName.toLowerCase() === 'pgsql' && profile.databaseName.toLowerCase() === 'postgres');
 }

src/sql/workbench/parts/connection/common/connectionContextKey.ts

@@ -5,6 +5,7 @@
 
 import { RawContextKey, IContextKeyService, IContextKey } from 'vs/platform/contextkey/common/contextkey';
 import { IConnectionProfile } from 'azdata';
+import { IQueryManagementService } from 'sql/platform/query/common/queryManagement';
 
 export class ConnectionContextKey implements IContextKey<IConnectionProfile> {
 
@@ -12,26 +13,32 @@ export class ConnectionContextKey implements IContextKey<IConnectionProfile> {
 	static Server = new RawContextKey<string>('serverName', undefined);
 	static Database = new RawContextKey<string>('databaseName', undefined);
 	static Connection = new RawContextKey<IConnectionProfile>('connection', undefined);
+	static IsQueryProvider = new RawContextKey<boolean>('isQueryProvider', false);
 
 	private _providerKey: IContextKey<string>;
 	private _serverKey: IContextKey<string>;
 	private _databaseKey: IContextKey<string>;
 	private _connectionKey: IContextKey<IConnectionProfile>;
+	private _isQueryProviderKey: IContextKey<boolean>;
 
 	constructor(
-		@IContextKeyService contextKeyService: IContextKeyService
+		@IContextKeyService contextKeyService: IContextKeyService,
+		@IQueryManagementService private queryManagementService: IQueryManagementService
 	) {
 		this._providerKey = ConnectionContextKey.Provider.bindTo(contextKeyService);
 		this._serverKey = ConnectionContextKey.Server.bindTo(contextKeyService);
 		this._databaseKey = ConnectionContextKey.Database.bindTo(contextKeyService);
 		this._connectionKey = ConnectionContextKey.Connection.bindTo(contextKeyService);
+		this._isQueryProviderKey = ConnectionContextKey.IsQueryProvider.bindTo(contextKeyService);
 	}
 
 	set(value: IConnectionProfile) {
+		let queryProviders = this.queryManagementService.getRegisteredProviders();
 		this._connectionKey.set(value);
 		this._providerKey.set(value && value.providerName);
 		this._serverKey.set(value && value.serverName);
 		this._databaseKey.set(value && value.databaseName);
+		this._isQueryProviderKey.set(value && value.providerName && queryProviders.indexOf(value.providerName) !== -1);
 	}
 
 	reset(): void {
@@ -39,6 +46,7 @@ export class ConnectionContextKey implements IContextKey<IConnectionProfile> {
 		this._serverKey.reset();
 		this._databaseKey.reset();
 		this._connectionKey.reset();
+		this._isQueryProviderKey.reset();
 	}
 
 	public get(): IConnectionProfile {

src/sql/workbench/parts/dashboard/dashboardEditor.ts

@@ -23,6 +23,7 @@ import { IConnectionProfile } from 'sql/platform/connection/common/interfaces';
 import { IConnectionManagementService } from 'sql/platform/connection/common/connectionManagement';
 import { CancellationToken } from 'vs/base/common/cancellation';
 import { IStorageService } from 'vs/platform/storage/common/storage';
+import { IQueryManagementService } from 'sql/platform/query/common/queryManagement';
 
 export class DashboardEditor extends BaseEditor {
 
@@ -37,7 +38,8 @@ export class DashboardEditor extends BaseEditor {
 		@IContextKeyService private _contextKeyService: IContextKeyService,
 		@IDashboardService private _dashboardService: IDashboardService,
 		@IConnectionManagementService private _connMan: IConnectionManagementService,
-		@IStorageService storageService: IStorageService
+		@IStorageService storageService: IStorageService,
+		@IQueryManagementService private queryManagementService: IQueryManagementService
 	) {
 		super(DashboardEditor.ID, telemetryService, themeService, storageService);
 	}
@@ -112,7 +114,7 @@ export class DashboardEditor extends BaseEditor {
 		const serverInfo = this._connMan.getConnectionInfo(this.input.uri).serverInfo;
 		this._dashboardService.changeToDashboard({ profile, serverInfo });
 		const scopedContextService = this._contextKeyService.createScoped(input.container);
-		const connectionContextKey = new ConnectionContextKey(scopedContextService);
+		const connectionContextKey = new ConnectionContextKey(scopedContextService, this.queryManagementService);
 		connectionContextKey.set(input.connectionProfile);
 
 		const params: IDashboardComponentParams = {

src/sql/workbench/parts/dataExplorer/common/nodeContext.ts

@@ -8,6 +8,7 @@ import { IOEShimService } from 'sql/workbench/parts/objectExplorer/common/object
 import { ITreeItem } from 'sql/workbench/common/views';
 import { Disposable } from 'vs/base/common/lifecycle';
 import { IContextKey, IContextKeyService, RawContextKey } from 'vs/platform/contextkey/common/contextkey';
+import { IQueryManagementService } from 'sql/platform/query/common/queryManagement';
 
 export interface INodeContextValue {
 	node: ITreeItem;
@@ -31,7 +32,8 @@ export class NodeContextKey extends Disposable implements IContextKey<INodeConte
 
 	constructor(
 		@IContextKeyService contextKeyService: IContextKeyService,
-		@IOEShimService private oeService: IOEShimService
+		@IOEShimService private oeService: IOEShimService,
+		@IQueryManagementService queryManagementService: IQueryManagementService
 	) {
 		super();
 
@@ -40,7 +42,7 @@ export class NodeContextKey extends Disposable implements IContextKey<INodeConte
 		this._viewIdKey = NodeContextKey.ViewId.bindTo(contextKeyService);
 		this._viewItemKey = NodeContextKey.ViewItem.bindTo(contextKeyService);
 		this._nodeContextKey = NodeContextKey.Node.bindTo(contextKeyService);
-		this._connectionContextKey = new ConnectionContextKey(contextKeyService);
+		this._connectionContextKey = new ConnectionContextKey(contextKeyService, queryManagementService);
 	}
 
 	set(value: INodeContextValue) {

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

@@ -124,7 +124,7 @@ export class ServerTreeActionProvider extends ContributableActionProvider {
 
 	private getContextKeyService(context: ObjectExplorerContext): IContextKeyService {
 		let scopedContextService = this._contextKeyService.createScoped();
-		let connectionContextKey = new ConnectionContextKey(scopedContextService);
+		let connectionContextKey = new ConnectionContextKey(scopedContextService, this._queryManagementService);
 		let connectionProfile = context && context.profile;
 		connectionContextKey.set(connectionProfile);
 		let serverInfoContextKey = new ServerInfoContextKey(scopedContextService);