Progressive Results Part 2: Result Completion Event (#134)

The main change in this pull request is to add a new event that will be fired upon completion of a resultset but before the completion of a batch. This event will only fire if a resultset is available and generated.

Changes:
* ConnectionService - Slight changes to enable mocking, cleanup 
* Batch - Moving summary generation into ResultSet class, adding generation of ordinals for resultset and locking of result set list (which needs further refinement, but would be outside scope of this change)
* Adding new event and associated parameters for completion of a resultset. Params return the resultset summary
* Adding logic for assigning the event a handler in the query execution service
* Adding unit tests for testing the new event /making sure the existing tests work
* Refactoring some private properties into member variables

* Refactor to remove SectionData class in favor of BufferRange

* Adding callback for batch completion that will let the extension know that a batch has completed execution

* Refactoring to make progressive results work as per async query execution

* Allowing retrieval of batch results while query is in progress

* reverting global.json, whoops

* Adding a few missing comments, and fixing a couple code style bugs

* Using SelectionData everywhere again

* One more missing comment

* Adding new notification type for result set completion

* Plumbing event for result set completion

* Unit tests for result set events

This includes a fairly substantial change to create a mock of the
ConnectionService and to create separate memorystream storage arrays. It
preserves more correct behavior with a integration test, fixes an issue
where the test db reader will return n-1 rows because the Reliable
Connection Helper steals a record.

* Adding locking to ResultSets for thread safety

* Adding/fixing unit tests

* Adding batch ID to result set summary

test/Microsoft.SqlTools.ServiceLayer.Test/QueryExecution/Execution/BatchTests.cs

@@ -42,6 +42,16 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution.Execution
 
             // ... It's ordinal ID should be what I set it to
             Assert.Equal(Common.Ordinal, batch.Id);
+
+            // ... The summary should have the same info
+            Assert.False(batch.Summary.HasError);
+            Assert.Equal(Common.Ordinal, batch.Summary.Id);
+            Assert.Empty(batch.Summary.ResultSetSummaries);
+            Assert.Empty(batch.Summary.Messages);
+            Assert.Equal(0, batch.Summary.Selection.StartLine);
+            Assert.NotEqual(default(DateTime).ToString("o"), batch.Summary.ExecutionStart); // Should have been set at construction
+            Assert.Equal(default(DateTime).ToString("o"), batch.Summary.ExecutionEnd);
+            Assert.Equal((default(DateTime) - DateTime.Parse(batch.Summary.ExecutionStart)).ToString(), batch.Summary.ExecutionElapsed);
         }
 
         [Fact]
@@ -49,18 +59,25 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution.Execution
         {
             // Setup: Create a callback for batch completion
             BatchSummary batchSummaryFromCallback = null;
-            bool completionCallbackFired = false;
-            Batch.BatchAsyncEventHandler callback = b =>
+            Batch.BatchAsyncEventHandler batchCallback = b =>
             {
-                completionCallbackFired = true;
                 batchSummaryFromCallback = b.Summary;
                 return Task.FromResult(0);
             };
 
+            // ... Create a callback for result completion
+            bool resultCallbackFired = false;
+            ResultSet.ResultSetAsyncEventHandler resultSetCallback = r =>
+            {
+                resultCallbackFired = true;
+                return Task.FromResult(0);
+            };
+
             // If I execute a query that should get no result sets
             var fileStreamFactory = Common.GetFileStreamFactory(new Dictionary<string, byte[]>());
             Batch batch = new Batch(Common.StandardQuery, Common.SubsectionDocument, Common.Ordinal, fileStreamFactory);
-            batch.BatchCompletion += callback;
+            batch.BatchCompletion += batchCallback;
+            batch.ResultSetCompletion += resultSetCallback;
             batch.Execute(GetConnection(Common.CreateTestConnectionInfo(null, false)), CancellationToken.None).Wait();
 
             // Then:
@@ -80,8 +97,20 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution.Execution
             Assert.Equal(1, batch.ResultMessages.Count());
 
             // ... The callback for batch completion should have been fired
-            Assert.True(completionCallbackFired);
             Assert.NotNull(batchSummaryFromCallback);
+
+            // ... The callback for the result set should NOT have been fired
+            Assert.False(resultCallbackFired);
+
+            // ... The summary should have the same info
+            Assert.False(batch.Summary.HasError);
+            Assert.Equal(Common.Ordinal, batch.Summary.Id);
+            Assert.Equal(0, batch.Summary.ResultSetSummaries.Length);
+            Assert.Equal(1, batch.Summary.Messages.Length);
+            Assert.Equal(0, batch.Summary.Selection.StartLine);
+            Assert.True(DateTime.Parse(batch.Summary.ExecutionStart) > default(DateTime));
+            Assert.True(DateTime.Parse(batch.Summary.ExecutionEnd) > default(DateTime));
+            Assert.NotNull(batch.Summary.ExecutionElapsed);
         }
 
         [Fact]
@@ -92,18 +121,25 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution.Execution
 
             // Setup: Create a callback for batch completion
             BatchSummary batchSummaryFromCallback = null;
-            bool completionCallbackFired = false;
-            Batch.BatchAsyncEventHandler callback = b =>
+            Batch.BatchAsyncEventHandler batchCallback = b =>
             {
-                completionCallbackFired = true;
                 batchSummaryFromCallback = b.Summary;
                 return Task.FromResult(0);
             };
 
