Adding caching in OE service (#22539)

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

@@ -70,7 +70,7 @@ export class AsyncServerTree extends WorkbenchAsyncDataTree<ConnectionProfileGro
 	 * This method overrides the original implementation to find the node by comparing the ids of the elements.
 	 * If the node is not found in the original implementation, we search for the node in the nodes map by ids.
 	 */
-	public override getDataNode(element: ConnectionProfileGroup | ServerTreeElement): IAsyncDataTreeNode<ConnectionProfileGroup, ServerTreeElement> {
+	public override getDataNode(element: ServerTreeElement, throwError: boolean = true): IAsyncDataTreeNode<ConnectionProfileGroup, ServerTreeElement> | undefined {
 		try {
 			const node = super.getDataNode(element);
 			return node;
@@ -79,7 +79,10 @@ export class AsyncServerTree extends WorkbenchAsyncDataTree<ConnectionProfileGro
 			if (node) {
 				return node;
 			}
-			throw e;
+			if (throwError) {
+				throw e;
+			}
+			return undefined;
 		}
 	}
 
@@ -132,46 +135,13 @@ export class AsyncServerTree extends WorkbenchAsyncDataTree<ConnectionProfileGro
 
 	public async expandElements(elements: ServerTreeElement[]): Promise<void> {
 		for (let element of elements) {
-			const id = element.id;
+			const node = this.getDataNode(element, false);
-			const node = this.getDataNodeById(id);
 			if (node) {
 				await this.expand(node.element);
-			} else {
-				// If the node is not found in the nodes map, we search for the node by comparing the relative paths of the elements
-				if (element) {
-					const elementPath = this.getRelativePath(element);
-					for (let n of this.nodes.values()) {
-						if (this.getRelativePath(n.element) === elementPath) {
-							await this.expand(n.element);
-							break;
-						}
-					}
-				}
 			}
 		}
 	}
 
-	/**
-	 * Get the relative path of the element in the tree. For connection and group, the path is the id of the element.
-	 * For other elements, the path is the node path of the element and the id of the connection they belong to.
-	 */
-	private getRelativePath(element: ServerTreeElement): string {
-		let path = '';
-		if (element instanceof TreeNode) {
-			path = element.nodePath;
-			let parent = element.parent;
-			while (parent.parent) {
-				parent = parent.parent;
-			}
-			if (parent.connection) {
-				path = parent.connection.id + '/' + path;
-			}
-		} else if (element instanceof ConnectionProfile || element instanceof ConnectionProfileGroup) {
-			path = element.id;
-		}
-		return path;
-	}
-
 	/**
 	 * Mark the element as dirty so that it will be refreshed when it is expanded next time
 	 * @param element The element to mark as dirty

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

@@ -179,6 +179,9 @@ export class ObjectExplorerService implements IObjectExplorerService {
 
 	private _connectionsWaitingForSession: Map<string, boolean> = new Map<string, boolean>();
 
+	// Cache of tree nodes for each connection by session ids
+	private _treeNodeCache: Map<string, Map<string, TreeNode>> = new Map<string, Map<string, TreeNode>>();
+
 	constructor(
 		@IConnectionManagementService private _connectionManagementService: IConnectionManagementService,
 		@IAdsTelemetryService private _telemetryService: IAdsTelemetryService,
@@ -233,6 +236,7 @@ export class ObjectExplorerService implements IObjectExplorerService {
 			if (!session) {
 				return;
 			}
+			this._treeNodeCache.delete(session.sessionId);
 			await this.closeSession(connection.providerName, session);
 			delete this._activeObjectExplorerNodes[connectionUri];
 			delete this._sessions[session.sessionId!];
@@ -340,6 +344,8 @@ export class ObjectExplorerService implements IObjectExplorerService {
 		try {
 			if (session.success && session.rootNode) {
 				let server = this.toTreeNode(session.rootNode, undefined);
+				this._treeNodeCache.set(sessionId, new Map<string, TreeNode>());
+				this._treeNodeCache.get(sessionId)!.set(this.getTreeNodeCacheKey(server.toNodeInfo()), server);
 				server.connection = connection;
 				server.session = session;
 				this._activeObjectExplorerNodes[connection!.id] = server;
@@ -730,9 +736,31 @@ export class ObjectExplorerService implements IObjectExplorerService {
 			throw new Error('Failed to expand node - no provider name');
 		}
 		const expandResult = await this.callExpandOrRefreshFromService(providerName, session, parentTree, refresh);
+		const sessionTreeNodeCache = this._treeNodeCache.get(session.sessionId!);
 		if (expandResult && expandResult.nodes) {
+			// In case of refresh, we want to clear the cache of the descendants of the node being refreshed
+			if (refresh && parentTree?.children) {
+				const stack = [...parentTree.children];
+				while (stack.length > 0) {
+					const currentTreeNode = stack.pop();
+					if (currentTreeNode) {
+						sessionTreeNodeCache.delete(this.getTreeNodeCacheKey(currentTreeNode.toNodeInfo()));
+						if (currentTreeNode.children) {
+							stack.push(...currentTreeNode.children);
+						}
+					}
+				}
+			}
 			const children = expandResult.nodes.map(node => {
-				return this.toTreeNode(node, parentTree);
+				const cacheKey = this.getTreeNodeCacheKey(node);
+				// In case of refresh, we want to update the existing node in the cache
+				if (!refresh && sessionTreeNodeCache.has(cacheKey)) {
+					return sessionTreeNodeCache.get(cacheKey);
+				} else {
+					const treeNode = this.toTreeNode(node, parentTree);
+					sessionTreeNodeCache.set(cacheKey, treeNode);
+					return treeNode;
+				}
 			});
 			parentTree.children = children.filter(c => c !== undefined);
 			return children;
@@ -1004,4 +1032,8 @@ export class ObjectExplorerService implements IObjectExplorerService {
 	public getObjectExplorerTimeout(): number {
 		return this._configurationService.getValue<number>(NODE_EXPANSION_CONFIG);
 	}
+
+	private getTreeNodeCacheKey(node: azdata.NodeInfo): string {
+		return node.nodePath;
+	}
 }