Fixing new connections not expanding in object explorer. (#22307) (#22325)

* Fixing new connection node not expanding in OE

* Fixing new connections not expanding and fixing expand request not resolving because of some provider error.

* Fixing test

* Adding a setting for node expansion timeout

* Not saving when loading tree based connections

* Adding some logs

* Removing special casing for mssql provider

* Missing providers

* Adding user toast for node expansion timeout

* Adding notification service to test

* Fixing node type for mssql

* remove polling

* Fixing onNodeStatus

* Fixing stuff

* consolidating functions

* Consolidating resolve logic

* removing extra try catch

* code cleanup

* adding size checks

* Removing commented code

* Ignoring errors from other sessions and nodepaths.

src/sql/workbench/contrib/objectExplorer/browser/serverTreeView.ts

@@ -293,8 +293,9 @@ export class ServerTreeView extends Disposable implements IServerTreeView {
 			}
 			await this.refreshTree();
 			if (newProfile && !newProfileIsSelected) {
-				await this._tree!.reveal(newProfile);
+				await this._tree.reveal(newProfile);
-				this._tree.setFocus(newProfile);
+				await this._tree.select(newProfile);
+				await this._tree.expand(newProfile);
 			}
 		}
 

src/sql/workbench/contrib/objectExplorer/common/serverGroup.contribution.ts

@@ -50,6 +50,12 @@ const serverTreeConfig: IConfigurationNode = {
 			'type': 'boolean',
 			'default': false,
 			'description': localize('serverTree.useAsyncServerTree', "(Preview) Use the new async server tree for the Servers view and Connection Dialog with support for new features such as dynamic node filtering.")
+		},
+		'serverTree.nodeExpansionTimeout': {
+			'type': 'number',
+			'default': '45',
+			'description': localize('serverTree.nodeExpansionTimeout', "The timeout in seconds for expanding a node in the Servers view"),
+			'minimum': 1
 		}
 	}
 };

src/sql/workbench/contrib/objectExplorer/test/browser/serverTreeView.test.ts

