diff --git a/test/Microsoft.SqlTools.ServiceLayer.TestDriver/Tests/QueryExecutionTests.cs b/test/Microsoft.SqlTools.ServiceLayer.TestDriver/Tests/QueryExecutionTests.cs index 05146a05..a06e1e90 100644 --- a/test/Microsoft.SqlTools.ServiceLayer.TestDriver/Tests/QueryExecutionTests.cs +++ b/test/Microsoft.SqlTools.ServiceLayer.TestDriver/Tests/QueryExecutionTests.cs @@ -3,9 +3,7 @@ // Licensed under the MIT license. See LICENSE file in the project root for full license information. // -using System; using System.Collections.Generic; -using System.Diagnostics; using System.Linq; using System.Threading; using System.Threading.Tasks; @@ -18,6 +16,190 @@ namespace Microsoft.SqlTools.ServiceLayer.TestDriver.Tests { public class QueryExecutionTests : TestBase { + [Fact] + public async Task TestQueryCancelReliability() + { + string ownerUri = System.IO.Path.GetTempFileName(); + string query = "SELECT * FROM sys.objects a CROSS JOIN sys.objects b CROSS JOIN sys.objects c"; + + await Connect(ownerUri, ConnectionTestUtils.AzureTestServerConnection); + + // Run and cancel 100 queries + for (int i = 0; i < 100; i++) + { + var queryTask = RunQuery(ownerUri, query); + + var cancelResult = await CancelQuery(ownerUri); + Assert.NotNull(cancelResult); + Assert.True(string.IsNullOrEmpty(cancelResult.Messages)); + + await queryTask; + } + + await Disconnect(ownerUri); + } + + [Fact] + public async Task TestQueryDoesNotBlockOtherRequests() + { + string ownerUri = System.IO.Path.GetTempFileName(); + string query = "SELECT * FROM sys.objects a CROSS JOIN sys.objects b CROSS JOIN sys.objects c"; + + await Connect(ownerUri, ConnectionTestUtils.AzureTestServerConnection); + + // Start a long-running query + var queryTask = RunQuery(ownerUri, query, 60000); + + // Interact with the service. None of these requests should time out while waiting for the query to finish + for (int i = 0; i < 10; i++) + { + string ownerUri2 = System.IO.Path.GetTempFileName(); + + await Connect(ownerUri2, ConnectionTestUtils.AzureTestServerConnection); + Assert.NotNull(await RequestCompletion(ownerUri2, "SELECT * FROM sys.objects", 0, 10)); + await Disconnect(ownerUri2); + } + + await CancelQuery(ownerUri); + await Disconnect(ownerUri); + } + + [Fact] + public async Task TestParallelQueryExecution() + { + int queryCount = 10; + + // Create n connections + string[] ownerUris = new string[queryCount]; + for (int i = 0; i < queryCount; i++) + { + ownerUris[i] = System.IO.Path.GetTempFileName(); + Assert.NotNull(await Connect(ownerUris[i], ConnectionTestUtils.AzureTestServerConnection)); + } + + // Run n queries at once + string query = "SELECT * FROM sys.objects"; + var queryTasks = new Task[queryCount]; + for (int i = 0; i < queryCount; i++) + { + queryTasks[i] = RunQuery(ownerUris[i], query); + } + await Task.WhenAll(queryTasks); + + // Verify that they all completed with results and Disconnect + for (int i = 0; i < queryCount; i++) + { + Assert.NotNull(queryTasks[i].Result); + Assert.NotNull(queryTasks[i].Result.BatchSummaries); + await Disconnect(ownerUris[i]); + } + } + + [Fact] + public async Task TestSaveResultsDoesNotBlockOtherRequests() + { + string ownerUri = System.IO.Path.GetTempFileName(); + string query = "SELECT * FROM sys.objects"; + + await Connect(ownerUri, ConnectionTestUtils.AzureTestServerConnection); + + // Execute a query + await RunQuery(ownerUri, query); + + // Spawn several tasks to save results + var saveTasks = new Task[100]; + for (int i = 0; i < 100; i++) + { + if (i % 2 == 0) + { + saveTasks[i] = SaveAsCsv(ownerUri, System.IO.Path.GetTempFileName(), 0, 0); + } + else + { + saveTasks[i] = SaveAsJson(ownerUri, System.IO.Path.GetTempFileName(), 0, 0); + } + } + + // Interact with the service. None of these requests should time out while waiting for the save results tasks to finish + for (int i = 0; i < 10; i++) + { + string ownerUri2 = System.IO.Path.GetTempFileName(); + + await Connect(ownerUri2, ConnectionTestUtils.AzureTestServerConnection); + Assert.NotNull(await RequestCompletion(ownerUri2, "SELECT * FROM sys.objects", 0, 10)); + await Disconnect(ownerUri2); + } + + await Task.WhenAll(saveTasks); + + await Disconnect(ownerUri); + } + + [Fact] + public async Task TestQueryingSubsetDoesNotBlockOtherRequests() + { + string ownerUri = System.IO.Path.GetTempFileName(); + string query = "SELECT * FROM sys.objects"; + + await Connect(ownerUri, ConnectionTestUtils.AzureTestServerConnection); + + // Execute a query + await RunQuery(ownerUri, query); + + // Spawn several tasks for subset requests + var subsetTasks = new Task[100]; + for (int i = 0; i < 100; i++) + { + subsetTasks[i] = ExecuteSubset(ownerUri, 0, 0, 0, 100); + } + + // Interact with the service. None of these requests should time out while waiting for the subset tasks to finish + for (int i = 0; i < 10; i++) + { + string ownerUri2 = System.IO.Path.GetTempFileName(); + + await Connect(ownerUri2, ConnectionTestUtils.AzureTestServerConnection); + Assert.NotNull(await RequestCompletion(ownerUri2, "SELECT * FROM sys.objects", 0, 10)); + await Disconnect(ownerUri2); + } + + await Task.WhenAll(subsetTasks); + + await Disconnect(ownerUri); + } + + [Fact] + public async Task TestCancelQueryWhileOtherOperationsAreInProgress() + { + string ownerUri = System.IO.Path.GetTempFileName(); + string query = "SELECT * FROM sys.objects a CROSS JOIN sys.objects b"; + List tasks = new List(); + + await Connect(ownerUri, ConnectionTestUtils.AzureTestServerConnection); + + // Execute a long-running query + var queryTask = RunQuery(ownerUri, query, 60000); + + // Queue up some tasks that interact with the service + for (int i = 0; i < 10; i++) + { + string ownerUri2 = System.IO.Path.GetTempFileName(); + + tasks.Add(Task.Run(async () => + { + await Connect(ownerUri2, ConnectionTestUtils.AzureTestServerConnection); + await RequestCompletion(ownerUri2, "SELECT * FROM sys.objects", 0, 10); + await RunQuery(ownerUri2, "SELECT * FROM sys.objects"); + await Disconnect(ownerUri2); + })); + } + + // Cancel the long-running query + await CancelQuery(ownerUri); + + await Disconnect(ownerUri); + } + [Fact] public async Task ExecuteBasicQueryTest() { @@ -78,7 +260,7 @@ namespace Microsoft.SqlTools.ServiceLayer.TestDriver.Tests } } - //[Fact] + [Fact] public async Task TestQueryingAfterCompletionRequests() { try diff --git a/test/Microsoft.SqlTools.ServiceLayer.TestDriver/Tests/TestBase.cs b/test/Microsoft.SqlTools.ServiceLayer.TestDriver/Tests/TestBase.cs index c785c107..97be72e7 100644 --- a/test/Microsoft.SqlTools.ServiceLayer.TestDriver/Tests/TestBase.cs +++ b/test/Microsoft.SqlTools.ServiceLayer.TestDriver/Tests/TestBase.cs @@ -211,7 +211,7 @@ namespace Microsoft.SqlTools.ServiceLayer.TestDriver.Tests /// /// Run a query using a given connection bound to a URI /// - protected async Task RunQuery(string ownerUri, string query) + protected async Task RunQuery(string ownerUri, string query, int timeoutMilliseconds = 5000) { // Write the query text to a backing file WriteToFile(ownerUri, query); @@ -223,7 +223,7 @@ namespace Microsoft.SqlTools.ServiceLayer.TestDriver.Tests var result = await Driver.SendRequest(QueryExecuteRequest.Type, queryParams); if (result != null && string.IsNullOrEmpty(result.Messages)) { - var eventResult = await Driver.WaitForEvent(QueryExecuteCompleteEvent.Type); + var eventResult = await Driver.WaitForEvent(QueryExecuteCompleteEvent.Type, timeoutMilliseconds); return eventResult; } else @@ -231,6 +231,64 @@ namespace Microsoft.SqlTools.ServiceLayer.TestDriver.Tests return null; } } + + /// + /// Request to cancel an executing query + /// + protected async Task CancelQuery(string ownerUri) + { + var cancelParams = new QueryCancelParams(); + cancelParams.OwnerUri = ownerUri; + + var result = await Driver.SendRequest(QueryCancelRequest.Type, cancelParams); + return result; + } + + /// + /// Request to save query results as CSV + /// + protected async Task SaveAsCsv(string ownerUri, string filename, int batchIndex, int resultSetIndex) + { + var saveParams = new SaveResultsAsCsvRequestParams(); + saveParams.OwnerUri = ownerUri; + saveParams.BatchIndex = batchIndex; + saveParams.ResultSetIndex = resultSetIndex; + saveParams.FilePath = filename; + + var result = await Driver.SendRequest(SaveResultsAsCsvRequest.Type, saveParams); + return result; + } + + /// + /// Request to save query results as JSON + /// + protected async Task SaveAsJson(string ownerUri, string filename, int batchIndex, int resultSetIndex) + { + var saveParams = new SaveResultsAsJsonRequestParams(); + saveParams.OwnerUri = ownerUri; + saveParams.BatchIndex = batchIndex; + saveParams.ResultSetIndex = resultSetIndex; + saveParams.FilePath = filename; + + var result = await Driver.SendRequest(SaveResultsAsJsonRequest.Type, saveParams); + return result; + } + + /// + /// Request a subset of results from a query + /// + protected async Task ExecuteSubset(string ownerUri, int batchIndex, int resultSetIndex, int rowStartIndex, int rowCount) + { + var subsetParams = new QueryExecuteSubsetParams(); + subsetParams.OwnerUri = ownerUri; + subsetParams.BatchIndex = batchIndex; + subsetParams.ResultSetIndex = resultSetIndex; + subsetParams.RowsStartIndex = rowStartIndex; + subsetParams.RowsCount = rowCount; + + var result = await Driver.SendRequest(QueryExecuteSubsetRequest.Type, subsetParams); + return result; + } protected void WriteToFile(string ownerUri, string query) {