+            // ... Create a callback for result set completion
+            bool resultCallbackFired = false;
+            ResultSet.ResultSetAsyncEventHandler resultSetCallback = r =>
+            {
+                resultCallbackFired = true;
+                return Task.FromResult(0);
+            };
+
             // If I execute a query that should get one result set
             var fileStreamFactory = Common.GetFileStreamFactory(new Dictionary<string, byte[]>());
             Batch batch = new Batch(Common.StandardQuery, Common.SubsectionDocument, Common.Ordinal, fileStreamFactory);
-            batch.BatchCompletion += callback;
+            batch.BatchCompletion += batchCallback;
+            batch.ResultSetCompletion += resultSetCallback;
             batch.Execute(GetConnection(ci), CancellationToken.None).Wait();
 
             // Then:
@@ -112,7 +148,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution.Execution
             Assert.False(batch.HasError, "The batch should not have an error");
 
             // ... There should be exactly one result set
-            Assert.Equal(resultSets, batch.ResultSets.Count());
+            Assert.Equal(resultSets, batch.ResultSets.Count);
             Assert.Equal(resultSets, batch.ResultSummaries.Length);
 
             // ... Inside the result set should be with 5 rows
@@ -127,8 +163,11 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution.Execution
             Assert.Equal(resultSets, batch.ResultMessages.Count());
 
             // ... The callback for batch completion should have been fired
-            Assert.True(completionCallbackFired);
             Assert.NotNull(batchSummaryFromCallback);
+
+            // ... The callback for resultset completion should have been fired
+            Assert.True(resultCallbackFired);  // We only want to validate that it happened, validation of the 
+                                               // summary is done in result set tests
         }
 
         [Fact]
@@ -140,18 +179,25 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution.Execution
 
             // Setup: Create a callback for batch completion
             BatchSummary batchSummaryFromCallback = null;
-            bool completionCallbackFired = false;
-            Batch.BatchAsyncEventHandler callback = b =>
+            Batch.BatchAsyncEventHandler batchCallback = b =>
             {
-                completionCallbackFired = true;
                 batchSummaryFromCallback = b.Summary;
                 return Task.FromResult(0);
             };
 
+            // ... Create a callback for resultset completion
+            int resultSummaryCount  = 0;
+            ResultSet.ResultSetAsyncEventHandler resultSetCallback = r =>
+            {
+                resultSummaryCount++;
+                return Task.FromResult(0);
+            };
+
             // If I execute a query that should get two result sets
             var fileStreamFactory = Common.GetFileStreamFactory(new Dictionary<string, byte[]>());
             Batch batch = new Batch(Common.StandardQuery, Common.SubsectionDocument, Common.Ordinal, fileStreamFactory);
-            batch.BatchCompletion += callback;
+            batch.BatchCompletion += batchCallback;
+            batch.ResultSetCompletion += resultSetCallback;
             batch.Execute(GetConnection(ci), CancellationToken.None).Wait();
 
             // Then:
@@ -184,8 +230,10 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution.Execution
             }
 
             // ... The callback for batch completion should have been fired
-            Assert.True(completionCallbackFired);
             Assert.NotNull(batchSummaryFromCallback);
+
+            // ... The callback for result set completion should have been fired
+            Assert.Equal(2, resultSummaryCount);
         }
 
         [Fact]
@@ -193,20 +241,25 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution.Execution
         {
             // Setup: Create a callback for batch completion
             BatchSummary batchSummaryFromCallback = null;
-            bool completionCallbackFired = false;
-            Batch.BatchAsyncEventHandler callback = b =>
+            Batch.BatchAsyncEventHandler batchCallback = b =>
             {
-                completionCallbackFired = true;
                 batchSummaryFromCallback = b.Summary;
                 return Task.FromResult(0);
             };
 
+            // ... Create a callback that will fail the test if it's called
+            ResultSet.ResultSetAsyncEventHandler resultSetCallback = r =>
+            {
+                throw new Exception("ResultSet callback was called when it should not have been.");
+            };
+
             ConnectionInfo ci = Common.CreateTestConnectionInfo(null, true);
 
             // If I execute a batch that is invalid
             var fileStreamFactory = Common.GetFileStreamFactory(new Dictionary<string, byte[]>());
             Batch batch = new Batch(Common.StandardQuery, Common.SubsectionDocument, Common.Ordinal, fileStreamFactory);
-            batch.BatchCompletion += callback;
+            batch.BatchCompletion += batchCallback;
+            batch.ResultSetCompletion += resultSetCallback;
             batch.Execute(GetConnection(ci), CancellationToken.None).Wait();
 
             // Then:
@@ -222,14 +275,25 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution.Execution
             Assert.NotEmpty(batch.ResultMessages);
 
             // ... The callback for batch completion should have been fired
-            Assert.True(completionCallbackFired);
             Assert.NotNull(batchSummaryFromCallback);
         }
 
         [Fact]
         public async Task BatchExecuteExecuted()
         {
-            // Setup: Create a callback for batch completion
+            ConnectionInfo ci = Common.CreateTestConnectionInfo(new[] { Common.StandardTestData }, false);
+
+            // If I execute a batch
+            var fileStreamFactory = Common.GetFileStreamFactory(new Dictionary<string, byte[]>());
+            Batch batch = new Batch(Common.StandardQuery, Common.SubsectionDocument, Common.Ordinal, fileStreamFactory);
+            batch.Execute(GetConnection(ci), CancellationToken.None).Wait();
+
+            // Then:
+            // ... It should have executed without error
+            Assert.True(batch.HasExecuted, "The batch should have been marked executed.");
+            Assert.False(batch.HasError, "The batch should not have an error");
+
+            // Setup for part 2: Create a callback for batch completion
             BatchSummary batchSummaryFromCallback = null;
             bool completionCallbackFired = false;
             Batch.BatchAsyncEventHandler callback = b =>
@@ -239,26 +303,10 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution.Execution
                 return Task.FromResult(0);
             };
 
