Fixing async server tree (#22511)

* Changing look of new OE * Fixing styling * Fixing moving of connected profile * Fixing drag and drop where treenodes delete connection nodes * Fixing Deleting and disconnecting in AsyncServerTree * Adding constant for OE timeout * Updated interfaces * Removing test compilation errors * Fixing most events in async server tree * Fixing connection pane styling * Fixing find node function * Fixing some more operations * Fixing some ops * All operations done * Fixed active connections * Fixed data source * Adding support for setting parents * code cleanup * Fixing icon styling issues * Fix errors * Fixing comment * Fixing spacing * Adding explanation to OE service. * Reverting server delegate changes * Reverting styling * reverting more styling change * reverting more styling * Fixing has children * Fixing drag and drop to tree nodes * fixing drag and drop * reverting timing * fixing drag and drop * cleaning some code * Fixed server and group moving * spell check * consolidating some logic * Fixed whitespace * fixing moving to root group
2026-01-17 17:22:42 -05:00 · 2023-03-29 13:59:35 -07:00
parent 7ecbbdf398
commit 2a9705c495
8 changed files with 641 additions and 107 deletions
--- a/src/sql/workbench/contrib/objectExplorer/browser/serverTreeView.ts
+++ b/src/sql/workbench/contrib/objectExplorer/browser/serverTreeView.ts
@@ -173,7 +173,10 @@ export class ServerTreeView extends Disposable implements IServerTreeView {
 			this._register(this._tree.onContextMenu(e => this.onTreeNodeContextMenu(e)));
 			this._register(this._tree.onMouseDblClick(async e => { await this.onTreeNodeDoubleClick(e.element); }));
 			this._register(this._connectionManagementService.onConnectionChanged(() => {
-				this.refreshTree().catch(err => errors.onUnexpectedError);
+				// No need to refresh AsyncServerTree when a connection is edited or added
+				if (!(this._tree instanceof AsyncServerTree)) {
+					this.refreshTree().catch(err => errors.onUnexpectedError);
+				}
 			}));
 		}

@@ -185,11 +188,19 @@ export class ServerTreeView extends Disposable implements IServerTreeView {
 			this.handleAddConnectionProfile(newProfile).catch(errors.onUnexpectedError);
 		}));
 		this._register(this._connectionManagementService.onDeleteConnectionProfile(() => {
-			this.refreshTree().catch(errors.onUnexpectedError);
+			// No need to refresh AsyncServerTree when a connection is deleted
+			if (!(this._tree instanceof AsyncServerTree)) {
+				this.refreshTree().catch(errors.onUnexpectedError);
+			}
 		}));
