Adding handling for expand node failures

localization/l10n/bundle.l10n.json

@@ -747,6 +747,7 @@
     ]
   },
   "Enable Experiences & Reload": "Enable Experiences & Reload",
+  "Error loading; refresh to try again": "Error loading; refresh to try again",
   "Connection Dialog (Preview)": "Connection Dialog (Preview)",
   "Azure Account": "Azure Account",
   "Azure Account is required": "Azure Account is required",

localization/xliff/vscode-mssql.xlf

@@ -506,6 +506,9 @@
     <trans-unit id="++CODE++6d89c16c7cf3f832b29fe1d181d420d4a583f93bb02e63fdfa266c33b6b3fcaf">
       <source xml:lang="en">Error loading preview</source>
     </trans-unit>
+    <trans-unit id="++CODE++2b850d590d607436d4cbf751406a7626e01d0e5c775f26c880ab4fd21e70bb66">
+      <source xml:lang="en">Error loading; refresh to try again</source>
+    </trans-unit>
     <trans-unit id="++CODE++619d1a6dc49f043552f7b8d14e6483de4a2caf1c8270d6168a642c8eccc3a752">
       <source xml:lang="en">Error occurred opening content in editor.</source>
     </trans-unit>

src/constants/locConstants.ts

@@ -680,6 +680,12 @@ export function enableRichExperiencesPrompt(learnMoreUrl: string) {
 }
 export let enableRichExperiences = l10n.t("Enable Experiences & Reload");
 