-            ConnectionInfo ci = Common.CreateTestConnectionInfo(new[] { Common.StandardTestData }, false);
-
-            // If I execute a batch
-            var fileStreamFactory = Common.GetFileStreamFactory(new Dictionary<string, byte[]>());
-            Batch batch = new Batch(Common.StandardQuery, Common.SubsectionDocument, Common.Ordinal, fileStreamFactory);
-            batch.BatchCompletion += callback;
-            batch.Execute(GetConnection(ci), CancellationToken.None).Wait();
-
-            // Then:
-            // ... It should have executed without error
-            Assert.True(batch.HasExecuted, "The batch should have been marked executed.");
-            Assert.False(batch.HasError, "The batch should not have an error");
-
-            // ... The callback for batch completion should have been fired
-            Assert.True(completionCallbackFired);
-            Assert.NotNull(batchSummaryFromCallback);
-
             // If I execute it again
             // Then:
             // ... It should throw an invalid operation exception
+            batch.BatchCompletion += callback;
             await Assert.ThrowsAsync<InvalidOperationException>(() =>
                 batch.Execute(GetConnection(ci), CancellationToken.None));
 
@@ -267,6 +315,10 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution.Execution
             Assert.True(batch.HasExecuted, "The batch should still be marked executed.");
             Assert.NotEmpty(batch.ResultSets);
             Assert.NotEmpty(batch.ResultSummaries);
+
+            // ... The callback for batch completion should not have been fired for the second run
+            Assert.False(completionCallbackFired);
+            Assert.Null(batchSummaryFromCallback);
         }
 
         [Theory]

test/Microsoft.SqlTools.ServiceLayer.Test/QueryExecution/Execution/QueryTests.cs

@@ -5,6 +5,7 @@
 
 using System;
 using System.Collections.Generic;
+using System.Threading.Tasks;
 using Microsoft.SqlTools.ServiceLayer.Connection;
 using Microsoft.SqlTools.ServiceLayer.QueryExecution;
 using Microsoft.SqlTools.ServiceLayer.SqlContext;
