Added Unified connection support (#3785)

* Added Unified connection support

* Use generic way to do expandNode.
Cleanup the ported code and removed unreference code. Added as needed later.
Resolved PR comments.

* Minor fixes and removed timer for all expanders for now. If any providers can't response, the tree node will spin and wait. We may improve later.

* Change handSessionClose to not thenable.
Added a node to OE to show error message instead of reject. So we could show partial expanded result if get any.
Resolve PR comments

* Minor fixes of PR comments

src/sql/parts/objectExplorer/common/objectExplorerService.ts

@@ -26,6 +26,10 @@ export const SERVICE_ID = 'ObjectExplorerService';
 
 export const IObjectExplorerService = createDecorator<IObjectExplorerService>(SERVICE_ID);
 
+export interface NodeExpandInfoWithProviderId extends sqlops.ObjectExplorerExpandInfo {
+	providerId: string;
+}
+
 export interface IObjectExplorerService {
 	_serviceBrand: any;
 
@@ -45,13 +49,15 @@ export interface IObjectExplorerService {
 
 	onSessionDisconnected(handle: number, sessionResponse: sqlops.ObjectExplorerSession);
 
-	onNodeExpanded(handle: number, sessionResponse: sqlops.ObjectExplorerExpandInfo);
+	onNodeExpanded(sessionResponse: NodeExpandInfoWithProviderId);
 
 	/**
 	 * Register a ObjectExplorer provider
 	 */
 	registerProvider(providerId: string, provider: sqlops.ObjectExplorerProvider): void;
 
+	registerNodeProvider(expander: sqlops.ObjectExplorerNodeProvider): void;
+
 	getObjectExplorerNode(connection: IConnectionProfile): TreeNode;
 
 	updateObjectExplorerNodes(connectionProfile: IConnectionProfile): Promise<void>;
@@ -82,15 +88,18 @@ export interface IObjectExplorerService {
 	 * For Testing purpose only. Get the context menu actions for an object explorer node.
 	*/
 	getNodeActions(connectionId: string, nodePath: string): Thenable<string[]>;
+
+	getSessionConnectionProfile(sessionId: string): sqlops.IConnectionProfile;
 }
 
 interface SessionStatus {
 	nodes: { [nodePath: string]: NodeStatus };
 	connection: ConnectionProfile;
+	expandNodeTimer?: number;
 }
 
 interface NodeStatus {
-	expandEmitter: Emitter<sqlops.ObjectExplorerExpandInfo>;
+	expandEmitter: Emitter<NodeExpandInfoWithProviderId>;
 }
 
 export interface ObjectExplorerNodeEventArgs {
@@ -103,6 +112,14 @@ export interface NodeInfoWithConnection {
 	nodeInfo: sqlops.NodeInfo;
 }
 
+export interface TopLevelChildrenPath {
+	providerId: string;
+	supportedProviderId: string;
+	groupingId: number;
+	path: string[];
+	providerObject: sqlops.ObjectExplorerNodeProvider | sqlops.ObjectExplorerProvider;
+}
+
 export class ObjectExplorerService implements IObjectExplorerService {
 
 	public _serviceBrand: any;
@@ -111,6 +128,8 @@ export class ObjectExplorerService implements IObjectExplorerService {
 
 	private _providers: { [handle: string]: sqlops.ObjectExplorerProvider; } = Object.create(null);
 
+	private _nodeProviders: { [handle: string]: sqlops.ObjectExplorerNodeProvider[]; } = Object.create(null);
+
 	private _activeObjectExplorerNodes: { [id: string]: TreeNode };
 	private _sessions: { [sessionId: string]: SessionStatus };
 
@@ -129,6 +148,7 @@ export class ObjectExplorerService implements IObjectExplorerService {
 		this._activeObjectExplorerNodes = {};
 		this._sessions = {};
 		this._providers = {};
+		this._nodeProviders = {};
 		this._onSelectionOrFocusChange = new Emitter<void>();
 	}
 
@@ -166,7 +186,7 @@ export class ObjectExplorerService implements IObjectExplorerService {
 	/**
 	 * Gets called when expanded node response is ready
 	 */
-	public onNodeExpanded(handle: number, expandResponse: sqlops.ObjectExplorerExpandInfo) {
+	public onNodeExpanded(expandResponse: NodeExpandInfoWithProviderId) {
 
 		if (expandResponse.errorMessage) {
 			error(expandResponse.errorMessage);
@@ -189,28 +209,43 @@ export class ObjectExplorerService implements IObjectExplorerService {
 	/**
 	 * Gets called when session is created
 	 */
-	public onSessionCreated(handle: number, session: sqlops.ObjectExplorerSession) {
-		let connection: ConnectionProfile = undefined;
-		let errorMessage: string = undefined;
-		if (this._sessions[session.sessionId]) {
-			connection = this._sessions[session.sessionId].connection;
+	public onSessionCreated(handle: number, session: sqlops.ObjectExplorerSession): void {
+		this.handleSessionCreated(session);
+	}
 
-			if (session && session.success && session.rootNode) {
-				let server = this.toTreeNode(session.rootNode, null);
-				server.connection = connection;
-				server.session = session;
-				this._activeObjectExplorerNodes[connection.id] = server;
-			} else {
-				errorMessage = session && session.errorMessage ? session.errorMessage :
-					nls.localize('OeSessionFailedError', 'Failed to create Object Explorer session');
-				error(errorMessage);
+	private async handleSessionCreated(session: sqlops.ObjectExplorerSession): Promise<void> {
+		try {
+			let connection: ConnectionProfile = undefined;
+			let errorMessage: string = undefined;
+			if (this._sessions[session.sessionId]) {
+				connection = this._sessions[session.sessionId].connection;
+
+				if (session && session.success && session.rootNode) {
+					let server = this.toTreeNode(session.rootNode, null);
+					server.connection = connection;
+					server.session = session;
+					this._activeObjectExplorerNodes[connection.id] = server;
+				}
+				else {
+					errorMessage = session && session.errorMessage ? session.errorMessage :
+						nls.localize('OeSessionFailedError', 'Failed to create Object Explorer session');
+					error(errorMessage);
+				}
+				// Send on session created about the session to all node providers so they can prepare for node expansion
+				let nodeProviders = this._nodeProviders[connection.providerName];
+				if (nodeProviders) {
+					let promises: Thenable<boolean>[] = nodeProviders.map(p => p.handleSessionOpen(session));
+					await Promise.all(promises);
+				}
+			}
+			else {
+				warn(`cannot find session ${session.sessionId}`);
 			}
 
-		} else {
-			warn(`cannot find session ${session.sessionId}`);
+			this.sendUpdateNodeEvent(connection, errorMessage);
+		} catch (error) {
+			warn(`cannot handle the session ${session.sessionId} in all nodeProviders`);
 		}
-
-		this.sendUpdateNodeEvent(connection, errorMessage);
 	}
 
 	/**
@@ -291,7 +326,7 @@ export class ObjectExplorerService implements IObjectExplorerService {
 			let provider = this._providers[providerId];
 			if (provider) {
 				TelemetryUtils.addTelemetry(this._telemetryService, TelemetryKeys.ObjectExplorerExpand, { refresh: 0, provider: providerId });
-				this.expandOrRefreshNode(provider, session, nodePath).then(result => {
+				this.expandOrRefreshNode(providerId, session, nodePath).then(result => {
 					resolve(result);
 				}, error => {
 					reject(error);
@@ -301,7 +336,8 @@ export class ObjectExplorerService implements IObjectExplorerService {
 			}
 		});
 	}
-	private callExpandOrRefreshFromProvider(provider: sqlops.ObjectExplorerProvider, nodeInfo: sqlops.ExpandNodeInfo, refresh: boolean = false) {
+
+	private callExpandOrRefreshFromProvider(provider: sqlops.ObjectExplorerProviderBase, nodeInfo: sqlops.ExpandNodeInfo, refresh: boolean = false) {
 		if (refresh) {
 			return provider.refreshNode(nodeInfo);
 		} else {
@@ -310,7 +346,7 @@ export class ObjectExplorerService implements IObjectExplorerService {
 	}
 
 	private expandOrRefreshNode(
-		provider: sqlops.ObjectExplorerProvider,
+		providerId: string,
 		session: sqlops.ObjectExplorerSession,
 		nodePath: string,
 		refresh: boolean = false): Thenable<sqlops.ObjectExplorerExpandInfo> {
@@ -320,29 +356,60 @@ export class ObjectExplorerService implements IObjectExplorerService {
 				let newRequest = false;
 				if (!self._sessions[session.sessionId].nodes[nodePath]) {
 					self._sessions[session.sessionId].nodes[nodePath] = {
-						expandEmitter: new Emitter<sqlops.ObjectExplorerExpandInfo>()
+						expandEmitter: new Emitter<NodeExpandInfoWithProviderId>()
 					};
 					newRequest = true;
 				}
-				self._sessions[session.sessionId].nodes[nodePath].expandEmitter.event(((expandResult) => {
-					if (expandResult && !expandResult.errorMessage) {
-						resolve(expandResult);
+				let provider = this._providers[providerId];
+				if (provider) {
+					let resultMap: Map<string, sqlops.ObjectExplorerExpandInfo> = new Map<string, sqlops.ObjectExplorerExpandInfo>();
+					let allProviders: sqlops.ObjectExplorerProviderBase[] = [provider];
+
+					let nodeProviders = this._nodeProviders[providerId];
+					if (nodeProviders) {
+						nodeProviders = nodeProviders.sort((a, b) => a.group.toLowerCase().localeCompare(b.group.toLowerCase()));
+						allProviders.push(...nodeProviders);
 					}
-					else {
-						reject(expandResult ? expandResult.errorMessage : undefined);
-					}
-					if (newRequest) {
-						delete self._sessions[session.sessionId].nodes[nodePath];
-					}
-				}));
-				if (newRequest) {
-					self.callExpandOrRefreshFromProvider(provider, {
-						sessionId: session.sessionId,
-						nodePath: nodePath
-					}, refresh).then(result => {
-					}, error => {
-						reject(error);
+
+					self._sessions[session.sessionId].nodes[nodePath].expandEmitter.event((expandResult) => {
+						if (expandResult && expandResult.providerId) {
+							resultMap.set(expandResult.providerId, expandResult);
+						} else {
+							console.log('OE provider returns empty result or providerId');
+						}
+
+						// When get all responses from all providers, merge results
+						if (resultMap.size === allProviders.length) {
+							resolve(self.mergeResults(allProviders, resultMap, nodePath));
+
+							// Have to delete it after get all reponses otherwise couldn't find session for not the first response
+							if (newRequest) {
+								delete self._sessions[session.sessionId].nodes[nodePath];
+							}
+						}
 					});
+					if (newRequest) {
+						allProviders.forEach(provider => {
+							TelemetryUtils.addTelemetry(this._telemetryService, TelemetryKeys.ObjectExplorerExpand, { refresh: 0, provider: providerId });
+							self.callExpandOrRefreshFromProvider(provider, {
+								sessionId: session.sessionId,
+								nodePath: nodePath
+							}, refresh).then(isExpanding => {
+								if (!isExpanding) {
+									// The provider stated it's not going to expand the node, therefore do not need to track when merging results
+									let emptyResult: sqlops.ObjectExplorerExpandInfo = {
+										errorMessage: undefined,
+										nodePath: nodePath,
+										nodes: [],
+										sessionId: session.sessionId
+									};
+									resultMap.set(provider.providerId, emptyResult);
+								}
+							}, error => {
+								reject(error);
+							});
+						});
+					}
 				}
 			} else {
 				reject(`session cannot find to expand node. id: ${session.sessionId} nodePath: ${nodePath}`);
@@ -350,11 +417,48 @@ export class ObjectExplorerService implements IObjectExplorerService {
 		});
 	}
 
+	private mergeResults(allProviders: sqlops.ObjectExplorerProviderBase[], resultMap: Map<string, sqlops.ObjectExplorerExpandInfo>, nodePath: string): sqlops.ObjectExplorerExpandInfo {
+		let finalResult: sqlops.ObjectExplorerExpandInfo;
+		let allNodes: sqlops.NodeInfo[] = [];
+		let errorNode: sqlops.NodeInfo = {
+			nodePath: nodePath,
+			label: 'Error',
+			errorMessage: '',
+			nodeType: 'folder',
+			isLeaf: true,
+			nodeSubType: '',
+			nodeStatus: '',
+			metadata: null
+		};
+
+		for (let provider of allProviders) {
+			if (resultMap.has(provider.providerId)) {
+				let result = resultMap.get(provider.providerId);
+				if (result) {
+					if (!result.errorMessage) {
+						finalResult = result;
+						allNodes = allNodes.concat(result.nodes);
+					} else {
+						errorNode.errorMessage += provider.providerId + 'returns ' + result.errorMessage + ' ';
+					}
+				}
+			}
+		}
+		if (finalResult) {
+			if (errorNode.errorMessage && errorNode.errorMessage.length > 0) {
+				allNodes = allNodes.concat([errorNode]);
+			}
+
+			finalResult.nodes = allNodes;
+		}
+		return finalResult;
+	}
+
 	public refreshNode(providerId: string, session: sqlops.ObjectExplorerSession, nodePath: string): Thenable<sqlops.ObjectExplorerExpandInfo> {
 		let provider = this._providers[providerId];
 		if (provider) {
 			TelemetryUtils.addTelemetry(this._telemetryService, TelemetryKeys.ObjectExplorerExpand, { refresh: 1, provider: providerId });
-			return this.expandOrRefreshNode(provider, session, nodePath, true);
+			return this.expandOrRefreshNode(providerId, session, nodePath, true);
 		}
 		return Promise.resolve(undefined);
 	}
@@ -369,7 +473,8 @@ export class ObjectExplorerService implements IObjectExplorerService {
 						sessionId: session.sessionId,
 						nodes: [],
 						nodePath: nodePath,
-						errorMessage: undefined
+						errorMessage: undefined,
+						providerId: providerId
 					});
 				}
 			});
@@ -377,6 +482,14 @@ export class ObjectExplorerService implements IObjectExplorerService {
 
 		let provider = this._providers[providerId];
 		if (provider) {
+			let nodeProviders = this._nodeProviders[providerId];
+			if (nodeProviders) {
+				for (let nodeProvider of nodeProviders) {
+					nodeProvider.handleSessionClose({
+						sessionId: session ? session.sessionId : undefined
+					});
+				}
+			}
 			return provider.closeSession({
 				sessionId: session ? session.sessionId : undefined
 			});
@@ -392,6 +505,12 @@ export class ObjectExplorerService implements IObjectExplorerService {
 		this._providers[providerId] = provider;
 	}
 
+	public registerNodeProvider(nodeProvider: sqlops.ObjectExplorerNodeProvider): void {
+		let nodeProviders = this._nodeProviders[nodeProvider.supportedProviderId] || [];
+		nodeProviders.push(nodeProvider);
+		this._nodeProviders[nodeProvider.supportedProviderId] = nodeProviders;
+	}
+
 	public dispose(): void {
 		this._disposables = dispose(this._disposables);
 	}
@@ -528,7 +647,7 @@ export class ObjectExplorerService implements IObjectExplorerService {
 	}
 
 	/**
-	 * For Testing purpose only. Get the context menu actions for an object explorer node
+	* For Testing purpose only. Get the context menu actions for an object explorer node
 	*/
 	public getNodeActions(connectionId: string, nodePath: string): Thenable<string[]> {
 		return this.getTreeNode(connectionId, nodePath).then(node => {
@@ -552,6 +671,10 @@ export class ObjectExplorerService implements IObjectExplorerService {
 		return treeNode;
 	}
 
+	public getSessionConnectionProfile(sessionId: string): sqlops.IConnectionProfile {
+		return this._sessions[sessionId].connection.toIConnectionProfile();
+	}
+
 	private async setNodeExpandedState(treeNode: TreeNode, expandedState: TreeItemCollapsibleState): Promise<void> {
 		treeNode = await this.getUpdatedTreeNode(treeNode);
 		let expandNode = this.getTreeItem(treeNode);

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

@@ -256,6 +256,13 @@ export interface IConnectionManagementService {
 	 */
 	getActiveConnectionCredentials(profileId: string): { [name: string]: string };
 
+	/**
+	 * Get the ServerInfo for a connected connection profile
+	 * @param {string} profileId The id of the connection profile to get the password for
+	 * @returns ServerInfo
+	 */
+	getServerInfo(profileId: string): sqlops.ServerInfo;
+
 	/**
 	 * Get the connection string for the provided connection ID
 	 */

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

@@ -1367,6 +1367,17 @@ export class ConnectionManagementService extends Disposable implements IConnecti
 		return credentials;
 	}
 
+	public getServerInfo(profileId: string): sqlops.ServerInfo {
+		let profile = this._connectionStatusManager.findConnectionByProfileId(profileId);
+		if (!profile) {
+			return undefined;
+		}
+
+		let serverInfo = profile.serverInfo;
+
+		return serverInfo;
+	}
+
 	/**
 	 * Get the connection string for the provided connection ID
 	 */

src/sql/sqlops.d.ts

@@ -21,6 +21,8 @@ declare module 'sqlops' {
 
 		export function registerObjectExplorerProvider(provider: ObjectExplorerProvider): vscode.Disposable;
 
+		export function registerObjectExplorerNodeProvider(provider: ObjectExplorerNodeProvider): vscode.Disposable;
+
 		export function registerTaskServicesProvider(provider: TaskServicesProvider): vscode.Disposable;
 
 		export function registerFileBrowserProvider(provider: FileBrowserProvider): vscode.Disposable;
@@ -101,6 +103,13 @@ declare module 'sqlops' {
 		 */
 		export function getCredentials(connectionId: string): Thenable<{ [name: string]: string }>;
 
+		/**
+		 * Get ServerInfo for a connectionId
+		 * @param {string} connectionId The id of the connection
+		 * @returns ServerInfo
+		 */
+		export function getServerInfo(connectionId: string): Thenable<ServerInfo>;
+
 		/**
 		 * Interface for representing a connection when working with connection APIs
 		*/
@@ -150,6 +159,13 @@ declare module 'sqlops' {
 		 */
 		export function findNodes(connectionId: string, type: string, schema: string, name: string, database: string, parentObjectNames: string[]): Thenable<ObjectExplorerNode[]>;
 
+		/**
+		 * Get connectionProfile from sessionId
+		 * *@param {string} sessionId The id of the session that the node exists on
+		 * @returns {IConnectionProfile} The IConnecitonProfile for the session
+		 */
+		export function getSessionConnectionProfile(sessionId: string): Thenable<IConnectionProfile>;
+
 		/**
 		 * Interface for representing and interacting with items in Object Explorer
 		*/
@@ -347,6 +363,10 @@ declare module 'sqlops' {
 		 * The Operating System version string of the machine running the instance.
 		 */
 		osVersion: string;
+		/**
+		 * options for all new server properties.
+		 */
+		options: {};
 	}
 
 	export interface DataProvider {
@@ -1163,22 +1183,41 @@ declare module 'sqlops' {
 		nodes: NodeInfo[];
 	}
 
-	export interface ObjectExplorerProvider extends DataProvider {
-		createNewSession(connInfo: ConnectionInfo): Thenable<ObjectExplorerSessionResponse>;
-
+	export interface ObjectExplorerProviderBase extends DataProvider {
 		expandNode(nodeInfo: ExpandNodeInfo): Thenable<boolean>;
 
 		refreshNode(nodeInfo: ExpandNodeInfo): Thenable<boolean>;
 
-		closeSession(closeSessionInfo: ObjectExplorerCloseSessionInfo): Thenable<ObjectExplorerCloseSessionResponse>;
-
 		findNodes(findNodesInfo: FindNodesInfo): Thenable<ObjectExplorerFindNodesResponse>;
 
+		registerOnExpandCompleted(handler: (response: ObjectExplorerExpandInfo) => any): void;
+	}
+
+	export interface ObjectExplorerProvider extends ObjectExplorerProviderBase {
+		createNewSession(connInfo: ConnectionInfo): Thenable<ObjectExplorerSessionResponse>;
+
+		closeSession(closeSessionInfo: ObjectExplorerCloseSessionInfo): Thenable<ObjectExplorerCloseSessionResponse>;
+
 		registerOnSessionCreated(handler: (response: ObjectExplorerSession) => any): void;
 
 		registerOnSessionDisconnected?(handler: (response: ObjectExplorerSession) => any): void;
+	}
 
-		registerOnExpandCompleted(handler: (response: ObjectExplorerExpandInfo) => any): void;
+	export interface ObjectExplorerNodeProvider extends ObjectExplorerProviderBase {
+		/**
+		 * The providerId for whichever type of ObjectExplorer connection this can add folders and objects to
+		 */
+		readonly supportedProviderId: string;
+
+		/**
+		 * Optional group name used to sort nodes in the tree. If not defined, the node order will be added in order based on provider ID, with
+		 * nodes from the main ObjectExplorerProvider for this provider type added first
+		 */
+		readonly group?: string;
+
+		handleSessionOpen(session: ObjectExplorerSession): Thenable<boolean>;
+
+		handleSessionClose(closeSessionInfo: ObjectExplorerCloseSessionInfo): void;
 	}
 
 	// Admin Services interfaces  -----------------------------------------------------------------------

src/sql/sqlops.proposed.d.ts

@@ -1240,6 +1240,7 @@ declare module 'sqlops' {
 		AgentServicesProvider = 'AgentServicesProvider',
 		CapabilitiesProvider = 'CapabilitiesProvider',
 		DacFxServicesProvider = 'DacFxServicesProvider',
+		ObjectExplorerNodeProvider = 'ObjectExplorerNodeProvider',
 	}
 
 	export namespace dataprotocol {

src/sql/workbench/api/common/sqlExtHostTypes.ts

@@ -288,6 +288,7 @@ export enum DataProviderType {
 	AgentServicesProvider = 'AgentServicesProvider',
 	CapabilitiesProvider = 'CapabilitiesProvider',
 	DacFxServicesProvider = 'DacFxServicesProvider',
+	ObjectExplorerNodeProvider = 'ObjectExplorerNodeProvider',
 }
 
 export enum DeclarativeDataType {
@@ -320,6 +321,17 @@ export enum AzureResource {
 	Sql = 1
 }
 
+export interface ServerInfoOption {
+	isBigDataCluster: boolean;
+	clusterEndpoints: ClusterEndpoint;
+}
+
+export interface ClusterEndpoint {
+	serviceName: string;
+	ipAddress: string;
+	port: number;
+}
+
 export class SqlThemeIcon {
 
 	static readonly Folder = new SqlThemeIcon('Folder');

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

@@ -32,6 +32,10 @@ export class ExtHostConnectionManagement extends ExtHostConnectionManagementShap
 		return this._proxy.$getCredentials(connectionId);
 	}
 
+	public $getServerInfo(connectionId: string): Thenable<sqlops.ServerInfo> {
+		return this._proxy.$getServerInfo(connectionId);
+	}
+
 	public $openConnectionDialog(providers?: string[], initialConnectionProfile?: sqlops.IConnectionProfile, connectionCompletionOptions?: sqlops.IConnectionCompletionOptions): Thenable<sqlops.connection.Connection> {
 		return this._proxy.$openConnectionDialog(providers, initialConnectionProfile, connectionCompletionOptions);
 	}

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

@@ -132,6 +132,12 @@ export class ExtHostDataProtocol extends ExtHostDataProtocolShape {
 		return rt;
 	}
 
+	$registerObjectExplorerNodeProvider(provider: sqlops.ObjectExplorerNodeProvider): vscode.Disposable {
+		let rt = this.registerProvider(provider, DataProviderType.ObjectExplorerNodeProvider);
+		this._proxy.$registerObjectExplorerNodeProvider(provider.providerId, provider.supportedProviderId, provider.group, provider.handle);
+		return rt;
+	}
+
 	$registerProfilerProvider(provider: sqlops.ProfilerProvider): vscode.Disposable {
 		let rt = this.registerProvider(provider, DataProviderType.ProfilerProvider);
 		this._proxy.$registerProfilerProvider(provider.providerId, provider.handle);
@@ -342,20 +348,28 @@ export class ExtHostDataProtocol extends ExtHostDataProtocolShape {
 		return this._resolveProvider<sqlops.ObjectExplorerProvider>(handle).createNewSession(connInfo);
 	}
 
+	public $createObjectExplorerNodeProviderSession(handle: number, session: sqlops.ObjectExplorerSession): Thenable<boolean> {
+		return this._resolveProvider<sqlops.ObjectExplorerNodeProvider>(handle).handleSessionOpen(session);
+	}
+
 	public $expandObjectExplorerNode(handle: number, nodeInfo: sqlops.ExpandNodeInfo): Thenable<boolean> {
-		return this._resolveProvider<sqlops.ObjectExplorerProvider>(handle).expandNode(nodeInfo);
+		return this._resolveProvider<sqlops.ObjectExplorerProviderBase> (handle).expandNode(nodeInfo);
 	}
 
 	public $refreshObjectExplorerNode(handle: number, nodeInfo: sqlops.ExpandNodeInfo): Thenable<boolean> {
-		return this._resolveProvider<sqlops.ObjectExplorerProvider>(handle).refreshNode(nodeInfo);
+		return this._resolveProvider<sqlops.ObjectExplorerProviderBase> (handle).refreshNode(nodeInfo);
 	}
 
 	public $closeObjectExplorerSession(handle: number, closeSessionInfo: sqlops.ObjectExplorerCloseSessionInfo): Thenable<sqlops.ObjectExplorerCloseSessionResponse> {
 		return this._resolveProvider<sqlops.ObjectExplorerProvider>(handle).closeSession(closeSessionInfo);
 	}
 
+	public $handleSessionClose(handle: number, closeSessionInfo: sqlops.ObjectExplorerCloseSessionInfo): void {
+		return this._resolveProvider<sqlops.ObjectExplorerNodeProvider>(handle).handleSessionClose(closeSessionInfo);
+	}
+
 	public $findNodes(handle: number, findNodesInfo: sqlops.FindNodesInfo): Thenable<sqlops.ObjectExplorerFindNodesResponse> {
-		return this._resolveProvider<sqlops.ObjectExplorerProvider>(handle).findNodes(findNodesInfo);
+		return this._resolveProvider<sqlops.ObjectExplorerProviderBase>(handle).findNodes(findNodesInfo);
 	}
 
 	public $onObjectExplorerSessionCreated(handle: number, response: sqlops.ObjectExplorerSession): void {
@@ -366,8 +380,8 @@ export class ExtHostDataProtocol extends ExtHostDataProtocolShape {
 		this._proxy.$onObjectExplorerSessionDisconnected(handle, response);
 	}
 
-	public $onObjectExplorerNodeExpanded(handle: number, response: sqlops.ObjectExplorerExpandInfo): void {
-		this._proxy.$onObjectExplorerNodeExpanded(handle, response);
+	public $onObjectExplorerNodeExpanded(providerId: string, response: sqlops.ObjectExplorerExpandInfo): void {
+		this._proxy.$onObjectExplorerNodeExpanded(providerId, response);
 	}
 
 	// Task Service

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

@@ -34,6 +34,10 @@ export class ExtHostObjectExplorer implements ExtHostObjectExplorerShape {
 	public $getNodeActions(connectionId: string, nodePath: string): Thenable<string[]> {
 		return this._proxy.$getNodeActions(connectionId, nodePath);
 	}
+
+	public $getSessionConnectionProfile(sessionId: string): Thenable<sqlops.IConnectionProfile> {
+		return this._proxy.$getSessionConnectionProfile(sessionId);
+	}
 }
 
 class ExtHostObjectExplorerNode implements sqlops.objectexplorer.ObjectExplorerNode {

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

@@ -55,6 +55,9 @@ export class MainThreadConnectionManagement implements MainThreadConnectionManag
 		return Promise.resolve(this._connectionManagementService.getActiveConnectionCredentials(connectionId));
 	}
 
+	public $getServerInfo(connectionId: string): Thenable<sqlops.ServerInfo> {
+		return Promise.resolve(this._connectionManagementService.getServerInfo(connectionId));
+	}
 
 	public async $openConnectionDialog(providers: string[], initialConnectionProfile?: IConnectionProfile, connectionCompletionOptions?: sqlops.IConnectionCompletionOptions): Promise<sqlops.connection.Connection> {
 		let connectionProfile = await this._connectionDialogService.openDialogAndWait(this._connectionManagementService, { connectionType: 1, providers: providers }, initialConnectionProfile);

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

@@ -15,7 +15,7 @@ import { ICapabilitiesService } from 'sql/platform/capabilities/common/capabilit
 import { IQueryManagementService } from 'sql/platform/query/common/queryManagement';
 import * as sqlops from 'sqlops';
 import { IMetadataService } from 'sql/platform/metadata/common/metadataService';
-import { IObjectExplorerService } from 'sql/parts/objectExplorer/common/objectExplorerService';
+import { IObjectExplorerService, NodeExpandInfoWithProviderId } from 'sql/parts/objectExplorer/common/objectExplorerService';
 import { IScriptingService } from 'sql/platform/scripting/common/scriptingService';
 import { IAdminService } from 'sql/workbench/services/admin/common/adminService';
 import { IJobManagementService } from 'sql/platform/jobManagement/common/interfaces';
@@ -231,6 +231,7 @@ export class MainThreadDataProtocol implements MainThreadDataProtocolShape {
 	public $registerObjectExplorerProvider(providerId: string, handle: number): TPromise<any> {
 		const self = this;
 		this._objectExplorerService.registerProvider(providerId, <sqlops.ObjectExplorerProvider>{
+			providerId: providerId,
 			createNewSession(connection: sqlops.ConnectionInfo): Thenable<sqlops.ObjectExplorerSessionResponse> {
 				return self._proxy.$createObjectExplorerSession(handle, connection);
 			},
@@ -251,6 +252,32 @@ export class MainThreadDataProtocol implements MainThreadDataProtocolShape {
 		return undefined;
 	}
 
+	public $registerObjectExplorerNodeProvider(providerId: string, supportedProviderId: string, group: string, handle: number): TPromise<any> {
+		const self = this;
+		this._objectExplorerService.registerNodeProvider(<sqlops.ObjectExplorerNodeProvider> {
+			supportedProviderId: supportedProviderId,
+			providerId: providerId,
+			group: group,
+			expandNode(nodeInfo: sqlops.ExpandNodeInfo): Thenable<boolean> {
+				return self._proxy.$expandObjectExplorerNode(handle, nodeInfo);
+			},
+			refreshNode(nodeInfo: sqlops.ExpandNodeInfo): Thenable<boolean> {
+				return self._proxy.$refreshObjectExplorerNode(handle, nodeInfo);
+			},
+			findNodes(findNodesInfo: sqlops.FindNodesInfo): Thenable<sqlops.ObjectExplorerFindNodesResponse> {
+				return self._proxy.$findNodes(handle, findNodesInfo);
+			},
+			handleSessionOpen(session: sqlops.ObjectExplorerSession): Thenable<boolean> {
+				return self._proxy.$createObjectExplorerNodeProviderSession(handle, session);
+			},
+			handleSessionClose(closeSessionInfo: sqlops.ObjectExplorerCloseSessionInfo): void {
+				return self._proxy.$handleSessionClose(handle, closeSessionInfo);
+			}
+		});
+
+		return undefined;
+	}
+
 	public $registerTaskServicesProvider(providerId: string, handle: number): TPromise<any> {
 		const self = this;
 		this._taskService.registerProvider(providerId, <sqlops.TaskServicesProvider>{
@@ -474,8 +501,9 @@ export class MainThreadDataProtocol implements MainThreadDataProtocolShape {
 		this._objectExplorerService.onSessionDisconnected(handle, sessionResponse);
 	}
 
-	public $onObjectExplorerNodeExpanded(handle: number, expandResponse: sqlops.ObjectExplorerExpandInfo): void {
-		this._objectExplorerService.onNodeExpanded(handle, expandResponse);
+	public $onObjectExplorerNodeExpanded(providerId: string, expandResponse: sqlops.ObjectExplorerExpandInfo): void {
+		let expandInfo: NodeExpandInfoWithProviderId = Object.assign({ providerId: providerId }, expandResponse);
+		this._objectExplorerService.onNodeExpanded(expandInfo);
 	}
 
 	//Tasks handlers

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

@@ -82,4 +82,8 @@ export class MainThreadObjectExplorer implements MainThreadObjectExplorerShape {
 	public $getNodeActions(connectionId: string, nodePath: string): Thenable<string[]> {
 		return this._objectExplorerService.getNodeActions(connectionId, nodePath);
 	}
+	
+	public $getSessionConnectionProfile(sessionId: string): Thenable<sqlops.IConnectionProfile> {
+		return Promise.resolve(this._objectExplorerService.getSessionConnectionProfile(sessionId));
+	}
 }

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

@@ -120,6 +120,9 @@ export function createApiFactory(
 				getCredentials(connectionId: string): Thenable<{ [name: string]: string }> {
 					return extHostConnectionManagement.$getCredentials(connectionId);
 				},
+				getServerInfo(connectionId: string): Thenable<sqlops.ServerInfo> {
+					return extHostConnectionManagement.$getServerInfo(connectionId);
+				},
 				openConnectionDialog(providers?: string[], initialConnectionProfile?: sqlops.IConnectionProfile, connectionCompletionOptions?: sqlops.IConnectionCompletionOptions): Thenable<sqlops.connection.Connection> {
 					return extHostConnectionManagement.$openConnectionDialog(providers, initialConnectionProfile, connectionCompletionOptions);
 				},
@@ -160,6 +163,9 @@ export function createApiFactory(
 				},
 				getNodeActions(connectionId: string, nodePath: string): Thenable<string[]> {
 					return extHostObjectExplorer.$getNodeActions(connectionId, nodePath);
+				},
+				getSessionConnectionProfile(sessionId: string): Thenable<sqlops.IConnectionProfile> {
+					return extHostObjectExplorer.$getSessionConnectionProfile(sessionId);
 				}
 			};
 
@@ -238,12 +244,20 @@ export function createApiFactory(
 				}
 
 				provider.registerOnExpandCompleted((response: sqlops.ObjectExplorerExpandInfo) => {
-					extHostDataProvider.$onObjectExplorerNodeExpanded(provider.handle, response);
+					extHostDataProvider.$onObjectExplorerNodeExpanded(provider.providerId, response);
 				});
 
 				return extHostDataProvider.$registerObjectExplorerProvider(provider);
 			};
 
+			let registerObjectExplorerNodeProvider = (provider: sqlops.ObjectExplorerNodeProvider): vscode.Disposable => {
+				provider.registerOnExpandCompleted((response: sqlops.ObjectExplorerExpandInfo) => {
+					extHostDataProvider.$onObjectExplorerNodeExpanded(provider.providerId, response);
+				});
+
+				return extHostDataProvider.$registerObjectExplorerNodeProvider(provider);
+			};
+
 			let registerTaskServicesProvider = (provider: sqlops.TaskServicesProvider): vscode.Disposable => {
 				provider.registerOnTaskCreated((response: sqlops.TaskInfo) => {
 					extHostDataProvider.$onTaskCreated(provider.handle, response);
@@ -335,6 +349,7 @@ export function createApiFactory(
 				registerFileBrowserProvider,
 				registerMetadataProvider,
 				registerObjectExplorerProvider,
+				registerObjectExplorerNodeProvider,
 				registerProfilerProvider,
 				registerRestoreProvider,
 				registerScriptingProvider,

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

@@ -128,6 +128,10 @@ export abstract class ExtHostDataProtocolShape {
 
 	$findNodes(handle: number, findNodesInfo: sqlops.FindNodesInfo): Thenable<sqlops.ObjectExplorerFindNodesResponse> { throw ni(); }
 
+	$createObjectExplorerNodeProviderSession(handle: number, sessionInfo: sqlops.ObjectExplorerSession): Thenable<boolean> { throw ni(); }
+
+	$handleSessionClose(handle: number, closeSessionInfo: sqlops.ObjectExplorerCloseSessionInfo): void { throw ni(); }
+
 	/**
 	 * Tasks
 	 */
@@ -507,6 +511,7 @@ export interface MainThreadDataProtocolShape extends IDisposable {
 	$registerQueryProvider(providerId: string, handle: number): TPromise<any>;
 	$registerProfilerProvider(providerId: string, handle: number): TPromise<any>;
 	$registerObjectExplorerProvider(providerId: string, handle: number): TPromise<any>;
+	$registerObjectExplorerNodeProvider(providerId: string, supportedProviderId: string, group: string, handle: number): TPromise<any>;
 	$registerMetadataProvider(providerId: string, handle: number): TPromise<any>;
 	$registerTaskServicesProvider(providerId: string, handle: number): TPromise<any>;
 	$registerFileBrowserProvider(providerId: string, handle: number): TPromise<any>;
@@ -526,7 +531,7 @@ export interface MainThreadDataProtocolShape extends IDisposable {
 	$onQueryMessage(handle: number, message: sqlops.QueryExecuteMessageParams): void;
 	$onObjectExplorerSessionCreated(handle: number, message: sqlops.ObjectExplorerSession): void;
 	$onObjectExplorerSessionDisconnected(handle: number, message: sqlops.ObjectExplorerSession): void;
-	$onObjectExplorerNodeExpanded(handle: number, message: sqlops.ObjectExplorerExpandInfo): void;
+	$onObjectExplorerNodeExpanded(providerId: string, message: sqlops.ObjectExplorerExpandInfo): void;
 	$onTaskCreated(handle: number, sessionResponse: sqlops.TaskInfo): void;
 	$onTaskStatusChanged(handle: number, sessionResponse: sqlops.TaskProgressInfo): void;
 	$onFileBrowserOpened(handle: number, response: sqlops.FileBrowserOpenedParams): void;
@@ -548,6 +553,7 @@ export interface MainThreadConnectionManagementShape extends IDisposable {
 	$getActiveConnections(): Thenable<sqlops.connection.Connection[]>;
 	$getCurrentConnection(): Thenable<sqlops.connection.Connection>;
 	$getCredentials(connectionId: string): Thenable<{ [name: string]: string }>;
+	$getServerInfo(connectedId: string): Thenable<sqlops.ServerInfo>;
 	$openConnectionDialog(providers: string[], initialConnectionProfile?: sqlops.IConnectionProfile, connectionCompletionOptions?: sqlops.IConnectionCompletionOptions): Thenable<sqlops.connection.Connection>;
 	$listDatabases(connectionId: string): Thenable<string[]>;
 	$getConnectionString(connectionId: string, includePassword: boolean): Thenable<string>;
@@ -717,6 +723,7 @@ export interface MainThreadObjectExplorerShape extends IDisposable {
 	$findNodes(connectionId: string, type: string, schema: string, name: string, database: string, parentObjectNames: string[]): Thenable<sqlops.NodeInfo[]>;
 	$refresh(connectionId: string, nodePath: string): Thenable<sqlops.NodeInfo>;
 	$getNodeActions(connectionId: string, nodePath: string): Thenable<string[]>;
+	$getSessionConnectionProfile(sessionId: string): Thenable<sqlops.IConnectionProfile>;
 }
 
 export interface ExtHostModelViewDialogShape {