+export class ObjectExplorer {
+    public static ErrorLoadingRefreshToTryAgain = l10n.t(
+        "Error loading; refresh to try again",
+    );
+}
+
 export class ConnectionDialog {
     public static connectionDialog = l10n.t("Connection Dialog (Preview)");
     public static azureAccount = l10n.t("Azure Account");

src/objectExplorer/objectExplorerService.ts

@@ -357,10 +357,19 @@ export class ObjectExplorerService {
         return undefined;
     }
 
-    private handleExpandSessionNotification(): NotificationHandler<ExpandResponse> {
+    /**
+     * Handler for async response from SQL Tools Service.
+     * Public only for testing
+     */
+    public handleExpandSessionNotification(): NotificationHandler<ExpandResponse> {
         const self = this;
         const handler = (result: ExpandResponse) => {
-            if (result && result.nodes) {
+            if (!result) {
+                return undefined;
+            }
+
+            if (result.nodes && !result.errorMessage) {
+                // successfully received children from SQL Tools Service
                 const credentials =
                     self._sessionIdToConnectionCredentialsMap.get(
                         result.sessionId,
@@ -402,6 +411,42 @@ export class ObjectExplorerService {
                         return;
                     }
                 }
+            } else {
+                // failure to expand node; display error
+
+                if (result.errorMessage) {
+                    self._connectionManager.vscodeWrapper.showErrorMessage(
+                        result.errorMessage,
+                    );
+                }
+
+                const expandParams: ExpandParams = {
+                    sessionId: result.sessionId,
+                    nodePath: result.nodePath,
+                };
+                const parentNode = self.getParentFromExpandParams(expandParams);
+
+                const errorNode = new vscode.TreeItem(
+                    LocalizedConstants.ObjectExplorer.ErrorLoadingRefreshToTryAgain,
+                    TreeItemCollapsibleState.None,
+                );
+
+                errorNode.tooltip = result.errorMessage;
+
+                self._treeNodeToChildrenMap.set(parentNode, [errorNode]);
+
+                for (let key of self._expandParamsToPromiseMap.keys()) {
+                    if (
+                        key.sessionId === expandParams.sessionId &&
+                        key.nodePath === expandParams.nodePath
+                    ) {
+                        let promise = self._expandParamsToPromiseMap.get(key);
+                        promise.resolve([errorNode as TreeNodeInfo]);
+                        self._expandParamsToPromiseMap.delete(key);
+                        self._expandParamsToTreeNodeInfoMap.delete(key);
+                        return;
+                    }
+                }
             }
         };
         return handler;

test/unit/objectExplorerProvider.test.ts

@@ -18,12 +18,18 @@ import { AccountSignInTreeNode } from "../../src/objectExplorer/accountSignInTre
 import { ConnectTreeNode } from "../../src/objectExplorer/connectTreeNode";
 import { NodeInfo } from "../../src/models/contracts/objectExplorer/nodeInfo";
 import { Deferred } from "../../src/protocol";
+import {
+    ExpandParams,
+    ExpandResponse,
+} from "../../src/models/contracts/objectExplorer/expandNodeRequest";
+import VscodeWrapper from "../../src/controllers/vscodeWrapper";
 
-suite("Object Explorer Provider Tests", () => {
+suite("Object Explorer Provider Tests", function () {
     let objectExplorerService: TypeMoq.IMock<ObjectExplorerService>;
     let connectionManager: TypeMoq.IMock<ConnectionManager>;
     let client: TypeMoq.IMock<SqlToolsServiceClient>;
     let objectExplorerProvider: ObjectExplorerProvider;
+    let vscodeWrapper: TypeMoq.IMock<VscodeWrapper>;
 
     setup(() => {
         let mockContext: TypeMoq.IMock<vscode.ExtensionContext> =
@@ -42,13 +48,29 @@ suite("Object Explorer Provider Tests", () => {
             c.onNotification(TypeMoq.It.isAny(), TypeMoq.It.isAny()),
         );
         connectionManager.object.client = client.object;
+
+        vscodeWrapper = TypeMoq.Mock.ofType(
+            VscodeWrapper,
+            TypeMoq.MockBehavior.Loose,
+        );
+
+        vscodeWrapper.setup((v) =>
+            v.showErrorMessage(TypeMoq.It.isAnyString()),
+        );
+
+        connectionManager
+            .setup((c) => c.vscodeWrapper)
+            .returns(() => vscodeWrapper.object);
+        connectionManager.object.vscodeWrapper = vscodeWrapper.object;
+
         objectExplorerProvider = new ObjectExplorerProvider(
             connectionManager.object,
         );
         expect(
             objectExplorerProvider,
             "Object Explorer Provider is initialzied properly",
         ).is.not.equal(undefined);
+
         objectExplorerService = TypeMoq.Mock.ofType(
             ObjectExplorerService,
             TypeMoq.MockBehavior.Loose,
@@ -334,6 +356,146 @@ suite("Object Explorer Provider Tests", () => {
         assert.equal(treeItem, node);
     });
 
+    const mockParentTreeNode = new TreeNodeInfo(
+        "Parent Node",
+        undefined,
+        undefined,
+        "parentNodePath",
+        undefined,
+        undefined,
+        undefined,
+        undefined,
+        undefined,
+        undefined,
+    );
+
+    test("Test handleExpandSessionNotification returns child nodes upon success", async function () {
+        const childNodeInfo: NodeInfo = {
+            nodePath: `${mockParentTreeNode.nodePath}/childNodePath`,
+            nodeStatus: undefined,
+            nodeSubType: undefined,
+            nodeType: undefined,
+            label: "Child Node",
+            isLeaf: true,
+            errorMessage: undefined,
+            metadata: undefined,
+        };
+
+        const mockExpandResponse: ExpandResponse = {
+            sessionId: "test_session",
+            nodePath: mockParentTreeNode.nodePath,
+            nodes: [childNodeInfo],
+            errorMessage: undefined,
+        };
+
+        const testOeService = new ObjectExplorerService(
+            connectionManager.object,
+            objectExplorerProvider,
+        );
+
+        let notificationObject =
+            testOeService.handleExpandSessionNotification();
+
+        const expandParams: ExpandParams = {
+            sessionId: mockExpandResponse.sessionId,
+            nodePath: mockExpandResponse.nodePath,
+        };
+
+        testOeService["_expandParamsToTreeNodeInfoMap"].set(
+            expandParams,
+            mockParentTreeNode,
+        );
+
+        testOeService["_sessionIdToConnectionCredentialsMap"].set(
+            mockExpandResponse.sessionId,
+            undefined,
+        );
+
+        const outputPromise = new Deferred<TreeNodeInfo[]>();
+
+        testOeService["_expandParamsToPromiseMap"].set(
+            expandParams,
+            outputPromise,
+        );
+
+        notificationObject.call(testOeService, mockExpandResponse);
+
+        const childNodes = await outputPromise;
+        assert.equal(childNodes.length, 1, "Child nodes length");
+        assert.equal(
+            childNodes[0].label,
+            childNodeInfo.label,
+            "Child node label",
+        );
+        assert.equal(
+            childNodes[0].nodePath,
+            childNodeInfo.nodePath,
+            "Child node path",
+        );
+    });
+
+    test("Test handleExpandSessionNotification returns message node upon failure", async function () {
+        this.timeout(0);
+
+        const mockExpandResponse: ExpandResponse = {
+            sessionId: "test_session",
+            nodePath: mockParentTreeNode.nodePath,
+            nodes: [],
+            errorMessage: "Error occurred when expanding node",
+        };
+
+        const testOeService = new ObjectExplorerService(
+            connectionManager.object,
+            objectExplorerProvider,
+        );
+
+        let notificationObject =
+            testOeService.handleExpandSessionNotification();
+
+        const expandParams: ExpandParams = {
+            sessionId: mockExpandResponse.sessionId,
+            nodePath: mockExpandResponse.nodePath,
+        };
+
+        testOeService["_expandParamsToTreeNodeInfoMap"].set(
+            expandParams,
+            mockParentTreeNode,
+        );
+
+        testOeService["_sessionIdToConnectionCredentialsMap"].set(
+            mockExpandResponse.sessionId,
+            undefined,
+        );
+
+        const outputPromise = new Deferred<TreeNodeInfo[]>();
+
+        testOeService["_expandParamsToPromiseMap"].set(
+            expandParams,
+            outputPromise,
+        );
+
+        notificationObject.call(testOeService, mockExpandResponse);
+
+        const childNodes = await outputPromise;
+
+        vscodeWrapper.verify(
+            (x) => x.showErrorMessage(mockExpandResponse.errorMessage),
+            TypeMoq.Times.once(),
+        );
+
+        assert.equal(childNodes.length, 1, "Child nodes length");
+        assert.equal(
+            childNodes[0].label,
+            "Error loading; refresh to try again",
+            "Error node label",
+        );
+        assert.equal(
+            childNodes[0].tooltip,
+            mockExpandResponse.errorMessage,
+            "Error node tooltip",
+        );
+    });
+
     test("Test signInNode function", () => {
         objectExplorerService.setup((s) =>
             s.signInNodeServer(TypeMoq.It.isAny()),