diff --git a/src/Microsoft.Kusto.ServiceLayer/Admin/AdminService.cs b/src/Microsoft.Kusto.ServiceLayer/Admin/AdminService.cs index ef7015b8..76a3f595 100644 --- a/src/Microsoft.Kusto.ServiceLayer/Admin/AdminService.cs +++ b/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 _instance = new Lazy(() => new AdminService()); - private static ConnectionService _connectionService; + private IConnectionManager _connectionManager; /// /// Gets the singleton instance object @@ -33,23 +29,23 @@ namespace Microsoft.Kusto.ServiceLayer.Admin /// /// Initializes the service instance /// - public void InitializeService(ServiceHost serviceHost, ConnectionService connectionService) + public void InitializeService(IProtocolEndpoint serviceHost, IConnectionManager connectionManager) { + _connectionManager = connectionManager; serviceHost.SetRequestHandler(GetDatabaseInfoRequest.Type, HandleGetDatabaseInfoRequest); - _connectionService = connectionService; } /// /// Handle get database info request /// - private async Task HandleGetDatabaseInfoRequest(GetDatabaseInfoParams databaseParams, RequestContext requestContext) + public async Task HandleGetDatabaseInfoRequest(GetDatabaseInfoParams databaseParams, RequestContext 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 /// /// /// - public DatabaseInfo GetDatabaseInfo(ConnectionInfo connInfo) + private DatabaseInfo GetDatabaseInfo(ConnectionInfo connInfo) { if (string.IsNullOrEmpty(connInfo.ConnectionDetails.DatabaseName)) { diff --git a/src/Microsoft.Kusto.ServiceLayer/Connection/ConnectionManager.cs b/src/Microsoft.Kusto.ServiceLayer/Connection/ConnectionManager.cs new file mode 100644 index 00000000..ec477935 --- /dev/null +++ b/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 + { + /// + /// Map from script URIs to ConnectionInfo objects + /// + private ConcurrentDictionary _ownerToConnectionMap; + + public ConnectionManager() + { + _ownerToConnectionMap = new ConcurrentDictionary(); + } + + 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 _); + } + } +} \ No newline at end of file diff --git a/src/Microsoft.Kusto.ServiceLayer/Connection/ConnectionService.cs b/src/Microsoft.Kusto.ServiceLayer/Connection/ConnectionService.cs index 73dd2e53..f24a7c19 100644 --- a/src/Microsoft.Kusto.ServiceLayer/Connection/ConnectionService.cs +++ b/src/Microsoft.Kusto.ServiceLayer/Connection/ConnectionService.cs @@ -25,7 +25,7 @@ namespace Microsoft.Kusto.ServiceLayer.Connection /// /// Main class for the Connection Management services /// - public class ConnectionService + public class ConnectionService : IConnectionService { private const string PasswordPlaceholder = "******"; @@ -53,11 +53,6 @@ namespace Microsoft.Kusto.ServiceLayer.Connection private ConcurrentDictionary connectedQueues = new ConcurrentDictionary(); - /// - /// Map from script URIs to ConnectionInfo objects - /// - private Dictionary OwnerToConnectionMap { get; } = new Dictionary(); - /// /// Database Lock manager instance /// @@ -122,8 +117,7 @@ namespace Microsoft.Kusto.ServiceLayer.Connection /// private IDataSourceConnectionFactory _dataSourceConnectionFactory; - // Attempts to link a URI to an actively used connection for this URI - public virtual bool TryFindConnection(string ownerUri, out ConnectionInfo connectionInfo) => this.OwnerToConnectionMap.TryGetValue(ownerUri, out connectionInfo); + private IConnectionManager _connectionManager; /// /// 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 /// Name of the database to change the connection to 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; } diff --git a/src/Microsoft.Kusto.ServiceLayer/Connection/IConnectionManager.cs b/src/Microsoft.Kusto.ServiceLayer/Connection/IConnectionManager.cs new file mode 100644 index 00000000..fdcdfd5b --- /dev/null +++ b/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); + } +} \ No newline at end of file diff --git a/src/Microsoft.Kusto.ServiceLayer/Connection/IConnectionService.cs b/src/Microsoft.Kusto.ServiceLayer/Connection/IConnectionService.cs new file mode 100644 index 00000000..419621c7 --- /dev/null +++ b/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 + { + /// + /// Register a new connection queue if not already registered + /// + /// + /// + void RegisterConnectedQueue(string type, IConnectedBindingQueue connectedQueue); + + /// + /// Open a connection with the specified ConnectParams + /// + Task Connect(ConnectParams connectionParams); + + /// + /// 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. + /// + /// URI identifying the resource mapped to this connection + /// + /// 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 + /// + /// + /// 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 + /// + /// A DB connection for the connection type requested + Task GetOrOpenConnection(string ownerUri, string connectionType, bool alwaysPersistSecurity = false); + + /// + /// Cancel a connection that is in the process of opening. + /// + bool CancelConnect(CancelConnectParams cancelParams); + + /// + /// Close a connection with the specified connection details. + /// + bool Disconnect(DisconnectParams disconnectParams); + + void InitializeService(IProtocolEndpoint serviceHost, IDataSourceConnectionFactory dataSourceConnectionFactory, + IConnectedBindingQueue connectedBindingQueue, IConnectionManager connectionManager); + + /// + /// Add a new method to be called when the onconnection request is submitted + /// + /// + void RegisterOnConnectionTask(ConnectionService.OnConnectionHandler activity); + + /// + /// Add a new method to be called when the ondisconnect request is submitted + /// + void RegisterOnDisconnectTask(ConnectionService.OnDisconnectHandler activity); + + /// + /// Handles a request to get a connection string for the provided connection + /// + Task HandleGetConnectionStringRequest( + GetConnectionStringParams connStringParams, + RequestContext requestContext); + + /// + /// Handles a request to serialize a connection string + /// + Task HandleBuildConnectionInfoRequest( + string connectionString, + RequestContext requestContext); + + ConnectionDetails ParseConnectionString(string connectionString); + } +} \ No newline at end of file diff --git a/src/Microsoft.Kusto.ServiceLayer/HostLoader.cs b/src/Microsoft.Kusto.ServiceLayer/HostLoader.cs index cd6e40df..e92d8dfa 100644 --- a/src/Microsoft.Kusto.ServiceLayer/HostLoader.cs +++ b/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(); + var scripter = serviceProvider.GetService(); var dataSourceConnectionFactory = serviceProvider.GetService(); var connectedBindingQueue = serviceProvider.GetService(); @@ -77,25 +77,25 @@ namespace Microsoft.Kusto.ServiceLayer WorkspaceService.Instance.InitializeService(serviceHost); serviceProvider.RegisterSingleService(WorkspaceService.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); diff --git a/src/Microsoft.Kusto.ServiceLayer/LanguageServices/LanguageService.cs b/src/Microsoft.Kusto.ServiceLayer/LanguageServices/LanguageService.cs index 1d00b0e8..3403c2d1 100644 --- a/src/Microsoft.Kusto.ServiceLayer/LanguageServices/LanguageService.cs +++ b/src/Microsoft.Kusto.ServiceLayer/LanguageServices/LanguageService.cs @@ -141,6 +141,7 @@ namespace Microsoft.Kusto.ServiceLayer.LanguageServices } private CancellationTokenSource existingRequestCancellation; + private IConnectionManager _connectionManager; /// /// Gets or sets the current workspace service instance @@ -205,9 +206,10 @@ namespace Microsoft.Kusto.ServiceLayer.LanguageServices /// /// /// - 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( - scriptFile.ClientUri, - out connInfo); + _connectionManager.TryGetValue(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( - scriptFile.ClientUri, - out connInfo); + _connectionManager.TryGetValue(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( - scriptFile.ClientUri, - out connInfo); + _connectionManager.TryGetValue(scriptFile.ClientUri, out connInfo); if (connInfo != null) { diff --git a/src/Microsoft.Kusto.ServiceLayer/Metadata/MetadataService.cs b/src/Microsoft.Kusto.ServiceLayer/Metadata/MetadataService.cs index 2898cf26..4ff70613 100644 --- a/src/Microsoft.Kusto.ServiceLayer/Metadata/MetadataService.cs +++ b/src/Microsoft.Kusto.ServiceLayer/Metadata/MetadataService.cs @@ -19,7 +19,7 @@ namespace Microsoft.Kusto.ServiceLayer.Metadata /// public sealed class MetadataService { - private static ConnectionService _connectionService; + private static IConnectionManager _connectionManager; private static readonly Lazy LazyInstance = new Lazy(); public static MetadataService Instance => LazyInstance.Value; @@ -28,9 +28,9 @@ namespace Microsoft.Kusto.ServiceLayer.Metadata /// /// /// - 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 LoadMetadata(MetadataQueryParams metadataParams) { - _connectionService.TryFindConnection(metadataParams.OwnerUri, out ConnectionInfo connInfo); + _connectionManager.TryGetValue(metadataParams.OwnerUri, out ConnectionInfo connInfo); if (connInfo == null) { diff --git a/src/Microsoft.Kusto.ServiceLayer/ObjectExplorer/ObjectExplorerService.cs b/src/Microsoft.Kusto.ServiceLayer/ObjectExplorer/ObjectExplorerService.cs index 131ff685..6f178a93 100644 --- a/src/Microsoft.Kusto.ServiceLayer/ObjectExplorer/ObjectExplorerService.cs +++ b/src/Microsoft.Kusto.ServiceLayer/ObjectExplorer/ObjectExplorerService.cs @@ -47,6 +47,8 @@ namespace Microsoft.Kusto.ServiceLayer.ObjectExplorer /// private ObjectExplorerSettings _settings; + private IConnectionManager _connectionManager; + /// /// Singleton constructor /// @@ -78,6 +80,7 @@ namespace Microsoft.Kusto.ServiceLayer.ObjectExplorer Validate.IsNotNull(nameof(provider), provider); _serviceProvider = provider; _connectionService = provider.GetService(); + _connectionManager = provider.GetService(); 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; } diff --git a/src/Microsoft.Kusto.ServiceLayer/QueryExecution/QueryExecutionService.cs b/src/Microsoft.Kusto.ServiceLayer/QueryExecution/QueryExecutionService.cs index 82eafe0a..6c7e7a40 100644 --- a/src/Microsoft.Kusto.ServiceLayer/QueryExecution/QueryExecutionService.cs +++ b/src/Microsoft.Kusto.ServiceLayer/QueryExecution/QueryExecutionService.cs @@ -123,6 +123,8 @@ namespace Microsoft.Kusto.ServiceLayer.QueryExecution /// private ConnectionService ConnectionService { get; } + private IConnectionManager _connectionManager; + private WorkspaceService WorkspaceService { get; } /// @@ -160,8 +162,10 @@ namespace Microsoft.Kusto.ServiceLayer.QueryExecution /// event handler. /// /// The service host instance to register with - 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 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); } diff --git a/src/Microsoft.Kusto.ServiceLayer/QueryExecution/SerializationService.cs b/src/Microsoft.Kusto.ServiceLayer/QueryExecution/SerializationService.cs index eb7c2a59..34192c2e 100644 --- a/src/Microsoft.Kusto.ServiceLayer/QueryExecution/SerializationService.cs +++ b/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, IComposableService { private ConcurrentDictionary inProgressSerializations; diff --git a/src/Microsoft.Kusto.ServiceLayer/Scripting/ScriptingService.cs b/src/Microsoft.Kusto.ServiceLayer/Scripting/ScriptingService.cs index eed37441..a86aa594 100644 --- a/src/Microsoft.Kusto.ServiceLayer/Scripting/ScriptingService.cs +++ b/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> operations = new Lazy>(() => new ConcurrentDictionary()); @@ -45,10 +47,11 @@ namespace Microsoft.Kusto.ServiceLayer.Scripting /// /// /// - 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; diff --git a/test/Microsoft.Kusto.ServiceLayer.UnitTests/Admin/AdminServiceTests.cs b/test/Microsoft.Kusto.ServiceLayer.UnitTests/Admin/AdminServiceTests.cs new file mode 100644 index 00000000..82382015 --- /dev/null +++ b/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(); + var mockDataSourceFactory = new Mock(); + var connectionInfo = new ConnectionInfo(mockDataSourceFactory.Object, null, new ConnectionDetails { DatabaseName = databaseName }); + var mockConnectionManager = new Mock(); + mockConnectionManager + .Setup(x => x.TryGetValue(It.IsAny(), out connectionInfo)) + .Returns(true); + + var mockRequestContext = new Mock>(); + var actualResponse = new GetDatabaseInfoResponse(); + mockRequestContext.Setup(x => x.SendResult(It.IsAny())) + .Callback(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 + { + { "FakeDatabaseName", "FakeSizeInMB" } + } + }; + + var mockDatasource = new Mock(); + mockDatasource + .Setup(x => x.GetDatabaseInfo(It.IsAny(), It.IsAny())) + .Returns(expectedDatabaseInfo); + + var mockDataSourceFactory = new Mock(); + mockDataSourceFactory + .Setup(x => x.Create(It.IsAny(), It.IsAny())) + .Returns(mockDatasource.Object); + + var expectedConnectionDetails = new ConnectionDetails + { + DatabaseName = "FakeDatabaseName" + }; + + var mockConnectionFactory = new Mock(); + 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(); + mockConnectionManager + .Setup(x => x.TryGetValue(It.IsAny(), out connectionInfo)) + .Returns(true); + + var mockRequestContext = new Mock>(); + var actualResponse = new GetDatabaseInfoResponse(); + mockRequestContext.Setup(x => x.SendResult(It.IsAny())) + .Callback(actual => actualResponse = actual) + .Returns(Task.CompletedTask); + + var mockServiceHost = new Mock(); + 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(); + var mockConnectionManager = new Mock(); + var mockRequestContext = new Mock>(); + var actualResponse = new GetDatabaseInfoResponse(); + mockRequestContext.Setup(x => x.SendResult(It.IsAny())) + .Callback(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(); + var mockConnectionManager = new Mock(); + ConnectionInfo connectionInfo; + var expectedException = new Exception("Fake Error Message"); + var actualException = new Exception(); + mockConnectionManager + .Setup(x => x.TryGetValue(It.IsAny(), out connectionInfo)) + .Throws(expectedException); + + var mockRequestContext = new Mock>(); + + mockRequestContext + .Setup(x => x.SendError(It.IsAny())) + .Callback(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); + } + } +} \ No newline at end of file diff --git a/test/Microsoft.Kusto.ServiceLayer.UnitTests/Admin/Contracts/AdminServiceTests.cs b/test/Microsoft.Kusto.ServiceLayer.UnitTests/Admin/Contracts/AdminServiceTests.cs deleted file mode 100644 index aa9094a4..00000000 --- a/test/Microsoft.Kusto.ServiceLayer.UnitTests/Admin/Contracts/AdminServiceTests.cs +++ /dev/null @@ -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(); - 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); - } - } -} \ No newline at end of file diff --git a/test/Microsoft.Kusto.ServiceLayer.UnitTests/Metadata/MetadataServiceTests.cs b/test/Microsoft.Kusto.ServiceLayer.UnitTests/Metadata/MetadataServiceTests.cs index 658ec2e2..1c8ac24d 100644 --- a/test/Microsoft.Kusto.ServiceLayer.UnitTests/Metadata/MetadataServiceTests.cs +++ b/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(); - var connectionServiceMock = new Mock(); + var connectionManagerMock = new Mock(); var connectionFactoryMock = new Mock(); var requestContextMock = new Mock>(); requestContextMock.Setup(x => x.SendResult(It.IsAny())).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(), out connectionInfo)); + connectionManagerMock.Setup(x => x.TryGetValue(It.IsAny(), 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);