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

src/Microsoft.SqlTools.ServiceLayer/Connection/ConnectionService.cs

@@ -29,7 +29,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection
         /// <summary>
         /// Singleton service instance
         /// </summary>
-        private static Lazy<ConnectionService> instance 
+        private static readonly Lazy<ConnectionService> instance 
             = new Lazy<ConnectionService>(() => new ConnectionService());
 
         /// <summary>
@@ -48,11 +48,11 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection
         /// </summary>
         private ISqlConnectionFactory connectionFactory;
            
-        private Dictionary<string, ConnectionInfo> ownerToConnectionMap = new Dictionary<string, ConnectionInfo>();
+        private readonly Dictionary<string, ConnectionInfo> ownerToConnectionMap = new Dictionary<string, ConnectionInfo>();
 
-        private ConcurrentDictionary<string, CancellationTokenSource> ownerToCancellationTokenSourceMap = new ConcurrentDictionary<string, CancellationTokenSource>();
+        private readonly ConcurrentDictionary<string, CancellationTokenSource> ownerToCancellationTokenSourceMap = new ConcurrentDictionary<string, CancellationTokenSource>();
 
-        private Object cancellationTokenSourceLock = new Object();
+        private readonly object cancellationTokenSourceLock = new object();
 
         /// <summary>
         /// Map from script URIs to ConnectionInfo objects
@@ -173,13 +173,12 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection
             connectionInfo = new ConnectionInfo(ConnectionFactory, connectionParams.OwnerUri, connectionParams.Connection);
 
             // try to connect
