Object explorer notifications (#354)

* changing oe protocols to send notifications

src/Microsoft.SqlTools.ServiceLayer/ObjectExplorer/Contracts/CreateSessionRequest.cs

@@ -14,6 +14,21 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer.Contracts
     /// requesting expansion of nodes, and a root node to display for this area.
     /// </summary>
     public class CreateSessionResponse
+    {
+        /// <summary>
+        /// Unique ID to use when sending any requests for objects in the
+        /// tree under the node
+        /// </summary>
+        public string SessionId { get; set; }
+
+    }
+
+    /// <summary>
+    /// Information returned from a <see cref="CreateSessionRequest"/>.
+    /// Contains success information, a <see cref="SessionId"/> to be used when
+    /// requesting expansion of nodes, and a root node to display for this area.
+    /// </summary>
+    public class SessionCreatedParameters
     {
         /// <summary>
         /// Boolean indicating if the connection was successful
@@ -30,6 +45,12 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer.Contracts
         /// Information describing the base node in the tree
         /// </summary>
         public NodeInfo RootNode { get; set; }
+
+
+        /// <summary>
+        /// Error message returned from the engine for a object explorer session failure reason, if any.
+        /// </summary>
+        public string ErrorMessage { get; set; }
     }
     /// <summary>
     /// Establishes an Object Explorer tree session for a specific connection.
@@ -42,4 +63,14 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer.Contracts
             RequestType<ConnectionDetails, CreateSessionResponse> Type =
             RequestType<ConnectionDetails, CreateSessionResponse>.Create("objectexplorer/createsession");
     }
+
+    /// <summary>
+    /// Session notification mapping entry 
+    /// </summary>
+    public class CreateSessionCompleteNotification
+    {
+        public static readonly
+            EventType<SessionCreatedParameters> Type =
+            EventType<SessionCreatedParameters>.Create("objectexplorer/sessionCreated");
+    }
 }

src/Microsoft.SqlTools.ServiceLayer/ObjectExplorer/Contracts/ExpandRequest.cs

@@ -22,6 +22,16 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer.Contracts
         /// Information describing the expanded nodes in the tree
         /// </summary>
         public NodeInfo[] Nodes { get; set; }
+
+        /// <summary>
+        /// Path identifying the node to expand. See <see cref="NodeInfo.NodePath"/> for details
+        /// </summary>
+        public string NodePath { get; set; }
+
+        /// <summary>
+        /// Error message returned from the engine for a object explorer expand failure reason, if any.
+        /// </summary>
+        public string ErrorMessage { get; set; }
     }
 
     /// <summary>
@@ -50,7 +60,17 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer.Contracts
         /// Returns children of a given node as a <see cref="NodeInfo"/> array.
         /// </summary>
         public static readonly
-            RequestType<ExpandParams, ExpandResponse> Type =
+            RequestType<ExpandParams, bool> Type =
-            RequestType<ExpandParams, ExpandResponse>.Create("objectexplorer/expand");
+            RequestType<ExpandParams, bool>.Create("objectexplorer/expand");
+    }
+
+    /// <summary>
+    /// Expand notification mapping entry 
+    /// </summary>
+    public class ExpandCompleteNotification
+    {
+        public static readonly
+            EventType<ExpandResponse> Type =
+            EventType<ExpandResponse>.Create("objectexplorer/expandCompleted");
     }
 }

src/Microsoft.SqlTools.ServiceLayer/ObjectExplorer/Contracts/NodeInfo.cs

@@ -53,5 +53,10 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer.Contracts
         /// Object Metadata for smo objects to be used for scripting
         /// </summary>
         public ObjectMetadata Metadata { get; set; }
+
+        /// <summary>
+        /// Error message returned from the engine for a object explorer node failure reason, if any.
+        /// </summary>
+        public string ErrorMessage { get; set; }
     }
 }

src/Microsoft.SqlTools.ServiceLayer/ObjectExplorer/Contracts/RefreshRequest.cs

@@ -7,38 +7,11 @@ using Microsoft.SqlTools.Hosting.Protocol.Contracts;
 
 namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer.Contracts
 {
-    /// <summary>
-    /// Information returned from a <see cref="RefreshRequest"/>.
-    /// </summary>
-    public class RefreshResponse
-    {
-        /// <summary>
-        /// Unique ID to use when sending any requests for objects in the
-        /// tree under the node
-        /// </summary>
-        public string SessionId { get; set; }
-
-        /// <summary>
-        /// Information describing the expanded nodes in the tree
-        /// </summary>
-        public NodeInfo[] Nodes { get; set; }
-    }
-
     /// <summary>
     /// Parameters to the <see cref="ExpandRequest"/>.
     /// </summary>
-    public class RefreshParams
+    public class RefreshParams: ExpandParams
     {
-        /// <summary>
-        /// The Id returned from a <see cref="CreateSessionRequest"/>. This
-        /// is used to disambiguate between different trees. 
-        /// </summary>
-        public string SessionId { get; set; }
-
-        /// <summary>
-        /// Path identifying the node to refresh. See <see cref="NodeInfo.NodePath"/> for details
-        /// </summary>
-        public string NodePath { get; set; }
     }
 
     /// <summary>
@@ -50,7 +23,7 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer.Contracts
         /// Returns children of a given node as a <see cref="NodeInfo"/> array.
         /// </summary>
         public static readonly
-            RequestType<RefreshParams, RefreshResponse> Type =
+            RequestType<RefreshParams, bool> Type =
-            RequestType<RefreshParams, RefreshResponse>.Create("objectexplorer/refresh");
+            RequestType<RefreshParams, bool>.Create("objectexplorer/refresh");
     }
 }

