From ec09ad21d291b33ac04b9e5ac4714ae958fbcfdb Mon Sep 17 00:00:00 2001
From: Charles Gagnon <chgagnon@microsoft.com>
Date: Mon, 29 Jul 2019 16:55:37 -0700
Subject: [PATCH] Revert "Default to Server level context for all OE nodes
 (#829)" (#838)

This reverts commit a4b6c300ac91875028ba4ed74f3304afaf45aa96.

(cherry picked from commit 7512a2bbe3154e4dd4b63927fcba52db78424994)
---
 .../ObjectExplorer/ObjectExplorerService.cs   | 29 ++++++++++++-------
 1 file changed, 19 insertions(+), 10 deletions(-)

diff --git a/src/Microsoft.SqlTools.ServiceLayer/ObjectExplorer/ObjectExplorerService.cs b/src/Microsoft.SqlTools.ServiceLayer/ObjectExplorer/ObjectExplorerService.cs
index 1c5fbc6c..4c773fa8 100644
--- a/src/Microsoft.SqlTools.ServiceLayer/ObjectExplorer/ObjectExplorerService.cs
+++ b/src/Microsoft.SqlTools.ServiceLayer/ObjectExplorer/ObjectExplorerService.cs
@@ -73,7 +73,7 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer
             {
                 this.bindingQueue = value;
             }
-        }
+        }        
 
         /// <summary>
         /// Internal for testing only
@@ -485,6 +485,7 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer
                 ObjectExplorerSession session = null;
                 connectionDetails.PersistSecurityInfo = true;
                 ConnectParams connectParams = new ConnectParams() { OwnerUri = uri, Connection = connectionDetails, Type = Connection.ConnectionType.ObjectExplorer };
+                bool isDefaultOrSystemDatabase = DatabaseUtils.IsSystemDatabaseConnection(connectionDetails.DatabaseName) || string.IsNullOrWhiteSpace(connectionDetails.DatabaseDisplayName);
 
                 ConnectionInfo connectionInfo;
                 ConnectionCompleteParams connectionResult = await Connect(connectParams, uri);
@@ -506,7 +507,7 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer
                            waitForLockTimeout: timeout,
                            bindOperation: (bindingContext, cancelToken) =>
                            {
-                               session = ObjectExplorerSession.CreateSession(connectionResult, serviceProvider, bindingContext.ServerConnection);
+                               session = ObjectExplorerSession.CreateSession(connectionResult, serviceProvider, bindingContext.ServerConnection, isDefaultOrSystemDatabase);
                                session.ConnectionInfo = connectionInfo;
 
                                sessionMap.AddOrUpdate(uri, session, (key, oldSession) => session);
@@ -525,7 +526,7 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer
                 await SendSessionFailedNotification(uri, ex.Message);
                 return null;
             }
-        }
+        }        
 
         private async Task<ConnectionCompleteParams> Connect(ConnectParams connectParams, string uri)
         {
@@ -544,7 +545,7 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer
                     await SendSessionFailedNotification(uri, result.ErrorMessage);
                     return null;
                 }
-
+               
             }
             catch (Exception ex)
             {
@@ -584,7 +585,7 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer
             ExpandTask = task;
             Task.Run(async () =>
             {
-                ObjectExplorerTaskResult result =  await RunTaskWithTimeout(task,
+                ObjectExplorerTaskResult result =  await RunTaskWithTimeout(task, 
                     settings?.ExpandTimeout ?? ObjectExplorerSettings.DefaultExpandTimeout);
 
                 if (result != null && !result.IsCompleted)
@@ -653,7 +654,7 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer
         internal static string GenerateUri(ConnectionDetails details)
         {
             return ConnectedBindingQueue.GetConnectionContextKey(details);
-        }
+        }      
 
         public IEnumerable<ChildFactory> GetApplicableChildFactories(TreeNode item)
         {
@@ -754,7 +755,7 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer
             {
                 bindingQueue.OnUnhandledException -= OnUnhandledException;
                 bindingQueue.Dispose();
-            }
+            }            
         }
 
         private async void OnUnhandledException(string queueKey, Exception ex)
@@ -809,12 +810,20 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer
 
             public string ErrorMessage { get; set; }
 
-            public static ObjectExplorerSession CreateSession(ConnectionCompleteParams response, IMultiServiceProvider serviceProvider, ServerConnection serverConnection)
+            public static ObjectExplorerSession CreateSession(ConnectionCompleteParams response, IMultiServiceProvider serviceProvider, ServerConnection serverConnection, bool isDefaultOrSystemDatabase)
             {
                 ServerNode rootNode = new ServerNode(response, serviceProvider, serverConnection);
-                return new ObjectExplorerSession(response.OwnerUri, rootNode, serviceProvider, serviceProvider.GetService<ConnectionService>());
-            }
+                var session = new ObjectExplorerSession(response.OwnerUri, rootNode, serviceProvider, serviceProvider.GetService<ConnectionService>());
+                if (!isDefaultOrSystemDatabase)
+                {
+                    // Assuming the databases are in a folder under server node
+                    DatabaseTreeNode databaseNode = new DatabaseTreeNode(rootNode, response.ConnectionSummary.DatabaseName);
+                    session.Root = databaseNode;
+                }
 
+                return session;
+            }
+            
         }
     }