Adding filtering to OE Service (#22900)

* Adding filtering to OE Service

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

Co-authored-by: Cheena Malhotra <13396919+cheenamalhotra@users.noreply.github.com>

* Update src/sql/workbench/services/objectExplorer/browser/serverTreeRenderer.ts

Co-authored-by: Cheena Malhotra <13396919+cheenamalhotra@users.noreply.github.com>

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

Co-authored-by: Cheena Malhotra <13396919+cheenamalhotra@users.noreply.github.com>

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

Co-authored-by: Cheena Malhotra <13396919+cheenamalhotra@users.noreply.github.com>

* Fixing compile errors

* Fixing circular dependency

---------

Co-authored-by: Cheena Malhotra <13396919+cheenamalhotra@users.noreply.github.com>

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

@@ -18,14 +18,14 @@ import { Disposable } from 'vs/base/common/lifecycle';
 import { IInstantiationService } from 'vs/platform/instantiation/common/instantiation';
 import { IKeyboardNavigationLabelProvider } from 'vs/base/browser/ui/list/list';
 import { IListAccessibilityProvider } from 'vs/base/browser/ui/list/listWidget';
-import { ServerTreeRenderer } from 'sql/workbench/services/objectExplorer/browser/serverTreeRenderer';
+import { ServerTreeRenderer, getLabelWithFilteredSuffix } from 'sql/workbench/services/objectExplorer/browser/serverTreeRenderer';
 import { ServerTreeElement } from 'sql/workbench/services/objectExplorer/browser/asyncServerTree';
 import { DefaultServerGroupColor } from 'sql/workbench/services/serverGroup/common/serverGroupViewModel';
 import { withNullAsUndefined } from 'vs/base/common/types';
 import { instanceOfSqlThemeIcon } from 'sql/workbench/services/objectExplorer/common/nodeType';
+import { IObjectExplorerService } from 'sql/workbench/services/objectExplorer/browser/objectExplorerService';
 
 const DefaultConnectionIconClass = 'server-page';
-
 export interface ConnectionProfileGroupDisplayOptions {
 	showColor: boolean;
 }
@@ -96,7 +96,8 @@ class ConnectionProfileTemplate extends Disposable {
 	constructor(
 		container: HTMLElement,
 		private _isCompact: boolean,
-		@IConnectionManagementService private _connectionManagementService: IConnectionManagementService
+		@IConnectionManagementService private _connectionManagementService: IConnectionManagementService,
+		@IObjectExplorerService private _objectExplorerService: IObjectExplorerService
 	) {
 		super();
 		container.parentElement!.classList.add('connection-profile');
@@ -122,6 +123,11 @@ class ConnectionProfileTemplate extends Disposable {
 		let label = element.title;
 		this._label.textContent = label;
 		this._root.title = element.serverInfo;
+
+		const treeNode = this._objectExplorerService.getObjectExplorerNode(element);
+		if (treeNode?.filters?.length > 0) {
+			this._label.textContent = getLabelWithFilteredSuffix(this._label.textContent);
+		}
 	}
 }
 
@@ -192,7 +198,8 @@ class TreeNodeTemplate extends Disposable {
 			iconRenderer.putIcon(this._icon, element.icon);
 		}
 
-		this._label.textContent = element.label;
+		this._label.textContent = element.filters.length > 0 ? getLabelWithFilteredSuffix(element.label) :
+			element.label;
 		this._root.title = element.label;
 	}
 }

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