-            var response = new ConnectionCompleteParams();
+            var response = new ConnectionCompleteParams {OwnerUri = connectionParams.OwnerUri};
-            response.OwnerUri = connectionParams.OwnerUri;
             CancellationTokenSource source = null;
             try
             {
                 // build the connection string from the input parameters
-                string connectionString = ConnectionService.BuildConnectionString(connectionInfo.ConnectionDetails);
+                string connectionString = BuildConnectionString(connectionInfo.ConnectionDetails);
 
                 // create a sql connection instance
                 connectionInfo.SqlConnection = connectionInfo.Factory.CreateSqlConnection(connectionString);
@@ -261,7 +260,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection
             // Update with the actual database name in connectionInfo and result
             // Doing this here as we know the connection is open - expect to do this only on connecting
             connectionInfo.ConnectionDetails.DatabaseName = connectionInfo.SqlConnection.Database;
-            response.ConnectionSummary = new ConnectionSummary()
+            response.ConnectionSummary = new ConnectionSummary
             {
                 ServerName = connectionInfo.ConnectionDetails.ServerName,
                 DatabaseName = connectionInfo.ConnectionDetails.DatabaseName,
@@ -269,7 +268,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection
             };
 
             // invoke callback notifications
-            invokeOnConnectionActivities(connectionInfo);
+            InvokeOnConnectionActivities(connectionInfo);
 
             // try to get information about the connected SQL Server instance
             try
@@ -278,7 +277,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection
                 DbConnection connection = reliableConnection != null ? reliableConnection.GetUnderlyingConnection() : connectionInfo.SqlConnection;
                 
                 ReliableConnectionHelper.ServerInfo serverInfo = ReliableConnectionHelper.GetServerVersion(connection);
-                response.ServerInfo = new Contracts.ServerInfo()
+                response.ServerInfo = new ServerInfo
                 {
                     ServerMajorVersion = serverInfo.ServerMajorVersion,
                     ServerMinorVersion = serverInfo.ServerMinorVersion,
@@ -399,7 +398,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection
             connection.Open();
             
             List<string> results = new List<string>();
-            var systemDatabases = new string[] {"master", "model", "msdb", "tempdb"};
+            var systemDatabases = new[] {"master", "model", "msdb", "tempdb"};
             using (DbCommand command = connection.CreateCommand())
             {
                 command.CommandText = "SELECT name FROM sys.databases ORDER BY name ASC";
@@ -473,7 +472,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection
 
             try
             {
-                RunConnectRequestHandlerTask(connectParams, requestContext);
+                RunConnectRequestHandlerTask(connectParams);
                 await requestContext.SendResult(true);
             }
             catch
@@ -482,7 +481,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection
             }
         }
 
-        private void RunConnectRequestHandlerTask(ConnectParams connectParams, RequestContext<bool> requestContext)
+        private void RunConnectRequestHandlerTask(ConnectParams connectParams)
         {
             // create a task to connect asynchronously so that other requests are not blocked in the meantime
             Task.Run(async () => 
@@ -490,7 +489,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection
                 try
                 {
                     // open connection based on request details
-                    ConnectionCompleteParams result = await ConnectionService.Instance.Connect(connectParams);
+                    ConnectionCompleteParams result = await Instance.Connect(connectParams);
                     await ServiceHost.SendEvent(ConnectionCompleteNotification.Type, result);
                 }
                 catch (Exception ex)
@@ -515,7 +514,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection
 
             try
             {
-                bool result = ConnectionService.Instance.CancelConnect(cancelParams);
+                bool result = Instance.CancelConnect(cancelParams);
                 await requestContext.SendResult(result);
             }
             catch(Exception ex)
@@ -535,7 +534,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection
 
             try
             {
-                bool result = ConnectionService.Instance.Disconnect(disconnectParams);
+                bool result = Instance.Disconnect(disconnectParams);
                 await requestContext.SendResult(result);
             }
             catch(Exception ex)
@@ -556,7 +555,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection
 
             try
             {
-                ListDatabasesResponse result = ConnectionService.Instance.ListDatabases(listDatabasesParams);
+                ListDatabasesResponse result = Instance.ListDatabases(listDatabasesParams);
                 await requestContext.SendResult(result);
             }
             catch(Exception ex)
@@ -579,10 +578,12 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection
         /// <param name="connectionDetails"></param>
         public static string BuildConnectionString(ConnectionDetails connectionDetails)
         {
-            SqlConnectionStringBuilder connectionBuilder = new SqlConnectionStringBuilder();
+            SqlConnectionStringBuilder connectionBuilder = new SqlConnectionStringBuilder
-            connectionBuilder["Data Source"] = connectionDetails.ServerName;
+            {
-            connectionBuilder["User Id"] = connectionDetails.UserName;
+                ["Data Source"] = connectionDetails.ServerName,
-            connectionBuilder["Password"] = connectionDetails.Password;
+                ["User Id"] = connectionDetails.UserName,
+                ["Password"] = connectionDetails.Password
+            };
 
             // Check for any optional parameters
             if (!string.IsNullOrEmpty(connectionDetails.DatabaseName))
@@ -722,7 +723,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection
 
                     // Fire a connection changed event
                     ConnectionChangedParams parameters = new ConnectionChangedParams();
-                    ConnectionSummary summary = (ConnectionSummary)(info.ConnectionDetails);
+                    ConnectionSummary summary = info.ConnectionDetails;
                     parameters.Connection = summary.Clone();
                     parameters.OwnerUri = ownerUri;
                     ServiceHost.SendEvent(ConnectionChangedNotification.Type, parameters);
@@ -741,7 +742,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection
             }
         }
 
-        private void invokeOnConnectionActivities(ConnectionInfo connectionInfo)
+        private void InvokeOnConnectionActivities(ConnectionInfo connectionInfo)
         {
             foreach (var activity in this.onConnectionActivities)
             {

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

@@ -88,6 +88,12 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
         /// </summary>
         public event BatchAsyncEventHandler BatchCompletion;
 
+        /// <summary>
+        /// Event that will be called when the resultset has completed execution. It will not be
+        /// called from the Batch but from the ResultSet instance
+        /// </summary>
+        public event ResultSet.ResultSetAsyncEventHandler ResultSetCompletion;
+
         /// <summary>
         /// The text of batch that will be executed
         /// </summary>
@@ -155,12 +161,10 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
         {
             get
             {
-                return ResultSets.Select((set, index) => new ResultSetSummary()
+                lock (resultSets)
                 {
-                    ColumnInfo = set.Columns,
+                    return resultSets.Select(set => set.Summary).ToArray();
-                    Id = index,
+                }
-                    RowCount = set.RowCount
-                }).ToArray();
             }
         }
 
@@ -221,7 +225,7 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
                     command = sqlConn.GetUnderlyingConnection().CreateCommand();
 
                     // Add a handler for when the command completes
-                    SqlCommand sqlCommand = (SqlCommand) command;
+                    SqlCommand sqlCommand = (SqlCommand)command;
                     sqlCommand.StatementCompleted += StatementCompletedHandler;
                 }
                 else
@@ -244,6 +248,7 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
                     // Execute the command to get back a reader
                     using (DbDataReader reader = await command.ExecuteReaderAsync(cancellationToken))
                     {
+                        int resultSetOrdinal = 0;
                         do
                         {
                             // Skip this result set if there aren't any rows (ie, UPDATE/DELETE/etc queries)
@@ -253,10 +258,15 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
                             }
 
                             // This resultset has results (ie, SELECT/etc queries)
-                            ResultSet resultSet = new ResultSet(reader, outputFileFactory);
+                            ResultSet resultSet = new ResultSet(reader, resultSetOrdinal, Id, outputFileFactory);
+                            resultSet.ResultCompletion += ResultSetCompletion;
 
                             // Add the result set to the results of the query
-                            resultSets.Add(resultSet);
+                            lock (resultSets)
+                            {
+                                resultSets.Add(resultSet);
+                                resultSetOrdinal++;
+                            }
 
                             // Read until we hit the end of the result set
                             await resultSet.ReadResultToEnd(cancellationToken).ConfigureAwait(false);
@@ -318,20 +328,21 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
         /// <returns>A subset of results</returns>
         public Task<ResultSetSubset> GetSubset(int resultSetIndex, int startRow, int rowCount)
         {
-            // Sanity check to make sure that the batch has finished
+            ResultSet targetResultSet;
-            if (!HasExecuted)
+            lock (resultSets)
             {
-                throw new InvalidOperationException(SR.QueryServiceSubsetBatchNotCompleted);
+                // Sanity check to make sure we have valid numbers
-            }
+                if (resultSetIndex < 0 || resultSetIndex >= resultSets.Count)
+                {
+                    throw new ArgumentOutOfRangeException(nameof(resultSetIndex),
+                        SR.QueryServiceSubsetResultSetOutOfRange);
+                }
 
-            // Sanity check to make sure we have valid numbers
+                targetResultSet = resultSets[resultSetIndex];
-            if (resultSetIndex < 0 || resultSetIndex >= resultSets.Count)
-            {
-                throw new ArgumentOutOfRangeException(nameof(resultSetIndex), SR.QueryServiceSubsetResultSetOutOfRange);
             }
 
             // Retrieve the result set
-            return resultSets[resultSetIndex].GetSubset(startRow, rowCount);
+            return targetResultSet.GetSubset(startRow, rowCount);
         }
 
         #endregion
@@ -431,9 +442,12 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
 
             if (disposing)
             {
-                foreach (ResultSet r in ResultSets)
+                lock (resultSets)
                 {
-                    r.Dispose();
+                    foreach (ResultSet r in resultSets)
+                    {
+                        r.Dispose();
+                    }
                 }
             }
 

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

@@ -0,0 +1,22 @@
+// 
+// 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
+{
+    public class QueryExecuteResultSetCompleteParams
+    {
+        public ResultSetSummary ResultSetSummary { get; set; }
+
+        public string OwnerUri { get; set; }
+    }
+
+    public class QueryExecuteResultSetCompleteEvent
+    {
+        public static readonly
+            EventType<QueryExecuteResultSetCompleteParams> Type =
+            EventType<QueryExecuteResultSetCompleteParams>.Create("query/resultSetComplete");
+    }
+}

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

@@ -15,6 +15,11 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts
         /// </summary>
         public int Id { get; set; }
 
+        /// <summary>
+        /// The ID of the batch set within the query
+        /// </summary>
+        public int BatchId { get; set; }
+
         /// <summary>
         /// The number of rows that was returned with the resultset
         /// </summary>

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

@@ -92,25 +92,13 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
             Batches = batchSelection.ToArray();
         }
 
-        #region Properties
+        #region Events
-
-        /// <summary>
-        /// Delegate type for callback when a query completes or fails
-        /// </summary>
-        /// <param name="q">The query that completed</param>
-        public delegate Task QueryAsyncEventHandler(Query q);
 
         /// <summary>
         /// Event to be called when a batch is completed.
         /// </summary>
         public event Batch.BatchAsyncEventHandler BatchCompleted;
 
-        /// <summary>
-        /// Delegate type for callback when a query connection fails
-        /// </summary>
-        /// <param name="message">Message to return</param>
-        public delegate Task QueryAsyncErrorEventHandler(string message);
-
         /// <summary>
         /// Callback for when the query has completed successfully
         /// </summary>
@@ -126,6 +114,27 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
         /// </summary>
         public event QueryAsyncErrorEventHandler QueryConnectionException;
 
+        /// <summary>
+        /// Event to be called when a resultset has completed.
+        /// </summary>
+        public event ResultSet.ResultSetAsyncEventHandler ResultSetCompleted;
+
+        #endregion
+
+        #region Properties
+
+        /// <summary>
+        /// Delegate type for callback when a query completes or fails
+        /// </summary>
+        /// <param name="q">The query that completed</param>
+        public delegate Task QueryAsyncEventHandler(Query q);
+
+        /// <summary>
+        /// Delegate type for callback when a query connection fails
+        /// </summary>
+        /// <param name="message">Message to return</param>
+        public delegate Task QueryAsyncErrorEventHandler(string message);
+
         /// <summary>
         /// The batches underneath this query
         /// </summary>
@@ -146,6 +155,10 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
             }
         }
 
