Add changeConnectionUri action for query (#1227)

* added wip rename query

* added first rename test

* added missing query handling and two new tests

* removed ServiceDispose

* added change to owneruri.

* made fields private, used accessor methods

* added readonly flag back

* added didSave to sqltools.servicelayer

* fix for docsavecallbacks

* removed tryremove

* test disconnect in queryExecutionService

* removed debug message working query.

* added comment to handlerenamerequest

* added uri replacement function

* removed duplicate removemap function

* added minor fixes

* changed connectionOwnerUri to getter setter

* removed handledidsave and added setters

* removed additional save parts

* fixed space issues

* more fixes for spaces

* restored spaces

* fixed workspaceservice

* Update TextDocument.cs

* changed renameRequest into notification WIP

* restored textdocument

* restored textdocument.cs

* added rightwards arrow

* renamed rename to ChangeConnectionUri

* added more renames

* rename changeconnectionuri to changeuri

* renamed file names

* Revert "renamed file names"

This reverts commit 55228e65025b5179b15ae9a0adc095d95538723d.

* Revert "rename changeconnectionuri to changeuri"

This reverts commit 23f3813f609e3947f103e057f7c0919184075bdc.

* removed logging message and other small changes

* renamed class and method

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; }
     }
 ```
+### <a name="query_connectionUriChanged"></a>`query/connectionUriChanged`
+
+Change the uri associated with a query.
+
+#### Notification
+```csharp
+    public class ConnectionUriChangedParams
+    {
+        public string NewOwnerUri { get; set; }
+        public string OriginalOwnerUri { get; set; 
+    }
+```
 
 ### <a name="query_saveCsv"></a>`query/saveCsv`
 

src/Microsoft.SqlTools.ServiceLayer/Connection/ConnectionInfo.cs

@@ -38,7 +38,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection
         /// <summary>
         /// URI identifying the owner/user of the connection. Could be a file, service, resource, etc.
         /// </summary>
-        public string OwnerUri { get; private set; }
+        public string OwnerUri { get; set; }
 
         /// <summary>
         /// Factory used for creating the SQL connection associated with the connection info.

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

@@ -766,6 +766,22 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection
             return false;
         }
 
+        /// <summary>
+        /// Reassign the uri associated with a connection info with a new uri.
+        /// </summary>
+        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;
+        }
+
         /// <summary>
         /// Close a connection with the specified connection details.
         /// </summary>

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
+{
+    /// <summary>
+    /// Parameters for the connection uri changed notification.
+    /// </summary>
+    public class ConnectionUriChangedParams
+    {
+        public string NewOwnerUri { get; set; }
+        public string OriginalOwnerUri { get; set; }
+    }
+    public class ConnectionUriChangedNotification
+    {
+        public static readonly 
+            EventType<ConnectionUriChangedParams> Type =
+            EventType<ConnectionUriChangedParams>.Create("query/connectionUriChanged");
+    }
+}

src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Query.cs

@@ -385,6 +385,20 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
             Batches[saveParams.BatchIndex].SaveAs(saveParams, fileFactory, successHandler, failureHandler);
         }
 
+        /// <summary>
+        /// Changes the OwnerURI for the editor connection.
+        /// </summary>
+        public String ConnectionOwnerURI { 
+            get
+            {
+                return this.editorConnection.OwnerUri;
+            }
+            set
+            {
+                this.editorConnection.OwnerUri = value;
+            }
+        }
+
         #endregion
 
         #region Private Helpers

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
             }
         }
 
+
+        /// <summary>
+        /// Handles a request to change the uri associated with an active query and connection info.
+        /// </summary>
+        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);
+            }
+        }
+
         /// <summary>
         /// Handles a request to get a subset of the results of this query
         /// </summary>

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<ExecuteRequestResult>(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<WorkspaceService<SqlToolsSettings>>();
+            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<Exception>(async () => await queryService.HandleConnectionUriChangedNotification(changeUriParams, new TestEventContext()));
+
+            Query query;
+            Assert.False(queryService.ActiveQueries.TryGetValue(Constants.OwnerUri, out query), "Query was removed from Active Queries");
+        }
+    }
+}