src/Microsoft.SqlTools.ServiceLayer/ObjectExplorer/Nodes/TreeNode.cs

@@ -77,6 +77,11 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer.Nodes
         /// </summary>
         public string NodeSubType { get; set; }
 
+        /// <summary>
+        /// Error message returned from the engine for a object explorer node failure reason, if any.
+        /// </summary>
+        public string ErrorMessage { get; set; }
+
         /// <summary>
         /// Node status - for example login can be disabled/enabled
         /// </summary>
@@ -193,7 +198,8 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer.Nodes
                 NodeType = this.NodeType,
                 Metadata = this.ObjectMetadata,
                 NodeStatus = this.NodeStatus,
-                NodeSubType = this.NodeSubType
+                NodeSubType = this.NodeSubType,
+                ErrorMessage = this.ErrorMessage
             };
         }
 
@@ -289,8 +295,8 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer.Nodes
             bool includeSystemObjects = context != null && context.Database != null ? ObjectExplorerUtils.IsSystemDatabaseConnection(context.Database.Name) : true;
             
 
-                if (children.IsPopulating || context == null)
+            if (children.IsPopulating || context == null)
-                return;
+            return;
 
             children.Clear();
             BeginChildrenInit();
@@ -302,16 +308,26 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer.Nodes
                 {
                     foreach (var factory in childFactories)
                     {
-                        IEnumerable<TreeNode> items = factory.Expand(this, refresh, name, includeSystemObjects);
+                        try
-                        if (items != null)
                         {
-                            foreach (TreeNode item in items)
+                            IEnumerable<TreeNode> items = factory.Expand(this, refresh, name, includeSystemObjects);
+                            if (items != null)
                             {
-                                children.Add(item);
+                                foreach (TreeNode item in items)
-                                item.Parent = this;
+                                {
+                                    children.Add(item);
+                                    item.Parent = this;
 
+                                }
                             }
                         }
+                        catch (Exception ex)
+                        {
+                            string error = string.Format(CultureInfo.InvariantCulture, "Failed populating oe children. error:{0} inner:{1} stacktrace:{2}",
+                            ex.Message, ex.InnerException != null ? ex.InnerException.Message : "", ex.StackTrace);
+                            Logger.Write(LogLevel.Error, error);
+                            ErrorMessage = ex.Message;
+                        }
                     }
                 }
             }
@@ -320,6 +336,7 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer.Nodes
                 string error = string.Format(CultureInfo.InvariantCulture, "Failed populating oe children. error:{0} inner:{1} stacktrace:{2}",
                     ex.Message, ex.InnerException != null ? ex.InnerException.Message : "", ex.StackTrace);
                 Logger.Write(LogLevel.Error, error);
+                ErrorMessage = ex.Message;
             }
             finally
             {

src/Microsoft.SqlTools.ServiceLayer/ObjectExplorer/ObjectExplorerService.cs

@@ -103,12 +103,130 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer
             serviceHost.SetRequestHandler(CloseSessionRequest.Type, HandleCloseSessionRequest);
         }
 
