Handle connect button press with firewall error (#17118)

* handle connect button press with firewall error

* add types

* created helper method getParentNode, localize

* display error message to user, move function outside of class

* reword error message

* cleanup

* change error message

localization/xliff/enu/constants/localizedConstants.enu.xlf

@@ -521,6 +521,9 @@
         <trans-unit id="notStarted">
             <source xml:lang="en">Not started</source>
         </trans-unit>
+        <trans-unit id="nodeErrorMessage">
+            <source xml:lang="en">Parent node was not TreeNodeInfo.</source>
+        </trans-unit>
     </body>
   </file>
 </xliff>

src/objectExplorer/connectTreeNode.ts

@@ -26,3 +26,8 @@ export class ConnectTreeNode extends vscode.TreeItem {
         return this._parentNode;
     }
 }
+
+export type TreeNodeType =
+    | TreeNodeInfo
+    | ConnectTreeNode;
+

src/objectExplorer/objectExplorerService.ts

@@ -17,7 +17,7 @@ import { AuthenticationTypes, IConnectionProfile } from '../models/interfaces';
 import * as LocalizedConstants from '../constants/localizedConstants';
 import { AddConnectionTreeNode } from './addConnectionTreeNode';
 import { AccountSignInTreeNode } from './accountSignInTreeNode';
-import { ConnectTreeNode } from './connectTreeNode';
+import { ConnectTreeNode, TreeNodeType } from './connectTreeNode';
 import { Deferred } from '../protocol';
 import * as Constants from '../constants/constants';
 import { ObjectExplorerUtils } from './objectExplorerUtils';
@@ -27,6 +27,15 @@ import { ConnectionProfile } from '../models/connectionProfile';
 import providerSettings from '../azure/providerSettings';
 import { IConnectionInfo } from 'vscode-mssql';
 
+function getParentNode(node: TreeNodeType): TreeNodeInfo {
+    node = node.parentNode;
+    if (!(node instanceof TreeNodeInfo)) {
+        vscode.window.showErrorMessage(LocalizedConstants.nodeErrorMessage);
+        throw new Error(`Parent node was not TreeNodeInfo.`);
+    }
+    return node;
+}
+
 export class ObjectExplorerService {
 
     private _client: SqlToolsServiceClient;
@@ -61,6 +70,9 @@ export class ObjectExplorerService {
     private handleSessionCreatedNotification(): NotificationHandler<SessionCreatedParameters> {
         const self = this;
         const handler = async (result: SessionCreatedParameters) => {
+            if (self._currentNode instanceof ConnectTreeNode) {
+                self.currentNode = getParentNode(self.currentNode);
+            }
             if (result.success) {
                 let nodeLabel = this._nodePathToNodeLabelMap.get(result.rootNode.nodePath);
                 // if no node label, check if it has a name in saved profiles
@@ -78,6 +90,7 @@ export class ObjectExplorerService {
                 }
                 // set connection and other things
                 let node: TreeNodeInfo;
+
                 if (self._currentNode && (self._currentNode.sessionId === result.sessionId)) {
                     nodeLabel = !nodeLabel ? self.createNodeLabel(self._currentNode.connectionInfo) : nodeLabel;
                     node = TreeNodeInfo.fromNodeInfo(result.rootNode, result.sessionId,
@@ -191,7 +204,10 @@ export class ObjectExplorerService {
         }
     }
 
-    public updateNode(node: TreeNodeInfo): void {
+    public updateNode(node: TreeNodeType): void {
+        if (node instanceof ConnectTreeNode) {
+            node = getParentNode(node);
+        }
         for (let rootTreeNode of this._rootTreeNodeArray) {
             if (Utils.isSameConnection(node.connectionInfo, rootTreeNode.connectionInfo) &&
                 rootTreeNode.label === node.label) {

src/objectExplorer/objectExplorerUtils.ts

@@ -8,6 +8,7 @@ import { IConnectionProfile } from '../models/interfaces';
 import * as Constants from '../constants/constants';
 import * as LocalizedConstants from '../constants/localizedConstants';
 import * as vscodeMssql from 'vscode-mssql';
+import { TreeNodeType } from './connectTreeNode';
 
 export class ObjectExplorerUtils {
 
@@ -26,8 +27,13 @@ export class ObjectExplorerUtils {
         }
     }
 
-    public static getNodeUri(node: TreeNodeInfo): string {
-        const profile = <IConnectionProfile>node.connectionInfo;
+    public static getNodeUri(node: TreeNodeType): string {
+        let profile: IConnectionProfile;
+        if (node instanceof TreeNodeInfo) {
+            profile = <IConnectionProfile>node.connectionInfo;
+        } else {
+            profile = <IConnectionProfile>node.parentNode.connectionInfo;
+        }
         return ObjectExplorerUtils.getNodeUriFromProfile(profile);
     }