Kusto AdminService Unit Tests (#1257)

* Added unit tests for AdminService. Added interface for IConnectionService.

* Kusto connection map to ConnectionManager (#1258)

* Created ConnectionManager and moved connectionMap. Changed references from ConnectionService to ConnectionManager.

* Changed creation of ConnectionManager to MEF instead of Instance creation

src/Microsoft.Kusto.ServiceLayer/Admin/AdminService.cs

@@ -4,15 +4,11 @@
 //
 
 using System;
-using System.Collections.Generic;
-using System.Linq;
 using System.Threading.Tasks;
 using Microsoft.SqlTools.Hosting.Protocol;
 using Microsoft.Kusto.ServiceLayer.Admin.Contracts;
 using Microsoft.Kusto.ServiceLayer.Connection;
 using Microsoft.Kusto.ServiceLayer.DataSource;
-using Microsoft.Kusto.ServiceLayer.DataSource.Metadata;
-using Microsoft.Kusto.ServiceLayer.Utility;
 
 namespace Microsoft.Kusto.ServiceLayer.Admin
 {
@@ -23,7 +19,7 @@ namespace Microsoft.Kusto.ServiceLayer.Admin
     {
         private static readonly Lazy<AdminService> _instance = new Lazy<AdminService>(() => new AdminService());
 
-        private static ConnectionService _connectionService;
+        private IConnectionManager _connectionManager;
 
         /// <summary>
         /// Gets the singleton instance object
@@ -33,23 +29,23 @@ namespace Microsoft.Kusto.ServiceLayer.Admin
         /// <summary>
         /// Initializes the service instance
         /// </summary>
-        public void InitializeService(ServiceHost serviceHost, ConnectionService connectionService)
+        public void InitializeService(IProtocolEndpoint serviceHost, IConnectionManager connectionManager)
         {
+            _connectionManager = connectionManager;
             serviceHost.SetRequestHandler(GetDatabaseInfoRequest.Type, HandleGetDatabaseInfoRequest);
-            _connectionService = connectionService;
         }
 
         /// <summary>
         /// Handle get database info request
         /// </summary>
-        private async Task HandleGetDatabaseInfoRequest(GetDatabaseInfoParams databaseParams, RequestContext<GetDatabaseInfoResponse> requestContext)
+        public async Task HandleGetDatabaseInfoRequest(GetDatabaseInfoParams databaseParams, RequestContext<GetDatabaseInfoResponse> requestContext)
         {
             try
             {
                 var infoResponse = await Task.Run(() =>
                 {
                     DatabaseInfo info = null;
-                    if (_connectionService.TryFindConnection(databaseParams.OwnerUri, out var connInfo))
+                    if (_connectionManager.TryGetValue(databaseParams.OwnerUri, out var connInfo))
                     {
                         info = GetDatabaseInfo(connInfo);
                     }
@@ -61,7 +57,7 @@ namespace Microsoft.Kusto.ServiceLayer.Admin
             }
             catch (Exception ex)
             {
-                await requestContext.SendError(ex.ToString());
+                await requestContext.SendError(ex);
             }
         }
 
@@ -70,7 +66,7 @@ namespace Microsoft.Kusto.ServiceLayer.Admin
         /// </summary>
         /// <param name="connInfo"></param>
         /// <returns></returns>
-        public DatabaseInfo GetDatabaseInfo(ConnectionInfo connInfo)
+        private DatabaseInfo GetDatabaseInfo(ConnectionInfo connInfo)
         {
             if (string.IsNullOrEmpty(connInfo.ConnectionDetails.DatabaseName))
             {

src/Microsoft.Kusto.ServiceLayer/Connection/ConnectionManager.cs

@@ -0,0 +1,39 @@
+using System.Collections.Concurrent;
+using System.Composition;
+
+namespace Microsoft.Kusto.ServiceLayer.Connection
+{
+    [Export(typeof(IConnectionManager))]
+    public class ConnectionManager : IConnectionManager
+    {
+        /// <summary>
+        /// Map from script URIs to ConnectionInfo objects
+        /// </summary>
+        private ConcurrentDictionary<string, ConnectionInfo> _ownerToConnectionMap;
+
+        public ConnectionManager()
+        {
+            _ownerToConnectionMap = new ConcurrentDictionary<string, ConnectionInfo>();
+        }
+
+        public bool TryGetValue(string ownerUri, out ConnectionInfo info)
+        {
+            return _ownerToConnectionMap.TryGetValue(ownerUri, out info);
+        }
+
+        public bool ContainsKey(string ownerUri)
+        {
+            return _ownerToConnectionMap.ContainsKey(ownerUri);
+        }
+
+        public bool TryAdd(string ownerUri, ConnectionInfo connectionInfo)
+        {
+            return _ownerToConnectionMap.TryAdd(ownerUri, connectionInfo);
+        }
+
+        public bool TryRemove(string ownerUri)
+        {
+            return _ownerToConnectionMap.TryRemove(ownerUri, out _);
+        }
+    }
+}

src/Microsoft.Kusto.ServiceLayer/Connection/ConnectionService.cs

@@ -25,7 +25,7 @@ namespace Microsoft.Kusto.ServiceLayer.Connection
     /// <summary>
     /// Main class for the Connection Management services
     /// </summary>
-    public class ConnectionService
+    public class ConnectionService : IConnectionService
     {
         private const string PasswordPlaceholder = "******";
 
@@ -53,11 +53,6 @@ namespace Microsoft.Kusto.ServiceLayer.Connection
 
         private ConcurrentDictionary<string, IConnectedBindingQueue> connectedQueues = new ConcurrentDictionary<string, IConnectedBindingQueue>();
 
-        /// <summary>
-        /// Map from script URIs to ConnectionInfo objects
-        /// </summary>
-        private Dictionary<string, ConnectionInfo> OwnerToConnectionMap { get; } = new Dictionary<string, ConnectionInfo>();
-
         /// <summary>
         /// Database Lock manager instance
         /// </summary>
@@ -122,8 +117,7 @@ namespace Microsoft.Kusto.ServiceLayer.Connection
         /// </summary>
         private IDataSourceConnectionFactory _dataSourceConnectionFactory;
 
-        // Attempts to link a URI to an actively used connection for this URI
+        private IConnectionManager _connectionManager;
-        public virtual bool TryFindConnection(string ownerUri, out ConnectionInfo connectionInfo) => this.OwnerToConnectionMap.TryGetValue(ownerUri, out connectionInfo);
 
         /// <summary>
         /// Validates the given ConnectParams object. 
@@ -171,7 +165,7 @@ namespace Microsoft.Kusto.ServiceLayer.Connection
             // If there is no ConnectionInfo in the map, create a new ConnectionInfo, 
             // but wait until later when we are connected to add it to the map.
             bool connectionChanged = false;
-            if (!OwnerToConnectionMap.TryGetValue(connectionParams.OwnerUri, out ConnectionInfo connectionInfo))
+            if (!_connectionManager.TryGetValue(connectionParams.OwnerUri, out ConnectionInfo connectionInfo))
             {
                 connectionInfo = new ConnectionInfo(_dataSourceConnectionFactory, connectionParams.OwnerUri, connectionParams.Connection);
             }
@@ -197,10 +191,10 @@ namespace Microsoft.Kusto.ServiceLayer.Connection
             }
 
             // If this is the first connection for this URI, add the ConnectionInfo to the map
-            bool addToMap = connectionChanged || !OwnerToConnectionMap.ContainsKey(connectionParams.OwnerUri);
+            bool addToMap = connectionChanged || !_connectionManager.ContainsKey(connectionParams.OwnerUri);
             if (addToMap)
             {
-                OwnerToConnectionMap[connectionParams.OwnerUri] = connectionInfo;
+                _connectionManager.TryAdd(connectionParams.OwnerUri, connectionInfo);
             }
 
             // Return information about the connected SQL Server instance
@@ -216,7 +210,7 @@ namespace Microsoft.Kusto.ServiceLayer.Connection
         internal bool TryRefreshAuthToken(string ownerUri, out string token)
         {
             token = string.Empty;
-            if (!TryFindConnection(ownerUri, out ConnectionInfo connection))
+            if (!_connectionManager.TryGetValue(ownerUri, out ConnectionInfo connection))
             {
                 return false;
             }
@@ -485,7 +479,7 @@ namespace Microsoft.Kusto.ServiceLayer.Connection
 
             // Try to get the ConnectionInfo, if it exists
             ConnectionInfo connectionInfo;
-            if (!OwnerToConnectionMap.TryGetValue(ownerUri, out connectionInfo))
+            if (!_connectionManager.TryGetValue(ownerUri, out connectionInfo))
             {
                 throw new ArgumentOutOfRangeException(SR.ConnectionServiceListDbErrorNotConnected(ownerUri));
             }
@@ -592,7 +586,7 @@ namespace Microsoft.Kusto.ServiceLayer.Connection
 
             // Lookup the ConnectionInfo owned by the URI
             ConnectionInfo info;
-            if (!OwnerToConnectionMap.TryGetValue(disconnectParams.OwnerUri, out info))
+            if (!_connectionManager.TryGetValue(disconnectParams.OwnerUri, out info))
             {
                 return false;
             }
@@ -617,7 +611,7 @@ namespace Microsoft.Kusto.ServiceLayer.Connection
             // If the ConnectionInfo has no more connections, remove the ConnectionInfo
             if (info.CountConnections == 0)
             {
-                OwnerToConnectionMap.Remove(disconnectParams.OwnerUri);
+                _connectionManager.TryRemove(disconnectParams.OwnerUri);
             }
 
             // Handle Telemetry disconnect events if we are disconnecting the default connection
@@ -719,7 +713,7 @@ namespace Microsoft.Kusto.ServiceLayer.Connection
             }
 
             // Use the existing connection as a base for the search
-            if (!TryFindConnection(listDatabasesParams.OwnerUri, out ConnectionInfo info))
+            if (!_connectionManager.TryGetValue(listDatabasesParams.OwnerUri, out ConnectionInfo info))
             {
                 throw new Exception(SR.ConnectionServiceListDbErrorNotConnected(listDatabasesParams.OwnerUri));
             }
@@ -731,10 +725,11 @@ namespace Microsoft.Kusto.ServiceLayer.Connection
         }
 
         public void InitializeService(IProtocolEndpoint serviceHost, IDataSourceConnectionFactory dataSourceConnectionFactory, 
-            IConnectedBindingQueue connectedBindingQueue)
+            IConnectedBindingQueue connectedBindingQueue, IConnectionManager connectionManager)
         {
             _serviceHost = serviceHost;
             _dataSourceConnectionFactory = dataSourceConnectionFactory;
+            _connectionManager = connectionManager;
             connectedQueues.AddOrUpdate("Default", connectedBindingQueue, (key, old) => connectedBindingQueue);
             LockedDatabaseManager.ConnectionService = this;
 
@@ -872,7 +867,7 @@ namespace Microsoft.Kusto.ServiceLayer.Connection
             {
                 string connectionString = string.Empty;
                 ConnectionInfo info;
-                if (TryFindConnection(connStringParams.OwnerUri, out info))
+                if (_connectionManager.TryGetValue(connStringParams.OwnerUri, out info))
                 {
                     try
                     {
@@ -946,7 +941,7 @@ namespace Microsoft.Kusto.ServiceLayer.Connection
         /// <param name="newDatabaseName">Name of the database to change the connection to</param>
         private bool ChangeConnectionDatabaseContext(string ownerUri, string newDatabaseName, bool force = false)
         {
-            if (!TryFindConnection(ownerUri, out ConnectionInfo info))
+            if (!_connectionManager.TryGetValue(ownerUri, out ConnectionInfo info))
             {
                 return false;
             }

src/Microsoft.Kusto.ServiceLayer/Connection/IConnectionManager.cs

@@ -0,0 +1,10 @@
+namespace Microsoft.Kusto.ServiceLayer.Connection
+{
+    public interface IConnectionManager
+    {
+        bool TryGetValue(string ownerUri, out ConnectionInfo info);
+        bool ContainsKey(string ownerUri);
+        bool TryAdd(string ownerUri, ConnectionInfo connectionInfo);
+        bool TryRemove(string ownerUri);
+    }
+}

src/Microsoft.Kusto.ServiceLayer/Connection/IConnectionService.cs

@@ -0,0 +1,80 @@
+using System.Threading.Tasks;
+using Microsoft.Kusto.ServiceLayer.Connection.Contracts;
+using Microsoft.Kusto.ServiceLayer.LanguageServices;
+using Microsoft.SqlTools.Hosting.Protocol;
+
+namespace Microsoft.Kusto.ServiceLayer.Connection
+{
+    public interface IConnectionService
+    {
+        /// <summary>
+        /// Register a new connection queue if not already registered
+        /// </summary>
+        /// <param name="type"></param>
+        /// <param name="connectedQueue"></param>
+        void RegisterConnectedQueue(string type, IConnectedBindingQueue connectedQueue);
+
+        /// <summary>
+        /// Open a connection with the specified ConnectParams
+        /// </summary>
+        Task<ConnectionCompleteParams> Connect(ConnectParams connectionParams);
+
+        /// <summary>
+        /// Gets the existing connection with the given URI and connection type string. If none exists, 
+        /// creates a new connection. This cannot be used to create a default connection or to create a 
+        /// connection if a default connection does not exist.
+        /// </summary>
+        /// <param name="ownerUri">URI identifying the resource mapped to this connection</param>
+        /// <param name="connectionType">
+        /// What the purpose for this connection is. A single resource
+        /// such as a SQL file may have multiple connections - one for Intellisense, another for query execution
+        /// </param>
+        /// <param name="alwaysPersistSecurity">
+        /// Workaround for .Net Core clone connection issues: should persist security be used so that
+        /// when SMO clones connections it can do so without breaking on SQL Password connections.
+        /// This should be removed once the core issue is resolved and clone works as expected
+        /// </param>
+        /// <returns>A DB connection for the connection type requested</returns>
+        Task<ReliableDataSourceConnection> GetOrOpenConnection(string ownerUri, string connectionType, bool alwaysPersistSecurity = false);
+
+        /// <summary>
+        /// Cancel a connection that is in the process of opening.
+        /// </summary>
+        bool CancelConnect(CancelConnectParams cancelParams);
+
+        /// <summary>
+        /// Close a connection with the specified connection details.
+        /// </summary>
+        bool Disconnect(DisconnectParams disconnectParams);
+
+        void InitializeService(IProtocolEndpoint serviceHost, IDataSourceConnectionFactory dataSourceConnectionFactory, 
+            IConnectedBindingQueue connectedBindingQueue, IConnectionManager connectionManager);
+
+        /// <summary> 
+        /// Add a new method to be called when the onconnection request is submitted 
+        /// </summary> 
+        /// <param name="activity"></param> 
+        void RegisterOnConnectionTask(ConnectionService.OnConnectionHandler activity);
+
+        /// <summary>
+        /// Add a new method to be called when the ondisconnect request is submitted
+        /// </summary>
+        void RegisterOnDisconnectTask(ConnectionService.OnDisconnectHandler activity);
+
+        /// <summary>
+        /// Handles a request to get a connection string for the provided connection
+        /// </summary>
+        Task HandleGetConnectionStringRequest(
+            GetConnectionStringParams connStringParams,
+            RequestContext<string> requestContext);
+
+        /// <summary>
+        /// Handles a request to serialize a connection string
+        /// </summary>
+        Task HandleBuildConnectionInfoRequest(
+            string connectionString,
+            RequestContext<ConnectionDetails> requestContext);
+
+        ConnectionDetails ParseConnectionString(string connectionString);
+    }
+}

src/Microsoft.Kusto.ServiceLayer/HostLoader.cs

@@ -11,7 +11,6 @@ using Microsoft.SqlTools.Hosting.Protocol;
 using Microsoft.Kusto.ServiceLayer.Admin;
 using Microsoft.Kusto.ServiceLayer.Metadata;
 using Microsoft.Kusto.ServiceLayer.Connection;
-using Microsoft.Kusto.ServiceLayer.DataSource;
 using Microsoft.Kusto.ServiceLayer.LanguageServices;
 using Microsoft.Kusto.ServiceLayer.QueryExecution;
 using Microsoft.Kusto.ServiceLayer.Scripting;
@@ -66,7 +65,8 @@ namespace Microsoft.Kusto.ServiceLayer
             ExtensionServiceProvider serviceProvider = ExtensionServiceProvider.CreateDefaultServiceProvider(inclusionList);
             serviceProvider.RegisterSingleService(sqlToolsContext);
             serviceProvider.RegisterSingleService(serviceHost);
-            
+            var connectionManager = serviceProvider.GetService<IConnectionManager>();
+
             var scripter = serviceProvider.GetService<IScripter>();
             var dataSourceConnectionFactory = serviceProvider.GetService<IDataSourceConnectionFactory>();
             var connectedBindingQueue = serviceProvider.GetService<IConnectedBindingQueue>();
@@ -77,25 +77,25 @@ namespace Microsoft.Kusto.ServiceLayer
             WorkspaceService<SqlToolsSettings>.Instance.InitializeService(serviceHost);
             serviceProvider.RegisterSingleService(WorkspaceService<SqlToolsSettings>.Instance);
 
-            LanguageService.Instance.InitializeService(serviceHost, connectedBindingQueue);
+            LanguageService.Instance.InitializeService(serviceHost, connectedBindingQueue, connectionManager);
             serviceProvider.RegisterSingleService(LanguageService.Instance);
 
-            ConnectionService.Instance.InitializeService(serviceHost, dataSourceConnectionFactory, connectedBindingQueue);
+            ConnectionService.Instance.InitializeService(serviceHost, dataSourceConnectionFactory, connectedBindingQueue, connectionManager);
             serviceProvider.RegisterSingleService(ConnectionService.Instance);
 
             CredentialService.Instance.InitializeService(serviceHost);
             serviceProvider.RegisterSingleService(CredentialService.Instance);
 
-            QueryExecutionService.Instance.InitializeService(serviceHost);
+            QueryExecutionService.Instance.InitializeService(serviceHost, connectionManager);
             serviceProvider.RegisterSingleService(QueryExecutionService.Instance);
 
-            ScriptingService.Instance.InitializeService(serviceHost, scripter, ConnectionService.Instance);
+            ScriptingService.Instance.InitializeService(serviceHost, scripter, ConnectionService.Instance, connectionManager);
             serviceProvider.RegisterSingleService(ScriptingService.Instance);
 
-            AdminService.Instance.InitializeService(serviceHost, ConnectionService.Instance);
+            AdminService.Instance.InitializeService(serviceHost, connectionManager);
             serviceProvider.RegisterSingleService(AdminService.Instance);
 
-            MetadataService.Instance.InitializeService(serviceHost, ConnectionService.Instance);
+            MetadataService.Instance.InitializeService(serviceHost, connectionManager);
             serviceProvider.RegisterSingleService(MetadataService.Instance);
             
             InitializeHostedServices(serviceProvider, serviceHost);

src/Microsoft.Kusto.ServiceLayer/LanguageServices/LanguageService.cs

@@ -141,6 +141,7 @@ namespace Microsoft.Kusto.ServiceLayer.LanguageServices
         }
 
         private CancellationTokenSource existingRequestCancellation;
+        private IConnectionManager _connectionManager;
 
         /// <summary>
         /// Gets or sets the current workspace service instance
@@ -205,9 +206,10 @@ namespace Microsoft.Kusto.ServiceLayer.LanguageServices
         /// <param name="context"></param>
         /// <param name="dataSourceFactory"></param>
         /// <param name="connectedBindingQueue"></param>
-        public void InitializeService(ServiceHost serviceHost, IConnectedBindingQueue connectedBindingQueue)
+        public void InitializeService(ServiceHost serviceHost, IConnectedBindingQueue connectedBindingQueue, IConnectionManager connectionManager)
         {
             _bindingQueue = connectedBindingQueue;
+            _connectionManager = connectionManager;
             // Register the requests that this service will handle
 
             //serviceHost.SetRequestHandler(SignatureHelpRequest.Type, HandleSignatureHelpRequest);     // Kusto api doesnt support this as of now. Implement it wherever applicable. Hover help is closest to signature help
@@ -315,9 +317,7 @@ namespace Microsoft.Kusto.ServiceLayer.LanguageServices
                     }
 
                     ConnectionInfo connInfo;
-                    ConnectionServiceInstance.TryFindConnection(
+                    _connectionManager.TryGetValue(scriptFile.ClientUri, out connInfo);
-                        scriptFile.ClientUri,
-                        out connInfo);
 
                     var completionItems = GetCompletionItems(
                         textDocumentPosition, scriptFile, connInfo);
@@ -378,7 +378,7 @@ namespace Microsoft.Kusto.ServiceLayer.LanguageServices
                     DefinitionResult definitionResult = null;
                     if (scriptFile != null)
                     {
-                        isConnected = ConnectionServiceInstance.TryFindConnection(scriptFile.ClientUri, out connInfo);
+                        isConnected = _connectionManager.TryGetValue(scriptFile.ClientUri, out connInfo);
                         definitionResult = GetDefinition(textDocumentPosition, scriptFile, connInfo);
                     }
 
@@ -549,9 +549,7 @@ namespace Microsoft.Kusto.ServiceLayer.LanguageServices
                 }
 
                 ConnectionInfo connInfo;
-                ConnectionServiceInstance.TryFindConnection(
+                _connectionManager.TryGetValue(scriptFile.ClientUri, out connInfo);
-                    scriptFile.ClientUri,
-                    out connInfo);
 
                 // check that there is an active connection for the current editor
                 if (connInfo != null)
@@ -782,7 +780,7 @@ namespace Microsoft.Kusto.ServiceLayer.LanguageServices
         internal Hover GetHoverItem(TextDocumentPosition textDocumentPosition, ScriptFile scriptFile)
         {
             ScriptParseInfo scriptParseInfo = GetScriptParseInfo(scriptFile.ClientUri);
-            if (!ConnectionServiceInstance.TryFindConnection(scriptFile.ClientUri, out var connInfo))
+            if (!_connectionManager.TryGetValue(scriptFile.ClientUri, out var connInfo))
             {
                 return null;
             }
@@ -993,9 +991,7 @@ namespace Microsoft.Kusto.ServiceLayer.LanguageServices
 
                 ScriptFileMarker[] semanticMarkers = null;
                 ConnectionInfo connInfo;
-                ConnectionServiceInstance.TryFindConnection(
+                _connectionManager.TryGetValue(scriptFile.ClientUri, out connInfo);
-                    scriptFile.ClientUri,
-                    out connInfo);
 
                 if (connInfo != null)
                 {

src/Microsoft.Kusto.ServiceLayer/Metadata/MetadataService.cs

@@ -19,7 +19,7 @@ namespace Microsoft.Kusto.ServiceLayer.Metadata
     /// </summary>
     public sealed class MetadataService
     {
-        private static ConnectionService _connectionService;
+        private static IConnectionManager _connectionManager;
         private static readonly Lazy<MetadataService> LazyInstance = new Lazy<MetadataService>();
         public static MetadataService Instance => LazyInstance.Value;
 
@@ -28,9 +28,9 @@ namespace Microsoft.Kusto.ServiceLayer.Metadata
         /// </summary>
         /// <param name="serviceHost"></param>
         /// <param name="connectionService"></param>
-        public void InitializeService(IProtocolEndpoint serviceHost, ConnectionService connectionService)
+        public void InitializeService(IProtocolEndpoint serviceHost, IConnectionManager connectionManager)
         {
-            _connectionService = connectionService;
+            _connectionManager = connectionManager;
             serviceHost.SetRequestHandler(MetadataListRequest.Type, HandleMetadataListRequest);
         }
 
@@ -57,7 +57,7 @@ namespace Microsoft.Kusto.ServiceLayer.Metadata
 
         private List<ObjectMetadata> LoadMetadata(MetadataQueryParams metadataParams)
         {
-            _connectionService.TryFindConnection(metadataParams.OwnerUri, out ConnectionInfo connInfo);
+            _connectionManager.TryGetValue(metadataParams.OwnerUri, out ConnectionInfo connInfo);
 
             if (connInfo == null)
             {

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

@@ -47,6 +47,8 @@ namespace Microsoft.Kusto.ServiceLayer.ObjectExplorer
         /// </summary>
         private ObjectExplorerSettings _settings;
 
+        private IConnectionManager _connectionManager;
+
         /// <summary>
         /// Singleton constructor
         /// </summary>
@@ -78,6 +80,7 @@ namespace Microsoft.Kusto.ServiceLayer.ObjectExplorer
             Validate.IsNotNull(nameof(provider), provider);
             _serviceProvider = provider;
             _connectionService = provider.GetService<ConnectionService>();
+            _connectionManager = provider.GetService<IConnectionManager>();
             
             try
             {
@@ -425,7 +428,7 @@ namespace Microsoft.Kusto.ServiceLayer.ObjectExplorer
 
                 ConnectionInfo connectionInfo;
                 ConnectionCompleteParams connectionResult = await Connect(connectParams, uri);
-                if (!_connectionService.TryFindConnection(uri, out connectionInfo))
+                if (!_connectionManager.TryGetValue(uri, out connectionInfo))
                 {
                     return null;
                 }

src/Microsoft.Kusto.ServiceLayer/QueryExecution/QueryExecutionService.cs

@@ -123,6 +123,8 @@ namespace Microsoft.Kusto.ServiceLayer.QueryExecution
         /// </summary>
         private ConnectionService ConnectionService { get; }
 
+        private IConnectionManager _connectionManager;
+
         private WorkspaceService<SqlToolsSettings> WorkspaceService { get; }
 
         /// <summary>
@@ -160,8 +162,10 @@ namespace Microsoft.Kusto.ServiceLayer.QueryExecution
         /// event handler.
         /// </summary>
         /// <param name="serviceHost">The service host instance to register with</param>
-        public void InitializeService(ServiceHost serviceHost)
+        public void InitializeService(ServiceHost serviceHost, IConnectionManager connectionManager)
         {
+            _connectionManager = connectionManager;
+            
             // Register handlers for requests
             serviceHost.SetRequestHandler(ExecuteDocumentSelectionRequest.Type, HandleExecuteRequest);
             serviceHost.SetRequestHandler(ExecuteDocumentStatementRequest.Type, HandleExecuteRequest);
@@ -252,7 +256,7 @@ namespace Microsoft.Kusto.ServiceLayer.QueryExecution
 
                 // get connection
                 ConnectionInfo connInfo;
-                if (!ConnectionService.TryFindConnection(executeParams.OwnerUri, out connInfo))
+                if (!_connectionManager.TryGetValue(executeParams.OwnerUri, out connInfo))
                 {
                     await requestContext.SendError(SR.QueryServiceQueryInvalidOwnerUri);
                     return;
@@ -269,7 +273,7 @@ namespace Microsoft.Kusto.ServiceLayer.QueryExecution
                     await ConnectionService.Connect(connectParams);
 
                     ConnectionInfo newConn;
-                    ConnectionService.TryFindConnection(randomUri, out newConn);
+                    _connectionManager.TryGetValue(randomUri, out newConn);
 
                     Func<string, Task> queryCreateFailureAction = message => requestContext.SendError(message);
 
@@ -707,7 +711,7 @@ namespace Microsoft.Kusto.ServiceLayer.QueryExecution
             {
                 connectionInfo = connInfo;
             } 
-            else if (!ConnectionService.TryFindConnection(executeParams.OwnerUri, out connectionInfo))
+            else if (!_connectionManager.TryGetValue(executeParams.OwnerUri, out connectionInfo))
             {
                 throw new ArgumentOutOfRangeException(nameof(executeParams.OwnerUri), SR.QueryServiceQueryInvalidOwnerUri);
             }

src/Microsoft.Kusto.ServiceLayer/QueryExecution/SerializationService.cs

@@ -20,8 +20,7 @@ using Microsoft.SqlTools.Utility;
 
 namespace Microsoft.Kusto.ServiceLayer.QueryExecution
 {
-
+    
-    [Export(typeof(IHostedService))]
     public class SerializationService : HostedService<SerializationService>, IComposableService
     {
         private ConcurrentDictionary<string, DataSerializer> inProgressSerializations;

src/Microsoft.Kusto.ServiceLayer/Scripting/ScriptingService.cs

@@ -25,7 +25,9 @@ namespace Microsoft.Kusto.ServiceLayer.Scripting
 
         public static ScriptingService Instance => LazyInstance.Value;
 
-        private static ConnectionService _connectionService;
+        private IConnectionService _connectionService;
+
+        private IConnectionManager _connectionManager;
 
         private readonly Lazy<ConcurrentDictionary<string, ScriptingOperation>> operations =
             new Lazy<ConcurrentDictionary<string, ScriptingOperation>>(() => new ConcurrentDictionary<string, ScriptingOperation>());
@@ -45,10 +47,11 @@ namespace Microsoft.Kusto.ServiceLayer.Scripting
         /// <param name="serviceHost"></param>
         /// <param name="scripter"></param>
         /// <param name="connectionService"></param>
-        public void InitializeService(ServiceHost serviceHost, IScripter scripter, ConnectionService connectionService)
+        public void InitializeService(ServiceHost serviceHost, IScripter scripter, IConnectionService connectionService, IConnectionManager connectionManager)
         {
             _scripter = scripter;
             _connectionService = connectionService;
+            _connectionManager = connectionManager;
             
             serviceHost.SetRequestHandler(ScriptingRequest.Type, this.HandleScriptExecuteRequest);
             serviceHost.SetRequestHandler(ScriptingCancelRequest.Type, this.HandleScriptCancelRequest);
@@ -95,7 +98,7 @@ namespace Microsoft.Kusto.ServiceLayer.Scripting
                 if (parameters.DatabaseName == null)
                 {
                     ConnectionInfo connInfo;
-                    _connectionService.TryFindConnection(parameters.OwnerUri, out connInfo);
+                    _connectionManager.TryGetValue(parameters.OwnerUri, out connInfo);
                     if (connInfo != null)
                     {
                         parameters.DatabaseName = connInfo.ConnectionDetails.DatabaseName;

test/Microsoft.Kusto.ServiceLayer.UnitTests/Admin/AdminServiceTests.cs

@@ -0,0 +1,141 @@
+using System;
+using System.Collections.Generic;
+using System.Linq;
+using System.Threading.Tasks;
+using Microsoft.Kusto.ServiceLayer.Admin;
+using Microsoft.Kusto.ServiceLayer.Admin.Contracts;
+using Microsoft.Kusto.ServiceLayer.Connection;
+using Microsoft.Kusto.ServiceLayer.Connection.Contracts;
+using Microsoft.Kusto.ServiceLayer.DataSource;
+using Microsoft.SqlTools.Hosting.Protocol;
+using Microsoft.SqlTools.ServiceLayer.Connection.ReliableConnection;
+using Moq;
+using NUnit.Framework;
+
+namespace Microsoft.Kusto.ServiceLayer.UnitTests.Admin
+{
+    public class AdminServiceTests
+    {
+        [TestCase(null)]
+        [TestCase("")]
+        public async Task HandleGetDatabaseInfoRequest_NoDatabase_Returns_Null(string databaseName)
+        {
+            var mockServiceHost = new Mock<IProtocolEndpoint>();
+            var mockDataSourceFactory = new Mock<IDataSourceConnectionFactory>();
+            var connectionInfo = new ConnectionInfo(mockDataSourceFactory.Object, null, new ConnectionDetails { DatabaseName = databaseName });
+            var mockConnectionManager = new Mock<IConnectionManager>();
+            mockConnectionManager
+                .Setup(x => x.TryGetValue(It.IsAny<string>(), out connectionInfo))
+                .Returns(true);
+            
+            var mockRequestContext = new Mock<RequestContext<GetDatabaseInfoResponse>>();
+            var actualResponse = new GetDatabaseInfoResponse();
+            mockRequestContext.Setup(x => x.SendResult(It.IsAny<GetDatabaseInfoResponse>()))
+                .Callback<GetDatabaseInfoResponse>(actual => actualResponse = actual)
+                .Returns(Task.CompletedTask);
+            
+            var adminService = new AdminService();
+            adminService.InitializeService(mockServiceHost.Object, mockConnectionManager.Object);
+            await adminService.HandleGetDatabaseInfoRequest(new GetDatabaseInfoParams(), mockRequestContext.Object);
+
+            Assert.AreEqual(null, actualResponse.DatabaseInfo);
+        }
+        
+        [Test]
+        public async Task HandleGetDatabaseInfoRequest_Returns_DatabaseName()
+        {
+            var expectedDatabaseInfo = new DatabaseInfo()
+            {
+                Options = new Dictionary<string, object>
+                {
+                    { "FakeDatabaseName", "FakeSizeInMB" }
+                }
+            };
+            
+            var mockDatasource = new Mock<IDataSource>();
+            mockDatasource
+                .Setup(x => x.GetDatabaseInfo(It.IsAny<string>(), It.IsAny<string>()))
+                .Returns(expectedDatabaseInfo);
+
+            var mockDataSourceFactory = new Mock<IDataSourceFactory>();
+            mockDataSourceFactory
+                .Setup(x => x.Create(It.IsAny<ConnectionDetails>(), It.IsAny<string>()))
+                .Returns(mockDatasource.Object);
+
+            var expectedConnectionDetails = new ConnectionDetails
+            {
+                DatabaseName = "FakeDatabaseName"
+            };
+            
+            var mockConnectionFactory = new Mock<IDataSourceConnectionFactory>();
+            var connectionInfo = new ConnectionInfo(mockConnectionFactory.Object, null, expectedConnectionDetails);
+            var connection = new ReliableDataSourceConnection(expectedConnectionDetails, RetryPolicyFactory.NoRetryPolicy,
+                RetryPolicyFactory.NoRetryPolicy, mockDataSourceFactory.Object, "");
+            connectionInfo.AddConnection(ConnectionType.Default, connection);
+            
+            var mockConnectionManager = new Mock<IConnectionManager>();
+            mockConnectionManager
+                .Setup(x => x.TryGetValue(It.IsAny<string>(), out connectionInfo))
+                .Returns(true);
+            
+            var mockRequestContext = new Mock<RequestContext<GetDatabaseInfoResponse>>();
+            var actualResponse = new GetDatabaseInfoResponse();
+            mockRequestContext.Setup(x => x.SendResult(It.IsAny<GetDatabaseInfoResponse>()))
+                .Callback<GetDatabaseInfoResponse>(actual => actualResponse = actual)
+                .Returns(Task.CompletedTask);
+            
+            var mockServiceHost = new Mock<IProtocolEndpoint>();
+            var adminService = new AdminService();
+            adminService.InitializeService(mockServiceHost.Object, mockConnectionManager.Object);
+            await adminService.HandleGetDatabaseInfoRequest(new GetDatabaseInfoParams(), mockRequestContext.Object);
+
+            Assert.AreEqual("FakeDatabaseName", actualResponse.DatabaseInfo.Options.First().Key);
+            Assert.AreEqual("FakeSizeInMB", actualResponse.DatabaseInfo.Options.First().Value);
+        }
+        
+        [Test]
+        public async Task HandleGetDatabaseInfoRequest_NoConnection_Returns_Null()
+        {
+            var mockServiceHost = new Mock<IProtocolEndpoint>();
+            var mockConnectionManager = new Mock<IConnectionManager>();
+            var mockRequestContext = new Mock<RequestContext<GetDatabaseInfoResponse>>();
+            var actualResponse = new GetDatabaseInfoResponse();
+            mockRequestContext.Setup(x => x.SendResult(It.IsAny<GetDatabaseInfoResponse>()))
+                .Callback<GetDatabaseInfoResponse>(actual => actualResponse = actual)
+                .Returns(Task.CompletedTask);
+
+            var adminService = new AdminService();
+            adminService.InitializeService(mockServiceHost.Object, mockConnectionManager.Object);
+            await adminService.HandleGetDatabaseInfoRequest(new GetDatabaseInfoParams(), mockRequestContext.Object);
+
+            Assert.AreEqual(null, actualResponse.DatabaseInfo);
+        }
+
+        [Test]
+        public async Task HandleDatabaseInfoRequest_ThrowsException_Returns_Error()
+        {
+            var mockServiceHost = new Mock<IProtocolEndpoint>();
+            var mockConnectionManager = new Mock<IConnectionManager>();
+            ConnectionInfo connectionInfo;
+            var expectedException = new Exception("Fake Error Message");
+            var actualException = new Exception();
+            mockConnectionManager
+                .Setup(x => x.TryGetValue(It.IsAny<string>(), out connectionInfo))
+                .Throws(expectedException);
+
+            var mockRequestContext = new Mock<RequestContext<GetDatabaseInfoResponse>>();
+
+            mockRequestContext
+                .Setup(x => x.SendError(It.IsAny<Exception>()))
+                .Callback<Exception>(ex => actualException = ex)
+                .Returns(Task.CompletedTask);
+
+            var adminService = new AdminService();
+            adminService.InitializeService(mockServiceHost.Object, mockConnectionManager.Object);
+            await adminService.HandleGetDatabaseInfoRequest(new GetDatabaseInfoParams(), mockRequestContext.Object);
+
+            Assert.AreEqual(expectedException.GetType(), actualException.GetType());
+            Assert.AreEqual(expectedException.Message, actualException.Message);
+        }
+    }
+}

test/Microsoft.Kusto.ServiceLayer.UnitTests/Admin/Contracts/AdminServiceTests.cs

@@ -1,27 +0,0 @@
-using Microsoft.Kusto.ServiceLayer.Admin;
-using Microsoft.Kusto.ServiceLayer.Connection;
-using Microsoft.Kusto.ServiceLayer.Connection.Contracts;
-using Moq;
-using NUnit.Framework;
-
-namespace Microsoft.Kusto.ServiceLayer.UnitTests.Admin.Contracts
-{
-    public class AdminServiceTests
-    {
-        [TestCase(null)]
-        [TestCase("")]
-        public void GetDatabaseInfo_Returns_Null_For_Invalid_DatabaseName(string databaseName)
-        {
-            var dataSourceConnectionFactory = new Mock<IDataSourceConnectionFactory>();
-            var connectionDetails = new ConnectionDetails
-            {
-                DatabaseName = databaseName 
-            };
-            var connectionInfo = new ConnectionInfo(dataSourceConnectionFactory.Object, "", connectionDetails); 
-            
-            var adminService = new AdminService();
-            var databaseInfo = adminService.GetDatabaseInfo(connectionInfo);
-            Assert.IsNull(databaseInfo);
-        }
-    }
-}

test/Microsoft.Kusto.ServiceLayer.UnitTests/Metadata/MetadataServiceTests.cs

@@ -20,7 +20,7 @@ namespace Microsoft.Kusto.ServiceLayer.UnitTests.Metadata
         public async Task HandleMetadataListRequest_Sets_MetadataListTask()
         {
             var serviceHostMock = new Mock<IProtocolEndpoint>();
-            var connectionServiceMock = new Mock<ConnectionService>();
+            var connectionManagerMock = new Mock<IConnectionManager>();
             var connectionFactoryMock = new Mock<IDataSourceConnectionFactory>();
             var requestContextMock = new Mock<RequestContext<MetadataQueryResult>>();
             requestContextMock.Setup(x => x.SendResult(It.IsAny<MetadataQueryResult>())).Returns(Task.CompletedTask);
@@ -44,10 +44,10 @@ namespace Microsoft.Kusto.ServiceLayer.UnitTests.Metadata
             var connectionInfo = new ConnectionInfo(connectionFactoryMock.Object, "", connectionDetails);
             connectionInfo.AddConnection(ConnectionType.Default, reliableDataSource);
             
-            connectionServiceMock.Setup(x => x.TryFindConnection(It.IsAny<string>(), out connectionInfo));
+            connectionManagerMock.Setup(x => x.TryGetValue(It.IsAny<string>(), out connectionInfo));
             
             var metadataService = new MetadataService();
-            metadataService.InitializeService(serviceHostMock.Object, connectionServiceMock.Object);
+            metadataService.InitializeService(serviceHostMock.Object, connectionManagerMock.Object);
 
             await metadataService.HandleMetadataListRequest(new MetadataQueryParams(), requestContextMock.Object);