diff --git a/docs/guide/jsonrpc_protocol.md b/docs/guide/jsonrpc_protocol.md index 423f8c0a..9820cb51 100644 --- a/docs/guide/jsonrpc_protocol.md +++ b/docs/guide/jsonrpc_protocol.md @@ -50,6 +50,7 @@ This document provides the protocol specification for all the service's JSON-RPC * :leftwards_arrow_with_hook: [query/subset](#query_subset) * :leftwards_arrow_with_hook: [query/dispose](#query_dispose) * :leftwards_arrow_with_hook: [query/cancel](#query_cancel) +* :arrow_right: [query/connectionUriChanged](#query_connectionUriChanged) * :leftwards_arrow_with_hook: [query/saveCsv](#query_saveCsv) * :leftwards_arrow_with_hook: [query/saveExcel](#query_saveExcel) * :leftwards_arrow_with_hook: [query/saveJson](#query_saveJson) @@ -889,6 +890,18 @@ Cancel the query in progress for the owner uri. public string Messages { get; set; } } ``` +### `query/connectionUriChanged` + +Change the uri associated with a query. + +#### Notification +```csharp + public class ConnectionUriChangedParams + { + public string NewOwnerUri { get; set; } + public string OriginalOwnerUri { get; set; + } +``` ### `query/saveCsv` diff --git a/src/Microsoft.SqlTools.ServiceLayer/Connection/ConnectionInfo.cs b/src/Microsoft.SqlTools.ServiceLayer/Connection/ConnectionInfo.cs index 59445a12..9e8334dc 100644 --- a/src/Microsoft.SqlTools.ServiceLayer/Connection/ConnectionInfo.cs +++ b/src/Microsoft.SqlTools.ServiceLayer/Connection/ConnectionInfo.cs @@ -38,7 +38,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection /// /// URI identifying the owner/user of the connection. Could be a file, service, resource, etc. /// - public string OwnerUri { get; private set; } + public string OwnerUri { get; set; } /// /// Factory used for creating the SQL connection associated with the connection info. diff --git a/src/Microsoft.SqlTools.ServiceLayer/Connection/ConnectionService.cs b/src/Microsoft.SqlTools.ServiceLayer/Connection/ConnectionService.cs index 3b13c2bc..23284e29 100644 --- a/src/Microsoft.SqlTools.ServiceLayer/Connection/ConnectionService.cs +++ b/src/Microsoft.SqlTools.ServiceLayer/Connection/ConnectionService.cs @@ -766,6 +766,22 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection return false; } + /// + /// Reassign the uri associated with a connection info with a new uri. + /// + public bool ReplaceUri(string originalOwnerUri, string newOwnerUri) + { + // Lookup the ConnectionInfo owned by the URI + ConnectionInfo info; + if (!OwnerToConnectionMap.TryGetValue(originalOwnerUri, out info)) + { + return false; + } + OwnerToConnectionMap.Remove(originalOwnerUri); + OwnerToConnectionMap.Add(newOwnerUri, info); + return true; + } + /// /// Close a connection with the specified connection details. /// diff --git a/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Contracts/ConnectionUriChangedNotification.cs b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Contracts/ConnectionUriChangedNotification.cs new file mode 100644 index 00000000..a1e8a28b --- /dev/null +++ b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Contracts/ConnectionUriChangedNotification.cs @@ -0,0 +1,24 @@ +// +// Copyright (c) Microsoft. All rights reserved. +// Licensed under the MIT license. See LICENSE file in the project root for full license information. +// + +using Microsoft.SqlTools.Hosting.Protocol.Contracts; + +namespace Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts +{ + /// + /// Parameters for the connection uri changed notification. + /// + public class ConnectionUriChangedParams + { + public string NewOwnerUri { get; set; } + public string OriginalOwnerUri { get; set; } + } + public class ConnectionUriChangedNotification + { + public static readonly + EventType Type = + EventType.Create("query/connectionUriChanged"); + } +} diff --git a/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Query.cs b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Query.cs index 05ff9f3d..b803e941 100644 --- a/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Query.cs +++ b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Query.cs @@ -385,6 +385,20 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution Batches[saveParams.BatchIndex].SaveAs(saveParams, fileFactory, successHandler, failureHandler); } + /// + /// Changes the OwnerURI for the editor connection. + /// + public String ConnectionOwnerURI { + get + { + return this.editorConnection.OwnerUri; + } + set + { + this.editorConnection.OwnerUri = value; + } + } + #endregion #region Private Helpers diff --git a/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/QueryExecutionService.cs b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/QueryExecutionService.cs index 6a411987..1a4eed92 100644 --- a/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/QueryExecutionService.cs +++ b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/QueryExecutionService.cs @@ -170,6 +170,7 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution serviceHost.SetRequestHandler(SubsetRequest.Type, HandleResultSubsetRequest); serviceHost.SetRequestHandler(QueryDisposeRequest.Type, HandleDisposeRequest); serviceHost.SetRequestHandler(QueryCancelRequest.Type, HandleCancelRequest); + serviceHost.SetEventHandler(ConnectionUriChangedNotification.Type, HandleConnectionUriChangedNotification); serviceHost.SetRequestHandler(SaveResultsAsCsvRequest.Type, HandleSaveResultsAsCsvRequest); serviceHost.SetRequestHandler(SaveResultsAsExcelRequest.Type, HandleSaveResultsAsExcelRequest); serviceHost.SetRequestHandler(SaveResultsAsJsonRequest.Type, HandleSaveResultsAsJsonRequest); @@ -351,6 +352,33 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution } } + + /// + /// Handles a request to change the uri associated with an active query and connection info. + /// + internal Task HandleConnectionUriChangedNotification(ConnectionUriChangedParams changeUriParams, + EventContext eventContext) + { + try { + string OriginalOwnerUri = changeUriParams.OriginalOwnerUri; + string NewOwnerUri = changeUriParams.NewOwnerUri; + // Attempt to load the query + Query query; + if(!ActiveQueries.TryRemove(OriginalOwnerUri, out query)){ + throw new Exception("Uri: " + OriginalOwnerUri + " is not associated with an active query."); + } + ConnectionService.ReplaceUri(OriginalOwnerUri, NewOwnerUri); + query.ConnectionOwnerURI = NewOwnerUri; + ActiveQueries.TryAdd(NewOwnerUri, query); + return Task.FromResult(true); + } + catch (Exception ex) + { + Logger.Write(TraceEventType.Error, "Error encountered " + ex.ToString()); + return Task.FromException(ex); + } + } + /// /// Handles a request to get a subset of the results of this query /// diff --git a/test/Microsoft.SqlTools.ServiceLayer.UnitTests/QueryExecution/ChangeConnectionUriTests.cs b/test/Microsoft.SqlTools.ServiceLayer.UnitTests/QueryExecution/ChangeConnectionUriTests.cs new file mode 100644 index 00000000..1a02b843 --- /dev/null +++ b/test/Microsoft.SqlTools.ServiceLayer.UnitTests/QueryExecution/ChangeConnectionUriTests.cs @@ -0,0 +1,73 @@ +// +// Copyright (c) Microsoft. All rights reserved. +// Licensed under the MIT license. See LICENSE file in the project root for full license information. +// + +using System; +using System.Threading.Tasks; +using Microsoft.SqlTools.ServiceLayer.QueryExecution; +using Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts; +using Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts.ExecuteRequests; +using Microsoft.SqlTools.ServiceLayer.SqlContext; +using Microsoft.SqlTools.ServiceLayer.Test.Common; +using Microsoft.SqlTools.ServiceLayer.Test.Common.RequestContextMocking; +using Microsoft.SqlTools.ServiceLayer.Workspace; +using Moq; +using NUnit.Framework; + +namespace Microsoft.SqlTools.ServiceLayer.UnitTests.QueryExecution +{ + public class ConnectionUriChangedTests + { + [Test] + public async Task ChangeUriForExecutedQuery() + { + // If: + // ... I request a query (doesn't matter what kind) + var workspaceService = Common.GetPrimedWorkspaceService(Constants.StandardQuery); + var queryService = Common.GetPrimedExecutionService(null, true, false, false, workspaceService); + var executeParams = new ExecuteDocumentSelectionParams {QuerySelection = null, OwnerUri = Constants.OwnerUri}; + var executeRequest = RequestContextMocks.Create(null); + await queryService.HandleExecuteRequest(executeParams, executeRequest.Object); + await queryService.WorkTask; + await queryService.ActiveQueries[Constants.OwnerUri].ExecutionTask; + + const string newOwnerUri = "newTestFile"; + Query query; + queryService.ActiveQueries.TryGetValue(Constants.OwnerUri, out query); + + // ... And then I change the uri for the query + var changeUriParams = new ConnectionUriChangedParams { + OriginalOwnerUri = Constants.OwnerUri, + NewOwnerUri = newOwnerUri + }; + + + await queryService.HandleConnectionUriChangedNotification(changeUriParams, new TestEventContext()); + + // Then: + // ... And the active queries should have the new query. + Assert.That(queryService.ActiveQueries.TryGetValue(newOwnerUri, out query), "Query with newOwnerUri not found."); + Assert.That(Equals(query.ConnectionOwnerURI, newOwnerUri), "OwnerUri was not changed!"); + } + + [Test] + public void ChangeUriForMissingQuery() + { + // If: + // ... I attempt to change the uri a query that doesn't exist + var workspaceService = new Mock>(); + var queryService = Common.GetPrimedExecutionService(null, false, false, false, workspaceService.Object); + const string newOwnerUri = "newTestFile"; + var changeUriParams = new ConnectionUriChangedParams { + OriginalOwnerUri = Constants.OwnerUri, + NewOwnerUri = newOwnerUri + }; + + Assert.ThrowsAsync(async () => await queryService.HandleConnectionUriChangedNotification(changeUriParams, new TestEventContext())); + + Query query; + Assert.False(queryService.ActiveQueries.TryGetValue(Constants.OwnerUri, out query), "Query was removed from Active Queries"); + } + } +}