@@ -182,6 +182,9 @@ export class ObjectExplorerService implements IObjectExplorerService {
 	// Cache of tree nodes for each connection by session ids
 	private _treeNodeCache: Map<string, Map<string, TreeNode>> = new Map<string, Map<string, TreeNode>>();
 
+	// Cache of node filters for each connection by session ids
+	private _nodeFilterCache: Map<string, Map<string, azdata.NodeFilter[]>> = new Map<string, Map<string, azdata.NodeFilter[]>>();
+
 	constructor(
 		@IConnectionManagementService private _connectionManagementService: IConnectionManagementService,
 		@IAdsTelemetryService private _telemetryService: IAdsTelemetryService,
@@ -237,6 +240,7 @@ export class ObjectExplorerService implements IObjectExplorerService {
 				return;
 			}
 			this._treeNodeCache.delete(session.sessionId);
+			this._nodeFilterCache.delete(session.sessionId);
 			await this.closeSession(connection.providerName, session);
 			delete this._activeObjectExplorerNodes[connectionUri];
 			delete this._sessions[session.sessionId!];
@@ -346,6 +350,7 @@ export class ObjectExplorerService implements IObjectExplorerService {
 				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);
+				this._nodeFilterCache.set(sessionId, new Map<string, azdata.NodeFilter[]>());
 				server.connection = connection;
 				server.session = session;
 				this._activeObjectExplorerNodes[connection!.id] = server;
@@ -475,6 +480,12 @@ export class ObjectExplorerService implements IObjectExplorerService {
 		refresh: boolean = false): Promise<azdata.ObjectExplorerExpandInfo> {
 		let self = this;
 		return new Promise<azdata.ObjectExplorerExpandInfo>((resolve, reject) => {
+			const sessionFilterCache = this._nodeFilterCache.get(session.sessionId!);
+			// If the node has filters we need to cache them so that we can reapply them when the node is refreshed.
+			if (node.filters && sessionFilterCache) {
+				sessionFilterCache.set(this.getTreeNodeCacheKey(node), node.filters);
+			}
+
 			if (session.sessionId! in self._sessions && self._sessions[session.sessionId!]) {
 				let newRequest = false;
 				if (!self._sessions[session.sessionId!].nodes[node.nodePath]) {
@@ -561,7 +572,8 @@ export class ObjectExplorerService implements IObjectExplorerService {
 							self.callExpandOrRefreshFromProvider(provider, {
 								sessionId: session.sessionId!,
 								nodePath: node.nodePath,
-								securityToken: session.securityToken
+								securityToken: session.securityToken,
+								filters: node.filters
 							}, 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
@@ -699,8 +711,9 @@ export class ObjectExplorerService implements IObjectExplorerService {
 	}
 
 	public resolveTreeNodeChildren(session: azdata.ObjectExplorerSession, parentTree: TreeNode): Promise<TreeNode[]> {
-		// Always refresh the node if it has an error, otherwise expand it normally
+		// Always refresh the node if it has an error or forceRefresh is set to true, otherwise expand it normally
-		let needsRefresh = !!parentTree.errorStateMessage;
+		let needsRefresh = !!parentTree.errorStateMessage || parentTree.forceRefresh;
+		parentTree.forceRefresh = false;
 		return this.expandOrRefreshTreeNode(session, parentTree, needsRefresh);
 	}
 
@@ -752,15 +765,25 @@ export class ObjectExplorerService implements IObjectExplorerService {
 				}
 			}
 			const children = expandResult.nodes.map(node => {
+				let treeNode;
 				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);
+					treeNode = sessionTreeNodeCache.get(cacheKey);
 				} else {
-					const treeNode = this.toTreeNode(node, parentTree);
+					treeNode = this.toTreeNode(node, parentTree);
 					sessionTreeNodeCache.set(cacheKey, treeNode);
-					return treeNode;
 				}
+
+				const filterCacheKey = this.getTreeNodeCacheKey(treeNode);
+				const sessionFilterCache = this._nodeFilterCache.get(session.sessionId!);
+				// Making sure we retain the filters for the node.
+				if (sessionFilterCache?.has(filterCacheKey)) {
+					treeNode.filters = sessionFilterCache.get(filterCacheKey) ?? [];
+				} else {
+					treeNode.filters = [];
+				}
+				return treeNode;
 			});
 			parentTree.children = children.filter(c => c !== undefined);
 			return children;
@@ -785,7 +808,7 @@ export class ObjectExplorerService implements IObjectExplorerService {
 		}
 
 		let node = new TreeNode(nodeInfo.nodeType, nodeInfo.objectType, nodeInfo.label, isLeaf, nodeInfo.nodePath, nodeInfo.parentNodePath,
-			nodeInfo.nodeSubType!, nodeInfo.nodeStatus, parent, nodeInfo.metadata, nodeInfo.iconType, nodeInfo.icon, {
+			nodeInfo.nodeSubType!, nodeInfo.nodeStatus, parent, nodeInfo.metadata, nodeInfo.iconType, nodeInfo.icon, nodeInfo.filterableProperties, {
 			getChildren: (treeNode?: TreeNode) => this.getChildren(treeNode),
 			isExpanded: treeNode => this.isExpanded(treeNode),
 			setNodeExpandedState: async (treeNode, expandedState) => await this.setNodeExpandedState(treeNode, expandedState),
@@ -1033,7 +1056,7 @@ export class ObjectExplorerService implements IObjectExplorerService {
 		return this._configurationService.getValue<number>(NODE_EXPANSION_CONFIG);
 	}
 
-	private getTreeNodeCacheKey(node: azdata.NodeInfo): string {
+	private getTreeNodeCacheKey(node: azdata.NodeInfo | TreeNode): string {
 		return node.nodePath;
 	}
 }

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

@@ -17,6 +17,8 @@ import { URI } from 'vs/base/common/uri';
 import { DefaultServerGroupColor } from 'sql/workbench/services/serverGroup/common/serverGroupViewModel';
 import { withNullAsUndefined } from 'vs/base/common/types';
 import { instanceOfSqlThemeIcon } from 'sql/workbench/services/objectExplorer/common/nodeType';
+import { IObjectExplorerService } from 'sql/workbench/services/objectExplorer/browser/objectExplorerService';
+import { localize } from 'vs/nls';
 
 export interface IConnectionTemplateData {
 	root: HTMLElement;
@@ -38,6 +40,10 @@ export interface IObjectExplorerTemplateData {
 	treeNode: TreeNode;
 }
 
+export function getLabelWithFilteredSuffix(label: string): string {
+	return localize('filteredTreeElementName', "{0} (filtered)", label);
+}
+
 /**
  * Renders the tree items.
  * Uses the dom template to render connection groups and connections.
@@ -59,7 +65,8 @@ export class ServerTreeRenderer implements IRenderer {
 
 	constructor(
 		isCompact: boolean,
-		@IConnectionManagementService private _connectionManagementService: IConnectionManagementService
+		@IConnectionManagementService private _connectionManagementService: IConnectionManagementService,
+		@IObjectExplorerService private _objectExplorerService: IObjectExplorerService
 	) {
 		// isCompact defaults to false unless explicitly set by instantiation call.
 		if (isCompact) {
@@ -163,9 +170,9 @@ export class ServerTreeRenderer implements IRenderer {
 		if (treeNode.icon && !instanceOfSqlThemeIcon(treeNode.icon)) {
 			iconRenderer.putIcon(templateData.icon, treeNode.icon);
 		}
-
+		const nodeLabel = treeNode.filters?.length > 0 ? getLabelWithFilteredSuffix(treeNode.label) : treeNode.label;
-		templateData.label.textContent = treeNode.label;
+		templateData.label.textContent = nodeLabel;
-		templateData.root.title = treeNode.label;
+		templateData.root.title = nodeLabel;
 	}
 
 	private getIconPath(connection: ConnectionProfile): IconPath | undefined {
@@ -224,10 +231,10 @@ export class ServerTreeRenderer implements IRenderer {
 
 		let iconPath = this.getIconPath(connection);
 		this.renderServerIcon(templateData.icon, iconPath, isConnected);
-
+		const treeNode = this._objectExplorerService.getObjectExplorerNode(connection);
-		let label = connection.title;
+		let label = treeNode?.filters?.length > 0 ? getLabelWithFilteredSuffix(connection.title) : connection.title;
 		templateData.label.textContent = label;
-		templateData.root.title = connection.serverInfo;
+		templateData.root.title = treeNode?.filters?.length > 0 ? getLabelWithFilteredSuffix(connection.title) : connection.serverInfo;
 		templateData.connectionProfile = connection;
 	}
 

src/sql/workbench/services/objectExplorer/common/treeNode.ts

@@ -93,6 +93,15 @@ export class TreeNode {
 	 */
 	public children?: TreeNode[];
 
+	/**
+	 * Filterable properties that this node supports
+	 */
+	public filterProperties?: azdata.NodeFilterProperty[];
+
+	/**
+	 * Filters that are currently applied to this node children.
+	 */
+	public filters?: azdata.NodeFilter[];
 
 	public connection?: ConnectionProfile;
 
@@ -104,10 +113,13 @@ export class TreeNode {
 
 	public icon?: IconPath | SqlThemeIcon;
 
+	public forceRefresh: boolean = false;
+
 	constructor(nodeTypeId: string, objectType: string, label: string, isAlwaysLeaf: boolean, nodePath: string, parentNodePath: string,
 		nodeSubType: string, nodeStatus?: string, parent?: TreeNode, metadata?: azdata.ObjectMetadata,
 		iconType?: string | SqlThemeIcon,
 		icon?: IconPath | SqlThemeIcon,
+		filterProperties?: azdata.NodeFilterProperty[],
 		private _objectExplorerCallbacks?: ObjectExplorerCallbacks) {
 		this.nodeTypeId = nodeTypeId;
 		this.objectType = objectType;
@@ -122,6 +134,7 @@ export class TreeNode {
 		this.nodeSubType = nodeSubType;
 		this.nodeStatus = nodeStatus;
 		this.icon = icon;
+		this.filterProperties = filterProperties;
 	}
 	public getConnectionProfile(): ConnectionProfile | undefined {
 		let currentNode: TreeNode = this;