-        public void CloseSession(string uri)
+       
+        
+        internal async Task HandleCreateSessionRequest(ConnectionDetails connectionDetails, RequestContext<CreateSessionResponse> context)
+        {
+            Logger.Write(LogLevel.Verbose, "HandleCreateSessionRequest");
+
+            Func<Task<CreateSessionResponse>> doCreateSession = async () =>
+            {
+                Validate.IsNotNull(nameof(connectionDetails), connectionDetails);
+                Validate.IsNotNull(nameof(context), context);
+                return await Task.Factory.StartNew(() =>
+                {
+                    string uri = GenerateUri(connectionDetails);
+
+                    return new CreateSessionResponse { SessionId = uri };
+                });
+            };
+
+            CreateSessionResponse response = await HandleRequestAsync(doCreateSession, context, "HandleCreateSessionRequest");
+            if (response != null)
+            {
+                RunCreateSessionTask(connectionDetails, response.SessionId);
+            }
+
+        }
+
+        internal async Task HandleExpandRequest(ExpandParams expandParams, RequestContext<bool> context)
+        {
+            Logger.Write(LogLevel.Verbose, "HandleExpandRequest");
+
+            Func<Task<bool>> expandNode = async () =>
+            {
+                Validate.IsNotNull(nameof(expandParams), expandParams);
+                Validate.IsNotNull(nameof(context), context);
+
+                string uri = expandParams.SessionId;
+                ObjectExplorerSession session = null;
+                NodeInfo[] nodes = null;
+                ExpandResponse response;
+                if (!sessionMap.TryGetValue(uri, out session))
+                {
+                    Logger.Write(LogLevel.Verbose, $"Cannot expand object explorer node. Couldn't find session for uri. {uri} ");
+                    await serviceHost.SendEvent(ExpandCompleteNotification.Type, new ExpandResponse
+                    {
+                        SessionId = expandParams.SessionId,
+                        NodePath = expandParams.NodePath,
+                        ErrorMessage = $"Couldn't find session for session: {uri}"
+                    });
+                    return false;
+                }
+                else
+                {
+                    RunExpandTask(session, expandParams);
+                    return true;
+                }
+            };
+            await HandleRequestAsync(expandNode, context, "HandleExpandRequest");
+        }
+
+        internal async Task HandleRefreshRequest(RefreshParams refreshParams, RequestContext<bool> context)
+        {
+            Logger.Write(LogLevel.Verbose, "HandleRefreshRequest");
+            Validate.IsNotNull(nameof(refreshParams), refreshParams);
+            Validate.IsNotNull(nameof(context), context);
+
+            string uri = refreshParams.SessionId;
+            ObjectExplorerSession session = null;
+            ExpandResponse response;
+            if (!sessionMap.TryGetValue(uri, out session))
+            {
+                Logger.Write(LogLevel.Verbose, $"Cannot expand object explorer node. Couldn't find session for uri. {uri} ");
+                await serviceHost.SendEvent(ExpandCompleteNotification.Type, new ExpandResponse
+                {
+                    SessionId = refreshParams.SessionId,
+                    NodePath = refreshParams.NodePath,
+                    ErrorMessage = $"Couldn't find session for session: {uri}"
+                });
+            }
+            else
+            {
+                RunExpandTask(session, refreshParams, true);
+            }
+            await context.SendResult(true);
+        }
+
+        internal async Task HandleCloseSessionRequest(CloseSessionParams closeSessionParams, RequestContext<CloseSessionResponse> context)
+        {
+
+            Logger.Write(LogLevel.Verbose, "HandleCloseSessionRequest");
+            Func<Task<CloseSessionResponse>> closeSession = () =>
+            {
+                Validate.IsNotNull(nameof(closeSessionParams), closeSessionParams);
+                Validate.IsNotNull(nameof(context), context);
+                return Task.Factory.StartNew(() =>
+                {
+                    string uri = closeSessionParams.SessionId;
+                    ObjectExplorerSession session = null;
+                    bool success = false;
+                    if (!sessionMap.TryGetValue(uri, out session))
+                    {
+                        Logger.Write(LogLevel.Verbose, $"Cannot close object explorer session. Couldn't find session for uri. {uri} ");
+                    }
+
+                    if (session != null)
+                    {
+                        // refresh the nodes for given node path
+                        CloseSession(uri);
+                        success = true;
+                    }
+
+                    var response = new CloseSessionResponse() { Success = success, SessionId = uri };
+                    return response;
+                });
+            };
+
+            await HandleRequestAsync(closeSession, context, "HandleCloseSessionRequest");
+        }
+
+        internal void CloseSession(string uri)
         {
             ObjectExplorerSession session;
             if (sessionMap.TryGetValue(uri, out session))
             {
-                // Establish a connection to the specified server/database
+                // Remove the session from active sessions and disconnect
                 sessionMap.TryRemove(session.Uri, out session);
                 connectionService.Disconnect(new DisconnectParams()
                 {
@@ -117,42 +235,51 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer
             }
         }
 
-        internal async Task HandleCreateSessionRequest(ConnectionDetails connectionDetails, RequestContext<CreateSessionResponse> context)
+        private void  RunCreateSessionTask(ConnectionDetails connectionDetails, string uri)
         {
-            Logger.Write(LogLevel.Verbose, "HandleCreateSessionRequest");
+            if (connectionDetails != null && !string.IsNullOrEmpty(uri))
-            Func<Task<CreateSessionResponse>> doCreateSession = async () =>
             {
-                Validate.IsNotNull(nameof(connectionDetails), connectionDetails);
+                Task task = CreateSessionAsync(connectionDetails, uri);
-                Validate.IsNotNull(nameof(context), context);
+                CreateSessionTask = task;
+                Task.Run(async () => await task);
+            }
+        }
 
-                string uri = GenerateUri(connectionDetails);
+        /// <summary>
+        /// For tests only
+        /// </summary>
+        internal Task CreateSessionTask
+        {
+            get;
+            private set;
+        }
 
-                ObjectExplorerSession session;
+        private async Task<SessionCreatedParameters> CreateSessionAsync(ConnectionDetails connectionDetails, string uri)
-                if (!sessionMap.TryGetValue(uri, out session))
+        {
-                {
+            ObjectExplorerSession session;
-                    // Establish a connection to the specified server/database
+            if (!sessionMap.TryGetValue(uri, out session))
-                    session = await DoCreateSession(connectionDetails, uri);
+            {
-                }
+                // Establish a connection to the specified server/database
+                session = await DoCreateSession(connectionDetails, uri);
+            }
 
-                CreateSessionResponse response;
+            SessionCreatedParameters response;
-                if (session == null)
+            if (session != null)
+            {
+                // Else we have a session available, response with existing session information
+                response = new SessionCreatedParameters
                 {
-                    response = new CreateSessionResponse() { Success = false };
+                    Success = true,
-                }
+                    RootNode = session.Root.ToNodeInfo(),
-                else
+                    SessionId = uri,
-                {
+                    ErrorMessage = session.ErrorMessage
-                    // Else we have a session available, response with existing session information
+                    
-                    response = new CreateSessionResponse()
+                };
-                    {
+                 await serviceHost.SendEvent(CreateSessionCompleteNotification.Type, response);
-                        Success = true,
-                        RootNode = session.Root.ToNodeInfo(),
-                        SessionId = session.Uri
-                    };
-                }
                 return response;
-            };
+            }
+            return null;
 
-            await HandleRequestAsync(doCreateSession, context, "HandleCreateSessionRequest");
         }
 
         internal async Task<NodeInfo[]> ExpandNode(ObjectExplorerSession session, string nodePath, bool forceRefresh = false)