@@ -61,11 +62,21 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution.Execution
         [Fact]
         public void QueryExecuteSingleBatch()
         {
+            // Setup:
+            // ... Create a callback for batch completion
+            int batchCallbacksReceived = 0;
+            Batch.BatchAsyncEventHandler batchCallback = summary =>
+            {
+                batchCallbacksReceived++;
+                return Task.CompletedTask;
+            };
+
             // If:
             // ... I create a query from a single batch (without separator)
             ConnectionInfo ci = Common.CreateTestConnectionInfo(null, false);
             var fileStreamFactory = Common.GetFileStreamFactory(new Dictionary<string, byte[]>());
             Query query = new Query(Common.StandardQuery, ci, new QueryExecutionSettings(), fileStreamFactory);
+            query.BatchCompleted += batchCallback;
 
             // Then:
             // ... I should get a single batch to execute that hasn't been executed
@@ -85,16 +96,27 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution.Execution
             Assert.True(query.HasExecuted);
             Assert.NotEmpty(query.BatchSummaries);
             Assert.Equal(1, query.BatchSummaries.Length);
+
+            // ... The batch callback should have been called precisely 1 time
+            Assert.Equal(1, batchCallbacksReceived);
         }
 
         [Fact]
         public void QueryExecuteNoOpBatch()
         {
+            // Setup:
+            // ... Create a callback for batch completion
+            Batch.BatchAsyncEventHandler batchCallback = summary =>
+            {
+                throw new Exception("Batch completion callback was called");
+            };
+
             // If:
             // ... I create a query from a single batch that does nothing
             ConnectionInfo ci = Common.CreateTestConnectionInfo(null, false);
             var fileStreamFactory = Common.GetFileStreamFactory(new Dictionary<string, byte[]>());
             Query query = new Query(Common.NoOpQuery, ci, new QueryExecutionSettings(), fileStreamFactory);
+            query.BatchCompleted += batchCallback;
 
             // Then:
             // ... I should get no batches back
@@ -117,12 +139,22 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution.Execution
         [Fact]
         public void QueryExecuteMultipleBatches()
         {
+            // Setup:
+            // ... Create a callback for batch completion
+            int batchCallbacksReceived = 0;
+            Batch.BatchAsyncEventHandler batchCallback = summary =>
+            {
+                batchCallbacksReceived++;
+                return Task.CompletedTask;
+            };
+
             // If:
             // ... I create a query from two batches (with separator)
             ConnectionInfo ci = Common.CreateTestConnectionInfo(null, false);
             string queryText = string.Format("{0}\r\nGO\r\n{0}", Common.StandardQuery);
             var fileStreamFactory = Common.GetFileStreamFactory(new Dictionary<string, byte[]>());
             Query query = new Query(queryText, ci, new QueryExecutionSettings(), fileStreamFactory);
+            query.BatchCompleted += batchCallback;
 
             // Then:
             // ... I should get back two batches to execute that haven't been executed
@@ -142,17 +174,30 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution.Execution
             Assert.True(query.HasExecuted);
             Assert.NotEmpty(query.BatchSummaries);
             Assert.Equal(2, query.BatchSummaries.Length);
+
+            // ... The batch callback should have been called precisely 2 times
+            Assert.Equal(2, batchCallbacksReceived);
         }
 
         [Fact]
         public void QueryExecuteMultipleBatchesWithNoOp()
         {
+            // Setup:
+            // ... Create a callback for batch completion
+            int batchCallbacksReceived = 0;
+            Batch.BatchAsyncEventHandler batchCallback = summary =>
+            {
+                batchCallbacksReceived++;
+                return Task.CompletedTask;
+            };
+
             // If:
             // ... I create a query from a two batches (with separator)
             ConnectionInfo ci = Common.CreateTestConnectionInfo(null, false);
             string queryText = string.Format("{0}\r\nGO\r\n{1}", Common.StandardQuery, Common.NoOpQuery);
             var fileStreamFactory = Common.GetFileStreamFactory(new Dictionary<string, byte[]>());
             Query query = new Query(queryText, ci, new QueryExecutionSettings(), fileStreamFactory);
+            query.BatchCompleted += batchCallback;
 
             // Then:
             // ... I should get back one batch to execute that hasn't been executed
@@ -171,16 +216,29 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution.Execution
             Assert.True(query.HasExecuted);
             Assert.NotEmpty(query.BatchSummaries);
             Assert.Equal(1, query.BatchSummaries.Length);
+
+            // ... The batch callback should have been called precisely 1 time
+            Assert.Equal(1, batchCallbacksReceived);
         }
 
         [Fact]
         public void QueryExecuteInvalidBatch()
         {
+            // Setup:
+            // ... Create a callback for batch completion
+            int batchCallbacksReceived = 0;
+            Batch.BatchAsyncEventHandler batchCallback = summary =>
+            {
+                batchCallbacksReceived++;
+                return Task.CompletedTask;
+            };
+
             // If:
             // ... I create a query from an invalid batch
             ConnectionInfo ci = Common.CreateTestConnectionInfo(null, true);
             var fileStreamFactory = Common.GetFileStreamFactory(new Dictionary<string, byte[]>());
             Query query = new Query(Common.InvalidQuery, ci, new QueryExecutionSettings(), fileStreamFactory);
+            query.BatchCompleted += batchCallback;
 
             // Then:
             // ... I should get back a query with one batch not executed
@@ -202,6 +260,9 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution.Execution
             Assert.Equal(1, query.BatchSummaries.Length);
             Assert.True(query.BatchSummaries[0].HasError);
             Assert.NotEmpty(query.BatchSummaries[0].Messages);
+
+            // ... The batch callback should have been called once
+            Assert.Equal(1, batchCallbacksReceived);
         }
 
     }

test/Microsoft.SqlTools.ServiceLayer.Test/QueryExecution/Execution/ResultSetTests.cs

@@ -22,14 +22,20 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution.Execution
         {
             // If:
             // ... I create a new result set with a valid db data reader
-
             DbDataReader mockReader = GetReader(null, false, string.Empty);
-            ResultSet resultSet = new ResultSet(mockReader, Common.GetFileStreamFactory(new Dictionary<string, byte[]>()));
+            ResultSet resultSet = new ResultSet(mockReader, Common.Ordinal, Common.Ordinal, Common.GetFileStreamFactory(new Dictionary<string, byte[]>()));
 
             // Then:
             // ... There should not be any data read yet
             Assert.Null(resultSet.Columns);
             Assert.Equal(0, resultSet.RowCount);
+            Assert.Equal(Common.Ordinal, resultSet.Id);
+
+            // ... The summary should include the same info
+            Assert.Null(resultSet.Summary.ColumnInfo);
+            Assert.Equal(0, resultSet.Summary.RowCount);
+            Assert.Equal(Common.Ordinal, resultSet.Summary.Id);
+            Assert.Equal(Common.Ordinal, resultSet.Summary.BatchId);
         }
 
         [Fact]
@@ -39,29 +45,45 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution.Execution
             // ... I create a new result set without a reader
             // Then:
             // ... It should throw an exception
