From d9220c809c0b2e32c4a6a60dfeca442aa233901b Mon Sep 17 00:00:00 2001
From: Aasim Khan <aasimkhan30@gmail.com>
Date: Tue, 16 May 2023 09:20:51 -0700
Subject: [PATCH] Adding collapse all to OE (#23132)

---
 .../browser/connection.contribution.ts        | 26 +++++++++++++++++--
 .../objectExplorer/browser/serverTreeView.ts  |  8 ++++++
 .../browser/objectExplorerService.ts          |  1 +
 3 files changed, 33 insertions(+), 2 deletions(-)

diff --git a/src/sql/workbench/contrib/connection/browser/connection.contribution.ts b/src/sql/workbench/contrib/connection/browser/connection.contribution.ts
index 7b2d14191f..420d613997 100644
--- a/src/sql/workbench/contrib/connection/browser/connection.contribution.ts
+++ b/src/sql/workbench/contrib/connection/browser/connection.contribution.ts
@@ -8,7 +8,7 @@ import { Registry } from 'vs/platform/registry/common/platform';
 import { ClearRecentConnectionsAction, GetCurrentConnectionStringAction } from 'sql/workbench/services/connection/browser/connectionActions';
 import * as azdata from 'azdata';
 import { IWorkbenchActionRegistry, Extensions } from 'vs/workbench/common/actions';
-import { MenuId, MenuRegistry, SyncActionDescriptor } from 'vs/platform/actions/common/actions';
+import { Action2, MenuId, MenuRegistry, SyncActionDescriptor, registerAction2 } from 'vs/platform/actions/common/actions';
 import { localize } from 'vs/nls';
 import { ConnectionStatusbarItem } from 'sql/workbench/contrib/connection/browser/connectionStatus';
 import { CommandsRegistry } from 'vs/platform/commands/common/commands';
@@ -28,8 +28,10 @@ const workbenchRegistry = Registry.as<IWorkbenchContributionsRegistry>(Workbench
 workbenchRegistry.registerWorkbenchContribution(ConnectionStatusbarItem, LifecyclePhase.Restored);
 
 import 'sql/workbench/contrib/connection/common/connectionTreeProviderExentionPoint';
-import { ServerTreeViewView } from 'sql/workbench/services/objectExplorer/browser/objectExplorerService';
+import { IObjectExplorerService, ServerTreeViewView } from 'sql/workbench/services/objectExplorer/browser/objectExplorerService';
 import { AuthenticationType } from 'sql/platform/connection/common/constants';
+import { Codicon } from 'vs/base/common/codicons';
+import { ServicesAccessor } from 'vs/editor/browser/editorExtensions';
 
 // Connection Dashboard registration
 
@@ -111,6 +113,26 @@ MenuRegistry.appendMenuItem(MenuId.ViewTitle, {
 	when: ContextKeyEqualsExpr.create('view', ConnectionViewletPanel.ID),
 });
 
+registerAction2(class extends Action2 {
+	constructor() {
+		super({
+			id: 'registeredServers.collapseAll',
+			title: localize('registeredServers.collapseAll', "Collapse All Connections"),
+			menu: {
+				id: MenuId.ViewTitle,
+				when: ContextKeyEqualsExpr.create('view', ConnectionViewletPanel.ID),
+				group: 'navigation',
+				order: Number.MAX_SAFE_INTEGER - 1,
+			},
+			icon: Codicon.collapseAll
+		});
+	}
+	async run(accessor: ServicesAccessor): Promise<void> {
+		const objectExplorerService = accessor.get(IObjectExplorerService);
+		await objectExplorerService.getServerTreeView().collapseAllConnections();
+	}
+});
+
 CommandsRegistry.registerCommand('azdata.connect',
 	function (accessor, args: {
 		serverName: string,
diff --git a/src/sql/workbench/contrib/objectExplorer/browser/serverTreeView.ts b/src/sql/workbench/contrib/objectExplorer/browser/serverTreeView.ts
index 2ac8e65896..1045c50608 100644
--- a/src/sql/workbench/contrib/objectExplorer/browser/serverTreeView.ts
+++ b/src/sql/workbench/contrib/objectExplorer/browser/serverTreeView.ts
@@ -952,4 +952,12 @@ export class ServerTreeView extends Disposable implements IServerTreeView {
 		}
 		return actionContext;
 	}
+
+	public collapseAllConnections(): void {
+		const root = TreeUpdateUtils.getTreeInput(this._connectionManagementService)!;
+		const connections = ConnectionProfileGroup.getConnectionsInGroup(root);
+		connections.forEach(con => {
+			this._tree!.collapse(con, true);
+		});
+	}
 }
diff --git a/src/sql/workbench/services/objectExplorer/browser/objectExplorerService.ts b/src/sql/workbench/services/objectExplorer/browser/objectExplorerService.ts
index 31a49c9faa..6c61fa8450 100644
--- a/src/sql/workbench/services/objectExplorer/browser/objectExplorerService.ts
+++ b/src/sql/workbench/services/objectExplorer/browser/objectExplorerService.ts
@@ -57,6 +57,7 @@ export interface IServerTreeView {
 	showFilteredTree(view: ServerTreeViewView): void;
 	filterElementChildren(node: TreeNode): Promise<void>;
 	getActionContext(element: ServerTreeElement): any;
+	collapseAllConnections(): void;
 	view: ServerTreeViewView;
 }