Adding support for query cancellation

Query cancellation support is added via CancellationToken mechanisms that
were implemented previously. This change adds a new request type
"query/cancel" that will issue the cancellation token. Unit tests were
also added.

src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Contracts/QueryCancelRequest.cs

@@ -0,0 +1,36 @@
+//
+// 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.ServiceLayer.Hosting.Protocol.Contracts;
+
+namespace Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts
+{
+    /// <summary>
+    /// Parameters for the query cancellation request
+    /// </summary>
+    public class QueryCancelParams
+    {
+        public string OwnerUri { get; set; }
+    }
+
+    /// <summary>
+    /// Parameters to return as the result of a query dispose request
+    /// </summary>
+    public class QueryCancelResult
+    {
+        /// <summary>
+        /// Any error messages that occurred during disposing the result set. Optional, can be set
+        /// to null if there were no errors.
+        /// </summary>
+        public string Messages { get; set; }
+    }
+
+    public class QueryCancelRequest
+    {
+        public static readonly 
+            RequestType<QueryCancelParams, QueryCancelResult> Type =
+            RequestType<QueryCancelParams, QueryCancelResult>.Create("query/cancel");
+    }
+}

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

@@ -233,6 +233,21 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
             };
         }
 
+        /// <summary>
+        /// Cancels the query by issuing the cancellation token
+        /// </summary>
+        public void Cancel()
+        {
+            // Make sure that the query hasn't completed execution
+            if (HasExecuted)
+            {
+                throw new InvalidOperationException("The query has already completed, it cannot be cancelled.");
+            }
+
+            // Issue the cancellation token for the query
+            cancellationSource.Cancel();
+        }
+
         /// <summary>
         /// Delegate handler for storing messages that are returned from the server
         /// NOTE: Only messages that are below a certain severity will be returned via this

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

@@ -73,6 +73,7 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
             serviceHost.SetRequestHandler(QueryExecuteRequest.Type, HandleExecuteRequest);
             serviceHost.SetRequestHandler(QueryExecuteSubsetRequest.Type, HandleResultSubsetRequest);
             serviceHost.SetRequestHandler(QueryDisposeRequest.Type, HandleDisposeRequest);
+            serviceHost.SetRequestHandler(QueryCancelRequest.Type, HandleCancelRequest);
 
             // Register handler for shutdown event
             serviceHost.RegisterShutdownTask((shutdownParams, requestContext) =>
@@ -178,6 +179,51 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
             }
         }
 
+        public async Task HandleCancelRequest(QueryCancelParams cancelParams,
+            RequestContext<QueryCancelResult> requestContext)
+        {
+            try
+            {
+                // Attempt to find the query for the owner uri
+                Query result;
+                if (!ActiveQueries.TryGetValue(cancelParams.OwnerUri, out result))
+                {
+                    await requestContext.SendResult(new QueryCancelResult
+                    {
+                        Messages = "Failed to cancel query, ID not found."
+                    });
+                    return;
+                }
+
+                // Cancel the query
+                result.Cancel();
+
+                // Attempt to dispose the query
+                if (!ActiveQueries.TryRemove(cancelParams.OwnerUri, out result))
+                {
+                    // It really shouldn't be possible to get to this scenario, but we'll cover it anyhow
+                    await requestContext.SendResult(new QueryCancelResult
+                    {
+                        Messages = "Query successfully cancelled, failed to dispose query. ID not found."
+                    });
+                    return;
+                }
+
+                await requestContext.SendResult(new QueryCancelResult());
+            }
+            catch (InvalidOperationException e)
+            {
+                await requestContext.SendResult(new QueryCancelResult
+                {
+                    Messages = e.Message
+                });
+            }
+            catch (Exception e)
+            {
+                await requestContext.SendError(e.Message);
+            }
+        }
+
         #endregion
 
         #region Private Helpers