+        /// <summary>
+        /// Storage for the async task for execution. Set as internal in order to await completion
+        /// in unit tests.
+        /// </summary>
         internal Task ExecutionTask { get; private set; }
 
         /// <summary>
@@ -242,7 +255,7 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
                 {
                     await conn.OpenAsync();
                 }
-                catch(Exception exception)
+                catch (Exception exception)
                 {
                     this.HasExecuted = true;
                     if (QueryConnectionException != null)
@@ -265,6 +278,7 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
                     foreach (Batch b in Batches)
                     {
                         b.BatchCompletion += BatchCompleted;
+                        b.ResultSetCompletion += ResultSetCompleted;
                         await b.Execute(conn, cancellationSource.Token);
                     }
 

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

@@ -437,7 +437,6 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
                 };
                 await requestContext.SendEvent(QueryExecuteCompleteEvent.Type, eventParams);
             };
-
             query.QueryCompleted += callback;
             query.QueryFailed += callback;
             query.QueryConnectionException += errorCallback;
@@ -454,6 +453,18 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
             };
             query.BatchCompleted += batchCallback;
 
+            // Setup the ResultSet completion callback
+            ResultSet.ResultSetAsyncEventHandler resultCallback = async r =>
+            {
+                QueryExecuteResultSetCompleteParams eventParams = new QueryExecuteResultSetCompleteParams
+                {
+                    ResultSetSummary = r.Summary,
+                    OwnerUri = executeParams.OwnerUri
+                };
+                await requestContext.SendEvent(QueryExecuteResultSetCompleteEvent.Type, eventParams);
+            };
+            query.ResultSetCompleted += resultCallback;
+
             // Launch this as an asynchronous task
             query.Execute();
 

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

