diff --git a/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/QueryExecutionService.cs b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/QueryExecutionService.cs
index d8222972..d0ff8d1a 100644
--- a/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/QueryExecutionService.cs
+++ b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/QueryExecutionService.cs
@@ -13,6 +13,7 @@ using Microsoft.SqlTools.ServiceLayer.Hosting.Protocol;
using Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts;
using Microsoft.SqlTools.ServiceLayer.QueryExecution.DataStorage;
using Microsoft.SqlTools.ServiceLayer.SqlContext;
+using Microsoft.SqlTools.ServiceLayer.Utility;
using Microsoft.SqlTools.ServiceLayer.Workspace;
using Microsoft.SqlTools.ServiceLayer.Workspace.Contracts;
using Newtonsoft.Json;
@@ -207,6 +208,9 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
return;
}
+ // Cleanup the query
+ result.Dispose();
+
// Success
await requestContext.SendResult(new QueryDisposeResult
{
@@ -237,6 +241,7 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
// Cancel the query
result.Cancel();
+ result.Dispose();
// Attempt to dispose the query
if (!ActiveQueries.TryRemove(cancelParams.OwnerUri, out result))
@@ -268,7 +273,7 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
///
/// Process request to save a resultSet to a file in CSV format
///
- public async Task HandleSaveResultsAsCsvRequest( SaveResultsAsCsvRequestParams saveParams,
+ public async Task HandleSaveResultsAsCsvRequest(SaveResultsAsCsvRequestParams saveParams,
RequestContext requestContext)
{
// retrieve query for OwnerUri
@@ -341,7 +346,7 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
///
/// Process request to save a resultSet to a file in JSON format
///
- public async Task HandleSaveResultsAsJsonRequest( SaveResultsAsJsonRequestParams saveParams,
+ public async Task HandleSaveResultsAsJsonRequest(SaveResultsAsJsonRequestParams saveParams,
RequestContext requestContext)
{
// retrieve query for OwnerUri
@@ -422,6 +427,7 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
await requestContext.SendError(ex.Message);
}
}
+
#endregion
#region Private Helpers
@@ -445,6 +451,7 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
Query oldQuery;
if (ActiveQueries.TryGetValue(executeParams.OwnerUri, out oldQuery) && oldQuery.HasExecuted)
{
+ oldQuery.Dispose();
ActiveQueries.TryRemove(executeParams.OwnerUri, out oldQuery);
}
@@ -546,8 +553,22 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
{
foreach (var query in ActiveQueries)
{
+ if (!query.Value.HasExecuted)
+ {
+ try
+ {
+ query.Value.Cancel();
+ }
+ catch (Exception e)
+ {
+ // We don't particularly care if we fail to cancel during shutdown
+ string message = string.Format("Failed to cancel query {0} during query service disposal: {1}", query.Key, e);
+ Logger.Write(LogLevel.Warning, message);
+ }
+ }
query.Value.Dispose();
}
+ ActiveQueries.Clear();
}
disposed = true;
diff --git a/test/Microsoft.SqlTools.ServiceLayer.Test/QueryExecution/DisposeTests.cs b/test/Microsoft.SqlTools.ServiceLayer.Test/QueryExecution/DisposeTests.cs
index 226afcc8..2837e892 100644
--- a/test/Microsoft.SqlTools.ServiceLayer.Test/QueryExecution/DisposeTests.cs
+++ b/test/Microsoft.SqlTools.ServiceLayer.Test/QueryExecution/DisposeTests.cs
@@ -4,9 +4,12 @@
//
using System;
+using System.Data.Common;
using System.Threading.Tasks;
using Microsoft.SqlTools.ServiceLayer.Hosting.Protocol;
+using Microsoft.SqlTools.ServiceLayer.QueryExecution;
using Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts;
+using Microsoft.SqlTools.ServiceLayer.QueryExecution.DataStorage;
using Microsoft.SqlTools.ServiceLayer.SqlContext;
using Microsoft.SqlTools.ServiceLayer.Test.Utility;
using Microsoft.SqlTools.ServiceLayer.Workspace;
@@ -18,6 +21,21 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
{
public class DisposeTests
{
+ [Fact]
+ public void DisposeResultSet()
+ {
+ // Setup: Mock file stream factory, mock db reader
+ var mockFileStreamFactory = new Mock();
+ var mockDataReader = Common.CreateTestConnection(null, false).CreateCommand().ExecuteReaderAsync().Result;
+
+ // If: I setup a single resultset and then dispose it
+ ResultSet rs = new ResultSet(mockDataReader, mockFileStreamFactory.Object);
+ rs.Dispose();
+
+ // Then: The file that was created should have been deleted
+ mockFileStreamFactory.Verify(fsf => fsf.DisposeFile(It.IsAny()), Times.Once);
+ }
+
[Fact]
public void DisposeExecutedQuery()
{
@@ -70,6 +88,37 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
Assert.NotEmpty(result.Messages);
}
+ [Fact]
+ public async Task ServiceDispose()
+ {
+ // Setup:
+ // ... We need a workspace service that returns a file
+ var fileMock = new Mock();
+ fileMock.SetupGet(file => file.Contents).Returns(Common.StandardQuery);
+ var workspaceService = new Mock>();
+ workspaceService.Setup(service => service.Workspace.GetFile(It.IsAny()))
+ .Returns(fileMock.Object);
+ // ... We need a query service
+ var queryService = Common.GetPrimedExecutionService(Common.CreateMockFactory(null, false), true,
+ workspaceService.Object);
+
+ // If:
+ // ... I execute some bogus query
+ var queryParams = new QueryExecuteParams { QuerySelection = Common.WholeDocument, OwnerUri = Common.OwnerUri };
+ var requestContext = RequestContextMocks.Create(null);
+ await queryService.HandleExecuteRequest(queryParams, requestContext.Object);
+
+ // ... And it sticks around as an active query
+ Assert.Equal(1, queryService.ActiveQueries.Count);
+
+ // ... The query execution service is disposed, like when the service is shutdown
+ queryService.Dispose();
+
+ // Then:
+ // ... There should no longer be an active query
+ Assert.Empty(queryService.ActiveQueries);
+ }
+
#region Mocking
private Mock> GetQueryDisposeResultContextMock(