-            Assert.Throws<ArgumentNullException>(() => new ResultSet(null, null));
+            Assert.Throws<ArgumentNullException>(() => new ResultSet(null, Common.Ordinal, Common.Ordinal, null));
 
         }
 
         [Fact]
         public async Task ReadToEndSuccess()
         {
+            // Setup: Create a callback for resultset completion
+            ResultSetSummary resultSummaryFromCallback = null;
+            ResultSet.ResultSetAsyncEventHandler callback = r =>
+            {
+                resultSummaryFromCallback = r.Summary;
+                return Task.FromResult(0);
+            };
+
             // If:
             // ... I create a new resultset with a valid db data reader that has data
             // ... and I read it to the end
             DbDataReader mockReader = GetReader(new [] {Common.StandardTestData}, false, Common.StandardQuery);
             var fileStreamFactory = Common.GetFileStreamFactory(new Dictionary<string, byte[]>());
-            ResultSet resultSet = new ResultSet(mockReader, fileStreamFactory);
+            ResultSet resultSet = new ResultSet(mockReader, Common.Ordinal, Common.Ordinal, fileStreamFactory);
+            resultSet.ResultCompletion += callback;
             await resultSet.ReadResultToEnd(CancellationToken.None);
 
             // Then:
             // ... The columns should be set
             // ... There should be rows to read back
             Assert.NotNull(resultSet.Columns);
-            Assert.NotEmpty(resultSet.Columns);
+            Assert.Equal(Common.StandardColumns, resultSet.Columns.Length);
             Assert.Equal(Common.StandardRows, resultSet.RowCount);
-        }
 
+            // ... The summary should have the same info
+            Assert.NotNull(resultSet.Summary.ColumnInfo);
+            Assert.Equal(Common.StandardColumns, resultSet.Summary.ColumnInfo.Length);
+            Assert.Equal(Common.StandardRows, resultSet.Summary.RowCount);
+
+            // ... The callback for result set completion should have been fired
+            Assert.NotNull(resultSummaryFromCallback);
+        }
         
         [Theory]
         [InlineData("JSON")]
@@ -78,12 +100,21 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution.Execution
             }
             Dictionary<string, string>[][] dataSets = {data.ToArray()};
 
+            // ... Create a callback for resultset completion
+            ResultSetSummary resultSummary = null;
+            ResultSet.ResultSetAsyncEventHandler callback = r =>
+            {
+                resultSummary = r.Summary;
+                return Task.FromResult(0);
+            };
+
             // If:
             // ... I create a new resultset with a valid db data reader that is FOR XML/JSON
             // ... and I read it to the end
             DbDataReader mockReader = GetReader(dataSets, false, Common.StandardQuery);
             var fileStreamFactory = Common.GetFileStreamFactory(new Dictionary<string, byte[]>());
-            ResultSet resultSet = new ResultSet(mockReader, fileStreamFactory);
+            ResultSet resultSet = new ResultSet(mockReader, Common.Ordinal, Common.Ordinal, fileStreamFactory);
+            resultSet.ResultCompletion += callback;
             await resultSet.ReadResultToEnd(CancellationToken.None);
 
             // Then:
@@ -93,6 +124,9 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution.Execution
             Assert.Equal(1, resultSet.Columns.Length);
             Assert.Equal(1, resultSet.RowCount);
 
+            // ... The callback should have been called
+            Assert.NotNull(resultSummary);
+
             // If:
             // ... I attempt to read back the results
             // Then: 
@@ -108,7 +142,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution.Execution
             // ... I create a new result set with a valid db data reader without executing it
             DbDataReader mockReader = GetReader(null, false, string.Empty);
             var fileStreamFactory = Common.GetFileStreamFactory(new Dictionary<string, byte[]>());
-            ResultSet resultSet = new ResultSet(mockReader, fileStreamFactory);
+            ResultSet resultSet = new ResultSet(mockReader, Common.Ordinal, Common.Ordinal, fileStreamFactory);
 
             // Then:
             // ... Attempting to read a subset should fail miserably
@@ -126,7 +160,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution.Execution
             // ... And execute the result
             DbDataReader mockReader = GetReader(new[] {Common.StandardTestData}, false, Common.StandardQuery);
             var fileStreamFactory = Common.GetFileStreamFactory(new Dictionary<string, byte[]>());
-            ResultSet resultSet = new ResultSet(mockReader, fileStreamFactory);
+            ResultSet resultSet = new ResultSet(mockReader, Common.Ordinal, Common.Ordinal, fileStreamFactory);
             await resultSet.ReadResultToEnd(CancellationToken.None);
 
             // ... And attempt to get a subset with invalid parameters
@@ -147,7 +181,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution.Execution
             // ... And execute the result set
             DbDataReader mockReader = GetReader(new[] { Common.StandardTestData }, false, Common.StandardQuery);
             var fileStreamFactory = Common.GetFileStreamFactory(new Dictionary<string, byte[]>());
-            ResultSet resultSet = new ResultSet(mockReader, fileStreamFactory);
+            ResultSet resultSet = new ResultSet(mockReader, Common.Ordinal, Common.Ordinal, fileStreamFactory);
             await resultSet.ReadResultToEnd(CancellationToken.None);
 
             // ... And attempt to get a subset with valid number of rows

test/Microsoft.SqlTools.ServiceLayer.Test/QueryExecution/Execution/ServiceIntegrationTests.cs

@@ -3,6 +3,7 @@
 // Licensed under the MIT license. See LICENSE file in the project root for full license information.
 //
 
+using System.Collections.Generic;
 using System.Linq;
 using System.Threading.Tasks;
 using Microsoft.SqlTools.ServiceLayer.Hosting.Protocol;