@@ -122,17 +122,17 @@ suite('ServerTreeView onAddConnectionProfile handler tests', () => {
 
 
 
-	test('onAddConnectionProfile handler focuses the new profile when no profile is already selected', async () => {
+	test('onAddConnectionProfile handler selects the new profile when no profile is already selected', async () => {
 		let newProfile = <IConnectionProfile>{
 			id: 'test_connection'
 		};
 		await runAddConnectionProfileHandler(undefined, newProfile);
 		mockRefreshTreeMethod.verify(x => x(), TypeMoq.Times.once());
 		mockTree.verify(x => x.clearSelection(), TypeMoq.Times.never());
-		mockTree.verify(x => x.setFocus(TypeMoq.It.is(profile => profile === newProfile)), TypeMoq.Times.once());
+		mockTree.verify(x => x.select(TypeMoq.It.is(profile => profile === newProfile)), TypeMoq.Times.once());
 	});
 
-	test('onAddConnectionProfile handler focuses the new profile when a different profile is already selected', async () => {
+	test('onAddConnectionProfile handler selects the new profile when a different profile is already selected', async () => {
 		let oldProfile = <IConnectionProfile>{
 			id: 'old_connection'
 		};
@@ -142,7 +142,7 @@ suite('ServerTreeView onAddConnectionProfile handler tests', () => {
 		await runAddConnectionProfileHandler(oldProfile, newProfile);
 		mockRefreshTreeMethod.verify(x => x(), TypeMoq.Times.once());
 		mockTree.verify(x => x.clearSelection(), TypeMoq.Times.once());
-		mockTree.verify(x => x.setFocus(TypeMoq.It.is(profile => profile === newProfile)), TypeMoq.Times.once());
+		mockTree.verify(x => x.select(TypeMoq.It.is(profile => profile === newProfile)), TypeMoq.Times.once());
 	});
 
 	test('onAddConnectionProfile handler does not clear the selection when the new profile is already selected', async () => {

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

@@ -24,6 +24,8 @@ import { ITree } from 'sql/base/parts/tree/browser/tree';
 import { AsyncServerTree, ServerTreeElement } from 'sql/workbench/services/objectExplorer/browser/asyncServerTree';
 import { mssqlProviderName } from 'sql/platform/connection/common/constants';
 import { ObjectExplorerRequestStatus } from 'sql/workbench/services/objectExplorer/browser/treeSelectionHandler';
+import { IConfigurationService } from 'vs/platform/configuration/common/configuration';
+import { INotificationService } from 'vs/platform/notification/common/notification';
 
 export const SERVICE_ID = 'ObjectExplorerService';
 
@@ -170,12 +172,15 @@ export class ObjectExplorerService implements IObjectExplorerService {
 	private _serverTreeView?: IServerTreeView;
 
 	private _onSelectionOrFocusChange: Emitter<void>;
+	private _onNodeExpandedError: Emitter<NodeExpandInfoWithProviderId> = new Emitter<NodeExpandInfoWithProviderId>();
 
 	constructor(
 		@IConnectionManagementService private _connectionManagementService: IConnectionManagementService,
 		@IAdsTelemetryService private _telemetryService: IAdsTelemetryService,
 		@ICapabilitiesService private _capabilitiesService: ICapabilitiesService,
-		@ILogService private logService: ILogService
+		@ILogService private logService: ILogService,
+		@IConfigurationService private _configurationService: IConfigurationService,
+		@INotificationService private _notificationService: INotificationService
 	) {
 		this._onUpdateObjectExplorerNodes = new Emitter<ObjectExplorerNodeEventArgs>();
 		this._activeObjectExplorerNodes = {};
@@ -248,9 +253,11 @@ export class ObjectExplorerService implements IObjectExplorerService {
 				}
 			} else {
 				this.logService.warn(`Cannot find node status for session: ${expandResponse.sessionId} and node path: ${expandResponse.nodePath}`);
+				this._onNodeExpandedError.fire(expandResponse);
 			}
 		} else {
 			this.logService.warn(`Cannot find session ${expandResponse.sessionId} for node path: ${expandResponse.nodePath}`);
+			this._onNodeExpandedError.fire(expandResponse);
 		}
 	}
 
@@ -415,16 +422,53 @@ export class ObjectExplorerService implements IObjectExplorerService {
 						nodeProviders = nodeProviders.sort((a, b) => a.group!.toLowerCase().localeCompare(b.group!.toLowerCase()));
 						allProviders.push(...nodeProviders);
 					}
+					this.logService.trace(`${session.sessionId}: got providers for node expansion: ${allProviders.map(p => p.providerId).join(', ')}`);
+
+					const resolveExpansion = () => {
+						resolve(self.mergeResults(allProviders, resultMap, node.nodePath));
+						// Have to delete it after get all responses otherwise couldn't find session for not the first response
+						clearTimeout(expansionTimeout);
+						if (newRequest) {
+							delete self._sessions[session.sessionId!].nodes[node.nodePath];
+							this.logService.trace(`Deleted node ${node.nodePath} from session ${session.sessionId}`);
+						}
+					}
+
+					// Incase node status not found.
+					this._onNodeExpandedError.event(e => {
+						//Ignore errors from other sessions and nodepaths.
+						if (e.sessionId === session.sessionId && e.nodePath === node.nodePath) {
+							resultMap.set(e.providerId, e);
+							// When get all responses from all providers, merge results
+							if (resultMap.size === allProviders.length) {
+								resolveExpansion();
+							}
+						}
+					});
+
+					const expansionTimeoutValueSec = this._configurationService.getValue<number>('serverTree.nodeExpansionTimeout');
+					const expansionTimeout = setTimeout(() => {
+						/**
+						 * If we don't get a response back from all the providers in specified expansion timeout seconds then we assume
+						 * it's not going to respond and resolve the promise with the results we have so far
+						 */
+						if (resultMap.size !== allProviders.length) {
+							const missingProviders = allProviders.filter(p => !resultMap.has(p.providerId));
+							this.logService.warn(`${session.sessionId}: Node expansion timed out for node ${node.nodePath} for providers ${missingProviders.map(p => p.providerId).join(', ')}`);
+							this._notificationService.error(nls.localize('nodeExpansionTimeout', "Node expansion timed out for node {0} for providers {1}", node.nodePath, missingProviders.map(p => p.providerId).join(', ')));
+						}
+						resolveExpansion();
+					}, expansionTimeoutValueSec * 1000);
 
 					self._sessions[session.sessionId!].nodes[node.nodePath].expandEmitter.event((expandResult: NodeExpandInfoWithProviderId) => {
 						if (expandResult && expandResult.providerId) {
+							this.logService.trace(`${session.sessionId}: Received expand result for node ${node.nodePath} from provider ${expandResult.providerId}`);
 							resultMap.set(expandResult.providerId, expandResult);
 							// If we got an error result back then send error our error event
-							// We only do this for the MSSQL provider
+							if (expandResult.errorMessage) {
-							if (expandResult.errorMessage && expandResult.providerId === mssqlProviderName) {
 								const errorType = expandResult.errorMessage.indexOf('Object Explorer task didn\'t complete') !== -1 ? 'Timeout' : 'Other';
-								// For folders send the actual name of the folder (since the nodeTypeId isn't useful in this case and the names are controlled by us)
+								// For folders send the actual name of the folder for the MSSQL provider (since the nodeTypeId isn't useful in this case and the names are controlled by us)
-								const nodeType = node.nodeTypeId === NodeType.Folder ? node.label : node.nodeTypeId;
+								const nodeType = expandResult.providerId === mssqlProviderName && node.nodeTypeId === NodeType.Folder ? node.label : node.nodeTypeId;
 								this._telemetryService.createErrorEvent(TelemetryKeys.TelemetryView.Shell, TelemetryKeys.TelemetryError.ObjectExplorerExpandError, undefined, errorType)
 									.withAdditionalProperties({
 										nodeType,
@@ -438,13 +482,7 @@ export class ObjectExplorerService implements IObjectExplorerService {
 
 						// When get all responses from all providers, merge results
 						if (resultMap.size === allProviders.length) {
-							resolve(self.mergeResults(allProviders, resultMap, node.nodePath));
+							resolveExpansion();
-
-							// Have to delete it after get all responses otherwise couldn't find session for not the first response
-							if (newRequest) {
-								delete self._sessions[session.sessionId!].nodes[node.nodePath];
-								this.logService.trace(`Deleted node ${node.nodePath} from session ${session.sessionId}`);
-							}
 						}
 					});
 					if (newRequest) {
@@ -463,6 +501,9 @@ export class ObjectExplorerService implements IObjectExplorerService {
 										sessionId: session.sessionId
 									};
 									resultMap.set(provider.providerId, emptyResult);
+									if (resultMap.size === allProviders.length) {
+										resolveExpansion();
+									}
 								}
 							}, error => {
 								reject(error);

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

@@ -20,10 +20,14 @@ import { TestConnectionManagementService } from 'sql/platform/connection/test/co
 import { TestCapabilitiesService } from 'sql/platform/capabilities/test/common/testCapabilitiesService';
 import { NullAdsTelemetryService } from 'sql/platform/telemetry/common/adsTelemetryService';
 import { ConnectionOptionSpecialType, ServiceOptionType } from 'sql/platform/connection/common/interfaces';
+import { TestConfigurationService } from 'sql/platform/connection/test/common/testConfigurationService';
+import { TestNotificationService } from 'vs/platform/notification/test/common/testNotificationService';
 
 suite('SQL Object Explorer Service tests', () => {
 	let sqlOEProvider: TypeMoq.Mock<TestObjectExplorerProvider>;
 	let connectionManagementService: TypeMoq.Mock<TestConnectionManagementService>;
+	let notificationService: TypeMoq.Mock<TestNotificationService>;
+	let configurationService: TypeMoq.Mock<TestConfigurationService>;
 	let connection: ConnectionProfile;
 	let connectionToFail: ConnectionProfile;
 	let conProfGroup: ConnectionProfileGroup;
@@ -267,10 +271,23 @@ suite('SQL Object Explorer Service tests', () => {
 			resolve(connection);
 		}));
 
+		configurationService = TypeMoq.Mock.ofType(TestConfigurationService, TypeMoq.MockBehavior.Strict);
+		configurationService.setup(x => x.getValue('serverTree.nodeExpansionTimeout')).returns(() => 45);
+
 		connectionManagementService.setup(x => x.getCapabilities(mssqlProviderName)).returns(() => undefined);
 
+		notificationService = TypeMoq.Mock.ofType(TestNotificationService, TypeMoq.MockBehavior.Strict);
+		notificationService.setup(x => x.error(TypeMoq.It.isAny())).returns(() => undefined);
+
 		const logService = new NullLogService();
-		objectExplorerService = new ObjectExplorerService(connectionManagementService.object, new NullAdsTelemetryService(), capabilitiesService, logService);
+		objectExplorerService = new ObjectExplorerService(
+			connectionManagementService.object,
+			new NullAdsTelemetryService(),
+			capabilitiesService,
+			logService,
+			configurationService.object,
+			notificationService.object);
+
 		objectExplorerService.registerProvider(mssqlProviderName, sqlOEProvider.object);
 		sqlOEProvider.setup(x => x.createNewSession(TypeMoq.It.is<azdata.ConnectionInfo>(x => x.options['serverName'] === connection.serverName))).returns(() => new Promise<any>((resolve) => {
 			resolve(response);