Revert "Revert "OE clicking and awaiting the expansion calls"" (#6923)

* Revert "Revert "OE clicking and awaiting the expansion calls" (#6921)"

This reverts commit b9e3a468ae.

* Handle edge cases better

* Polling for OE load time

extensions/integration-tests/src/utils.ts

@@ -14,7 +14,7 @@ import { TestServerProfile } from './testConfig';
  * @param timeout optional timeout parameter
  * Returns connection id for a new connection
  */
-export async function connectToServer(server: TestServerProfile, timeout: number = 3000): Promise<string> {
+export async function connectToServer(server: TestServerProfile, timeout: number = 10000): Promise<string> {
 	let connectionProfile: azdata.IConnectionProfile = {
 		serverName: server.serverName,
 		databaseName: server.database,
@@ -36,10 +36,37 @@ export async function connectToServer(server: TestServerProfile, timeout: number
 
 	//workaround
 	//wait for OE to load
-	await new Promise(c => setTimeout(c, timeout));
+	await pollTimeout(async () => {
+		const nodes = await azdata.objectexplorer.getActiveConnectionNodes();
+		let found = nodes.some(node => {
+			return node.connectionId === result.connectionId;
+		});
+		if (found === undefined) {
+			found = false;
+		}
+		return found;
+	}, 1000, timeout);
+
 	return result.connectionId;
 }
 
+export async function pollTimeout(predicate: () => Thenable<boolean>, intervalDelay: number, timeoutTime: number): Promise<boolean> {
+	let interval: NodeJS.Timer;
+	return new Promise(pollOver => {
+		const complete = (success = false) => {
+			clearInterval(interval);
+			pollOver(success);
+		};
+		interval = setInterval(async () => {
+			const predResult = await predicate();
+			if (predResult) {
+				complete(true);
+			}
+		}, intervalDelay);
+		setTimeout(complete, timeoutTime);
+	});
+}
+
 export async function ensureConnectionViewOpened() {
 	await vscode.commands.executeCommand('dataExplorer.servers.focus');
 }

src/sql/workbench/parts/objectExplorer/browser/serverTreeDataSource.ts

@@ -75,8 +75,8 @@ export class ServerTreeDataSource implements IDataSource {
 					// It has been tested for connecting to the server in profile itself and things work fine there.
 					this._objectExplorerService.resolveTreeNodeChildren(node.getSession(), node).then(() => {
 						resolve(node.children);
-					}, expandError => {
+					}, async expandError => {
-						node.setExpandedState(TreeItemCollapsibleState.Collapsed);
+						await node.setExpandedState(TreeItemCollapsibleState.Collapsed);
 						node.errorStateMessage = expandError;
 						this.showError(expandError);
 						// collapse node and refresh in case of error so remove tree cache

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

@@ -171,13 +171,13 @@ export class ServerTreeView extends Disposable {
 			}));
 		}
 
-		return new Promise<void>((resolve, reject) => {
+		return new Promise<void>(async (resolve, reject) => {
 			this.refreshTree();
 			const root = <ConnectionProfileGroup>this._tree.getInput();
 
 			const expandGroups: boolean = this._configurationService.getValue(SERVER_GROUP_CONFIG)[SERVER_GROUP_AUTOEXPAND_CONFIG];
 			if (expandGroups) {
-				this._tree.expandAll(ConnectionProfileGroup.getSubgroups(root));
+				await this._tree.expandAll(ConnectionProfileGroup.getSubgroups(root));
 			}
 
 			if (root && !root.hasValidConnections) {
@@ -270,9 +270,9 @@ export class ServerTreeView extends Disposable {
 		if (connection) {
 			const conn = this.getConnectionInTreeInput(connection.id);
 			if (conn) {
-				return this._objectExplorerService.deleteObjectExplorerNode(conn).then(() => {
+				return this._objectExplorerService.deleteObjectExplorerNode(conn).then(async () => {
-					this._tree.collapse(conn);
+					await this._tree.collapse(conn);
-					this._tree.refresh(conn);
+					return this._tree.refresh(conn);
 				});
 			}
 		}
@@ -342,10 +342,10 @@ export class ServerTreeView extends Disposable {
 			} else {
 				treeInput = filteredResults[0];
 			}
-			this._tree.setInput(treeInput).then(() => {
+			this._tree.setInput(treeInput).then(async () => {
 				if (isHidden(this.messages)) {
 					this._tree.getFocus();
-					this._tree.expandAll(ConnectionProfileGroup.getSubgroups(treeInput));
+					await this._tree.expandAll(ConnectionProfileGroup.getSubgroups(treeInput));
 				} else {
 					this._tree.clearFocus();
 				}
@@ -374,10 +374,10 @@ export class ServerTreeView extends Disposable {
 		// Add all connections to tree root and set tree input
 		const treeInput = new ConnectionProfileGroup('searchroot', undefined, 'searchroot', undefined, undefined);
 		treeInput.addConnections(filteredResults);
-		this._tree.setInput(treeInput).then(() => {
+		this._tree.setInput(treeInput).then(async () => {
 			if (isHidden(this.messages)) {
 				this._tree.getFocus();
-				this._tree.expandAll(ConnectionProfileGroup.getSubgroups(treeInput));
+				await this._tree.expandAll(ConnectionProfileGroup.getSubgroups(treeInput));
 			} else {
 				this._tree.clearFocus();
 			}

src/sql/workbench/parts/objectExplorer/browser/treeSelectionHandler.ts

@@ -15,8 +15,9 @@ import { TreeUpdateUtils } from 'sql/workbench/parts/objectExplorer/browser/tree
 export class TreeSelectionHandler {
 	// progressRunner: IProgressRunner;
 
-	private _lastClicked: any;
+	private _lastClicked: any[];
 	private _clickTimer: any = undefined;
+	private _otherTimer: any = undefined;
 
 	// constructor(@IProgressService private _progressService: IProgressService) {
 
@@ -38,14 +39,21 @@ export class TreeSelectionHandler {
 		return event && event.payload && event.payload.origin === 'mouse';
 	}
 
+	private isKeyboardEvent(event: any): boolean {
+		return event && event.payload && event.payload.origin === 'keyboard';
+	}
+
 	/**
-	 * Handle selection of tree element
+	 * Handle select	ion of tree element
 	 */
 	public onTreeSelect(event: any, tree: ITree, connectionManagementService: IConnectionManagementService, objectExplorerService: IObjectExplorerService, connectionCompleteCallback: () => void) {
-		let sendSelectionEvent = ((event: any, selection: any, isDoubleClick: boolean) => {
+		let sendSelectionEvent = ((event: any, selection: any, isDoubleClick: boolean, userInteraction: boolean) => {
-			let isKeyboard = event && event.payload && event.payload.origin === 'keyboard';
+			// userInteraction: defensive - don't touch this something else is handling it.
+			if (userInteraction === true && this._lastClicked && this._lastClicked[0] === selection[0]) {
+				this._lastClicked = undefined;
+			}
 			if (!TreeUpdateUtils.isInDragAndDrop) {
-				this.handleTreeItemSelected(connectionManagementService, objectExplorerService, isDoubleClick, isKeyboard, selection, tree, connectionCompleteCallback);
+				this.handleTreeItemSelected(connectionManagementService, objectExplorerService, isDoubleClick, this.isKeyboardEvent(event), selection, tree, connectionCompleteCallback);
 			}
 		});
 
@@ -56,18 +64,29 @@ export class TreeSelectionHandler {
 		}
 		let specificSelection = selection[0];
 
-		if (this.isMouseEvent(event)) {
+		if (this.isMouseEvent(event) || this.isKeyboardEvent(event)) {
-			if (this._lastClicked === specificSelection) {
+			if (this._lastClicked !== undefined) {
 				clearTimeout(this._clickTimer);
-				sendSelectionEvent(event, selection, true);
+				let lastSpecificClick = this._lastClicked[0];
-				return;
-			}
-			this._lastClicked = specificSelection;
-		}
 
-		this._clickTimer = setTimeout(() => {
+				if (lastSpecificClick === specificSelection) {
-			sendSelectionEvent(event, selection, false);
+					sendSelectionEvent(event, selection, true, true);
-		}, 300);
+					return;
+				} else {
+					sendSelectionEvent(event, this._lastClicked, false, true);
+				}
+			}
+			this._lastClicked = selection;
+
+			this._clickTimer = setTimeout(() => {
+				sendSelectionEvent(event, selection, false, true);
+			}, 400);
+		} else {
+			clearTimeout(this._otherTimer);
+			this._otherTimer = setTimeout(() => {
+				sendSelectionEvent(event, selection, false, false);
+			}, 400);
+		}
 	}
 
 	/**

src/sql/workbench/parts/objectExplorer/browser/treeUpdateUtils.ts

@@ -74,13 +74,13 @@ export class TreeUpdateUtils {
 			treeInput = TreeUpdateUtils.getTreeInput(connectionManagementService, providers);
 		}
 		const previousTreeInput: any = tree.getInput();
-		return tree.setInput(treeInput).then(() => {
+		return tree.setInput(treeInput).then(async () => {
 			if (previousTreeInput instanceof Disposable) {
 				previousTreeInput.dispose();
 			}
 			// Make sure to expand all folders that where expanded in the previous session
 			if (targetsToExpand) {
-				tree.expandAll(targetsToExpand);
+				await tree.expandAll(targetsToExpand);
 			}
 			if (selectedElement) {
 				tree.select(selectedElement);
@@ -118,10 +118,10 @@ export class TreeUpdateUtils {
 		let treeInput = TreeUpdateUtils.getTreeInput(connectionManagementService);
 		if (treeInput) {
 			if (treeInput !== tree.getInput()) {
-				return tree.setInput(treeInput).then(() => {
+				return tree.setInput(treeInput).then(async () => {
 					// Make sure to expand all folders that where expanded in the previous session
 					if (targetsToExpand) {
-						tree.expandAll(targetsToExpand);
+						await tree.expandAll(targetsToExpand);
 					}
 					if (selectedElement) {
 						tree.select(selectedElement);
@@ -316,4 +316,4 @@ export class TreeUpdateUtils {
 		}
 		return connectionProfile;
 	}
-}
+}

src/sql/workbench/parts/tasks/browser/tasksView.ts

@@ -120,10 +120,10 @@ export class TaskHistoryView {
 		//Get the tree Input
 		let treeInput = this._taskService.getAllTasks();
 		if (treeInput) {
-			this._tree.setInput(treeInput).then(() => {
+			this._tree.setInput(treeInput).then(async () => {
 				// Make sure to expand all folders that where expanded in the previous session
 				if (targetsToExpand) {
-					this._tree.expandAll(targetsToExpand);
+					await this._tree.expandAll(targetsToExpand);
 				}
 				if (selectedElement) {
 					this._tree.select(selectedElement);

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

@@ -604,7 +604,7 @@ export class ObjectExplorerService implements IObjectExplorerService {
 			nodeInfo.nodeSubType, nodeInfo.nodeStatus, parent, nodeInfo.metadata, nodeInfo.iconType, {
 				getChildren: treeNode => this.getChildren(treeNode),
 				isExpanded: treeNode => this.isExpanded(treeNode),
-				setNodeExpandedState: (treeNode, expandedState) => this.setNodeExpandedState(treeNode, expandedState),
+				setNodeExpandedState: async (treeNode, expandedState) => await this.setNodeExpandedState(treeNode, expandedState),
 				setNodeSelected: (treeNode, selected, clearOtherSelections: boolean = undefined) => this.setNodeSelected(treeNode, selected, clearOtherSelections)
 			});
 		node.childProvider = nodeInfo.childProvider;