@@ -20,7 +21,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution.Execution
     {
 
         [Fact]
-        public async void QueryExecuteValidNoResultsTest()
+        public async void QueryExecuteSingleBatchNoResultsTest()
         {
             // Given:
             // ... Default settings are stored in the workspace service
@@ -38,7 +39,8 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution.Execution
             QueryExecuteBatchCompleteParams batchCompleteParams = null;
             var requestContext = RequestContextMocks.Create<QueryExecuteResult>(qer => result = qer)
                 .AddEventHandling(QueryExecuteCompleteEvent.Type, (et, p) => completeParams = p)
-                .AddEventHandling(QueryExecuteBatchCompleteEvent.Type, (et, p) => batchCompleteParams = p);
+                .AddEventHandling(QueryExecuteBatchCompleteEvent.Type, (et, p) => batchCompleteParams = p)
+                .AddEventHandling(QueryExecuteResultSetCompleteEvent.Type, null);
             await Common.AwaitExecution(queryService, queryParams, requestContext.Object);
 
             // Then:
@@ -46,7 +48,8 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution.Execution
             // ... A successful result should have been sent with messages on the first batch
             // ... A completion event should have been fired with empty results
             // ... A batch completion event should have been fired with empty results
-            VerifyQueryExecuteCallCount(requestContext, Times.Once(), Times.Once(), Times.Once(), Times.Never());
+            // ... A result set completion event should not have been fired
+            VerifyQueryExecuteCallCount(requestContext, Times.Once(), Times.Once(), Times.Once(), Times.Never(), Times.Never());
             Assert.Null(result.Messages);
 
             Assert.Equal(1, completeParams.BatchSummaries.Length);
@@ -56,7 +59,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution.Execution
             Assert.NotNull(batchCompleteParams);
             Assert.Empty(batchCompleteParams.BatchSummary.ResultSetSummaries);
             Assert.NotEmpty(batchCompleteParams.BatchSummary.Messages);
-            Assert.Equal(completeParams.OwnerUri, batchCompleteParams.OwnerUri);
+            Assert.Equal(Common.OwnerUri, batchCompleteParams.OwnerUri);
 
             // ... There should be one active query
             Assert.Equal(1, queryService.ActiveQueries.Count);
@@ -64,7 +67,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution.Execution
 
 
         [Fact]
-        public async void QueryExecuteValidResultsTest()
+        public async void QueryExecuteSingleBatchSingleResultTest()
         {
             // Given:
             // ... A workspace with a standard query is configured
@@ -78,16 +81,20 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution.Execution
             QueryExecuteResult result = null;
             QueryExecuteCompleteParams completeParams = null;
             QueryExecuteBatchCompleteParams batchCompleteParams = null;
+            QueryExecuteResultSetCompleteParams resultCompleteParams = null;
             var requestContext = RequestContextMocks.Create<QueryExecuteResult>(qer => result = qer)
                 .AddEventHandling(QueryExecuteCompleteEvent.Type, (et, p) => completeParams = p)
-                .AddEventHandling(QueryExecuteBatchCompleteEvent.Type, (et, p) => batchCompleteParams = p);
+                .AddEventHandling(QueryExecuteBatchCompleteEvent.Type, (et, p) => batchCompleteParams = p)
+                .AddEventHandling(QueryExecuteResultSetCompleteEvent.Type, (et, p) => resultCompleteParams = p);
             await Common.AwaitExecution(queryService, queryParams, requestContext.Object);
 
             // Then:
             // ... No errors should have been sent
-            // ... A successful result should have been sent with messages
+            // ... A successful result should have been sent without messages
             // ... A completion event should have been fired with one result
-            VerifyQueryExecuteCallCount(requestContext, Times.Once(), Times.Once(), Times.Once(), Times.Never());
+            // ... A batch completion event should have been fired
+            // ... A resultset completion event should have been fired
+            VerifyQueryExecuteCallCount(requestContext, Times.Once(), Times.Once(), Times.Once(), Times.Once(), Times.Never());
             Assert.Null(result.Messages);
 
             Assert.Equal(1, completeParams.BatchSummaries.Length);
@@ -98,7 +105,124 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution.Execution
             Assert.NotNull(batchCompleteParams);
             Assert.NotEmpty(batchCompleteParams.BatchSummary.ResultSetSummaries);
             Assert.NotEmpty(batchCompleteParams.BatchSummary.Messages);
-            Assert.Equal(completeParams.OwnerUri, batchCompleteParams.OwnerUri);
+            Assert.Equal(Common.OwnerUri, batchCompleteParams.OwnerUri);
+
+            Assert.NotNull(resultCompleteParams);
+            Assert.Equal(Common.StandardColumns, resultCompleteParams.ResultSetSummary.ColumnInfo.Length);
+            Assert.Equal(Common.StandardRows, resultCompleteParams.ResultSetSummary.RowCount);
+            Assert.Equal(Common.OwnerUri, resultCompleteParams.OwnerUri);
+
+            // ... There should be one active query
+            Assert.Equal(1, queryService.ActiveQueries.Count);
+        }
+
+        [Fact]
+        public async Task QueryExecuteSingleBatchMultipleResultTest()
+        {
+            // Given:
+            // ... A workspace with a standard query is configured
+            var workspaceService = Common.GetPrimedWorkspaceService(Common.StandardQuery);
+
+            // If:
+            // ... I request to execute a valid query with one batch and multiple result sets
+            var dataset = new[] { Common.StandardTestData, Common.StandardTestData };
+            var queryService = Common.GetPrimedExecutionService(dataset, true, false, workspaceService);
+            var queryParams = new QueryExecuteParams { OwnerUri = Common.OwnerUri, QuerySelection = Common.WholeDocument };
+
+            QueryExecuteResult result = null;
+            QueryExecuteCompleteParams completeParams = null;
+            QueryExecuteBatchCompleteParams batchCompleteParams = null;
+            List<QueryExecuteResultSetCompleteParams> resultCompleteParams = new List<QueryExecuteResultSetCompleteParams>();
+            var requestContext = RequestContextMocks.Create<QueryExecuteResult>(qer => result = qer)
+                .AddEventHandling(QueryExecuteCompleteEvent.Type, (et, p) => completeParams = p)
+                .AddEventHandling(QueryExecuteBatchCompleteEvent.Type, (et, p) => batchCompleteParams = p)
+                .AddEventHandling(QueryExecuteResultSetCompleteEvent.Type, (et, p) => resultCompleteParams.Add(p));
+            await Common.AwaitExecution(queryService, queryParams, requestContext.Object);
+
+            // Then:
+            // ... No errors should have been sent
+            // ... A successful result should have been sent without messages
+            // ... A completion event should have been fired with one result
+            // ... A batch completion event should have been fired
+            // ... Two resultset completion events should have been fired
+            VerifyQueryExecuteCallCount(requestContext, Times.Once(), Times.Once(), Times.Once(), Times.Exactly(2), Times.Never());
+            Assert.Null(result.Messages);
+
+            Assert.Equal(1, completeParams.BatchSummaries.Length);
+            Assert.NotEmpty(completeParams.BatchSummaries[0].ResultSetSummaries);
+            Assert.NotEmpty(completeParams.BatchSummaries[0].Messages);
+            Assert.False(completeParams.BatchSummaries[0].HasError);
+
+            Assert.NotNull(batchCompleteParams);
+            Assert.NotEmpty(batchCompleteParams.BatchSummary.ResultSetSummaries);
+            Assert.NotEmpty(batchCompleteParams.BatchSummary.Messages);
+            Assert.Equal(Common.OwnerUri, batchCompleteParams.OwnerUri);
+
+            Assert.Equal(2, resultCompleteParams.Count);
+            foreach (var resultParam in resultCompleteParams)
+            {
+                Assert.NotNull(resultCompleteParams);
+                Assert.Equal(Common.StandardColumns, resultParam.ResultSetSummary.ColumnInfo.Length);
+                Assert.Equal(Common.StandardRows, resultParam.ResultSetSummary.RowCount);
+                Assert.Equal(Common.OwnerUri, resultParam.OwnerUri);
+            }
+        }
+
+        [Fact]
+        public async Task QueryExecuteMultipleBatchSingleResultTest()
+        {
+            // Given:
+            // ... A workspace with a standard query is configured
+            var workspaceService = Common.GetPrimedWorkspaceService(string.Format("{0}\r\nGO\r\n{0}", Common.StandardQuery));
+
+            // If:
+            // ... I request a to execute a valid query with multiple batches
+            var dataSet = new[] { Common.StandardTestData };
+            var queryService = Common.GetPrimedExecutionService(dataSet, true, false, workspaceService);
+            var queryParams = new QueryExecuteParams { OwnerUri = Common.OwnerUri, QuerySelection = Common.WholeDocument };
+
+            QueryExecuteResult result = null;
+            QueryExecuteCompleteParams completeParams = null;
+            List<QueryExecuteBatchCompleteParams> batchCompleteParams = new List<QueryExecuteBatchCompleteParams>();
+            List<QueryExecuteResultSetCompleteParams> resultCompleteParams = new List<QueryExecuteResultSetCompleteParams>();
+            var requestContext = RequestContextMocks.Create<QueryExecuteResult>(qer => result = qer)
+                .AddEventHandling(QueryExecuteCompleteEvent.Type, (et, p) => completeParams = p)
+                .AddEventHandling(QueryExecuteBatchCompleteEvent.Type, (et, p) => batchCompleteParams.Add(p))
+                .AddEventHandling(QueryExecuteResultSetCompleteEvent.Type, (et, p) => resultCompleteParams.Add(p));
+            await Common.AwaitExecution(queryService, queryParams, requestContext.Object);
+
+            // Then:
+            // ... No errors should have been sent
+            // ... A successful result should have been sent without messages
+
+            VerifyQueryExecuteCallCount(requestContext, Times.Once(), Times.Once(), Times.Exactly(2), Times.Exactly(2), Times.Never());
+            Assert.Null(result.Messages);
+
+            // ... A completion event should have been fired with one two batch summaries, one result each
+            Assert.Equal(2, completeParams.BatchSummaries.Length);
+            Assert.Equal(1, completeParams.BatchSummaries[0].ResultSetSummaries.Length);
+            Assert.Equal(1, completeParams.BatchSummaries[1].ResultSetSummaries.Length);
+            Assert.NotEmpty(completeParams.BatchSummaries[0].Messages);
+            Assert.NotEmpty(completeParams.BatchSummaries[1].Messages);
+
+            // ... Two batch completion events should have been fired
+            Assert.Equal(2, batchCompleteParams.Count);
+            foreach (var batch in batchCompleteParams)
+            {
+                Assert.NotEmpty(batch.BatchSummary.ResultSetSummaries);
+                Assert.NotEmpty(batch.BatchSummary.Messages);
+                Assert.Equal(Common.OwnerUri, batch.OwnerUri);
+            }
+
+            // ... Two resultset completion events should have been fired
+            Assert.Equal(2, resultCompleteParams.Count);
+            foreach (var resultParam in resultCompleteParams)
+            {
+                Assert.NotNull(resultParam.ResultSetSummary);
+                Assert.Equal(Common.StandardColumns, resultParam.ResultSetSummary.ColumnInfo.Length);
+                Assert.Equal(Common.StandardRows, resultParam.ResultSetSummary.RowCount);
+                Assert.Equal(Common.OwnerUri, resultParam.OwnerUri);
+            }
 
             // ... There should be one active query
             Assert.Equal(1, queryService.ActiveQueries.Count);
@@ -126,7 +250,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution.Execution
             // ... No result should have been returned
             // ... No completion event should have been fired
             // ... There should be no active queries
-            VerifyQueryExecuteCallCount(requestContext, Times.Never(), Times.Never(), Times.Never(), Times.Once());
+            VerifyQueryExecuteCallCount(requestContext, Times.Never(), Times.Never(), Times.Never(), Times.Never(), Times.Once());
             Assert.IsType<string>(error);
             Assert.NotEmpty((string)error);
             Assert.Empty(queryService.ActiveQueries);
@@ -159,8 +283,9 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution.Execution
             // ... An error should have been sent
             // ... A result should have not have been sent
             // ... No completion event should have been fired
+            // ... A batch completion event should have fired, but not a resultset event
             // ... There should only be one active query
-            VerifyQueryExecuteCallCount(secondRequestContext, Times.Never(), Times.AtMostOnce(), Times.AtMostOnce(), Times.Once());
+            VerifyQueryExecuteCallCount(secondRequestContext, Times.Never(), Times.AtMostOnce(), Times.AtMostOnce(), Times.Never(), Times.Once());
             Assert.IsType<string>(error);
             Assert.NotEmpty((string)error);
             Assert.Equal(1, queryService.ActiveQueries.Count);
@@ -197,7 +322,8 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution.Execution
             // ... No errors should have been sent
             // ... A result should have been sent with no errors
             // ... There should only be one active query
-            VerifyQueryExecuteCallCount(secondRequestContext, Times.Once(), Times.Once(), Times.Once(), Times.Never());
+            // ... A batch completion event should have fired, but not a result set completion event
+            VerifyQueryExecuteCallCount(secondRequestContext, Times.Once(), Times.Once(), Times.Once(), Times.Never(), Times.Never());
             Assert.Null(result.Messages);
 
             Assert.False(complete.BatchSummaries.Any(b => b.HasError));
@@ -230,9 +356,9 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution.Execution
             // Then:
             // ... Am error should have been sent
             // ... No result should have been sent
-            // ... No completion event should have been fired
+            // ... No completion events should have been fired
             // ... An active query should not have been added
-            VerifyQueryExecuteCallCount(requestContext, Times.Never(), Times.Never(), Times.Never(), Times.Once());
+            VerifyQueryExecuteCallCount(requestContext, Times.Never(), Times.Never(), Times.Never(), Times.Never(), Times.Once());
             Assert.NotNull(errorResult);
             Assert.IsType<string>(errorResult);
             Assert.DoesNotContain(Common.OwnerUri, queryService.ActiveQueries.Keys);
@@ -264,8 +390,8 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution.Execution
             // Then:
             // ... No errors should have been sent
             // ... A result should have been sent with success (we successfully started the query)
-            // ... A completion event should have been sent with error
-            VerifyQueryExecuteCallCount(requestContext, Times.Once(), Times.Once(), Times.Once(), Times.Never());
+            // ... A completion event (query, batch, not resultset) should have been sent with error
+            VerifyQueryExecuteCallCount(requestContext, Times.Once(), Times.Once(), Times.Once(), Times.Never(), Times.Never());
             Assert.Null(result.Messages);
 
             Assert.Equal(1, complete.BatchSummaries.Length);
@@ -279,7 +405,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution.Execution
         }
 
         private static void VerifyQueryExecuteCallCount(Mock<RequestContext<QueryExecuteResult>> mock, Times sendResultCalls,
-            Times sendCompletionEventCalls, Times sendBatchCompletionEvent, Times sendErrorCalls)
+            Times sendCompletionEventCalls, Times sendBatchCompletionEvent, Times sendResultCompleteEvent, Times sendErrorCalls)
         {
             mock.Verify(rc => rc.SendResult(It.IsAny<QueryExecuteResult>()), sendResultCalls);
             mock.Verify(rc => rc.SendEvent(
@@ -288,6 +414,10 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution.Execution
             mock.Verify(rc => rc.SendEvent(
                 It.Is<EventType<QueryExecuteBatchCompleteParams>>(m => m == QueryExecuteBatchCompleteEvent.Type),
                 It.IsAny<QueryExecuteBatchCompleteParams>()), sendBatchCompletionEvent);
+            mock.Verify(rc => rc.SendEvent(
+                It.Is<EventType<QueryExecuteResultSetCompleteParams>>(m => m == QueryExecuteResultSetCompleteEvent.Type),
+                It.IsAny<QueryExecuteResultSetCompleteParams>()), sendResultCompleteEvent);
+
             mock.Verify(rc => rc.SendError(It.IsAny<object>()), sendErrorCalls);
         }
     }