avoid unnecessary connections in azure view (#20817)

* avoid unnecessary connections in azure view

* pr comments

src/sql/workbench/common/views.ts

@@ -29,7 +29,11 @@ export interface ITreeItem extends vsITreeItem {
 	payload?: IConnectionProfile; // its possible we will want this to be more generic
 	sqlIcon?: string;
 	type?: NodeType;
-	nodeInfo?: NodeInfo
+	nodeInfo?: NodeInfo;
+	/**
+	 * The Object Explorer session id that the tree item belongs to.
+	 */
+	sessionId?: string;
 }
 
 export interface ITreeView extends vsITreeView {

src/sql/workbench/contrib/dataExplorer/browser/nodeCommands.common.ts

@@ -19,7 +19,7 @@ CommandsRegistry.registerCommand({
 	handler: (accessor, args: TreeViewItemHandleArg) => {
 		if (args.$treeItem) {
 			const oeService = accessor.get(IOEShimService);
-			return oeService.disconnectNode(args.$treeViewId, args.$treeItem).then(() => {
+			return oeService.disconnectNode(args.$treeItem).then(() => {
 				const { treeView } = (<ICustomViewDescriptor>Registry.as<IViewsRegistry>(Extensions.ViewsRegistry).getView(args.$treeViewId));
 				// we need to collapse it then refresh it so that the tree doesn't try and use it's cache next time the user expands the node
 				treeView.collapse(args.$treeItem);

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

@@ -27,7 +27,7 @@ export const IOEShimService = createDecorator<IOEShimService>(SERVICE_ID);
 export interface IOEShimService {
 	_serviceBrand: undefined;
 	getChildren(node: ITreeItem, viewId: string): Promise<ITreeItem[]>;
-	disconnectNode(viewId: string, node: ITreeItem): Promise<boolean>;
+	disconnectNode(node: ITreeItem): Promise<boolean>;
 	providerExists(providerId: string): boolean;
 	isNodeConnected(viewId: string, node: ITreeItem): boolean;
 	getNodeInfoForTreeItem(treeItem: ITreeItem): azdata.NodeInfo | undefined;
@@ -36,7 +36,6 @@ export interface IOEShimService {
 export class OEShimService extends Disposable implements IOEShimService {
 	_serviceBrand: undefined;
 
-	private sessionMap = new Map<number, string>();
 	private nodeHandleMap = new Map<number, string>();
 	private nodeInfoMap = new Map<ITreeItem, azdata.NodeInfo>();
 
@@ -110,28 +109,12 @@ export class OEShimService extends Disposable implements IOEShimService {
 		return connProfile;
 	}
 
-	public async disconnectNode(viewId: string, node: ITreeItem): Promise<boolean> {
+	public async disconnectNode(node: ITreeItem): Promise<boolean> {
-		// we assume only nodes with payloads can be connected
+		if (node.sessionId !== undefined) {
-		// check to make sure we have an existing connection
+			const closed = (await this.oe.closeSession(node.childProvider!, this.oe.getSession(node.sessionId)!))!.success;
-		let key = generateSessionMapKey(viewId, node);
-		let session = this.sessionMap.get(key);
-		if (session) {
-			let closed = (await this.oe.closeSession(node.childProvider!, this.oe.getSession(session)!))!.success;
-			if (closed) {
-				this.sessionMap.delete(key);
-			}
 			return closed;
 		}
-		return Promise.resolve(false);
+		return false;
-	}
-
-	private async getOrCreateSession(viewId: string, node: ITreeItem): Promise<string> {
-		// verify the map is correct
-		let key = generateSessionMapKey(viewId, node);
-		if (!this.sessionMap.has(key)) {
-			this.sessionMap.set(key, await this.createSession(viewId, node.childProvider!, node));
-		}
-		return this.sessionMap.get(key)!;
 	}
 
 	public async getChildren(node: ITreeItem, viewId: string): Promise<ITreeItem[]> {
@@ -140,13 +123,15 @@ export class OEShimService extends Disposable implements IOEShimService {
 				node.payload.authenticationType = this.getDefaultAuthenticationType(this.configurationService);  // we need to set auth type here, because it's value is part of the session key
 			}
 
-			const sessionId = await this.getOrCreateSession(viewId, node);
+			if (node.sessionId === undefined) {
+				node.sessionId = await this.createSession(viewId, node.childProvider!, node);
+			}
 			const requestHandle = this.nodeHandleMap.get(generateNodeMapKey(viewId, node)) || node.handle;
 			const treeNode = new TreeNode(undefined!, undefined!, undefined!, undefined!, requestHandle, undefined!); // hack since this entire system is a hack anyways
 			treeNode.connection = new ConnectionProfile(this.capabilities, node.payload);
 			const childrenNodes = await this.oe.refreshTreeNode({
 				success: true,
-				sessionId,
+				sessionId: node.sessionId!,
 				rootNode: undefined!, // hack since this entire system is a hack anyways
 				errorMessage: undefined
 			}, treeNode);
@@ -209,7 +194,8 @@ export class OEShimService extends Disposable implements IOEShimService {
 			payload: node.payload || (databaseChanged ? updatedPayload : parentNode.payload),
 			contextValue: node.nodeTypeId,
 			sqlIcon: icon,
-			nodeInfo: nodeInfo
+			nodeInfo: nodeInfo,
+			sessionId: parentNode.sessionId
 		};
 		this.nodeHandleMap.set(generateNodeMapKey(viewId, newTreeItem), nodePath);
 		this.nodeInfoMap.set(newTreeItem, nodeInfo);
@@ -221,7 +207,7 @@ export class OEShimService extends Disposable implements IOEShimService {
 	}
 
 	public isNodeConnected(viewId: string, node: ITreeItem): boolean {
-		return this.sessionMap.has(generateSessionMapKey(viewId, node));
+		return node.sessionId !== undefined;
 	}
 
 	public getNodeInfoForTreeItem(treeItem: ITreeItem): azdata.NodeInfo | undefined {
@@ -236,10 +222,6 @@ export class OEShimService extends Disposable implements IOEShimService {
 	}
 }
 
-function generateSessionMapKey(viewId: string, node: ITreeItem): number {
-	return hash([viewId, node.childProvider, node.payload]);
-}
-
 function generateNodeMapKey(viewId: string, node: ITreeItem): number {
 	return hash([viewId, node.handle]);
 }