@@ -186,9 +313,10 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer
             connectionDetails.PersistSecurityInfo = true;
             ConnectParams connectParams = new ConnectParams() { OwnerUri = uri, Connection = connectionDetails };
 
-            ConnectionCompleteParams connectionResult = await Connect(connectParams);
+            ConnectionCompleteParams connectionResult = await Connect(connectParams, uri);
             if (connectionResult == null)
             {
+                // Connection failed and notification is already sent
                 return null;
             }
 
@@ -198,7 +326,7 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer
         }
         
 
-        private async Task<ConnectionCompleteParams> Connect(ConnectParams connectParams)
+        private async Task<ConnectionCompleteParams> Connect(ConnectParams connectParams, string uri)
         {
             string connectionErrorMessage = string.Empty;
             try
@@ -213,7 +341,11 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer
                 else
                 {
                     Logger.Write(LogLevel.Warning, $"Connection Failed for OE. connection error: {connectionErrorMessage}");
-                    await serviceHost.SendEvent(ConnectionCompleteNotification.Type, result);
+                    await serviceHost.SendEvent(CreateSessionCompleteNotification.Type, new SessionCreatedParameters
+                    {
+                        ErrorMessage = result.ErrorMessage,
+                        SessionId = uri
+                    });
                     return null;
                 }
                
@@ -222,108 +354,46 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer
             {
                 Logger.Write(LogLevel.Warning, $"Connection Failed for OE. connection error:{connectionErrorMessage} error: {ex.Message}");
                 // Send a connection failed error message in this case.
-                ConnectionCompleteParams result = new ConnectionCompleteParams()
+                SessionCreatedParameters result = new SessionCreatedParameters()
                 {
-                    Messages = ex.ToString()
+                    ErrorMessage = ex.ToString(),
+                    SessionId = uri
                 };
-                await serviceHost.SendEvent(ConnectionCompleteNotification.Type, result);
+                await serviceHost.SendEvent(CreateSessionCompleteNotification.Type, result);
                 return null;
             }
         }
 
-        internal async Task HandleExpandRequest(ExpandParams expandParams, RequestContext<ExpandResponse> context)
+        private void RunExpandTask(ObjectExplorerSession session, ExpandParams expandParams, bool forceRefresh = false)
         {
-            Logger.Write(LogLevel.Verbose, "HandleExpandRequest");
+            Task task = ExpandNodeAsync(session, expandParams, forceRefresh);
-            Func<Task<ExpandResponse>> expandNode = async () =>
+            ExpandTask = task;
+            Task.Run(async () =>
             {
-                Validate.IsNotNull(nameof(expandParams), expandParams);
+                await task;
-                Validate.IsNotNull(nameof(context), context);
+            });
-
-                string uri = expandParams.SessionId;
-                ObjectExplorerSession session = null;
-                NodeInfo[] nodes = null;
-                ExpandResponse response;
-                if (!sessionMap.TryGetValue(uri, out session))
-                {
-                    Logger.Write(LogLevel.Verbose, $"Cannot expand object explorer node. Couldn't find session for uri. {uri} ");
-                }
-
-                if (session != null)
-                {
-                    // expand the nodes for given node path
-                    nodes = await ExpandNode(session, expandParams.NodePath);
-                }
-
-                response = new ExpandResponse() { Nodes = nodes, SessionId = uri };
-                return response;
-            };
-
-            await HandleRequestAsync(expandNode, context, "HandleExpandRequest");
         }
 
