From ccaf5c45941d2f5c6549fda803a481d523e9df8d Mon Sep 17 00:00:00 2001
From: Charles Gagnon <chgagnon@microsoft.com>
Date: Thu, 19 Jan 2023 16:14:17 -0800
Subject: [PATCH] Fix null ref when expanding nodes from other providers
 (#1816)

* Fix null ref when expanding nodes from other providers

(cherry picked from commit 572a7b37b2a3651db5a101a8780c874c51b55b96)

* cleanup
---
 .../ObjectExplorer/Nodes/TreeNode.cs          |  4 ++--
 .../ObjectExplorer/ObjectExplorerService.cs   | 21 ++++++++++++-------
 .../ObjectExplorer/ObjectExplorerUtils.cs     |  6 +++---
 3 files changed, 19 insertions(+), 12 deletions(-)

diff --git a/src/Microsoft.SqlTools.ServiceLayer/ObjectExplorer/Nodes/TreeNode.cs b/src/Microsoft.SqlTools.ServiceLayer/ObjectExplorer/Nodes/TreeNode.cs
index 2a91d8b7..863e71f7 100644
--- a/src/Microsoft.SqlTools.ServiceLayer/ObjectExplorer/Nodes/TreeNode.cs
+++ b/src/Microsoft.SqlTools.ServiceLayer/ObjectExplorer/Nodes/TreeNode.cs
@@ -202,9 +202,9 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer.Nodes
             nodePath = path;
         }
 
-        public TreeNode FindNodeByPath(string path, bool expandIfNeeded = false)
+        public TreeNode? FindNodeByPath(string path, bool expandIfNeeded = false)
         {
-            TreeNode nodeForPath = ObjectExplorerUtils.FindNode(this, node =>
+            TreeNode? nodeForPath = ObjectExplorerUtils.FindNode(this, node =>
             {
                 return node.GetNodePath() == path;
             }, nodeToFilter =>
diff --git a/src/Microsoft.SqlTools.ServiceLayer/ObjectExplorer/ObjectExplorerService.cs b/src/Microsoft.SqlTools.ServiceLayer/ObjectExplorer/ObjectExplorerService.cs
index 00da316c..ab1915e6 100644
--- a/src/Microsoft.SqlTools.ServiceLayer/ObjectExplorer/ObjectExplorerService.cs
+++ b/src/Microsoft.SqlTools.ServiceLayer/ObjectExplorer/ObjectExplorerService.cs
@@ -321,7 +321,7 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer
                     ObjectExplorerTaskResult result = await RunTaskWithTimeout(task,
                         settings?.CreateSessionTimeout ?? ObjectExplorerSettings.DefaultCreateSessionTimeout);
 
-                    if (result != null && !result.IsCompleted)
+                    if (result != null && !result.IsSuccessful)
                     {
                         cancellationTokenSource.Cancel();
                         SessionCreatedParameters response = new SessionCreatedParameters
@@ -384,11 +384,11 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer
         internal ExpandResponse QueueExpandNodeRequest(ObjectExplorerSession session, string nodePath, bool forceRefresh = false, SecurityToken? securityToken = null)
         {
             NodeInfo[] nodes = null;
-            TreeNode node = session.Root.FindNodeByPath(nodePath);
+            TreeNode? node = session.Root.FindNodeByPath(nodePath);
             ExpandResponse response = null;
 
             // Performance Optimization for table designer to load the database model earlier based on user configuration.
-            if (node.NodeTypeId == NodeTypes.Database && TableDesignerService.Instance.Settings.PreloadDatabaseModel)
+            if (node?.NodeTypeId == NodeTypes.Database && TableDesignerService.Instance.Settings.PreloadDatabaseModel)
             {
                 // The operation below are not blocking, but just in case, wrapping it with a task run to make sure it has no impact on the node expansion time.
                 var _ = Task.Run(() =>
@@ -604,7 +604,7 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer
                 ObjectExplorerTaskResult result = await RunTaskWithTimeout(task,
                     settings?.ExpandTimeout ?? ObjectExplorerSettings.DefaultExpandTimeout);
 
-                if (result != null && !result.IsCompleted)
+                if (result != null && !result.IsSuccessful)
                 {
                     cancellationTokenSource.Cancel();
                     ExpandResponse response = CreateExpandResponse(session, expandParams);
@@ -620,9 +620,10 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer
             ObjectExplorerTaskResult result = new ObjectExplorerTaskResult();
             TimeSpan timeout = TimeSpan.FromSeconds(timeoutInSec);
             await Task.WhenAny(task, Task.Delay(timeout));
-            result.IsCompleted = task.IsCompleted;
+            result.IsSuccessful = task.IsCompleted;
             if (task.Exception != null)
             {
+                result.IsSuccessful = false;
                 result.Exception = task.Exception;
             }
             else if (!task.IsCompleted)
@@ -761,8 +762,14 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer
 
         internal class ObjectExplorerTaskResult
         {
-            public bool IsCompleted { get; set; }
-            public Exception Exception { get; set; }
+            /// <summary>
+            /// Whether the task was successfully completed. False if an error of any kind occurred during execution.
+            /// </summary>
+            public bool IsSuccessful { get; set; }
+            /// <summary>
+            /// The Exception that occurred during execution, if any. 
+            /// </summary>
+            public Exception? Exception { get; set; }
         }
 
         public void Dispose()
diff --git a/src/Microsoft.SqlTools.ServiceLayer/ObjectExplorer/ObjectExplorerUtils.cs b/src/Microsoft.SqlTools.ServiceLayer/ObjectExplorer/ObjectExplorerUtils.cs
index dd01a39a..29223af7 100644
--- a/src/Microsoft.SqlTools.ServiceLayer/ObjectExplorer/ObjectExplorerUtils.cs
+++ b/src/Microsoft.SqlTools.ServiceLayer/ObjectExplorer/ObjectExplorerUtils.cs
@@ -48,8 +48,8 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer
         /// <param name="condition">Predicate function that accesses the tree and
         /// determines whether to stop going further up the tree</param>
         /// <param name="filter">Predicate function to filter the children when traversing</param>
-        /// <returns>A Tree Node that matches the condition</returns>
-        public static TreeNode FindNode(TreeNode node, Predicate<TreeNode> condition, Predicate<TreeNode> filter, bool expandIfNeeded = false)
+        /// <returns>A Tree Node that matches the condition, or null if no matching node could be found</returns>
+        public static TreeNode? FindNode(TreeNode node, Predicate<TreeNode> condition, Predicate<TreeNode> filter, bool expandIfNeeded = false)
         {
             if(node == null)
             {
@@ -65,7 +65,7 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer
             {
                 if (filter != null && filter(child))
                 {
-                    TreeNode childNode = FindNode(child, condition, filter, expandIfNeeded);
+                    TreeNode? childNode = FindNode(child, condition, filter, expandIfNeeded);
                     if (childNode != null)
                     {
                         return childNode;