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

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>