-        internal async Task HandleRefreshRequest(RefreshParams refreshParams, RequestContext<RefreshResponse> context)
+        /// <summary>
+        /// For tests only
+        /// </summary>
+        internal Task ExpandTask
         {
-            Logger.Write(LogLevel.Verbose, "HandleRefreshRequest");
+            get;
-            Func<Task<RefreshResponse>> refreshNode = async () =>
+            set;
-            {
-                Validate.IsNotNull(nameof(refreshParams), refreshParams);
-                Validate.IsNotNull(nameof(context), context);
-
-                string uri = refreshParams.SessionId;
-                ObjectExplorerSession session = null;
-                NodeInfo[] nodes = null;
-                RefreshResponse response;
-                if (!sessionMap.TryGetValue(uri, out session))
-                {
-                    Logger.Write(LogLevel.Verbose, $"Cannot refresh object explorer node. Couldn't find session for uri. {uri} ");
-                }
-
-                if (session != null)
-                {
-                    // refresh the nodes for given node path
-                    nodes = await ExpandNode(session, refreshParams.NodePath, true);
-                }
-
-                response = new RefreshResponse() { Nodes = nodes, SessionId = uri };
-                return response;
-            };
-
-            await HandleRequestAsync(refreshNode, context, "HandleRefreshRequest");
         }
 
-        internal async Task HandleCloseSessionRequest(CloseSessionParams closeSessionParams, RequestContext<CloseSessionResponse> context)
+        private async Task ExpandNodeAsync(ObjectExplorerSession session, ExpandParams expandParams, bool forceRefresh = false)
         {
-            Validate.IsNotNull(nameof(closeSessionParams), closeSessionParams);
+            NodeInfo[] nodes = null;
-            Validate.IsNotNull(nameof(context), context);
+            nodes = await ExpandNode(session, expandParams.NodePath, forceRefresh);
-            Logger.Write(LogLevel.Verbose, "HandleCloseSessionRequest");
+            ExpandResponse response = new ExpandResponse() { Nodes = nodes, SessionId = session.Uri, NodePath = expandParams.NodePath };
-            Func<Task<CloseSessionResponse>> closeSession = () =>
+            await serviceHost.SendEvent(ExpandCompleteNotification.Type, response);
-            {
-                return Task.Factory.StartNew(() =>
-                {
-                    string uri = closeSessionParams.SessionId;
-                    ObjectExplorerSession session = null;
-                    bool success = false;
-                    if (!sessionMap.TryGetValue(uri, out session))
-                    {
-                        Logger.Write(LogLevel.Verbose, $"Cannot close object explorer session. Couldn't find session for uri. {uri} ");
-                    }
-
-                    if (session != null)
-                    {
-                        // refresh the nodes for given node path
-                        CloseSession(uri);
-                        success = true;
-                    }
-
-                    var response = new CloseSessionResponse() {Success = success, SessionId = uri};
-                    return response;
-                });
-            };
-
-            await HandleRequestAsync(closeSession, context, "HandleCloseSessionRequest");
         }
 
-        private async Task HandleRequestAsync<T>(Func<Task<T>> handler, RequestContext<T> requestContext, string requestType)
+       
+
+        private async Task<T> HandleRequestAsync<T>(Func<Task<T>> handler, RequestContext<T> requestContext, string requestType)
         {
             Logger.Write(LogLevel.Verbose, requestType);
 
@@ -331,11 +401,13 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer
             {
                 T result = await handler();
                 await requestContext.SendResult(result);
+                return result;
             }
             catch (Exception ex)
             {
                 await requestContext.SendError(ex.ToString());
             }
+            return default(T);
         }
 
         /// <summary>
@@ -441,21 +513,19 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer
             public string Uri { get; private set; }
             public TreeNode Root { get; private set; }
 
+            public string ErrorMessage { get; set; }
+
             public static ObjectExplorerSession CreateSession(ConnectionCompleteParams response, IMultiServiceProvider serviceProvider)
             {
-                TreeNode rootNode = new ServerNode(response, serviceProvider);
+                ServerNode rootNode = new ServerNode(response, serviceProvider);
                 var session = new ObjectExplorerSession(response.OwnerUri, rootNode, serviceProvider, serviceProvider.GetService<ConnectionService>());
                 if (!ObjectExplorerUtils.IsSystemDatabaseConnection(response.ConnectionSummary.DatabaseName))
                 {
                     // Assuming the databases are in a folder under server node
-                    var children = rootNode.Expand();
+                    DatabaseTreeNode databaseNode = new DatabaseTreeNode(rootNode, response.ConnectionSummary.DatabaseName);
-                    var databasesRoot = children.FirstOrDefault(x => x.NodeTypeId == NodeTypes.Databases);
-                    var databasesChildren = databasesRoot.Expand(response.ConnectionSummary.DatabaseName);
-                    var databases = databasesChildren.Where(x => x.NodeType == NodeTypes.Database.ToString());
-                    var databaseNode = databases.FirstOrDefault(d => d.Label == response.ConnectionSummary.DatabaseName);
-                    databaseNode.Label = rootNode.Label;
                     session.Root = databaseNode;
                 }
+               
                 return session;
             }
             