@@ -59,9 +59,10 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
         private readonly IFileStreamFactory fileStreamFactory;
 
         /// <summary>
-        /// Whether or not the result set has been read in from the database
+        /// Whether or not the result set has been read in from the database,
+        /// set as internal in order to fake value in unit tests
         /// </summary>
-        private bool hasBeenRead;
+        internal bool hasBeenRead;
 
         /// <summary>
         /// Whether resultSet is a 'for xml' or 'for json' result
@@ -74,7 +75,7 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
         private readonly string outputFileName;
 
         /// <summary>
-        /// Whether the resultSet is in the process of being disposed
+        /// All save tasks currently saving this ResultSet
         /// </summary>
         private readonly ConcurrentDictionary<string, Task> saveTasks;
 
@@ -84,13 +85,17 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
         /// Creates a new result set and initializes its state
         /// </summary>
         /// <param name="reader">The reader from executing a query</param>
+        /// <param name="ordinal">The ID of the resultset, the ordinal of the result within the batch</param>
+        /// <param name="batchOrdinal">The ID of the batch, the ordinal of the batch within the query</param>
         /// <param name="factory">Factory for creating a reader/writer</param>
-        public ResultSet(DbDataReader reader, IFileStreamFactory factory)
+        public ResultSet(DbDataReader reader, int ordinal, int batchOrdinal, IFileStreamFactory factory)
         {
             // Sanity check to make sure we got a reader
             Validate.IsNotNull(nameof(reader), SR.QueryServiceResultSetReaderNull);
 
             dataReader = new StorageDataReader(reader);
+            Id = ordinal;
+            BatchId = batchOrdinal;
 
             // Initialize the storage
             outputFileName = factory.CreateFile();
@@ -104,6 +109,17 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
 
         #region Properties
 
+        /// <summary>
+        /// Asynchronous handler for when a resultset has completed
+        /// </summary>
+        /// <param name="resultSet">The result set that completed</param>
+        public delegate Task ResultSetAsyncEventHandler(ResultSet resultSet);
+
+        /// <summary>
+        /// Event that will be called when the result set has completed execution
+        /// </summary>
+        public event ResultSetAsyncEventHandler ResultCompletion;
+
         /// <summary>
         /// Whether the resultSet is in the process of being disposed
         /// </summary>
@@ -115,6 +131,16 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
         /// </summary>
         public DbColumnWrapper[] Columns { get; private set; }
 
+        /// <summary>
+        /// ID of the result set, relative to the batch
+        /// </summary>
+        public int Id { get; private set; }
+
+        /// <summary>
+        /// ID of the batch set, relative to the query
+        /// </summary>
+        public int BatchId { get; private set; }
+
         /// <summary>
         /// Maximum number of characters to store for a field
         /// </summary>
@@ -130,6 +156,23 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
         /// </summary>
         public long RowCount { get; private set; }
 
+        /// <summary>
+        /// Generates a summary of this result set
+        /// </summary>
+        public ResultSetSummary Summary
+        {
+            get
+            {
+                return new ResultSetSummary
+                {
+                    ColumnInfo = Columns,
+                    Id = Id,
+                    BatchId = BatchId,
+                    RowCount = RowCount
+                };
+            }
+        }
+
         #endregion
 
         #region Public Methods
@@ -170,9 +213,9 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
                     if (isSingleColumnXmlJsonResultSet)
                     {
                         // Iterate over all the rows and process them into a list of string builders
+                        // ReSharper disable once AccessToDisposedClosure   The lambda is used immediately in string.Join call
                         IEnumerable<string> rowValues = fileOffsets.Select(rowOffset => fileStreamReader.ReadRow(rowOffset, Columns)[0].DisplayValue);
                         rows = new[] { new[] { string.Join(string.Empty, rowValues) } };
-
                     }
                     else
                     {
@@ -180,8 +223,9 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
                         IEnumerable<long> rowOffsets = fileOffsets.Skip(startRow).Take(rowCount);
 
                         // Iterate over the rows we need and process them into output
-                        rows = rowOffsets.Select(rowOffset =>
+                        // ReSharper disable once AccessToDisposedClosure   The lambda is used immediately in .ToArray call
-                            fileStreamReader.ReadRow(rowOffset, Columns).Select(cell => cell.DisplayValue).ToArray())
+                        rows = rowOffsets.Select(rowOffset => fileStreamReader.ReadRow(rowOffset, Columns)
+                            .Select(cell => cell.DisplayValue).ToArray())
                             .ToArray();
 
                     }
@@ -201,33 +245,41 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
         /// <param name="cancellationToken">Cancellation token for cancelling the query</param>
         public async Task ReadResultToEnd(CancellationToken cancellationToken)
         {
-            // Mark that result has been read
+            try
-            hasBeenRead = true;
-
-            // Open a writer for the file
-            using (IFileStreamWriter fileWriter = fileStreamFactory.GetWriter(outputFileName, MaxCharsToStore, MaxXmlCharsToStore))
             {
-                // If we can initialize the columns using the column schema, use that
+                // Mark that result has been read
-                if (!dataReader.DbDataReader.CanGetColumnSchema())
+                hasBeenRead = true;
+
+                // Open a writer for the file
+                var fileWriter = fileStreamFactory.GetWriter(outputFileName, MaxCharsToStore, MaxCharsToStore);
+                using (fileWriter)
                 {
-                    throw new InvalidOperationException(SR.QueryServiceResultSetNoColumnSchema);
+                    // If we can initialize the columns using the column schema, use that
+                    if (!dataReader.DbDataReader.CanGetColumnSchema())
+                    {
+                        throw new InvalidOperationException(SR.QueryServiceResultSetNoColumnSchema);
+                    }
+                    Columns = dataReader.Columns;
+                    long currentFileOffset = 0;
+
+                    while (await dataReader.ReadAsync(cancellationToken))
+                    {
+                        RowCount++;
+                        fileOffsets.Add(currentFileOffset);
+                        currentFileOffset += fileWriter.WriteRow(dataReader);
+                    }
                 }
-                Columns = dataReader.Columns;
+                // Check if resultset is 'for xml/json'. If it is, set isJson/isXml value in column metadata
-
+                SingleColumnXmlJsonResultSet();
-                long currentFileOffset = 0;
+            }
-                while (await dataReader.ReadAsync(cancellationToken))
+            finally
+            {
+                // Fire off a result set completion event if we have one
+                if (ResultCompletion != null)
                 {
-                    // Store the beginning of the row
+                    await ResultCompletion(this);
-                    long rowStart = currentFileOffset;
-                    currentFileOffset += fileWriter.WriteRow(dataReader);
-
-                    // Add the row to the list of rows we have only if the row was successfully written
-                    RowCount++;
-                    fileOffsets.Add(rowStart);
                 }
             }
-            // Check if resultset is 'for xml/json'. If it is, set isJson/isXml value in column metadata
-            SingleColumnXmlJsonResultSet();
         }
 
         #endregion
@@ -284,7 +336,8 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
         /// If the result set represented by this class corresponds to a single JSON
         /// column that contains results of "for json" query, set isJson = true
         /// </summary>
-        private void SingleColumnXmlJsonResultSet() {
+        private void SingleColumnXmlJsonResultSet()
+        {
 
             if (Columns?.Length == 1 && RowCount != 0)
             {

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

@@ -91,7 +91,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
         }
 
         [Fact]
-        public async void CancelNonExistantTest()
+        public async Task CancelNonExistantTest()
         {
 
             var workspaceService = new Mock<WorkspaceService<SqlToolsSettings>>();

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

@@ -32,7 +32,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
 
         public const string NoOpQuery = "-- No ops here, just us chickens.";
 
-        public const int Ordinal = 0;
+        public const int Ordinal = 100;     // We'll pick something other than default(int)
 
         public const string OwnerUri = "testFile";
 
@@ -161,7 +161,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
             var connectionMock = new Mock<DbConnection> { CallBase = true };
             connectionMock.Protected()
                 .Setup<DbCommand>("CreateDbCommand")
-                .Returns(CreateTestCommand(data, throwOnRead));
+                .Returns(() => CreateTestCommand(data, throwOnRead));
             connectionMock.Setup(dbc => dbc.Open())
                 .Callback(() => connectionMock.SetupGet(dbc => dbc.State).Returns(ConnectionState.Open));
             connectionMock.Setup(dbc => dbc.Close())
@@ -174,7 +174,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
         {
             var mockFactory = new Mock<ISqlConnectionFactory>();
             mockFactory.Setup(factory => factory.CreateSqlConnection(It.IsAny<string>()))
-                .Returns(CreateTestConnection(data, throwOnRead));
+                .Returns(() => CreateTestConnection(data, throwOnRead));
 
             return mockFactory.Object;
         }

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

@@ -28,7 +28,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
             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);
+            ResultSet rs = new ResultSet(mockDataReader, Common.Ordinal, Common.Ordinal, mockFileStreamFactory.Object);
             rs.Dispose();
 
             // Then: The file that was created should have been deleted
@@ -101,6 +101,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
             // ... We need a query service
             var queryService = Common.GetPrimedExecutionService(null, true, false, workspaceService.Object);
 
+
             // If:
             // ... I execute some bogus query
             var queryParams = new QueryExecuteParams { QuerySelection = Common.WholeDocument, OwnerUri = Common.OwnerUri };

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 batchCallback = b =>
-            Batch.BatchAsyncEventHandler callback = 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 batchCallback = b =>
-            Batch.BatchAsyncEventHandler callback = 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 batchCallback = b =>
-            Batch.BatchAsyncEventHandler callback = 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 batchCallback = b =>
-            Batch.BatchAsyncEventHandler callback = 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
+            // ... A completion event (query, batch, not resultset) should have been sent with error
-            VerifyQueryExecuteCallCount(requestContext, Times.Once(), Times.Once(), Times.Once(), Times.Never());
+            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);
         }
     }

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

