Merge pull request #18 from Microsoft/feature/queryCancellation

Adding support for query cancellation

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

@@ -177,12 +177,10 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
             {
                 HasError = true;
                 UnwrapDbException(dbe);
-                conn?.Dispose();
             }
             catch (Exception)
             {
                 HasError = true;
-                conn?.Dispose();
                 throw;
             }
             finally
@@ -233,6 +231,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,52 @@ 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)
+            {
+                // If this exception occurred, we most likely were trying to cancel a completed query
+                await requestContext.SendResult(new QueryCancelResult
+                {
+                    Messages = e.Message
+                });
+            }
+            catch (Exception e)
+            {
+                await requestContext.SendError(e.Message);
+            }
+        }
+
         #endregion
 
         #region Private Helpers

test/Microsoft.SqlTools.ServiceLayer.Test/QueryExecution/CancelTests.cs

@@ -0,0 +1,124 @@
+//
+// 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.Hosting.Protocol;
+using Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts;
+using Moq;
+using Xunit;
+
+namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
+{
+    public class CancelTests
+    {
+        [Fact]
+        public void CancelInProgressQueryTest()
+        {
+            // If:
+            // ... I request a query (doesn't matter what kind) and execute it
+            var queryService = Common.GetPrimedExecutionService(Common.CreateMockFactory(null, false), true);
+            var executeParams = new QueryExecuteParams { QueryText = "Doesn't Matter", OwnerUri = Common.OwnerUri };
+            var executeRequest = Common.GetQueryExecuteResultContextMock(null, null, null);
+            queryService.HandleExecuteRequest(executeParams, executeRequest.Object).Wait();
+            queryService.ActiveQueries[Common.OwnerUri].HasExecuted = false;    // Fake that it hasn't completed execution
+
+            // ... And then I request to cancel the query
+            var cancelParams = new QueryCancelParams {OwnerUri = Common.OwnerUri};
+            QueryCancelResult result = null;
+            var cancelRequest = GetQueryCancelResultContextMock(qcr => result = qcr, null);
+            queryService.HandleCancelRequest(cancelParams, cancelRequest.Object).Wait();
+
+            // Then:
+            // ... I should have seen a successful event (no messages)
+            VerifyQueryCancelCallCount(cancelRequest, Times.Once(), Times.Never());
+            Assert.Null(result.Messages);
+            
+            // ... The query should have been disposed as well
+            Assert.Empty(queryService.ActiveQueries);
+        }
+
+        [Fact]
+        public void CancelExecutedQueryTest()
+        {
+            // If:
+            // ... I request a query (doesn't matter what kind) and wait for execution
+            var queryService = Common.GetPrimedExecutionService(Common.CreateMockFactory(null, false), true);
+            var executeParams = new QueryExecuteParams {QueryText = "Doesn't Matter", OwnerUri = Common.OwnerUri};
+            var executeRequest = Common.GetQueryExecuteResultContextMock(null, null, null);
+            queryService.HandleExecuteRequest(executeParams, executeRequest.Object).Wait();
+
+            // ... And then I request to cancel the query
+            var cancelParams = new QueryCancelParams {OwnerUri = Common.OwnerUri};
+            QueryCancelResult result = null;
+            var cancelRequest = GetQueryCancelResultContextMock(qcr => result = qcr, null);
+            queryService.HandleCancelRequest(cancelParams, cancelRequest.Object).Wait();
+
+            // Then:
+            // ... I should have seen a result event with an error message
+            VerifyQueryCancelCallCount(cancelRequest, Times.Once(), Times.Never());
+            Assert.NotNull(result.Messages);
+
+            // ... The query should not have been disposed
+            Assert.NotEmpty(queryService.ActiveQueries);
+        }
+
+        [Fact]
+        public void CancelNonExistantTest()
+        {
+            // If:
+            // ... I request to cancel a query that doesn't exist
+            var queryService = Common.GetPrimedExecutionService(Common.CreateMockFactory(null, false), false);
+            var cancelParams = new QueryCancelParams {OwnerUri = "Doesn't Exist"};
+            QueryCancelResult result = null;
+            var cancelRequest = GetQueryCancelResultContextMock(qcr => result = qcr, null);
+            queryService.HandleCancelRequest(cancelParams, cancelRequest.Object).Wait();
+
+            // Then:
+            // ... I should have seen a result event with an error message
+            VerifyQueryCancelCallCount(cancelRequest, Times.Once(), Times.Never());
+            Assert.NotNull(result.Messages);
+        }
+
+        #region Mocking
+
+        private static Mock<RequestContext<QueryCancelResult>> GetQueryCancelResultContextMock(
+            Action<QueryCancelResult> resultCallback,
+            Action<object> errorCallback)
+        {
+            var requestContext = new Mock<RequestContext<QueryCancelResult>>();
+
+            // Setup the mock for SendResult
+            var sendResultFlow = requestContext
+                .Setup(rc => rc.SendResult(It.IsAny<QueryCancelResult>()))
+                .Returns(Task.FromResult(0));
+            if (resultCallback != null)
+            {
+                sendResultFlow.Callback(resultCallback);
+            }
+
+            // Setup the mock for SendError
+            var sendErrorFlow = requestContext
+                .Setup(rc => rc.SendError(It.IsAny<object>()))
+                .Returns(Task.FromResult(0));
+            if (errorCallback != null)
+            {
+                sendErrorFlow.Callback(errorCallback);
+            }
+
+            return requestContext;
+        }
+
+        private static void VerifyQueryCancelCallCount(Mock<RequestContext<QueryCancelResult>> mock,
+            Times sendResultCalls, Times sendErrorCalls)
+        {
+            mock.Verify(rc => rc.SendResult(It.IsAny<QueryCancelResult>()), sendResultCalls);
+            mock.Verify(rc => rc.SendError(It.IsAny<object>()), sendErrorCalls);
+        }
+
+        #endregion
+
+    }
+}