Spark features with dashboard are enabled (#3883)

* Spark features are enabled

* Fixed as PR comments

* minor change

* PR comments fixed

* minor fix

* change constant name to avoid conflicts with sqlopsextension

* sqlContext to context

* Changed tab name to SQL Server Big Data Cluster

* Added isCluster to ContextProvider to control display big data cluster dashboard tab
Ported New/open Notebook code to mssql extension and enable them in dashboard

* Fixed tslint

extensions/mssql/src/objectExplorerNodeProvider/connection.ts

@@ -6,214 +6,57 @@
 'use strict';
 
 import * as sqlops from 'sqlops';
-import * as UUID from 'vscode-languageclient/lib/utils/uuid';
 import * as nls from 'vscode-nls';
 const localize = nls.loadMessageBundle();
 
 import * as constants from '../constants';
-import * as LocalizedConstants from '../localizedConstants';
-import * as utils from '../utils';
 import { IFileSource, IHdfsOptions, IRequestParams, FileSourceFactory } from './fileSources';
-import { IEndpoint } from './objectExplorerNodeProvider';
 
-function appendIfExists(uri: string, propName: string, propValue: string): string {
-	if (propValue) {
-		uri = `${uri};${propName}=${propValue}`;
-	}
-	return uri;
-}
-
-interface IValidationResult {
-	isValid: boolean;
-	errors: string;
-}
-
-export class Connection {
+export class SqlClusterConnection {
+	private _connection: sqlops.connection.Connection;
+	private _profile: sqlops.IConnectionProfile;
 	private _host: string;
-	private _knoxPort: string;
+	private _port: string;
+	private _user: string;
+	private _password: string;
 
-	constructor(private connectionInfo: sqlops.ConnectionInfo, private connectionUri?: string, private _connectionId?: string) {
-		if (!this.connectionInfo) {
-			throw new Error(localize('connectionInfoMissing', 'connectionInfo is required'));
-		}
-
-		if (!this._connectionId) {
-			this._connectionId = UUID.generateUuid();
+	constructor(connectionInfo: sqlops.connection.Connection | sqlops.IConnectionProfile) {
+		this.validate(connectionInfo);
+
+		if ('id' in connectionInfo) {
+			this._profile = connectionInfo;
+			this._connection = this.toConnection(this._profile);
+		} else {
+			this._connection = connectionInfo;
+			this._profile = this.toConnectionProfile(this._connection);
 		}
+		this._host = this._connection.options[constants.hostPropName];
+		this._port = this._connection.options[constants.knoxPortPropName];
+		this._user = this._connection.options[constants.userPropName];
+		this._password = this._connection.options[constants.passwordPropName];
 	}
 
-	public get uri(): string {
-		return this.connectionUri;
+	public get connection(): sqlops.connection.Connection { return this._connection; }
+	public get profile(): sqlops.IConnectionProfile { return this._profile; }
+	public get host(): string { return this._host; }
+	public get port(): string { return this._port || constants.defaultKnoxPort; }
+	public get user(): string { return this._user; }
+	public get password(): string { return this._password; }
+
+	public isMatch(connection: SqlClusterConnection | sqlops.ConnectionInfo): boolean {
+		if (!connection) { return false; }
+		let options1 = connection instanceof SqlClusterConnection ?
+			connection._connection.options : connection.options;
+		let options2 = this._connection.options;
+		return [constants.hostPropName, constants.knoxPortPropName, constants.userPropName]
+			.every(e => options1[e] === options2[e]);
 	}
 
-	public saveUriWithPrefix(prefix: string): string {
-		let uri = `${prefix}${this.host}`;
-		uri = appendIfExists(uri, constants.knoxPortPropName, this.knoxport);
-		uri = appendIfExists(uri, constants.userPropName, this.user);
-		uri = appendIfExists(uri, constants.groupIdPropName, this.connectionInfo.options[constants.groupIdPropName]);
-		this.connectionUri = uri;
-		return this.connectionUri;
-	}
-
-	public async tryConnect(factory?: FileSourceFactory): Promise<sqlops.ConnectionInfoSummary> {
-		let fileSource = this.createHdfsFileSource(factory, {
-			timeout: this.connecttimeout
-		});
-		let summary: sqlops.ConnectionInfoSummary = undefined;
-		try {
-			await fileSource.enumerateFiles(constants.hdfsRootPath);
-			summary = {
-				ownerUri: this.connectionUri,
-				connectionId: this.connectionId,
-				connectionSummary: {
-					serverName: this.host,
-					databaseName: undefined,
-					userName: this.user
-				},
-				errorMessage: undefined,
-				errorNumber: undefined,
-				messages: undefined,
-				serverInfo: this.getEmptyServerInfo()
-			};
-		} catch (error) {
-			summary = {
-				ownerUri: this.connectionUri,
-				connectionId: undefined,
-				connectionSummary: undefined,
-				errorMessage: this.getConnectError(error),
-				errorNumber: undefined,
-				messages: undefined,
-				serverInfo: undefined
-			};
-		}
-		return summary;
-	}
-
-	private getConnectError(error: string | Error): string {
-		let errorMsg = utils.getErrorMessage(error);
-		if (errorMsg.indexOf('ETIMEDOUT') > -1) {
-			errorMsg = LocalizedConstants.msgTimeout;
-		} else if (errorMsg.indexOf('ENOTFOUND') > -1) {
-			errorMsg = LocalizedConstants.msgTimeout;
-		}
-		return localize('connectError', 'Connection failed with error: {0}', errorMsg);
-	}
-
-	private getEmptyServerInfo(): sqlops.ServerInfo {
-		let info: sqlops.ServerInfo = {
-			serverMajorVersion: 0,
-			serverMinorVersion: 0,
-			serverReleaseVersion: 0,
-			engineEditionId: 0,
-			serverVersion: '',
-			serverLevel: '',
-			serverEdition: '',
-			isCloud: false,
-			azureVersion: 0,
-			osVersion: '',
-			options: {}
-		};
-		return info;
-	}
-
-	public get connectionId(): string {
-		return this._connectionId;
-	}
-
-	public get host(): string {
-		if (!this._host) {
-			this.ensureHostAndPort();
-		}
-		return this._host;
-	}
-
-	/**
-	 * Sets host and port values, using any ',' or ':' delimited port in the hostname in
-	 * preference to the built in port.
-	 */
-	private ensureHostAndPort(): void {
-		this._host = this.connectionInfo.options[constants.hostPropName];
-		this._knoxPort = Connection.getKnoxPortOrDefault(this.connectionInfo);
-		// determine whether the host has either a ',' or ':' in it
-		this.setHostAndPort(',');
-		this.setHostAndPort(':');
-	}
-
-	// set port and host correctly after we've identified that a delimiter exists in the host name
-	private setHostAndPort(delimeter: string): void {
-		let originalHost = this._host;
-		let index = originalHost.indexOf(delimeter);
-		if (index > -1) {
-			this._host = originalHost.slice(0, index);
-			this._knoxPort = originalHost.slice(index + 1);
-		}
-	}
-
-	public get user(): string {
-		return this.connectionInfo.options[constants.userPropName];
-	}
-
-	public get password(): string {
-		return this.connectionInfo.options[constants.passwordPropName];
-	}
-
-	public get knoxport(): string {
-		if (!this._knoxPort) {
-			this.ensureHostAndPort();
-		}
-		return this._knoxPort;
-	}
-
-	private static getKnoxPortOrDefault(connInfo: sqlops.ConnectionInfo): string {
-		let port = connInfo.options[constants.knoxPortPropName];
-		if (!port) {
-			port = constants.defaultKnoxPort;
-		}
-		return port;
-	}
-
-	public get connecttimeout(): number {
-		let timeoutSeconds: number = this.connectionInfo.options['connecttimeout'];
-		if (!timeoutSeconds) {
-			timeoutSeconds = constants.hadoopConnectionTimeoutSeconds;
-		}
-		// connect timeout is in milliseconds
-		return timeoutSeconds * 1000;
-	}
-
-	public get sslverification(): string {
-		return this.connectionInfo.options['sslverification'];
-	}
-
-	public get groupId(): string {
-		return this.connectionInfo.options[constants.groupIdName];
-	}
-
-	public async isMatch(connectionInfo: sqlops.ConnectionInfo): Promise<boolean> {
-		if (!connectionInfo) {
-			return false;
-		}
-		let profile = connectionInfo as sqlops.IConnectionProfile;
-		if (profile) {
-			let result: IEndpoint = await utils.getClusterEndpoint(profile.id, constants.hadoopKnoxEndpointName);
-			if (result === undefined || !result.ipAddress || !result.port) {
-				return false;
-			}
-			return connectionInfo.options.groupId === this.groupId
-				&& result.ipAddress === this.host
-				&& String(result.port).startsWith(this.knoxport)
-				&& String(result.port).endsWith(this.knoxport);
-				// TODO: enable the user check when the unified user is used
-				//&& connectionInfo.options.user === this.user;
-		}
-	}
-
-	public createHdfsFileSource(factory?: FileSourceFactory, additionalRequestParams?: IRequestParams): IFileSource {
-		factory = factory || FileSourceFactory.instance;
+	public createHdfsFileSource(): IFileSource {
 		let options: IHdfsOptions = {
 			protocol: 'https',
 			host: this.host,
-			port: this.knoxport,
+			port: this.port,
 			user: this.user,
 			path: 'gateway/default/webhdfs/v1',
 			requestParams: {
@@ -223,9 +66,49 @@ export class Connection {
 				}
 			}
 		};
-		if (additionalRequestParams) {
-			options.requestParams = Object.assign(options.requestParams, additionalRequestParams);
+		return FileSourceFactory.instance.createHdfsFileSource(options);
+	}
+
+	private validate(connectionInfo: sqlops.ConnectionInfo): void {
+		if (!connectionInfo) {
+			throw new Error(localize('connectionInfoUndefined', 'ConnectionInfo is undefined.'));
 		}
-		return factory.createHdfsFileSource(options);
+		if (!connectionInfo.options) {
+			throw new Error(localize('connectionInfoOptionsUndefined', 'ConnectionInfo.options is undefined.'));
+		}
+		let missingProperties: string[] = this.getMissingProperties(connectionInfo);
+		if (missingProperties && missingProperties.length > 0) {
+			throw new Error(localize('connectionInfoOptionsMissingProperties',
+				'Some missing properties in connectionInfo.options: {0}',
+				missingProperties.join(', ')));
+		}
+	}
+
+	private getMissingProperties(connectionInfo: sqlops.ConnectionInfo): string[] {
+		if (!connectionInfo || !connectionInfo.options) { return undefined; }
+		return [
+			constants.hostPropName, constants.knoxPortPropName,
+			constants.userPropName, constants.passwordPropName
+		].filter(e => connectionInfo.options[e] === undefined);
+	}
+
+	private toConnection(connProfile: sqlops.IConnectionProfile): sqlops.connection.Connection {
+		let connection: sqlops.connection.Connection = Object.assign(connProfile,
+			{ connectionId: this._profile.id });
+		return connection;
+	}
+
+	private toConnectionProfile(connectionInfo: sqlops.connection.Connection): sqlops.IConnectionProfile {
+		let options = connectionInfo.options;
+		let connProfile: sqlops.IConnectionProfile = Object.assign(<sqlops.IConnectionProfile>{},
+			connectionInfo,
+			{
+				serverName: `${options[constants.hostPropName]},${options[constants.knoxPortPropName]}`,
+				userName: options[constants.userPropName],
+				password: options[constants.passwordPropName],
+				id: connectionInfo.connectionId,
+			}
+		);
+		return connProfile;
 	}
 }

extensions/mssql/src/objectExplorerNodeProvider/hdfsCommands.ts

@@ -21,7 +21,7 @@ import { IPrompter, IQuestion, QuestionTypes } from '../prompts/question';
 import * as constants from '../constants';
 import * as LocalizedConstants from '../localizedConstants';
 import * as utils from '../utils';
-import { Connection } from './connection';
+import { SqlClusterConnection } from './connection';
 import { AppContext } from '../appContext';
 import { TreeNode } from './treeNodes';
 import { MssqlObjectExplorerNodeProvider } from './objectExplorerNodeProvider';
@@ -45,14 +45,14 @@ function getSaveableUri(apiWrapper: ApiWrapper, fileName: string, isPreview?: bo
 	return vscode.Uri.file(fspath.join(root, fileName));
 }
 
-export async function getNode<T extends TreeNode>(context: ICommandViewContext |ICommandObjectExplorerContext, appContext: AppContext): Promise<T> {
+export async function getNode<T extends TreeNode>(context: ICommandViewContext | ICommandObjectExplorerContext, appContext: AppContext): Promise<T> {
 	let node: T = undefined;
 	if (context && context.type === constants.ViewType && context.node) {
 		node = context.node as T;
 	} else if (context && context.type === constants.ObjectExplorerService) {
-		let oeProvider = appContext.getService<MssqlObjectExplorerNodeProvider>(constants.ObjectExplorerService);
-		if (oeProvider) {
-			node = await oeProvider.findNodeForContext<T>(context.explorerContext);
+		let oeNodeProvider = appContext.getService<MssqlObjectExplorerNodeProvider>(constants.ObjectExplorerService);
+		if (oeNodeProvider) {
+			node = await oeNodeProvider.findSqlClusterNodeByContext<T>(context);
 		}
 	} else {
 		throw new Error(LocalizedConstants.msgMissingNodeContext);
@@ -73,7 +73,7 @@ export class UploadFilesCommand extends ProgressCommand {
 	async execute(context: ICommandViewContext | ICommandObjectExplorerContext, ...args: any[]): Promise<void> {
 		try {
 			let folderNode = await getNode<FolderNode>(context, this.appContext);
-			const allFilesFilter =  localize('allFiles', 'All Files');
+			const allFilesFilter = localize('allFiles', 'All Files');
 			let filter = {};
 			filter[allFilesFilter] = '*';
 			if (folderNode) {
@@ -180,11 +180,11 @@ export class DeleteFilesCommand extends Command {
 		super('mssqlCluster.deleteFiles', appContext);
 	}
 
-	protected async preExecute(context: ICommandViewContext |ICommandObjectExplorerContext, args: object = {}): Promise<any> {
+	protected async preExecute(context: ICommandViewContext | ICommandObjectExplorerContext, args: object = {}): Promise<any> {
 		return this.execute(context, args);
 	}
 
-	async execute(context: ICommandViewContext |ICommandObjectExplorerContext, ...args: any[]): Promise<void> {
+	async execute(context: ICommandViewContext | ICommandObjectExplorerContext, ...args: any[]): Promise<void> {
 		try {
 			let node = await getNode<TreeNode>(context, this.appContext);
 			if (node) {
@@ -282,6 +282,7 @@ export class SaveFileCommand extends ProgressCommand {
 		await this.apiWrapper.executeCommand('vscode.open', fileUri);
 	}
 }
+
 export class PreviewFileCommand extends ProgressCommand {
 	public static readonly DefaultMaxSize = 30 * 1024 * 1024;
 
@@ -334,6 +335,7 @@ export class PreviewFileCommand extends ProgressCommand {
 		}
 	}
 }
+
 export class CopyPathCommand extends Command {
 	public static readonly DefaultMaxSize = 30 * 1024 * 1024;
 
@@ -359,79 +361,3 @@ export class CopyPathCommand extends Command {
 		}
 	}
 }
-
-/**
- * The connect task is only expected to work in the file-tree based APIs, not Object Explorer
- */
-export class ConnectTask {
-	constructor(private hdfsProvider: HdfsProvider, private prompter: IPrompter, private apiWrapper: ApiWrapper) {
-
-	}
-
-	async execute(profile: sqlops.IConnectionProfile, ...args: any[]): Promise<void> {
-		if (profile) {
-			return this.createFromProfile(profile);
-		}
-		return this.createHdfsConnection();
-	}
-
-	private createFromProfile(profile: sqlops.IConnectionProfile): Promise<void> {
-		let connection = new Connection(profile);
-		if (profile.providerName === constants.mssqlClusterProviderName && connection.host) {
-			// TODO need to get the actual port and auth to be used since this will be non-default
-			// in future versions
-			this.hdfsProvider.addHdfsConnection(<IHdfsOptions> {
-				protocol: 'https',
-				host: connection.host,
-				port: connection.knoxport,
-				user: connection.user,
-				path: 'gateway/default/webhdfs/v1',
-				requestParams: {
-					auth: {
-						user: connection.user,
-						pass: connection.password
-					}
-				}
-			});
-		}
-		return Promise.resolve(undefined);
-	}
-
-	private addConnection(options: IHdfsOptions): void {
-		let display: string = `${options.user}@${options.host}:${options.port}`;
-		this.hdfsProvider.addConnection(display, FileSourceFactory.instance.createHdfsFileSource(options));
-	}
-
-	private async createHdfsConnection(profile?: sqlops.IConnectionProfile): Promise<void> {
-		let questions: IQuestion[] = [
-			{
-				type: QuestionTypes.input,
-				name: constants.hdfsHost,
-				message: localize('msgSetWebHdfsHost', 'HDFS URL and port'),
-				default: 'localhost:50070'
-			},
-			{
-				type: QuestionTypes.input,
-				name: constants.hdfsUser,
-				message: localize('msgSetWebHdfsUser', 'User Name'),
-				default: 'root'
-			}];
-
-		let answers = await this.prompter.prompt(questions);
-		if (answers) {
-			let hostAndPort: string = answers[constants.hdfsHost];
-			let parts = hostAndPort.split(':');
-			let host: string = parts[0];
-			let port: string = parts.length > 1 ? parts[1] : undefined;
-			let user: string = answers[constants.hdfsUser];
-
-
-			let options: IHdfsOptions = {
-				host: host,
-				port: port,
-				user: user
-			};
-			this.addConnection(options);
-		}
-	}
-}

extensions/mssql/src/objectExplorerNodeProvider/hdfsProvider.ts

@@ -130,7 +130,7 @@ export class FolderNode extends HdfsFileSourceNode {
 					// Note: for now, assuming HDFS-provided sorting is sufficient
 					this.children = files.map((file) => {
 						let node: TreeNode = file.isDirectory ? new FolderNode(this.context, file.path, this.fileSource)
-												: new FileNode(this.context, file.path, this.fileSource);
+							: new FileNode(this.context, file.path, this.fileSource);
 						node.parent = this;
 						return node;
 					});

extensions/mssql/src/objectExplorerNodeProvider/objectExplorerNodeProvider.ts

@@ -10,32 +10,27 @@ import * as vscode from 'vscode';
 import * as nls from 'vscode-nls';
 const localize = nls.loadMessageBundle();
 
-import * as UUID from 'vscode-languageclient/lib/utils/uuid';
 import { ProviderBase } from './providerBase';
-import { Connection } from './connection';
+import { SqlClusterConnection } from './connection';
 import * as utils from '../utils';
 import { TreeNode } from './treeNodes';
 import { ConnectionNode, TreeDataContext, ITreeChangeHandler } from './hdfsProvider';
 import { IFileSource } from './fileSources';
 import { AppContext } from '../appContext';
 import * as constants from '../constants';
+import * as SqlClusterLookUp from '../sqlClusterLookUp';
+import { ICommandObjectExplorerContext } from './command';
 
-const outputChannel = vscode.window.createOutputChannel(constants.providerId);
-export interface IEndpoint {
-	serviceName: string;
-	ipAddress: string;
-	port: number;
-}
+export const mssqlOutputChannel = vscode.window.createOutputChannel(constants.providerId);
 
 export class MssqlObjectExplorerNodeProvider extends ProviderBase implements sqlops.ObjectExplorerNodeProvider, ITreeChangeHandler {
 	public readonly supportedProviderId: string = constants.providerId;
-	private sessionMap: Map<string, Session>;
+	private sessionMap: Map<string, SqlClusterSession>;
 	private expandCompleteEmitter = new vscode.EventEmitter<sqlops.ObjectExplorerExpandInfo>();
 
 	constructor(private appContext: AppContext) {
 		super();
-
-		this.sessionMap = new Map();
+		this.sessionMap = new Map<string, SqlClusterSession>();
 		this.appContext.registerService<MssqlObjectExplorerNodeProvider>(constants.ObjectExplorerService, this);
 	}
 
@@ -49,44 +44,19 @@ export class MssqlObjectExplorerNodeProvider extends ProviderBase implements sql
 		});
 	}
 
-	private async doSessionOpen(sessionInfo: sqlops.ObjectExplorerSession): Promise<boolean> {
-		let connectionProfile = await sqlops.objectexplorer.getSessionConnectionProfile(sessionInfo.sessionId);
-		if (!connectionProfile) {
-			return false;
-		} else {
-			let credentials = await sqlops.connection.getCredentials(connectionProfile.id);
-			let serverInfo = await sqlops.connection.getServerInfo(connectionProfile.id);
-			if (!serverInfo || !credentials || !serverInfo.options) {
-				return false;
-			}
-			let endpoints: IEndpoint[] = serverInfo.options[constants.clusterEndpointsProperty];
-			if (!endpoints || endpoints.length === 0) {
-				return false;
-			}
-			let index = endpoints.findIndex(ep => ep.serviceName === constants.hadoopKnoxEndpointName);
-			if (index === -1) {
-				return false;
-			}
+	private async doSessionOpen(session: sqlops.ObjectExplorerSession): Promise<boolean> {
+		if (!session || !session.sessionId) { return false; }
 
-			let connInfo: sqlops.connection.Connection = {
-				options: {
-					'host': endpoints[index].ipAddress,
-					'groupId': connectionProfile.options.groupId,
-					'knoxport': endpoints[index].port,
-					'user': 'root', //connectionProfile.options.userName cluster setup has to have the same user for master and big data cluster
-					'password': credentials.password,
-				},
-				providerName: constants.mssqlClusterProviderName,
-				connectionId: UUID.generateUuid()
-			};
+		let sqlConnProfile = await sqlops.objectexplorer.getSessionConnectionProfile(session.sessionId);
+		if (!sqlConnProfile) { return false; }
 
-			let connection = new Connection(connInfo);
-			connection.saveUriWithPrefix(constants.objectExplorerPrefix);
-			let session = new Session(connection, sessionInfo.sessionId);
-			session.root = new RootNode(session, new TreeDataContext(this.appContext.extensionContext, this), sessionInfo.rootNode.nodePath);
-			this.sessionMap.set(sessionInfo.sessionId, session);
-			return true;
-		}
+		let clusterConnInfo = await SqlClusterLookUp.getSqlClusterConnection(sqlConnProfile);
+		if (!clusterConnInfo) { return false; }
+
+		let clusterConnection = new SqlClusterConnection(clusterConnInfo);
+		let clusterSession = new SqlClusterSession(clusterConnection, session, sqlConnProfile, this.appContext, this);
+		this.sessionMap.set(session.sessionId, clusterSession);
+		return true;
 	}
 
 	expandNode(nodeInfo: sqlops.ExpandNodeInfo, isRefresh: boolean = false): Thenable<boolean> {
@@ -125,15 +95,15 @@ export class MssqlObjectExplorerNodeProvider extends ProviderBase implements sql
 		return true;
 	}
 
-	private async startExpansion(session: Session, nodeInfo: sqlops.ExpandNodeInfo, isRefresh: boolean = false): Promise<void> {
+	private async startExpansion(session: SqlClusterSession, nodeInfo: sqlops.ExpandNodeInfo, isRefresh: boolean = false): Promise<void> {
 		let expandResult: sqlops.ObjectExplorerExpandInfo = {
-			sessionId: session.uri,
+			sessionId: session.sessionId,
 			nodePath: nodeInfo.nodePath,
 			errorMessage: undefined,
 			nodes: []
 		};
 		try {
-			let node = await session.root.findNodeByPath(nodeInfo.nodePath, true);
+			let node = await session.rootNode.findNodeByPath(nodeInfo.nodePath, true);
 			if (node) {
 				expandResult.errorMessage = node.getNodeInfo().errorMessage;
 				let children = await node.getChildren(true);
@@ -182,57 +152,55 @@ export class MssqlObjectExplorerNodeProvider extends ProviderBase implements sql
 
 	private async notifyNodeChangesAsync(node: TreeNode): Promise<void> {
 		try {
-			let session = this.getSessionForNode(node);
+			let session = this.getSqlClusterSessionForNode(node);
 			if (!session) {
 				this.appContext.apiWrapper.showErrorMessage(localize('sessionNotFound', 'Session for node {0} does not exist', node.nodePathValue));
 			} else {
 				let nodeInfo = node.getNodeInfo();
 				let expandInfo: sqlops.ExpandNodeInfo = {
 					nodePath: nodeInfo.nodePath,
-					sessionId: session.uri
+					sessionId: session.sessionId
 				};
 				await this.refreshNode(expandInfo);
 			}
 		} catch (err) {
-			outputChannel.appendLine(localize('notifyError', 'Error notifying of node change: {0}', err));
+			mssqlOutputChannel.appendLine(localize('notifyError', 'Error notifying of node change: {0}', err));
 		}
 	}
 
-	private getSessionForNode(node: TreeNode): Session {
-		let rootNode: DataServicesNode = undefined;
-		while (rootNode === undefined && node !== undefined) {
+	private getSqlClusterSessionForNode(node: TreeNode): SqlClusterSession {
+		let sqlClusterSession: SqlClusterSession = undefined;
+		while (node !== undefined) {
 			if (node instanceof DataServicesNode) {
-				rootNode = node;
+				sqlClusterSession = node.session;
 				break;
 			} else {
 				node = node.parent;
 			}
 		}
-		if (rootNode) {
-			return rootNode.session;
-		}
-		// Not found
-		return undefined;
+		return sqlClusterSession;
 	}
 
-	async findNodeForContext<T extends TreeNode>(explorerContext: sqlops.ObjectExplorerContext): Promise<T> {
+	async findSqlClusterNodeByContext<T extends TreeNode>(context: ICommandObjectExplorerContext | sqlops.ObjectExplorerContext): Promise<T> {
 		let node: T = undefined;
-		let session = await this.findSessionForConnection(explorerContext.connectionProfile);
+		let explorerContext = 'explorerContext' in context ? context.explorerContext : context;
+		let sqlConnProfile = explorerContext.connectionProfile;
+		let session = this.findSqlClusterSessionBySqlConnProfile(sqlConnProfile);
 		if (session) {
 			if (explorerContext.isConnectionNode) {
 				// Note: ideally fix so we verify T matches RootNode and go from there
-				node = <T><any>session.root;
+				node = <T><any>session.rootNode;
 			} else {
 				// Find the node under the session
-				node = <T><any>await session.root.findNodeByPath(explorerContext.nodeInfo.nodePath, true);
+				node = <T><any>await session.rootNode.findNodeByPath(explorerContext.nodeInfo.nodePath, true);
 			}
 		}
 		return node;
 	}
 
-	private async findSessionForConnection(connectionProfile: sqlops.IConnectionProfile): Promise<Session> {
+	public findSqlClusterSessionBySqlConnProfile(connectionProfile: sqlops.IConnectionProfile): SqlClusterSession {
 		for (let session of this.sessionMap.values()) {
-			if (session.connection && await session.connection.isMatch(connectionProfile)) {
+			if (session.isMatchedSqlConnection(connectionProfile)) {
 				return session;
 			}
 		}
@@ -240,50 +208,58 @@ export class MssqlObjectExplorerNodeProvider extends ProviderBase implements sql
 	}
 }
 
-export class Session {
-	private _root: RootNode;
-	constructor(private _connection: Connection, private sessionId?: string) {
+export class SqlClusterSession {
+	private _rootNode: SqlClusterRootNode;
+
+	constructor(
+		private _sqlClusterConnection: SqlClusterConnection,
+		private _sqlSession: sqlops.ObjectExplorerSession,
+		private _sqlConnectionProfile: sqlops.IConnectionProfile,
+		private _appContext: AppContext,
+		private _changeHandler: ITreeChangeHandler
+	) {
+		this._rootNode = new SqlClusterRootNode(this,
+			new TreeDataContext(this._appContext.extensionContext, this._changeHandler),
+			this._sqlSession.rootNode.nodePath);
 	}
 
-	public get uri(): string {
-		return this.sessionId || this._connection.uri;
-	}
+	public get sqlClusterConnection(): SqlClusterConnection { return this._sqlClusterConnection; }
+	public get sqlSession(): sqlops.ObjectExplorerSession { return this._sqlSession; }
+	public get sqlConnectionProfile(): sqlops.IConnectionProfile { return this._sqlConnectionProfile; }
+	public get sessionId(): string { return this._sqlSession.sessionId; }
+	public get rootNode(): SqlClusterRootNode { return this._rootNode; }
 
-	public get connection(): Connection {
-		return this._connection;
-	}
-
-	public set root(node: RootNode) {
-		this._root = node;
-	}
-
-	public get root(): RootNode {
-		return this._root;
+	public isMatchedSqlConnection(sqlConnProfile: sqlops.IConnectionProfile): boolean {
+		return this._sqlConnectionProfile.id === sqlConnProfile.id;
 	}
 }
 
-class RootNode extends TreeNode {
-	private children: TreeNode[];
-	constructor(private _session: Session, private context: TreeDataContext, private nodePath: string) {
+class SqlClusterRootNode extends TreeNode {
+	private _children: TreeNode[];
+	constructor(
+		private _session: SqlClusterSession,
+		private _treeDataContext: TreeDataContext,
+		private _nodePathValue: string
+	) {
 		super();
 	}
 
-	public get session(): Session {
+	public get session(): SqlClusterSession {
 		return this._session;
 	}
 
 	public get nodePathValue(): string {
-		return this.nodePath;
+		return this._nodePathValue;
 	}
 
 	public getChildren(refreshChildren: boolean): TreeNode[] | Promise<TreeNode[]> {
-		if (refreshChildren || !this.children) {
-			this.children = [];
-			let dataServicesNode = new DataServicesNode(this._session, this.context, this.nodePath);
+		if (refreshChildren || !this._children) {
+			this._children = [];
+			let dataServicesNode = new DataServicesNode(this._session, this._treeDataContext, this._nodePathValue);
 			dataServicesNode.parent = this;
-			this.children.push(dataServicesNode);
+			this._children.push(dataServicesNode);
 		}
-		return this.children;
+		return this._children;
 	}
 
 	getTreeItem(): vscode.TreeItem | Promise<vscode.TreeItem> {
@@ -307,31 +283,28 @@ class RootNode extends TreeNode {
 }
 
 class DataServicesNode extends TreeNode {
-	private children: TreeNode[];
-	constructor(private _session: Session, private context: TreeDataContext, private nodePath: string) {
+	private _children: TreeNode[];
+	constructor(private _session: SqlClusterSession, private _context: TreeDataContext, private _nodePath: string) {
 		super();
 	}
 
-	public get session(): Session {
+	public get session(): SqlClusterSession {
 		return this._session;
 	}
 
 	public get nodePathValue(): string {
-		return this.nodePath;
+		return this._nodePath;
 	}
 
 	public getChildren(refreshChildren: boolean): TreeNode[] | Promise<TreeNode[]> {
-		if (refreshChildren || !this.children) {
-			this.children = [];
-			let hdfsNode = new ConnectionNode(this.context, localize('hdfsFolder', 'HDFS'), this.createHdfsFileSource());
+		if (refreshChildren || !this._children) {
+			this._children = [];
+			let fileSource: IFileSource = this.session.sqlClusterConnection.createHdfsFileSource();
+			let hdfsNode = new ConnectionNode(this._context, localize('hdfsFolder', 'HDFS'), fileSource);
 			hdfsNode.parent = this;
-			this.children.push(hdfsNode);
+			this._children.push(hdfsNode);
 		}
-		return this.children;
-	}
-
-	private createHdfsFileSource(): IFileSource {
-		return this.session.connection.createHdfsFileSource();
+		return this._children;
 	}
 
 	getTreeItem(): vscode.TreeItem | Promise<vscode.TreeItem> {