@@ -149,7 +149,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
         /// Test saving results to CSV file when the requested result set is no longer active
         /// </summary>
         [Fact]
-        public async void SaveResultsAsCsvQueryNotFoundTest()
+        public async Task SaveResultsAsCsvQueryNotFoundTest()
         {
             // Create a query execution service
             var workspaceService = Common.GetPrimedWorkspaceService(Common.StandardQuery);
@@ -302,7 +302,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
         /// Test saving results to JSON file when the requested result set is no longer active
         /// </summary>
         [Fact]
-        public async void SaveResultsAsJsonQueryNotFoundTest()
+        public async Task SaveResultsAsJsonQueryNotFoundTest()
         {
             // Create a query service
             var workspaceService = Common.GetPrimedWorkspaceService(Common.StandardQuery);

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

@@ -4,6 +4,7 @@
 //
 
 using System;
+using System.Collections.Generic;
 using System.Linq;
 using System.Threading.Tasks;
 using Microsoft.SqlTools.ServiceLayer.Hosting.Protocol;
@@ -20,9 +21,9 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
         #region ResultSet Class Tests
 
         [Theory]
-        [InlineData(0,2)]
+        [InlineData(0, 2)]
-        [InlineData(0,20)]
+        [InlineData(0, 20)]
-        [InlineData(1,2)]
+        [InlineData(1, 2)]
         public void ResultSetValidTest(int startRow, int rowCount)
         {
             // Setup:
@@ -57,6 +58,17 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
             Assert.ThrowsAsync<ArgumentOutOfRangeException>(() => rs.GetSubset(rowStartIndex, rowCount)).Wait();
         }
 
+        [Fact]
+        public async Task ResultSetNotReadTest()
+        {
+            // If:
+            // ... I have a resultset that hasn't been executed and I request a valid result set from it
+            // Then:
+            // ... It should throw an exception for having not been read
+            ResultSet rs = new ResultSet(new TestDbDataReader(null), Common.Ordinal, Common.Ordinal, Common.GetFileStreamFactory(new Dictionary<string, byte[]>()));
+            await Assert.ThrowsAsync<InvalidOperationException>(() => rs.GetSubset(0, 1));
+        }
+
         #endregion
 
         #region Batch Class Tests
@@ -92,21 +104,6 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
             Assert.ThrowsAsync<ArgumentOutOfRangeException>(() => b.GetSubset(resultSetIndex, 0, 2)).Wait();
         }
 
-        [Fact]
-        public async Task BatchSubsetIncompleteTest()
-        {
-            // If:
-            // ... I have a batch that hasn't completed execution
-            Batch b = new Batch(Common.StandardQuery, Common.WholeDocument, Common.Ordinal, Common.GetFileStreamFactory(null));
-            Assert.False(b.HasExecuted);
-
-            // ... And I ask for a subset
-            // Then:
-            // ... It should throw an exception
-            await Assert.ThrowsAsync<InvalidOperationException>(() => b.GetSubset(Common.Ordinal, 0, 2));
-
-        }
-
         #endregion
 
         #region Query Class Tests
@@ -142,7 +139,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
             await queryService.ActiveQueries[Common.OwnerUri].ExecutionTask;
 
             // ... And I then ask for a valid set of results from it
-            var subsetParams = new QueryExecuteSubsetParams {OwnerUri = Common.OwnerUri, RowsCount = 1, ResultSetIndex = 0, RowsStartIndex = 0};
+            var subsetParams = new QueryExecuteSubsetParams { OwnerUri = Common.OwnerUri, RowsCount = 1, ResultSetIndex = 0, RowsStartIndex = 0 };
             QueryExecuteSubsetResult result = null;
             var subsetRequest = GetQuerySubsetResultContextMock(qesr => result = qesr, null);
             await queryService.HandleResultSubsetRequest(subsetParams, subsetRequest.Object);
@@ -157,7 +154,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
         }
 
         [Fact]