-		this._register(this._connectionManagementService.onDisconnect((connectionParams) => {
-			if (this.isObjectExplorerConnectionUri(connectionParams.connectionUri)) {
-				this.deleteObjectExplorerNodeAndRefreshTree(connectionParams.connectionProfile).catch(errors.onUnexpectedError);
+
+		this._register(this._connectionManagementService.onDisconnect(async (connectionParams) => {
+			if (!(this._tree instanceof AsyncServerTree)) {
+				if (this.isObjectExplorerConnectionUri(connectionParams.connectionUri)) {
+					this.deleteObjectExplorerNodeAndRefreshTree(connectionParams.connectionProfile).catch(errors.onUnexpectedError);
+				}
+			} else {
+				await this.disconnectConnection(<ConnectionProfile>connectionParams.connectionProfile);
 			}
 		}));
 		this._register(this._configurationService.onDidChangeConfiguration(e => {
@@ -202,13 +213,163 @@ export class ServerTreeView extends Disposable implements IServerTreeView {
 			this._register(this._objectExplorerService.onUpdateObjectExplorerNodes(args => {
 				if (args.errorMessage) {
 					this.showError(args.errorMessage);
-				}
-				if (args.connection) {
+				} else if (args.connection) {
+					if (this._tree instanceof AsyncServerTree) {
+						// Rerendering the node to update the badge
+						this._tree.rerender(<ConnectionProfile>args.connection);
+					}
 					this.onObjectExplorerSessionCreated(args.connection).catch(err => errors.onUnexpectedError);
 				}
 			}));
 		}

+		// Add connection profile to parent group and update group children. Then reveal and expand the new connection
+		this._register(this._connectionManagementService.onConnectionProfileCreated(async (newConnection) => {
+			if (this._tree instanceof AsyncServerTree) {
+				const connectionParentGroup = this._tree.getElementById(newConnection.groupId) as ConnectionProfileGroup;
+				if (connectionParentGroup) {
+					connectionParentGroup.connections.push(newConnection);
+					newConnection.parent = connectionParentGroup;
+					newConnection.groupId = connectionParentGroup.id;
+					await this._tree.updateChildren(connectionParentGroup);
+					await this._tree.revealSelectFocusElement(newConnection);
+					await this._tree.expand(newConnection);
+				}
+			}
+		}));
+
+		// Rerender the connection in the tree to update the badge and update the children of the connection.
+		this._register(this._connectionManagementService.onConnectionProfileConnected(async (connectedConnection) => {
+			if (this._tree instanceof AsyncServerTree) {
+				const connectionInTree = this._tree.getElementById(connectedConnection.id);
+				if (connectionInTree) {
+					await this._tree.rerender(connectionInTree);
+					await this._tree.revealSelectFocusElement(connectionInTree);
+					await this._tree.updateChildren(connectionInTree);
+					await this._tree.expand(connectionInTree);
+				}
+			}
+		}));
+
+		// Remove the connection from the parent group and update the parent's children.
+		this._register(this._connectionManagementService.onConnectionProfileDeleted(async (e) => {
+			if (this._tree instanceof AsyncServerTree) {
+				const parentGroup = <ConnectionProfileGroup>this._tree.getElementById(e.groupId);
+				if (parentGroup) {
+					parentGroup.connections = parentGroup.connections.filter(c => c.id !== e.id);
+					await this._tree.updateChildren(parentGroup);
+					await this._tree.revealSelectFocusElement(parentGroup);
+				}
+			}
+		}));
+
+
+		this._register(this._connectionManagementService.onConnectionProfileEdited(async (e) => {
+			if (this._tree instanceof AsyncServerTree) {
+				const oldProfile = <ConnectionProfile>this._tree.getElementById(e.oldProfileId);
+				const oldProfileParent = <ConnectionProfileGroup>this._tree.getElementById(oldProfile.groupId);
+				if (oldProfileParent.id !== e.profile.groupId) {
+					// If the profile was moved to a different group then remove it from the old group and add it to the new group.
+					oldProfileParent.connections = oldProfileParent.connections.filter(c => c.id !== oldProfile.id);
+					await this._tree.updateChildren(oldProfileParent);
+					const newProfileParent = <ConnectionProfileGroup>this._tree.getElementById(e.profile.groupId);
+					newProfileParent.connections.push(e.profile);
+					e.profile.parent = newProfileParent;
+					e.profile.groupId = newProfileParent.id;
+					await this._tree.updateChildren(newProfileParent);
+					await this._tree.revealSelectFocusElement(e.profile);
+					await this._tree.expand(e.profile);
+				} else {
+					// If the profile was not moved to a different group then just update the profile in the group.
+					oldProfileParent.connections[oldProfileParent.connections.findIndex(c => c.id === e.oldProfileId)] = e.profile;
+					e.profile.parent = oldProfileParent;
+					e.profile.groupId = oldProfileParent.id;
+					await this._tree.updateChildren(oldProfileParent)
+					await this._tree.revealSelectFocusElement(e.profile);
+					await this._tree.expand(e.profile);
+				}
+			}
+		}));
+
+		this._register(this._connectionManagementService.onConnectionProfileMoved(async (e) => {
+			if (this._tree instanceof AsyncServerTree) {
+				const movedConnection = <ConnectionProfile>e.source;
+				const oldParent = <ConnectionProfileGroup>this._tree.getElementById(e.oldGroupId);
+				const newParent = <ConnectionProfileGroup>this._tree.getElementById(e.newGroupId);
+				// Storing the expanded state of children of the moved connection so that they can be expanded after the move.
+				const profileExpandedState = this._tree.getExpandedState(movedConnection);
+				if (oldParent) {
+					oldParent.connections = oldParent.connections.filter(c => c.id !== e.source.id);
+					await this._tree.updateChildren(oldParent);
+				}
+				if (newParent) {
+					newParent.connections.push(movedConnection);
+					movedConnection.parent = newParent;
+					movedConnection.groupId = newParent.id;
+					await this._tree.updateChildren(newParent);
+				}
+				const newConnection = this._tree.getElementById(movedConnection.id);
+				if (newConnection) {
+					await this._tree.revealSelectFocusElement(newConnection);
+					// Expanding the previously expanded children of the moved connection after the move.
+					await this._tree.expandElements(profileExpandedState);
+				}
+			}
+		}));
+
+		this._register(this._connectionManagementService.onConnectionProfileGroupDeleted(async (e) => {
+			if (this._tree instanceof AsyncServerTree) {
+				const parent = <ConnectionProfileGroup>this._tree.getElementById(e.parentId);
+				parent.children = parent.children.filter(c => c.id !== e.id);
+				await this._tree.updateChildren(parent);
+				await this._tree.revealSelectFocusElement(parent);
+			}
+		}));
+
+		this._register(this._connectionManagementService.onConnectionProfileGroupCreated(async (e) => {
+			if (this._tree instanceof AsyncServerTree) {
+				let parent = <ConnectionProfileGroup>this._tree.getElementById(e.parentId);
+				if (!parent) {
+					parent = this._tree.getInput(); // If the parent is not found then add the group to the root.
+				}
+				parent.children.push(e);
+				e.parent = parent;
+				e.parentId = parent.id;
+				await this._tree.updateChildren(parent);
+				await this._tree.revealSelectFocusElement(e);
+			}
+		}));
+
+		this._register(this._connectionManagementService.onConnectionProfileGroupEdited(async (e) => {
+			if (this._tree instanceof AsyncServerTree) {
+				const newParent = <ConnectionProfileGroup>this._tree.getElementById(e.parentId);
+				if (newParent) {
+					newParent.children[newParent.children.findIndex(c => c.id === e.id)] = e;
+					await this._tree.updateChildren(newParent);
+					await this._tree.revealSelectFocusElement(e);
+				}
+			}
+		}));
+
+		this._register(this._connectionManagementService.onConnectionProfileGroupMoved(async (e) => {
+			if (this._tree instanceof AsyncServerTree) {
+				const movedGroup = <ConnectionProfileGroup>e.source;
+				const oldParent = <ConnectionProfileGroup>this._tree.getElementById(e.oldGroupId);
+				const newParent = <ConnectionProfileGroup>this._tree.getElementById(e.newGroupId);
+				// Storing the expanded state of children of the moved group so that they can be expanded after the move.
+				const profileExpandedState = this._tree.getExpandedState(movedGroup);
+				oldParent.children = oldParent.children.filter(c => c.id !== movedGroup.id);
+				await this._tree.updateChildren(oldParent);
+				newParent.children.push(movedGroup);
+				(<ConnectionProfileGroup>movedGroup).parent = newParent;
+				(<ConnectionProfileGroup>movedGroup).parentId = newParent.id;
+				await this._tree.updateChildren(newParent);
+				await this._tree.revealSelectFocusElement(movedGroup);
+				// Expanding the previously expanded children of the moved group after the move.
+				this._tree.expandElements(profileExpandedState);
+			}
+		}));
+
 		return new Promise<void>(async (resolve, reject) => {
 			await this.refreshTree();
 			const root = <ConnectionProfileGroup>this._tree!.getInput();
@@ -216,9 +377,10 @@ export class ServerTreeView extends Disposable implements IServerTreeView {
 			const expandGroups = this._configurationService.getValue<{ autoExpand: boolean }>(SERVER_GROUP_CONFIG).autoExpand;
 			if (expandGroups) {
 				if (this._tree instanceof AsyncServerTree) {
-					await Promise.all(ConnectionProfileGroup.getSubgroups(root).map(subgroup => {
-						return this._tree!.expand(subgroup);
-					}));
+					const subGroups = ConnectionProfileGroup.getSubgroups(root);
+					for (let group of subGroups) {
+						await this._tree.expand(group);
+					}
 				} else {
 					await this._tree!.expandAll(ConnectionProfileGroup.getSubgroups(root));
 				}
@@ -246,26 +408,6 @@ export class ServerTreeView extends Disposable implements IServerTreeView {
 		}

 		if (this._tree instanceof AsyncServerTree) {
-			// When new connection groups are added the event is fired with undefined so
-			// we still want to refresh the tree in that case to pick up the changes
-			await this.refreshTree();
-			if (newProfile) {
-				const currentSelections = this._tree.getSelection();
-				const currentSelectedElement = currentSelections && currentSelections.length >= 1 ? currentSelections[0] : undefined;
-				const newProfileIsSelected = currentSelectedElement && currentSelectedElement.id === newProfile.id;
-				// Clear any other selected elements first
-				if (currentSelectedElement && !newProfileIsSelected) {
-					this._tree.setSelection([]);
-				}
-				const newConnectionProfile = this.getConnectionInTreeInput(newProfile.id);
-				if (newConnectionProfile) {
-					// Re-render to update the connection status badge
-					this._tree.rerender(newConnectionProfile);
-					this._tree.setSelection([newConnectionProfile]);
-					this._tree.expand(newConnectionProfile);
-				}
-			}
-
 		} else {
 			if (newProfile) {
 				const groups = this._connectionManagementService.getConnectionGroups();
@@ -324,17 +466,39 @@ export class ServerTreeView extends Disposable implements IServerTreeView {
 		}
 	}

+	private async disconnectConnection(profile: ConnectionProfile, deleteConnFromConnectionService: boolean = false): Promise<void> {
+		if (this._tree instanceof AsyncServerTree) {
+			if (deleteConnFromConnectionService) {
+				await this._connectionManagementService.deleteConnection(profile);
+			}
+			const connectionProfile = this.getConnectionInTreeInput(profile.id);
+			// Delete the node from the tree
+			await this._objectExplorerService.deleteObjectExplorerNode(connectionProfile);
+			// Collapse the node
+			await this._tree.collapse(connectionProfile);
+			// Rerendering node to turn the badge red
+			await this._tree.rerender(connectionProfile);
+			connectionProfile.isDisconnecting = true;
+			await this._tree.updateChildren(connectionProfile);
+			connectionProfile.isDisconnecting = false;
+			// Make the connection dirty so that the next expansion will refresh the node
+			await this._tree.makeElementDirty(connectionProfile);
+			await this._tree.revealSelectFocusElement(connectionProfile);
+		}
+	}
+
 	private async onObjectExplorerSessionCreated(connection: IConnectionProfile): Promise<void> {
 		const element = this.getConnectionInTreeInput(connection.id);
 		if (element) {
 			if (this._tree instanceof AsyncServerTree) {
-				this._tree.rerender(element);
+				await this._tree.rerender(element);
+				await this._tree.revealSelectFocusElement(element);
 			} else {
 				await this._tree!.refresh(element);
+				await this._tree!.expand(element);
+				await this._tree!.reveal(element, 0.5);
+				this._treeSelectionHandler.onTreeActionStateChange(false);
 			}
-			await this._tree!.expand(element);
-			await this._tree!.reveal(element, 0.5);
-			this._treeSelectionHandler.onTreeActionStateChange(false);
 		}
 	}

@@ -354,7 +518,8 @@ export class ServerTreeView extends Disposable implements IServerTreeView {
 					// Collapse the node before refreshing so the refresh doesn't try to fetch
 					// the children again (which causes it to try and connect)
 					this._tree.collapse(conn);
-					await this.refreshTree();
+					this._tree.rerender(conn);
+					this._tree.makeElementDirty(conn);
 				} else {
 					await this._tree!.collapse(conn);
 					return this._tree!.refresh(conn);
@@ -416,7 +581,7 @@ export class ServerTreeView extends Disposable implements IServerTreeView {
 	/**
 	 * Set tree elements based on the view (recent/active)
 	 */
-	public showFilteredTree(view: ServerTreeViewView): void {
+	public async showFilteredTree(view: ServerTreeViewView): Promise<void> {
 		hide(this.messages!);
 		this._viewKey.set(view);
 		const root = TreeUpdateUtils.getTreeInput(this._connectionManagementService);
@@ -430,24 +595,29 @@ export class ServerTreeView extends Disposable implements IServerTreeView {
 			} else {
 				treeInput = filteredResults[0];
 			}
-			this._tree!.setInput(treeInput!).then(async () => {
-				if (isHidden(this.messages!)) {
-					this._tree!.getFocus();
-					if (this._tree instanceof AsyncServerTree) {
-						await Promise.all(ConnectionProfileGroup.getSubgroups(treeInput!).map(subgroup => {
-							this._tree!.expand(subgroup);
-						}));
-					} else {
-						await this._tree!.expandAll(ConnectionProfileGroup.getSubgroups(treeInput!));
+
+			if (this._tree instanceof AsyncServerTree) {
+				await this._tree.setInput(treeInput!);
+				await this._tree.updateChildren(treeInput!);
+				return;
+			}
+			await this._tree.setInput(treeInput!);
+			if (isHidden(this.messages!)) {
+				this._tree.getFocus();
+				if (this._tree instanceof AsyncServerTree) {
+					for (const subgroup of ConnectionProfileGroup.getSubgroups(treeInput)) {
+						await this._tree.expand(subgroup);
 					}
 				} else {
-					if (this._tree instanceof AsyncServerTree) {
-						this._tree.setFocus([]);
-					} else {
-						this._tree!.clearFocus();
-					}
+					await this._tree!.expandAll(ConnectionProfileGroup.getSubgroups(treeInput!));
 				}
-			}, errors.onUnexpectedError);
+			} else {
+				if (this._tree instanceof AsyncServerTree) {
+					this._tree.setFocus([]);
+				} else {
+					this._tree!.clearFocus();
+				}
+			}
 		} else {
 			//no op
 		}
--- a/src/sql/workbench/services/connection/browser/connectionManagementService.ts
+++ b/src/sql/workbench/services/connection/browser/connectionManagementService.ts
@@ -8,7 +8,7 @@ import * as WorkbenchUtils from 'sql/workbench/common/sqlWorkbenchUtils';
 import {
 	IConnectionManagementService, INewConnectionParams,
 	ConnectionType, IConnectableInput, IConnectionCompletionOptions, IConnectionCallbacks,
-	IConnectionParams, IConnectionResult, RunQueryOnConnectionMode
+	IConnectionParams, IConnectionResult, RunQueryOnConnectionMode, ConnectionElementMovedParams, ConnectionProfileEditedParams
 } from 'sql/platform/connection/common/connectionManagement';
 import { ConnectionStore } from 'sql/platform/connection/common/connectionStore';
 import { ConnectionManagementInfo } from 'sql/platform/connection/common/connectionManagementInfo';
@@ -76,6 +76,17 @@ export class ConnectionManagementService extends Disposable implements IConnecti
 	private _connectionGlobalStatus = new ConnectionGlobalStatus(this._notificationService);
 	private _uriToReconnectPromiseMap: { [uri: string]: Promise<IConnectionResult> } = {};

+	private _onConnectionProfileCreated = new Emitter<ConnectionProfile>();
+	private _onConnectionProfileDeleted = new Emitter<ConnectionProfile>();
+	private _onConnectionProfileEdited = new Emitter<ConnectionProfileEditedParams>();
+	private _onConnectionProfileMoved = new Emitter<ConnectionElementMovedParams>();
+	private _onConnectionProfileConnected = new Emitter<ConnectionProfile>();
+	private _onConnectionProfileDisconnected = new Emitter<ConnectionProfile>();
+	private _onConnectionProfileGroupCreated = new Emitter<ConnectionProfileGroup>();
+	private _onConnectionProfileGroupDeleted = new Emitter<ConnectionProfileGroup>();
+	private _onConnectionProfileGroupEdited = new Emitter<ConnectionProfileGroup>();
+	private _onConnectionProfileGroupMoved = new Emitter<ConnectionElementMovedParams>();
+
 	private _mementoContext: Memento;
 	private _mementoObj: MementoObject;
 	private _connectionStore: ConnectionStore;
@@ -191,6 +202,49 @@ export class ConnectionManagementService extends Disposable implements IConnecti
 		return this._onLanguageFlavorChanged.event;
 	}

+	/**
+	 * Async tree event emitters
+	 */
+	public get onConnectionProfileCreated(): Event<ConnectionProfile> {
+		return this._onConnectionProfileCreated.event;
+	}
+
+	public get onConnectionProfileEdited(): Event<ConnectionProfileEditedParams> {
+		return this._onConnectionProfileEdited.event;
+	}
+
+	public get onConnectionProfileDeleted(): Event<ConnectionProfile> {
+		return this._onConnectionProfileDeleted.event;
+	}
+
+	public get onConnectionProfileMoved(): Event<ConnectionElementMovedParams> {
+		return this._onConnectionProfileMoved.event;
+	}
+
+	public get onConnectionProfileConnected(): Event<ConnectionProfile> {
+		return this._onConnectionProfileConnected.event;
+	}
+
+	public get onConnectionProfileDisconnected(): Event<ConnectionProfile> {
+		return this._onConnectionProfileDisconnected.event;
+	}
+
+	public get onConnectionProfileGroupCreated(): Event<ConnectionProfileGroup> {
+		return this._onConnectionProfileGroupCreated.event;
+	}
+
+	public get onConnectionProfileGroupDeleted(): Event<ConnectionProfileGroup> {
+		return this._onConnectionProfileGroupDeleted.event;
+	}
+
+	public get onConnectionProfileGroupEdited(): Event<ConnectionProfileGroup> {
+		return this._onConnectionProfileGroupEdited.event;
+	}
+
+	public get onConnectionProfileGroupMoved(): Event<ConnectionElementMovedParams> {
+		return this._onConnectionProfileGroupMoved.event;
+	}
+
 	public get providerNameToDisplayNameMap(): { readonly [providerDisplayName: string]: string } {
 		return this._providerNameToDisplayNameMap;
 	}
@@ -539,6 +593,18 @@ export class ConnectionManagementService extends Disposable implements IConnecti

 					await this.saveToSettings(uri, connection, matcher).then(value => {
 						this._onAddConnectionProfile.fire(connection);
+						if (isEdit) {
+							this._onConnectionProfileEdited.fire({
+								oldProfileId: options.params.oldProfileId,
+								profile: <ConnectionProfile>connection
+							});
+						} else {
+							if (options.params === undefined) {
+								this._onConnectionProfileConnected.fire(<ConnectionProfile>connection);
+							} else {
+								this._onConnectionProfileCreated.fire(<ConnectionProfile>connection);
+							}
+						}
 						this.doActionsAfterConnectionComplete(value, options);
 					});
 				} else {
@@ -716,7 +782,22 @@ export class ConnectionManagementService extends Disposable implements IConnecti
 	}

 	public getConnectionGroups(providers?: string[]): ConnectionProfileGroup[] {
-		return this._connectionStore.getConnectionProfileGroups(false, providers);
+		const groups = this._connectionStore.getConnectionProfileGroups(false, providers);
+		return groups;
+	}
+
+	public getConnectionGroupById(id: string): ConnectionProfileGroup | undefined {
+		const groups = this.getConnectionGroups();
+		for (let group of groups) {
+			if (group.id === id) {
+				return group;
+			}
+			const subgroup = ConnectionProfileGroup.getSubgroups(group).find(g => g.id === id);
+			if (subgroup) {
+				return subgroup;
+			}
+		}
+		return undefined;
 	}

 	public getRecentConnections(providers?: string[]): ConnectionProfile[] {
@@ -745,10 +826,15 @@ export class ConnectionManagementService extends Disposable implements IConnecti
 		}
 	}

-	public saveProfileGroup(profile: IConnectionProfileGroup): Promise<string> {
+	public saveProfileGroup(group: IConnectionProfileGroup): Promise<string> {
 		this._telemetryService.sendActionEvent(TelemetryKeys.TelemetryView.Shell, TelemetryKeys.TelemetryAction.AddServerGroup);
-		return this._connectionStore.saveProfileGroup(profile).then(groupId => {
+		return this._connectionStore.saveProfileGroup(group).then(groupId => {
 			this._onAddConnectionProfile.fire(undefined);
+			//Getting id for the new profile group
+			group.id = groupId;
+			const parentGroup = this.getConnectionGroupById(group.parentId);
+			this._onConnectionProfileGroupCreated.fire(ConnectionProfileGroup.createConnectionProfileGroup(group, parentGroup));
+
 			return groupId;
 		});
 	}
@@ -1224,21 +1310,30 @@ export class ConnectionManagementService extends Disposable implements IConnecti
 	public onIntelliSenseCacheComplete(handle: number, connectionUri: string): void {
 	}

-	public changeGroupIdForConnectionGroup(source: ConnectionProfileGroup, target: ConnectionProfileGroup): Promise<void> {
+	public async changeGroupIdForConnectionGroup(source: ConnectionProfileGroup, target: ConnectionProfileGroup): Promise<void> {
 		this._telemetryService.sendActionEvent(TelemetryKeys.TelemetryView.Shell, TelemetryKeys.TelemetryAction.MoveServerConnection);
-		return this._connectionStore.changeGroupIdForConnectionGroup(source, target);
+		await this._connectionStore.changeGroupIdForConnectionGroup(source, target);
+		this._onConnectionProfileGroupMoved.fire({
+			source: source,
+			oldGroupId: source.parentId,
+			newGroupId: target.id
+		});
 	}

-	public changeGroupIdForConnection(source: ConnectionProfile, targetGroupId: string): Promise<void> {
+	public async changeGroupIdForConnection(source: ConnectionProfile, targetGroupId: string): Promise<void> {
+		const oldProfileId = source.groupId;
 		let id = Utils.generateUri(source);
 		this._telemetryService.sendActionEvent(TelemetryKeys.TelemetryView.Shell, TelemetryKeys.TelemetryAction.MoveServerGroup);
-		return this._connectionStore.changeGroupIdForConnection(source, targetGroupId).then(result => {
-			this._onAddConnectionProfile.fire(source);
-			if (id && targetGroupId) {
-				source.groupId = targetGroupId;
-			}
-			// change the connection uri with the new group so that connection does not appear as disconnected in OE.
-			this.changeConnectionUri(Utils.generateUri(source), id);
+		await this._connectionStore.changeGroupIdForConnection(source, targetGroupId)
+		this._onAddConnectionProfile.fire(source);
+		if (id && targetGroupId) {
+			source.groupId = targetGroupId;
+		}
+		this.changeConnectionUri(Utils.generateUri(source), id);
+		this._onConnectionProfileMoved.fire({
+			source: source,
+			oldGroupId: oldProfileId,
+			newGroupId: targetGroupId,
 		});
 	}

@@ -1508,6 +1603,7 @@ export class ConnectionManagementService extends Disposable implements IConnecti
 	public editGroup(group: ConnectionProfileGroup): Promise<void> {
 		return this._connectionStore.editGroup(group).then(groupId => {
 			this._onAddConnectionProfile.fire(undefined);
+			this._onConnectionProfileGroupEdited.fire(group);
 		});
 	}

@@ -1515,7 +1611,7 @@ export class ConnectionManagementService extends Disposable implements IConnecti
 	 * Deletes a connection from registered servers.
 	 * Disconnects a connection before removing from settings.
 	 */
-	public deleteConnection(connection: ConnectionProfile): Promise<boolean> {
+	public async deleteConnection(connection: ConnectionProfile): Promise<boolean> {
 		this._telemetryService.createActionEvent(TelemetryKeys.TelemetryView.Shell, TelemetryKeys.TelemetryAction.DeleteConnection)
 			.withAdditionalProperties({
 				provider: connection.providerName
@@ -1528,6 +1624,7 @@ export class ConnectionManagementService extends Disposable implements IConnecti
 					// Remove profile from configuration
 					return this._connectionStore.deleteConnectionFromConfiguration(connection).then(() => {
 						this._onDeleteConnectionProfile.fire();
+						this._onConnectionProfileDeleted.fire(connection);
 						return true;
 					});

@@ -1539,6 +1636,7 @@ export class ConnectionManagementService extends Disposable implements IConnecti
 			// Remove disconnected profile from settings
 			return this._connectionStore.deleteConnectionFromConfiguration(connection).then(() => {
 				this._onDeleteConnectionProfile.fire();
+				this._onConnectionProfileDeleted.fire(connection);
 				return true;
 			});
 		}
@@ -1567,6 +1665,7 @@ export class ConnectionManagementService extends Disposable implements IConnecti
 			// Remove profiles and groups from config
 			return this._connectionStore.deleteGroupFromConfiguration(group).then(() => {
 				this._onDeleteConnectionProfile.fire();
+				this._onConnectionProfileGroupDeleted.fire(group);
 				return true;
 			});
 		}).catch(() => false);
--- a/src/sql/workbench/services/objectExplorer/browser/asyncServerTree.ts
+++ b/src/sql/workbench/services/objectExplorer/browser/asyncServerTree.ts
@@ -4,18 +4,196 @@
 *--------------------------------------------------------------------------------------------*/

 import { ConnectionProfileGroup } from 'sql/platform/connection/common/connectionProfileGroup';
-import { WorkbenchAsyncDataTree } from 'vs/platform/list/browser/listService';
+import { IListService, IWorkbenchAsyncDataTreeOptions, WorkbenchAsyncDataTree } from 'vs/platform/list/browser/listService';
 import { FuzzyScore } from 'vs/base/common/filters';
 import { TreeNode } from 'sql/workbench/services/objectExplorer/common/treeNode';
 import { ConnectionProfile } from 'sql/platform/connection/common/connectionProfile';
-import { IAsyncDataTreeViewState } from 'vs/base/browser/ui/tree/asyncDataTree';
+import { IAsyncDataTreeNode, IAsyncDataTreeUpdateChildrenOptions, IAsyncDataTreeViewState } from 'vs/base/browser/ui/tree/asyncDataTree';
+import { IContextKeyService } from 'vs/platform/contextkey/common/contextkey';
+import { IThemeService } from 'vs/platform/theme/common/themeService';
+import { IConfigurationService } from 'vs/platform/configuration/common/configuration';
+import { IKeybindingService } from 'vs/platform/keybinding/common/keybinding';
+import { IAccessibilityService } from 'vs/platform/accessibility/common/accessibility';
+import { IListVirtualDelegate } from 'vs/base/browser/ui/list/list';
+import { IAsyncDataSource, ITreeRenderer } from 'vs/base/browser/ui/tree/tree';
+import { StandardKeyboardEvent } from 'vs/base/browser/keyboardEvent';
+import { KeyCode } from 'vs/base/common/keyCodes';

 export class AsyncServerTree extends WorkbenchAsyncDataTree<ConnectionProfileGroup, ServerTreeElement, FuzzyScore> {
+
+	constructor(
+		user: string,
+		container: HTMLElement,
+		delegate: IListVirtualDelegate<ServerTreeElement>,
+		renderers: ITreeRenderer<ServerTreeElement, FuzzyScore, any>[],
+		dataSource: IAsyncDataSource<ConnectionProfileGroup, ServerTreeElement>,
+		options: IWorkbenchAsyncDataTreeOptions<ServerTreeElement, FuzzyScore>,
+		@IContextKeyService contextKeyService: IContextKeyService,
+		@IListService listService: IListService,
+		@IThemeService themeService: IThemeService,
+		@IConfigurationService configurationService: IConfigurationService,
+		@IKeybindingService keybindingService: IKeybindingService,
+		@IAccessibilityService accessibilityService: IAccessibilityService,
+	) {
+		super(
+			user, container, delegate,
+			renderers, dataSource, options,
+			contextKeyService, listService,
+			themeService, configurationService, keybindingService, accessibilityService);
+
+		// Adding support for expand/collapse on enter/space
+		this.onKeyDown(e => {
+			const standardKeyboardEvent = new StandardKeyboardEvent(e);
+			if (standardKeyboardEvent.keyCode === KeyCode.Enter || standardKeyboardEvent.keyCode === KeyCode.Space) {
+				const selectedElement = this.getSelection()[0];
+				if (selectedElement) {
+					if (this.isCollapsed(selectedElement)) {
+						this.expand(selectedElement);
+					} else {
+						this.collapse(selectedElement);
+					}
+				}
+			}
+		})
+	}
+
+	// Overriding the setInput method to dispose the original input when a new input is set
 	override async setInput(input: ConnectionProfileGroup, viewState?: IAsyncDataTreeViewState): Promise<void> {
 		const originalInput = this.getInput();
 		await super.setInput(input, viewState);
 		originalInput?.dispose();
 	}
+
+	/**
+	 * The original implementation of getDataNode compares refrences of the elements to find the node.
+	 * This is not working for our case as we are creating new elements everytime we refresh the tree.
+	 * 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> {
+		try {
+			const node = super.getDataNode(element);
+			return node;
+		} catch (e) {
+			let node = this.getDataNodeById(element?.id);
+			if (node) {
+				return node;
+			}
+			throw e;
+		}
+	}
+
+	/**
+	 * Gets the element by id in the tree
+	 */
+	public getElementById(id: string): ServerTreeElement | undefined {
+		if (this.getInput().id === id) {
+			return this.getInput();
+		}
+		return this.getDataNodeById(id)?.element;
+	}
+
+	/**
+	 * Get the list of expanded elements in the tree
+	 */
+	public getExpandedState(element: ServerTreeElement): ServerTreeElement[] {
+		const node = this.getDataNode(element);
+		const stack = [node];
+		const expanded: ServerTreeElement[] = [];
+		while (stack.length > 0) {
+			const node = stack.pop();
+			if (node) {
+				if (!this.isCollapsed(node.element)) {
+					expanded.push(node.element);
+					if (node.children) {
+						node.children.forEach(child => stack.push(child));
+					}
+				}
+			}
+		}
+		return expanded;
+	}
+
+	private getDataNodeById(id: string): IAsyncDataTreeNode<ConnectionProfileGroup, ServerTreeElement> | undefined {
+		let node = undefined;
+		this.nodes.forEach((v, k) => {
+			if (id === v?.id) {
+				node = v;
+			}
+		});
+		return node;
+	}
+
+	public override async updateChildren(element?: ServerTreeElement, recursive?: boolean, rerender?: boolean, options?: IAsyncDataTreeUpdateChildrenOptions<ServerTreeElement>): Promise<void> {
+		const viewState = this.getViewState();
+		const expandedElementIds = viewState?.expanded;
+		await super.updateChildren(element, recursive, rerender, options);
+		if (expandedElementIds) {
+			for (let i = 0; i <= expandedElementIds.length; i++) {
+				const id = expandedElementIds[i];
+				const node = this.getDataNodeById(id);
+				if (node) {
+					await this.expand(node.element);
+				}
+			}
+		}
+	}
+
+	public async expandElements(elements: ServerTreeElement[]): Promise<void> {
+		for (let element of elements) {
+			const id = element.id;
+			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
+	 */
+	public async makeElementDirty(element: ServerTreeElement) {
+		this.getDataNode(element).stale = true;
+	}
+
+	public async revealSelectFocusElement(element: ServerTreeElement) {
+		await this.reveal(element);
+		await this.setSelection([element]);
+		this.setFocus([element]);
+	}
 }

 export type ServerTreeElement = ConnectionProfile | ConnectionProfileGroup | TreeNode;
--- a/src/sql/workbench/services/objectExplorer/browser/dragAndDropController.ts
+++ b/src/sql/workbench/services/objectExplorer/browser/dragAndDropController.ts
@@ -167,7 +167,6 @@ export class ServerTreeDragAndDrop implements IDragAndDrop {
 			// to avoid creating a circular structure.
 			canDragOver = source.id !== targetConnectionProfileGroup.id && !source.isAncestorOf(targetConnectionProfileGroup);
 		}
-
 		return canDragOver;
 	}

@@ -179,33 +178,41 @@ export class ServerTreeDragAndDrop implements IDragAndDrop {
 	public onDragOver(tree: AsyncServerTree | ITree, data: IDragAndDropData, targetElement: any, originalEvent: DragMouseEvent): IDragOverReaction {
 		let canDragOver: boolean = true;

-		if (targetElement instanceof ConnectionProfile || targetElement instanceof ConnectionProfileGroup) {
-			let targetConnectionProfileGroup = this.getTargetGroup(targetElement);
-			// Verify if the connection can be moved to the target group
-			const source = data.getData()[0];
-			if (source instanceof ConnectionProfile) {
-				if (!this._connectionManagementService.canChangeConnectionConfig(source, targetConnectionProfileGroup.id!)) {
-					canDragOver = false;
-				}
-			} else if (source instanceof ConnectionProfileGroup) {
-				// Dropping a group to itself or its descendants nodes is not allowed
-				// to avoid creating a circular structure.
-				canDragOver = source.id !== targetElement.id && !source.isAncestorOf(targetElement);
+		const source = data.getData()[0];
+		if (source instanceof ConnectionProfileGroup) {
+			if (targetElement instanceof ConnectionProfileGroup) {
+				// If target group is parent of the source connection, then don't allow drag over
+				canDragOver = this.canDragToConnectionProfileGroup(source, targetElement);
+			} else if (targetElement instanceof ConnectionProfile) {
+				canDragOver = source.parentId !== targetElement.groupId;
+			} else if (targetElement instanceof TreeNode) {
+				const treeNodeParentGroupId = this.getTreeNodeParentGroup(targetElement).id;
+				canDragOver = source.parentId !== treeNodeParentGroupId && source.id !== treeNodeParentGroupId;
 			}
-		} else {
-			canDragOver = true;
+		} else if (source instanceof ConnectionProfile) {
+			if (targetElement instanceof ConnectionProfileGroup) {
+				canDragOver = this.canDragToConnectionProfileGroup(source, targetElement);
+			} else if (targetElement instanceof ConnectionProfile) {
+				canDragOver = source.groupId !== targetElement.groupId &&
+					this._connectionManagementService.canChangeConnectionConfig(source, targetElement.groupId);
+			} else if (targetElement instanceof TreeNode) {
+				canDragOver = source.groupId !== this.getTreeNodeParentGroup(targetElement).id;
+			}
+		} else if (source instanceof TreeNode) {
+			canDragOver = false;
 		}
-
 		if (canDragOver) {
 			if (targetElement instanceof ConnectionProfile) {
 				const isConnected = this._connectionManagementService.isProfileConnected(targetElement);
 				// Don't auto-expand disconnected connections - doing so will try to connect the connection
 				// when expanded which is not something we want to support currently
 				return DRAG_OVER_ACCEPT_BUBBLE_DOWN(isConnected);
+			} else if (targetElement instanceof ConnectionProfileGroup) {
+				return DRAG_OVER_ACCEPT_BUBBLE_DOWN(true);
+			} else {
+				// Don't auto-expand treeNodes as we don't support drag and drop on them
+				return DRAG_OVER_ACCEPT_BUBBLE_DOWN(false);
 			}
-			// Auto-expand other elements (groups, tree nodes) so their children can be
-			// exposed for further dragging
-			return DRAG_OVER_ACCEPT_BUBBLE_DOWN(true);
 		} else {
 			return DRAG_OVER_REJECT;
 		}
@@ -224,23 +231,32 @@ export class ServerTreeDragAndDrop implements IDragAndDrop {
 			let oldParent: ConnectionProfileGroup = source.getParent();
 			const self = this;
 			if (this.isDropAllowed(targetConnectionProfileGroup, oldParent, source)) {
-
-				if (source instanceof ConnectionProfile) {
-					// Change group id of profile
-					this._connectionManagementService.changeGroupIdForConnection(source, targetConnectionProfileGroup.id!).then(() => {
-						if (tree) {
-							TreeUpdateUtils.registeredServerUpdate(tree, self._connectionManagementService, targetConnectionProfileGroup);
+				if (tree instanceof AsyncServerTree) {
+					if (oldParent && source && targetConnectionProfileGroup) {
+						if (source instanceof ConnectionProfileGroup) {
+							this._connectionManagementService.changeGroupIdForConnectionGroup(source, targetConnectionProfileGroup);
+						} else if (source instanceof ConnectionProfile) {
+							this._connectionManagementService.changeGroupIdForConnection(source, targetConnectionProfileGroup.id!);
 						}
+					}

-					});
-				} else if (source instanceof ConnectionProfileGroup) {
-					// Change parent id of group
-					this._connectionManagementService.changeGroupIdForConnectionGroup(source, targetConnectionProfileGroup).then(() => {
-						if (tree) {
-							TreeUpdateUtils.registeredServerUpdate(tree, self._connectionManagementService);
-						}
+				} else {
+					if (source instanceof ConnectionProfile) {
+						// Change group id of profile
+						this._connectionManagementService.changeGroupIdForConnection(source, targetConnectionProfileGroup.id!).then(async () => {
+							if (tree) {
+								TreeUpdateUtils.registeredServerUpdate(tree, self._connectionManagementService, targetConnectionProfileGroup);
+							}

-					});
+						});
+					} else if (source instanceof ConnectionProfileGroup) {
+						// Change parent id of group
+						this._connectionManagementService.changeGroupIdForConnectionGroup(source, targetConnectionProfileGroup).then(async () => {
+							if (tree) {
+								TreeUpdateUtils.registeredServerUpdate(tree, self._connectionManagementService);
+							}
+						});
+					}
 				}
 			}
 		}
@@ -250,13 +266,17 @@ export class ServerTreeDragAndDrop implements IDragAndDrop {
 		TreeUpdateUtils.isInDragAndDrop = false;
 	}

-	private getTargetGroup(targetElement: ConnectionProfileGroup | ConnectionProfile): ConnectionProfileGroup {
+	private getTargetGroup(targetElement: ConnectionProfileGroup | ConnectionProfile | TreeNode): ConnectionProfileGroup {
 		let targetConnectionProfileGroup: ConnectionProfileGroup;
 		if (targetElement instanceof ConnectionProfile) {
 			targetConnectionProfileGroup = targetElement.getParent()!;
-		}
-		else {
+		} else if (targetElement instanceof ConnectionProfileGroup) {
 			targetConnectionProfileGroup = targetElement;
+		} else if (targetElement instanceof TreeNode) {
+			targetConnectionProfileGroup = this.getTreeNodeParentGroup(targetElement);
+			if (!targetConnectionProfileGroup) {
+				throw new Error('Cannot find parent for the node');
+			}
 		}

 		return targetConnectionProfileGroup;
@@ -271,6 +291,20 @@ export class ServerTreeDragAndDrop implements IDragAndDrop {
 		let isUnsavedDrag = source && (source instanceof ConnectionProfileGroup) && (source.id === UNSAVED_GROUP_ID);
 		return (!isDropToSameLevel && !isDropToItself && !isUnsavedDrag);
 	}
+
+	private getTreeNodeParentGroup(element: TreeNode): ConnectionProfileGroup | undefined {
+		let treeNode = element;
+		while (!treeNode?.connection) {
+			treeNode = treeNode.parent;
+		}
+		if (treeNode) {
+			const groupId = treeNode.connection.groupId;
+			if (groupId) {
+				return this._connectionManagementService.getConnectionGroupById(groupId);
+			}
+		}
+		return undefined;
+	}
 }

 /**