src/Microsoft.SqlTools.ServiceLayer/ObjectExplorer/SmoModel/DatabaseTreeNode.cs

@@ -10,6 +10,14 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer.SmoModel
 {
     internal partial class DatabaseTreeNode
     {
+        public DatabaseTreeNode(ServerNode serverNode, string databaseName): this()
+        {
+            Parent = serverNode;
+            NodeValue = databaseName;
+            Database db = new Database(serverNode.GetContextAs<SmoQueryContext>().Server, this.NodeValue);
+            db.Refresh();
+            CacheInfoFromModel(db);
+        }
 
         /// <summary>
         /// Initializes the context and ensures that 

src/Microsoft.SqlTools.ServiceLayer/ObjectExplorer/SmoModel/SmoChildFactoryBase.cs

@@ -42,6 +42,7 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer.SmoModel
                 string error = string.Format(CultureInfo.InvariantCulture, "Failed expanding oe children. parent:{0} error:{1} inner:{2} stacktrace:{3}", 
                     parent != null ? parent.GetNodePath() : "", ex.Message, ex.InnerException != null ? ex.InnerException.Message : "", ex.StackTrace);
                 Logger.Write(LogLevel.Error, error);
+                throw ex;
             }
             finally
             {
@@ -97,7 +98,8 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer.SmoModel
                 string propertyFilter = GetProperyFilter(filters, querier.GetType(), validForFlag);
                 try
                 {
-                    foreach (var smoObject in querier.Query(context, propertyFilter, refresh))
+                    var smoObjectList = querier.Query(context, propertyFilter, refresh).ToList();
+                    foreach (var smoObject in smoObjectList)
                     {
                         if (smoObject == null)
                         {
@@ -116,6 +118,7 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer.SmoModel
                     string error = string.Format(CultureInfo.InvariantCulture, "Failed getting smo objects. parent:{0} querier: {1} error:{2} inner:{3} stacktrace:{4}",
                     parent != null ? parent.GetNodePath() : "", querier.GetType(), ex.Message, ex.InnerException != null ? ex.InnerException.Message : "", ex.StackTrace);
                     Logger.Write(LogLevel.Error, error);
+                    throw ex;
                 }
             }
         }

src/Microsoft.SqlTools.ServiceLayer/ObjectExplorer/SmoModel/SmoQuerier.cs

@@ -83,7 +83,7 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer.SmoModel
             string urn = string.Empty;
             try
             {
-                string parentUrn = smoObject.Urn;
+                string parentUrn = smoObject.Urn != null ? smoObject.Urn.Value : string.Empty;
                 urn = parentUrn != null ? $"{parentUrn.ToString()}/{objectName}" + filter : string.Empty;
 
                 if (!string.IsNullOrEmpty(urn))
@@ -102,6 +102,7 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer.SmoModel
                 string error = string.Format(CultureInfo.InvariantCulture, "Failed getting urns. error:{0} inner:{1} stacktrace:{2}",
                  ex.Message, ex.InnerException != null ? ex.InnerException.Message : "", ex.StackTrace);
                 Logger.Write(LogLevel.Error, error);
+                throw ex;
             }
 
             return urns;

test/Microsoft.SqlTools.ServiceLayer.UnitTests/ObjectExplorer/ObjectExplorerServiceTests.cs

@@ -50,8 +50,6 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.ObjectExplorer
         {
             // Given the connection service fails to connect
             ConnectionDetails details = TestObjects.GetTestConnectionDetails();
-            ConnectionCompleteParams completeParams = null;
-            serviceHostMock.AddEventHandling(ConnectionCompleteNotification.Type, (et, p) => completeParams = p);
 
             string expectedExceptionText = "Error!!!";
             connectionServiceMock.Setup(c => c.Connect(It.IsAny<ConnectParams>()))
@@ -59,19 +57,17 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.ObjectExplorer
 
             // when creating a new session
             // then expect the create session request to return false
-            await RunAndVerify<CreateSessionResponse>(
+            await RunAndVerify<CreateSessionResponse, SessionCreatedParameters>(
-                test: (requestContext) => service.HandleCreateSessionRequest(details, requestContext),
+                test: (requestContext) => CallCreateSession(details, requestContext),
                 verify: (actual =>
                 {
-                    Assert.False(actual.Success);
+                    Assert.NotNull(actual.SessionId);
-                    Assert.Null(actual.SessionId);
+                    Assert.NotNull(actual);
-                    Assert.Null(actual.RootNode);
+                    Assert.True(actual.ErrorMessage.Contains(expectedExceptionText));
                 }));
 
             // And expect error notification to be sent
-            serviceHostMock.Verify(x => x.SendEvent(ConnectionCompleteNotification.Type, It.IsAny<ConnectionCompleteParams>()), Times.Once());
+            serviceHostMock.Verify(x => x.SendEvent(CreateSessionCompleteNotification.Type, It.IsAny<SessionCreatedParameters>()), Times.Once());
-            Assert.NotNull(completeParams);
-            Assert.True(completeParams.Messages.Contains(expectedExceptionText));
         }
         
 
@@ -138,10 +134,11 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.ObjectExplorer
                 NodePath = "Any path"
             };
 
+
             // when expanding
             // then expect the nodes are server children 
-            await RunAndVerify<ExpandResponse>(
+            await RunAndVerify<bool, ExpandResponse>(
-                test: (requestContext) => service.HandleExpandRequest(expandParams, requestContext),
+                test: (requestContext) => CallServiceExpand(expandParams, requestContext),
                 verify: (actual =>
                 {
                     Assert.Equal(actual.SessionId, expandParams.SessionId);
@@ -160,8 +157,8 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.ObjectExplorer
 
             // when expanding
             // then expect the nodes are server children 
-            await RunAndVerify<RefreshResponse>(
+            await RunAndVerify<bool, ExpandResponse>(
-                test: (requestContext) => service.HandleRefreshRequest(expandParams, requestContext),
+                test: (requestContext) => CallServiceRefresh(expandParams, requestContext),
                 verify: (actual =>
                 {
                     Assert.Equal(actual.SessionId, expandParams.SessionId);
@@ -179,8 +176,8 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.ObjectExplorer
 
             // when expanding
             // then expect the nodes are server children 
-            await RunAndVerify<CloseSessionResponse>(
+            await RunAndVerify<CloseSessionResponse, CloseSessionResponse>(
-                test: (requestContext) => service.HandleCloseSessionRequest(closeSessionParamsparams, requestContext),
+                test: (requestContext) => CallCloseSession(closeSessionParamsparams, requestContext),
                 verify: (actual =>
                 {
                     Assert.Equal(actual.SessionId, closeSessionParamsparams.SessionId);
@@ -199,8 +196,8 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.ObjectExplorer
 
             // when expanding
             // then expect the nodes are server children 
-            await RunAndVerify<CloseSessionResponse>(
+            await RunAndVerify<CloseSessionResponse, CloseSessionResponse>(
-                test: (requestContext) => service.HandleCloseSessionRequest(closeSessionParamsparams, requestContext),
+                test: (requestContext) => CallCloseSession(closeSessionParamsparams, requestContext),
                 verify: (actual =>
                 {
                     Assert.Equal(actual.SessionId, closeSessionParamsparams.SessionId);
@@ -211,7 +208,7 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.ObjectExplorer
             connectionServiceMock.Verify(c => c.Disconnect(It.IsAny<DisconnectParams>()));
         }
 
-        private async Task<CreateSessionResponse> CreateSession()
+        private async Task<SessionCreatedParameters> CreateSession()
         {
             ConnectionDetails details = new ConnectionDetails()
             {
@@ -221,9 +218,11 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.ObjectExplorer
                 ServerName = "serverName"
             };
 
-            serviceHostMock.AddEventHandling(ConnectionCompleteNotification.Type, null);
+            SessionCreatedParameters sessionResult = null;
+            serviceHostMock.AddEventHandling(CreateSessionCompleteNotification.Type, (et, p) => sessionResult = p);
             CreateSessionResponse result = default(CreateSessionResponse);
             var contextMock = RequestContextMocks.Create<CreateSessionResponse>(r => result = r).AddErrorHandling(null);
+          
             connectionServiceMock.Setup(c => c.Connect(It.IsAny<ConnectParams>()))
                 .Returns((ConnectParams connectParams) => Task.FromResult(GetCompleteParamsForConnection(connectParams.OwnerUri, details)));
 
@@ -236,8 +235,9 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.ObjectExplorer
 
             connectionServiceMock.Setup(c => c.Disconnect(It.IsAny<DisconnectParams>())).Returns(true);
             await service.HandleCreateSessionRequest(details, contextMock.Object);
+            await service.CreateSessionTask;
 
-            return result;
+            return sessionResult;
         }
 
         private async Task ExpandAndVerifyServerNodes()
@@ -251,8 +251,8 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.ObjectExplorer
 
             // when expanding
             // then expect the nodes are server children 
-            await RunAndVerify<ExpandResponse>(
+            await RunAndVerify<bool, ExpandResponse>(
-                test: (requestContext) => service.HandleExpandRequest(expandParams, requestContext),
+                test: (requestContext) => CallServiceExpand(expandParams, requestContext),
                 verify: (actual =>
                 {
                     Assert.Equal(actual.SessionId, session.SessionId);
@@ -272,8 +272,8 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.ObjectExplorer
 
             // when expanding
             // then expect the nodes are server children 
-            await RunAndVerify<RefreshResponse>(
+            await RunAndVerify<bool, ExpandResponse>(
-                test: (requestContext) => service.HandleRefreshRequest(expandParams, requestContext),
+                test: (requestContext) => CallServiceRefresh(expandParams, requestContext),
                 verify: (actual =>
                 {
                     Assert.Equal(actual.SessionId, session.SessionId);
@@ -282,17 +282,75 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.ObjectExplorer
                 }));
         }
 
+        private async Task<ExpandResponse> CallServiceRefresh(RefreshParams expandParams, RequestContext<bool> requestContext)
+        {
+            ExpandResponse result = null;
+            serviceHostMock.AddEventHandling(ExpandCompleteNotification.Type, (et, p) => result = p);
+
+            await service.HandleRefreshRequest(expandParams, requestContext);
+            Task task = service.ExpandTask;
+            if (task != null)
+            {
+                await task;
+            }
+
+            return result;
+
+        }
+
+        private async Task<ExpandResponse> CallServiceExpand(ExpandParams expandParams, RequestContext<bool> requestContext)
+        {
+            ExpandResponse result = null;
+            serviceHostMock.AddEventHandling(ExpandCompleteNotification.Type, (et, p) => result = p);
+
+            await service.HandleExpandRequest(expandParams, requestContext);
+            Task task = service.ExpandTask;
+            if (task != null)
+            {
+                await task;
+            }
+            return result;
+        }
+
+        private async Task<SessionCreatedParameters> CallCreateSession(ConnectionDetails connectionDetails, RequestContext<CreateSessionResponse> context)
+        {
+            SessionCreatedParameters result = null;
+            serviceHostMock.AddEventHandling(CreateSessionCompleteNotification.Type, (et, p) => result = p);
+
+            await service.HandleCreateSessionRequest(connectionDetails, context);
+            Task task =  service.CreateSessionTask;
+            if (task != null)
+            {
+                await task;
+            }
+            return result;
+        }
+
+        private async Task<CloseSessionResponse> CallCloseSession(CloseSessionParams closeSessionParams, RequestContext<CloseSessionResponse> context)
+        {
+            SessionCreatedParameters result = null;
+            serviceHostMock.AddEventHandling(CreateSessionCompleteNotification.Type, (et, p) => result = p);
+
+            await service.HandleCloseSessionRequest(closeSessionParams, context);
+            return null;
+        }
+
         private async Task CreateSessionRequestAndVerifyServerNodeHelper(ConnectionDetails details)
         {
             serviceHostMock.AddEventHandling(ConnectionCompleteNotification.Type, null);
+            //SessionCreatedParameters sessionResult
+            
 
             connectionServiceMock.Setup(c => c.Connect(It.IsAny<ConnectParams>()))
                 .Returns((ConnectParams connectParams) => Task.FromResult(GetCompleteParamsForConnection(connectParams.OwnerUri, details)));
             
             // when creating a new session
             // then expect the create session request to return false
-            await RunAndVerify<CreateSessionResponse>(
+            await RunAndVerify<CreateSessionResponse, SessionCreatedParameters>(
-                test: (requestContext) => service.HandleCreateSessionRequest(details, requestContext),
+                test: (requestContext) => 
+                {
+                    return CallCreateSession(details, requestContext);
+                },
                 verify: (actual =>
                 {
                     Assert.True(actual.Success);
@@ -339,15 +397,19 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.ObjectExplorer
             };
         }
 
-        private async Task RunAndVerify<T>(Func<RequestContext<T>, Task> test, Action<T> verify)
+        private async Task RunAndVerify<T, TResult>(Func<RequestContext<T>, Task<TResult>> test, Action<TResult> verify)
         {
             T result = default(T);
             var contextMock = RequestContextMocks.Create<T>(r => result = r).AddErrorHandling(null);
-            await test(contextMock.Object);
+            TResult actualResult = await test(contextMock.Object);
-            VerifyResult(contextMock, verify, result);
+            if (actualResult == null && typeof(TResult) == typeof(T))
+            {
+                actualResult = (TResult)Convert.ChangeType(result, typeof(TResult));
+            }
+            VerifyResult(contextMock, verify, actualResult);
         }
         
-        private void VerifyResult<T>(Mock<RequestContext<T>> contextMock, Action<T> verify, T actual)
+        private void VerifyResult<T, TResult>(Mock<RequestContext<T>> contextMock, Action<TResult> verify, TResult actual)
         {
             contextMock.Verify(c => c.SendResult(It.IsAny<T>()), Times.Once);
             contextMock.Verify(c => c.SendError(It.IsAny<string>(), It.IsAny<int>()), Times.Never);
@@ -359,6 +421,5 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.ObjectExplorer
             contextMock.Verify(c => c.SendResult(It.IsAny<T>()), Times.Never);
             contextMock.Verify(c => c.SendError(It.IsAny<string>(), It.IsAny<int>()), Times.Once);
         }
-
     }
 }