-        public async void SubsetServiceMissingQueryTest()
+        public async Task SubsetServiceMissingQueryTest()
         {
             // If:
             // ... I ask for a set of results for a file that hasn't executed a query
@@ -178,7 +175,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
         }
 
         [Fact]
-        public async void SubsetServiceUnexecutedQueryTest()
+        public async Task SubsetServiceUnexecutedQueryTest()
         {
             // If:
             // ... I have a query that hasn't finished executing (doesn't matter what)
@@ -188,13 +185,13 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
             var executeRequest = RequestContextMocks.Create<QueryExecuteResult>(null);
             await queryService.HandleExecuteRequest(executeParams, executeRequest.Object);
             await queryService.ActiveQueries[Common.OwnerUri].ExecutionTask;
-            queryService.ActiveQueries[Common.OwnerUri].HasExecuted = false;
+            queryService.ActiveQueries[Common.OwnerUri].Batches[0].ResultSets[0].hasBeenRead = false;
 
             // ... And I then ask for a valid set of results from it
             var subsetParams = new QueryExecuteSubsetParams { OwnerUri = Common.OwnerUri, RowsCount = 1, ResultSetIndex = 0, RowsStartIndex = 0 };
             QueryExecuteSubsetResult result = null;
             var subsetRequest = GetQuerySubsetResultContextMock(qesr => result = qesr, null);
-            queryService.HandleResultSubsetRequest(subsetParams, subsetRequest.Object).Wait();
+            await queryService.HandleResultSubsetRequest(subsetParams, subsetRequest.Object);
 
             // Then:
             // ... I should get an error result