Progressive Results Part 1: Batch Completion Notification (#95)

The main feature of this pull request is a new callback that's added to the query class that is called when a batch has completed execution and retrieval of results. This callback will send an event to the extension with the batch summary information. After that, the extension can submit subset requests for the resultsets of the batch.
Other smaller changes in this pull request:
Refactor to assign a batch a id when its created instead of when returning the list of batch summaries
Passing the SelectionData around instead of extracting the values for it
Moving creation of BatchSummary into the Batch class
Retrieval of results is now permitted even if the entire query has not completed, as long as the batch requested has completed.
Also note, this does not break the protocol. It adds a new event that a queryRunner can listen to, but it doesn't require it to be listened to.

* 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

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

@@ -57,16 +57,19 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
 
         #endregion
 
-        internal Batch(string batchText, int startLine, int startColumn, int endLine, int endColumn, IFileStreamFactory outputFileFactory)
+        internal Batch(string batchText, SelectionData selection, int ordinalId, IFileStreamFactory outputFileFactory)
         {
             // Sanity check for input
             Validate.IsNotNullOrEmptyString(nameof(batchText), batchText);
             Validate.IsNotNull(nameof(outputFileFactory), outputFileFactory);
+            Validate.IsGreaterThan(nameof(ordinalId), ordinalId, 0);
 
             // Initialize the internal state
             BatchText = batchText;
-            Selection = new SelectionData(startLine, startColumn, endLine, endColumn);
+            Selection = selection;
+            executionStartTime = DateTime.Now;
             HasExecuted = false;
+            Id = ordinalId;
             resultSets = new List<ResultSet>();
             resultMessages = new List<ResultMessage>();
             this.outputFileFactory = outputFileFactory;
@@ -74,6 +77,17 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
 
         #region Properties
 
+        /// <summary>
+        /// Asynchronous handler for when batches are completed
+        /// </summary>
+        /// <param name="batch">The batch that completed</param>
+        public delegate Task BatchAsyncEventHandler(Batch batch);
+
+        /// <summary>
+        /// Event that will be called when the batch has completed execution
+        /// </summary>
+        public event BatchAsyncEventHandler BatchCompletion;
+
         /// <summary>
         /// The text of batch that will be executed
         /// </summary>
@@ -113,6 +127,11 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
         /// </summary>
         public bool HasExecuted { get; set; }
 
+        /// <summary>
+        /// Ordinal of the batch in the query
+        /// </summary>
+        public int Id { get; private set; }
+
         /// <summary>
         /// Messages that have come back from the server
         /// </summary>
@@ -145,6 +164,27 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
             }
         }
 
+        /// <summary>
+        /// Creates a <see cref="BatchSummary"/> based on the batch instance
+        /// </summary>
+        public BatchSummary Summary
+        {
+            get
+            {
+                return new BatchSummary
+                {
+                    HasError = HasError,
+                    Id = Id,
+                    ResultSetSummaries = ResultSummaries,
+                    Messages = ResultMessages.ToArray(),
+                    Selection = Selection,
+                    ExecutionElapsed = ExecutionElapsedTime,
+                    ExecutionStart = ExecutionStartTimeStamp,
+                    ExecutionEnd = ExecutionEndTimeStamp
+                };
+            }
+        }
+
         /// <summary>
         /// The range from the file that is this batch
         /// </summary>
@@ -169,8 +209,10 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
 
             try
             {
-                DbCommand command = null;
+                // Register the message listener to *this instance* of the batch
+                // Note: This is being done to associate messages with batches
                 ReliableSqlConnection sqlConn = conn as ReliableSqlConnection;
+                DbCommand command;
                 if (sqlConn != null)
                 {
                     // Register the message listener to *this instance* of the batch
@@ -258,6 +300,12 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
                 // Mark that we have executed
                 HasExecuted = true;
                 executionEndTime = DateTime.Now;
+
+                // Fire an event to signify that the batch has completed
+                if (BatchCompletion != null)
+                {
+                    await BatchCompletion(this);
+                }
             }
         }
 
@@ -270,6 +318,12 @@ 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
+            if (!HasExecuted)
+            {
+                throw new InvalidOperationException(SR.QueryServiceSubsetBatchNotCompleted);
+            }
+
             // Sanity check to make sure we have valid numbers
             if (resultSetIndex < 0 || resultSetIndex >= resultSets.Count)
             {

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

@@ -0,0 +1,32 @@
+// 
+// Copyright (c) Microsoft. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+
+using Microsoft.SqlTools.ServiceLayer.Hosting.Protocol.Contracts;
+
+namespace Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts
+{
+    /// <summary>
+    /// Parameters to be sent back as part of a QueryExecuteBatchCompleteEvent to indicate that a
+    /// batch of a query completed.
+    /// </summary>
+    public class QueryExecuteBatchCompleteParams
+    {
+        /// <summary>
+        /// Summary of the batch that just completed
+        /// </summary>
+        public BatchSummary BatchSummary { get; set; }
+
+        /// <summary>
+        /// URI for the editor that owns the query
+        /// </summary>
+        public string OwnerUri { get; set; }
+    }
+
+    public class QueryExecuteBatchCompleteEvent
+    {
+        public static readonly 
+            EventType<QueryExecuteBatchCompleteParams> Type =
+            EventType<QueryExecuteBatchCompleteParams>.Create("query/batchComplete");
+    }
+}

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

@@ -7,21 +7,6 @@ using Microsoft.SqlTools.ServiceLayer.Hosting.Protocol.Contracts;
 
 namespace Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts
 {
-    /// <summary>
-    /// Container class for a selection range from file
-    /// </summary>
-    public class SelectionData {
-        public int StartLine { get; set; }
-        public int StartColumn { get; set; }
-        public int EndLine { get; set; }
-        public int EndColumn { get; set; }
-        public SelectionData(int startLine, int startColumn, int endLine, int endColumn) {
-            StartLine = startLine;
-            StartColumn = startColumn;
-            EndLine = endLine;
-            EndColumn = endColumn;
-        }
-    }
     /// <summary>
     /// Parameters for the query execute request
     /// </summary>

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

@@ -0,0 +1,52 @@
+// 
+// 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.Workspace.Contracts;
+
+namespace Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts
+{
+    /// <summary> 
+    /// Container class for a selection range from file 
+    /// </summary>
+    /// TODO: Remove this in favor of buffer range end-to-end
+    public class SelectionData
+    {
+        public SelectionData() { }
+
+        public SelectionData(int startLine, int startColumn, int endLine, int endColumn)
+        {
+            StartLine = startLine;
+            StartColumn = startColumn;
+            EndLine = endLine;
+            EndColumn = endColumn;
+        }
+
+        #region Properties
+
+        public int EndColumn { get; set; }
+
+        public int EndLine { get; set; }
+
+        public int StartColumn { get; set; }
+        public int StartLine { get; set; }
+
+        #endregion
+
+        public BufferRange ToBufferRange()
+        {
+            return new BufferRange(StartLine, StartColumn, EndLine, EndColumn);
+        }
+
+        public static SelectionData FromBufferRange(BufferRange range)
+        {
+            return new SelectionData
+            {
+                StartLine = range.Start.Line,
+                StartColumn = range.Start.Column,
+                EndLine = range.End.Line,
+                EndColumn = range.End.Column
+            };
+        }
+    }
+}

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

@@ -51,11 +51,6 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
         /// </summary>
         private bool hasExecuteBeenCalled;
 
-        /// <summary>
-        /// The factory to use for outputting the results of this query
-        /// </summary>
-        private readonly IFileStreamFactory outputFileFactory;
-
         #endregion
 
         /// <summary>
@@ -77,7 +72,6 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
             QueryText = queryText;
             editorConnection = connection;
             cancellationSource = new CancellationTokenSource();
-            outputFileFactory = outputFactory;
 
             // Process the query into batches
             ParseResult parseResult = Parser.Parse(queryText, new ParseOptions
@@ -85,13 +79,17 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
                 BatchSeparator = settings.BatchSeparator
             });
             // NOTE: We only want to process batches that have statements (ie, ignore comments and empty lines)
-            Batches = parseResult.Script.Batches.Where(b => b.Statements.Count > 0)
+            var batchSelection = parseResult.Script.Batches
-                .Select(b => new Batch(b.Sql, 
+                .Where(batch => batch.Statements.Count > 0)
-                                       b.StartLocation.LineNumber - 1, 
+                .Select((batch, index) =>
-                                       b.StartLocation.ColumnNumber - 1, 
+                    new Batch(batch.Sql,
-                                       b.EndLocation.LineNumber - 1, 
+                        new SelectionData(
-                                       b.EndLocation.ColumnNumber - 1, 
+                            batch.StartLocation.LineNumber - 1,
-                                       outputFileFactory)).ToArray();
+                            batch.StartLocation.ColumnNumber - 1,
+                            batch.EndLocation.LineNumber - 1,
+                            batch.EndLocation.ColumnNumber - 1),
+                        index, outputFactory));
+            Batches = batchSelection.ToArray();
         }
 
         #region Properties
@@ -102,10 +100,15 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
         /// <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="q">The query that completed</param>
+        /// <param name="message">Message to return</param>
         public delegate Task QueryAsyncErrorEventHandler(string message);
 
         /// <summary>
@@ -139,18 +142,7 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
                 {
                     throw new InvalidOperationException("Query has not been executed.");
                 }
-
+                return Batches.Select(b => b.Summary).ToArray();
-                return Batches.Select((batch, index) => new BatchSummary
-                {
-                    Id = index,
-                    ExecutionStart = batch.ExecutionStartTimeStamp,
-                    ExecutionEnd = batch.ExecutionEndTimeStamp,
-                    ExecutionElapsed = batch.ExecutionElapsedTime,
-                    HasError = batch.HasError,
-                    Messages = batch.ResultMessages.ToArray(),
-                    ResultSetSummaries = batch.ResultSummaries,
-                    Selection = batch.Selection
-                }).ToArray();
             }
         }
 
@@ -214,12 +206,6 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
         /// <returns>A subset of results</returns>
         public Task<ResultSetSubset> GetSubset(int batchIndex, int resultSetIndex, int startRow, int rowCount)
         {
-            // Sanity check that the results are available
-            if (!HasExecuted)
-            {
-                throw new InvalidOperationException(SR.QueryServiceSubsetNotCompleted);
-            }
-
             // Sanity check to make sure that the batch is within bounds
             if (batchIndex < 0 || batchIndex >= Batches.Length)
             {
@@ -278,6 +264,7 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
                     // We need these to execute synchronously, otherwise the user will be very unhappy
                     foreach (Batch b in Batches)
                     {
+                        b.BatchCompletion += BatchCompleted;
                         await b.Execute(conn, cancellationSource.Token);
                     }
 

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

@@ -442,6 +442,18 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
             query.QueryFailed += callback;
             query.QueryConnectionException += errorCallback;
 
+            // Setup the batch completion callback
+            Batch.BatchAsyncEventHandler batchCallback = async b =>
+            {
+                QueryExecuteBatchCompleteParams eventParams = new QueryExecuteBatchCompleteParams
+                {
+                    BatchSummary = b.Summary,
+                    OwnerUri = executeParams.OwnerUri
+                };
+                await requestContext.SendEvent(QueryExecuteBatchCompleteEvent.Type, eventParams);
+            };
+            query.BatchCompleted += batchCallback;
+
             // Launch this as an asynchronous task
             query.Execute();
 

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

@@ -61,7 +61,7 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
         /// </summary>
         /// <param name="field">The field to encode</param>
         /// <returns>The CSV encoded version of the original field</returns>
-        internal static String EncodeCsvField(String field)
+        internal static string EncodeCsvField(string field)
         {
             StringBuilder sbField = new StringBuilder(field);
 
@@ -102,9 +102,8 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
 
             //Replace all quotes in the original field with double quotes
             sbField.Replace("\"", "\"\"");
-
+            string ret = sbField.ToString();
-            String ret = sbField.ToString();
+          
-
             if (embedInQuotes)
             {
                 ret = "\"" + ret + "\"";
@@ -121,7 +120,7 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
         internal static bool IsSaveSelection(SaveResultsRequestParams saveParams)
         {
             return (saveParams.ColumnStartIndex != null && saveParams.ColumnEndIndex != null
-                && saveParams.RowEndIndex != null && saveParams.RowEndIndex != null);
+                && saveParams.RowStartIndex != null && saveParams.RowEndIndex != null);
         }
 
         /// <summary>

src/Microsoft.SqlTools.ServiceLayer/sr.cs

@@ -173,11 +173,11 @@ namespace Microsoft.SqlTools.ServiceLayer
             }
         }
 
-        public static string QueryServiceSubsetNotCompleted
+        public static string QueryServiceSubsetBatchNotCompleted
         {
             get
             {
-                return Keys.GetString(Keys.QueryServiceSubsetNotCompleted);
+                return Keys.GetString(Keys.QueryServiceSubsetBatchNotCompleted);
             }
         }
 
@@ -468,7 +468,7 @@ namespace Microsoft.SqlTools.ServiceLayer
             public const string QueryServiceQueryCancelled = "QueryServiceQueryCancelled";
 
 
-            public const string QueryServiceSubsetNotCompleted = "QueryServiceSubsetNotCompleted";
+            public const string QueryServiceSubsetBatchNotCompleted = "QueryServiceSubsetBatchNotCompleted";
 
 
             public const string QueryServiceSubsetBatchOutOfRange = "QueryServiceSubsetBatchOutOfRange";

src/Microsoft.SqlTools.ServiceLayer/sr.resx

@@ -209,8 +209,8 @@
     <value>Query was canceled by user</value>
     <comment></comment>
   </data>
-  <data name="QueryServiceSubsetNotCompleted" xml:space="preserve">
+  <data name="QueryServiceSubsetBatchNotCompleted" xml:space="preserve">
-    <value>The query has not completed, yet</value>
+    <value>The batch has not completed, yet</value>
     <comment></comment>
   </data>
   <data name="QueryServiceSubsetBatchOutOfRange" xml:space="preserve">

src/Microsoft.SqlTools.ServiceLayer/sr.strings

@@ -83,7 +83,7 @@ QueryServiceQueryCancelled = Query was canceled by user
 
 ### Subset Request
 
-QueryServiceSubsetNotCompleted = The query has not completed, yet
+QueryServiceSubsetBatchNotCompleted = The batch has not completed, yet
 
 QueryServiceSubsetBatchOutOfRange = Batch index cannot be less than 0 or greater than the number of batches
 

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

@@ -33,7 +33,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
             // If:
             // ... I request a query (doesn't matter what kind) and execute it
             var queryService = await Common.GetPrimedExecutionService(Common.CreateMockFactory(null, false), true, workspaceService.Object);
-            var executeParams = new QueryExecuteParams { QuerySelection = Common.GetSubSectionDocument(), OwnerUri = Common.OwnerUri };
+            var executeParams = new QueryExecuteParams { QuerySelection = Common.SubsectionDocument, OwnerUri = Common.OwnerUri };
             var executeRequest = 
                 RequestContextMocks.SetupRequestContextMock<QueryExecuteResult, QueryExecuteCompleteParams>(null, QueryExecuteCompleteEvent.Type, null, null);
             await queryService.HandleExecuteRequest(executeParams, executeRequest.Object);

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

@@ -1,19 +1,20 @@
-//
+// 
 // Copyright (c) Microsoft. All rights reserved.
 // Licensed under the MIT license. See LICENSE file in the project root for full license information.
-//
 
 using System;
 using System.Collections.Generic;
 using System.Data;
 using System.Data.Common;
-using System.IO;
 using System.Data.SqlClient;
+using System.IO;
 using System.Threading;
 using System.Threading.Tasks;
+using Microsoft.SqlServer.Management.Common;
+using Microsoft.SqlServer.Management.SmoMetadataProvider;
+using Microsoft.SqlServer.Management.SqlParser.Binder;
 using Microsoft.SqlTools.ServiceLayer.Connection;
 using Microsoft.SqlTools.ServiceLayer.Connection.Contracts;
-using Microsoft.SqlServer.Management.Common;
 using Microsoft.SqlTools.ServiceLayer.LanguageServices;
 using Microsoft.SqlTools.ServiceLayer.QueryExecution;
 using Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts;
@@ -29,37 +30,61 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
 {
     public class Common
     {
-        public const SelectionData WholeDocument = null;
+        #region Constants
-        
-        public const string StandardQuery = "SELECT * FROM sys.objects";
 
         public const string InvalidQuery = "SELECT *** FROM sys.objects";
 
         public const string NoOpQuery = "-- No ops here, just us chickens.";
 
-        public const string UdtQuery = "SELECT hierarchyid::Parse('/')";
+        public const int Ordinal = 0;
 
         public const string OwnerUri = "testFile";
 
-        public const int StandardRows = 5;
-
         public const int StandardColumns = 5;
 
-        public static string TestServer { get; set; }
+        public const string StandardQuery = "SELECT * FROM sys.objects";
 
-        public static string TestDatabase { get; set; }
+        public const int StandardRows = 5;
 
-        static Common()
+        public const string UdtQuery = "SELECT hierarchyid::Parse('/')";
+
+        public const SelectionData WholeDocument = null;
+
+        public static readonly ConnectionDetails StandardConnectionDetails = new ConnectionDetails
         {
-            TestServer = "sqltools11";
+            DatabaseName = "123",
-            TestDatabase = "master";
+            Password = "456",
-        }
+            ServerName = "789",
+            UserName = "012"
+        };
+
+        public static readonly SelectionData SubsectionDocument = new SelectionData(0, 0, 2, 2);
+
+        #endregion
 
         public static Dictionary<string, string>[] StandardTestData
         {
             get { return GetTestData(StandardRows, StandardColumns); }
         }
 
+        #region Public Methods
+
+        public static Batch GetBasicExecutedBatch()
+        {
+            Batch batch = new Batch(StandardQuery, SubsectionDocument, 1, GetFileStreamFactory());
+            batch.Execute(CreateTestConnection(new[] {StandardTestData}, false), CancellationToken.None).Wait();
+            return batch;
+        }
+
+        public static Query GetBasicExecutedQuery()
+        {
+            ConnectionInfo ci = CreateTestConnectionInfo(new[] {StandardTestData}, false);
+            Query query = new Query(StandardQuery, ci, new QueryExecutionSettings(), GetFileStreamFactory());
+            query.Execute();
+            query.ExecutionTask.Wait();
+            return query;
+        }
+
         public static Dictionary<string, string>[] GetTestData(int columns, int rows)
         {
             Dictionary<string, string>[] output = new Dictionary<string, string>[rows];
@@ -76,26 +101,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
             return output;
         }
 
-        public static SelectionData GetSubSectionDocument() 
+        #endregion
-        {
-            return new SelectionData(0, 0, 2, 2);
-        }
-
-        public static Batch GetBasicExecutedBatch()
-        {
-            Batch batch = new Batch(StandardQuery, 0, 0, 2, 2, GetFileStreamFactory());
-            batch.Execute(CreateTestConnection(new[] {StandardTestData}, false), CancellationToken.None).Wait();
-            return batch;
-        }
-
-        public static Query GetBasicExecutedQuery()
-        {
-            ConnectionInfo ci = CreateTestConnectionInfo(new[] {StandardTestData}, false);
-            Query query = new Query(StandardQuery, ci, new QueryExecutionSettings(), GetFileStreamFactory());
-            query.Execute();
-            query.ExecutionTask.Wait();
-            return query;
-        }
 
         #region FileStreamWriteMocking 
 
@@ -113,20 +119,30 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
 
         public class InMemoryWrapper : IFileStreamWrapper
         {
-            private readonly byte[] storage = new byte[8192];
             private readonly MemoryStream memoryStream;
             private bool readingOnly;
+            private readonly byte[] storage = new byte[8192];
 
             public InMemoryWrapper()
             {
                 memoryStream = new MemoryStream(storage);
             }
 
+            public void Close()
+            {
+                memoryStream.Dispose();
+            }
+
             public void Dispose()
             {
                 // We'll dispose this via a special method
             }
 
+            public void Flush()
+            {
+                if (readingOnly) { throw new InvalidOperationException(); }
+            }
+
             public void Init(string fileName, int bufferSize, FileAccess fAccess)
             {
                 readingOnly = fAccess == FileAccess.Read;
@@ -150,16 +166,6 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
                 memoryStream.Flush();
                 return bytes;
             }
-
-            public void Flush()
-            {
-                if (readingOnly) { throw new InvalidOperationException(); }
-            }
-
-            public void Close()
-            {
-                memoryStream.Dispose();
-            }
         }
 
         #endregion
@@ -213,27 +219,18 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
 
         public static ConnectionInfo CreateTestConnectionInfo(Dictionary<string, string>[][] data, bool throwOnRead)
         {
-            // Create connection info
+            return new ConnectionInfo(CreateMockFactory(data, throwOnRead), OwnerUri, StandardConnectionDetails);
-            ConnectionDetails connDetails = new ConnectionDetails
-            {
-                UserName = "sa",
-                Password = "Yukon900",
-                DatabaseName = Common.TestDatabase,
-                ServerName = Common.TestServer
-            };
-
-            return new ConnectionInfo(CreateMockFactory(data, throwOnRead), OwnerUri, connDetails);
         }
 
         #endregion
 
         #region Service Mocking
-        
+
         public static void GetAutoCompleteTestObjects(
             out TextDocumentPosition textDocument,
             out ScriptFile scriptFile,
             out ConnectionInfo connInfo
-        )
+            )
         {
             textDocument = new TextDocumentPosition
             {
@@ -245,6 +242,12 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
                 }
             };
 
+
+            connInfo = CreateTestConnectionInfo(null, false);
+
+            var srvConn = GetServerConnection(connInfo);
+            var metadataProvider = SmoMetadataProvider.CreateConnectedProvider(srvConn);
+            var binder = BinderProvider.CreateBinder(metadataProvider);
             connInfo = Common.CreateTestConnectionInfo(null, false);
 
             LanguageService.Instance.ScriptParseInfoMap.Add(textDocument.TextDocument.Uri,  new ScriptParseInfo());
@@ -259,17 +262,6 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
             var sqlConnection = new SqlConnection(connectionString);
             return new ServerConnection(sqlConnection);
         }
-        
-        public static ConnectionDetails GetTestConnectionDetails()
-        {
-            return new ConnectionDetails
-            {
-                DatabaseName = "123",
-                Password = "456",
-                ServerName = "789",
-                UserName = "012"
-            };
-        }
 
         public static async Task<QueryExecutionService> GetPrimedExecutionService(ISqlConnectionFactory factory, bool isConnected, WorkspaceService<SqlToolsSettings> workspaceService)
         {
@@ -278,7 +270,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
             {
                 await connectionService.Connect(new ConnectParams
                 {
-                    Connection = GetTestConnectionDetails(),
+                    Connection = StandardConnectionDetails,
                     OwnerUri = OwnerUri
                 });
             }
@@ -300,6 +292,5 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
         }
 
         #endregion
-        
     }
 }

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

@@ -32,7 +32,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
         public void BatchCreationTest()
         {
             // If I create a new batch...
-            Batch batch = new Batch(Common.StandardQuery, 0, 0, 2, 2, Common.GetFileStreamFactory());
+            Batch batch = new Batch(Common.StandardQuery, Common.SubsectionDocument, Common.Ordinal, Common.GetFileStreamFactory());
 
             // Then: 
             // ... The text of the batch should be stored
@@ -49,13 +49,27 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
 
             // ... The start line of the batch should be 0
             Assert.Equal(0, batch.Selection.StartLine);
+
+            // ... It's ordinal ID should be what I set it to
+            Assert.Equal(Common.Ordinal, batch.Id);
         }
 
         [Fact]
         public void BatchExecuteNoResultSets()
         {
+            // Setup: Create a callback for batch completion
+            BatchSummary batchSummaryFromCallback = null;
+            bool completionCallbackFired = false;
+            Batch.BatchAsyncEventHandler callback = b =>
+            {
+                completionCallbackFired = true;
+                batchSummaryFromCallback = b.Summary;
+                return Task.FromResult(0);
+            };
+
             // If I execute a query that should get no result sets
-            Batch batch = new Batch(Common.StandardQuery, 0, 0, 2, 2, Common.GetFileStreamFactory());
+            Batch batch = new Batch(Common.StandardQuery, Common.SubsectionDocument, Common.Ordinal, Common.GetFileStreamFactory());
+            batch.BatchCompletion += callback;
             batch.Execute(GetConnection(Common.CreateTestConnectionInfo(null, false)), CancellationToken.None).Wait();
 
             // Then:
@@ -73,16 +87,31 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
 
             // ... There should be a message for how many rows were affected
             Assert.Equal(1, batch.ResultMessages.Count());
+
+            // ... The callback for batch completion should have been fired
+            Assert.True(completionCallbackFired);
+            Assert.NotNull(batchSummaryFromCallback);
         }
 
         [Fact]
         public void BatchExecuteOneResultSet()
         {
-            int resultSets = 1;
+            const int resultSets = 1;
             ConnectionInfo ci = Common.CreateTestConnectionInfo(new[] { Common.StandardTestData }, false);
 
+            // Setup: Create a callback for batch completion
+            BatchSummary batchSummaryFromCallback = null;
+            bool completionCallbackFired = false;
+            Batch.BatchAsyncEventHandler callback = b =>
+            {
+                completionCallbackFired = true;
+                batchSummaryFromCallback = b.Summary;
+                return Task.FromResult(0);
+            };
+
             // If I execute a query that should get one result set
-            Batch batch = new Batch(Common.StandardQuery, 0, 0, 2, 2, Common.GetFileStreamFactory());
+            Batch batch = new Batch(Common.StandardQuery, Common.SubsectionDocument, Common.Ordinal, Common.GetFileStreamFactory());
+            batch.BatchCompletion += callback;
             batch.Execute(GetConnection(ci), CancellationToken.None).Wait();
 
             // Then:
@@ -104,6 +133,10 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
 
             // ... There should be a message for how many rows were affected
             Assert.Equal(resultSets, batch.ResultMessages.Count());
+
+            // ... The callback for batch completion should have been fired
+            Assert.True(completionCallbackFired);
+            Assert.NotNull(batchSummaryFromCallback);
         }
 
         [Fact]
@@ -113,8 +146,19 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
             int resultSets = dataset.Length;
             ConnectionInfo ci = Common.CreateTestConnectionInfo(dataset, false);
 
+            // Setup: Create a callback for batch completion
+            BatchSummary batchSummaryFromCallback = null;
+            bool completionCallbackFired = false;
+            Batch.BatchAsyncEventHandler callback = b =>
+            {
+                completionCallbackFired = true;
+                batchSummaryFromCallback = b.Summary;
+                return Task.FromResult(0);
+            };
+
             // If I execute a query that should get two result sets
-            Batch batch = new Batch(Common.StandardQuery, 0, 0, 1, 1, Common.GetFileStreamFactory());
+            Batch batch = new Batch(Common.StandardQuery, Common.SubsectionDocument, Common.Ordinal, Common.GetFileStreamFactory());
+            batch.BatchCompletion += callback;
             batch.Execute(GetConnection(ci), CancellationToken.None).Wait();
 
             // Then:
@@ -145,15 +189,30 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
                 // ... Inside each result summary, there should be 5 column definitions
                 Assert.Equal(Common.StandardColumns, rs.ColumnInfo.Length);
             }
+
+            // ... The callback for batch completion should have been fired
+            Assert.True(completionCallbackFired);
+            Assert.NotNull(batchSummaryFromCallback);
         }
 
         [Fact]
         public void BatchExecuteInvalidQuery()
         {
+            // Setup: Create a callback for batch completion
+            BatchSummary batchSummaryFromCallback = null;
+            bool completionCallbackFired = false;
+            Batch.BatchAsyncEventHandler callback = b =>
+            {
+                completionCallbackFired = true;
+                batchSummaryFromCallback = b.Summary;
+                return Task.FromResult(0);
+            };
+
             ConnectionInfo ci = Common.CreateTestConnectionInfo(null, true);
 
             // If I execute a batch that is invalid
-            Batch batch = new Batch(Common.StandardQuery, 0, 0, 2, 2, Common.GetFileStreamFactory());
+            Batch batch = new Batch(Common.StandardQuery, Common.SubsectionDocument, Common.Ordinal, Common.GetFileStreamFactory());
+            batch.BatchCompletion += callback;
             batch.Execute(GetConnection(ci), CancellationToken.None).Wait();
 
             // Then:
@@ -167,15 +226,30 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
 
             // ... There should be plenty of messages for the error
             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
+            BatchSummary batchSummaryFromCallback = null;
+            bool completionCallbackFired = false;
+            Batch.BatchAsyncEventHandler callback = b =>
+            {
+                completionCallbackFired = true;
+                batchSummaryFromCallback = b.Summary;
+                return Task.FromResult(0);
+            };
+
             ConnectionInfo ci = Common.CreateTestConnectionInfo(new[] { Common.StandardTestData }, false);
 
             // If I execute a batch
-            Batch batch = new Batch(Common.StandardQuery, 0, 0, 2, 2, Common.GetFileStreamFactory());
+            Batch batch = new Batch(Common.StandardQuery, Common.SubsectionDocument, Common.Ordinal, Common.GetFileStreamFactory());
+            batch.BatchCompletion += callback;
             batch.Execute(GetConnection(ci), CancellationToken.None).Wait();
 
             // Then:
@@ -183,6 +257,10 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
             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
@@ -205,7 +283,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
             // ... I create a batch that has an empty query
             // Then:
             // ... It should throw an exception
-            Assert.Throws<ArgumentException>(() => new Batch(query, 0, 0, 2, 2, Common.GetFileStreamFactory()));
+            Assert.Throws<ArgumentException>(() => new Batch(query, Common.SubsectionDocument, Common.Ordinal, Common.GetFileStreamFactory()));
         }
 
         [Fact]
@@ -215,7 +293,17 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
             // ... I create a batch that has no file stream factory
             // Then:
             // ... It should throw an exception
-            Assert.Throws<ArgumentNullException>(() => new Batch("stuff", 0, 0, 2, 2, null));
+            Assert.Throws<ArgumentNullException>(() => new Batch("stuff", Common.SubsectionDocument, Common.Ordinal, null));
+        }
+
+        [Fact]
+        public void BatchInvalidOrdinal()
+        {
+            // If:
+            // ... I create a batch has has an ordinal less than 0
+            // Then:
+            // ... It should throw an exception
+            Assert.Throws<ArgumentOutOfRangeException>(() => new Batch("stuff", Common.SubsectionDocument, -1, Common.GetFileStreamFactory()));
         }
 
         #endregion
@@ -268,10 +356,20 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
         [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);
             Query query = new Query(Common.StandardQuery, ci, new QueryExecutionSettings(), Common.GetFileStreamFactory());
+            query.BatchCompleted += batchCallback;
 
             // Then:
             // ... I should get a single batch to execute that hasn't been executed
@@ -291,15 +389,26 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
             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);
             Query query = new Query(Common.NoOpQuery, ci, new QueryExecutionSettings(), Common.GetFileStreamFactory());
+            query.BatchCompleted += batchCallback;
 
             // Then:
             // ... I should get no batches back
@@ -322,11 +431,21 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
         [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);
             Query query = new Query(queryText, ci, new QueryExecutionSettings(), Common.GetFileStreamFactory());
+            query.BatchCompleted += batchCallback;
 
             // Then:
             // ... I should get back two batches to execute that haven't been executed
@@ -346,16 +465,29 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
             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);
             Query query = new Query(queryText, ci, new QueryExecutionSettings(), Common.GetFileStreamFactory());
+            query.BatchCompleted += batchCallback;
 
             // Then:
             // ... I should get back one batch to execute that hasn't been executed
@@ -374,15 +506,28 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
             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);
             Query query = new Query(Common.InvalidQuery, ci, new QueryExecutionSettings(), Common.GetFileStreamFactory());
+            query.BatchCompleted += batchCallback;
 
             // Then:
             // ... I should get back a query with one batch not executed
@@ -404,6 +549,9 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
             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);
         }
 
         #endregion
@@ -431,24 +579,29 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
 
             QueryExecuteResult result = null;
             QueryExecuteCompleteParams completeParams = null;
-            var requestContext = 
+            QueryExecuteBatchCompleteParams batchCompleteParams = null;
-                RequestContextMocks.SetupRequestContextMock<QueryExecuteResult, QueryExecuteCompleteParams>(
+            var requestContext = RequestContextMocks.Create<QueryExecuteResult>(qer => result = qer)
-                    resultCallback: qer => result = qer,
+                .AddEventHandling(QueryExecuteCompleteEvent.Type, (et, p) => completeParams = p)
-                    expectedEvent: QueryExecuteCompleteEvent.Type,
+                .AddEventHandling(QueryExecuteBatchCompleteEvent.Type, (et, p) => batchCompleteParams = p);
-                    eventCallback: (et, cp) => completeParams = cp,
+            queryService.HandleExecuteRequest(queryParams, requestContext.Object).Wait();
-                    errorCallback: null);
-            await AwaitExecution(queryService, queryParams, requestContext.Object);
 
             // Then:
             // ... No Errors should have been sent
             // ... A successful result should have been sent with messages on the first batch
             // ... A completion event should have been fired with empty results
-            VerifyQueryExecuteCallCount(requestContext, Times.Once(), Times.Once(), Times.Never());
+            // ... A batch completion event should have been fired with empty results
+            VerifyQueryExecuteCallCount(requestContext, Times.Once(), Times.Once(), Times.Once(), Times.Never());
             Assert.Null(result.Messages);
+
             Assert.Equal(1, completeParams.BatchSummaries.Length);
             Assert.Empty(completeParams.BatchSummaries[0].ResultSetSummaries);
             Assert.NotEmpty(completeParams.BatchSummaries[0].Messages);
 
+            Assert.NotNull(batchCompleteParams);
+            Assert.Empty(batchCompleteParams.BatchSummary.ResultSetSummaries);
+            Assert.NotEmpty(batchCompleteParams.BatchSummary.Messages);
+            Assert.Equal(completeParams.OwnerUri, batchCompleteParams.OwnerUri);
+
             // ... There should be one active query
             Assert.Equal(1, queryService.ActiveQueries.Count);
         }
@@ -472,25 +625,29 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
 
             QueryExecuteResult result = null;
             QueryExecuteCompleteParams completeParams = null;
-            var requestContext = 
+            QueryExecuteBatchCompleteParams batchCompleteParams = null;
-                RequestContextMocks.SetupRequestContextMock<QueryExecuteResult, QueryExecuteCompleteParams>(
+            var requestContext = RequestContextMocks.Create<QueryExecuteResult>(qer => result = qer)
-                    resultCallback: qer => result = qer, 
+                .AddEventHandling(QueryExecuteCompleteEvent.Type, (et, p) => completeParams = p)
-                    expectedEvent: QueryExecuteCompleteEvent.Type, 
+                .AddEventHandling(QueryExecuteBatchCompleteEvent.Type, (et, p) => batchCompleteParams = p);
-                    eventCallback: (et, cp) => completeParams = cp, 
+            queryService.HandleExecuteRequest(queryParams, requestContext.Object).Wait();
-                    errorCallback: null);
-            await AwaitExecution(queryService, queryParams, requestContext.Object);
 
             // Then:
             // ... No errors should have been sent
             // ... A successful result should have been sent with messages
             // ... A completion event should have been fired with one result
-            VerifyQueryExecuteCallCount(requestContext, Times.Once(), Times.Once(), Times.Never());
+            VerifyQueryExecuteCallCount(requestContext, Times.Once(), Times.Once(), Times.Once(), 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(completeParams.OwnerUri, batchCompleteParams.OwnerUri);
+
             // ... There should be one active query
             Assert.Equal(1, queryService.ActiveQueries.Count);
         }
@@ -508,14 +665,14 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
             object error = null;
             var requestContext = RequestContextMocks.Create<QueryExecuteResult>(null)
                 .AddErrorHandling(e => error = e);
-            await queryService.HandleExecuteRequest(queryParams, requestContext.Object);
+            queryService.HandleExecuteRequest(queryParams, requestContext.Object).Wait();
 
             // Then:
             // ... An error should have been returned
             // ... 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.Once());
+            VerifyQueryExecuteCallCount(requestContext, Times.Never(), Times.Never(), Times.Never(), Times.Once());
             Assert.IsType<string>(error);
             Assert.NotEmpty((string)error);
             Assert.Empty(queryService.ActiveQueries);
@@ -553,11 +710,11 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
             // ... An error should have been sent
             // ... A result should have not have been sent
             // ... No completion event should have been fired
-            // ... The original query should exist
+            // ... There should only be one active query
-            VerifyQueryExecuteCallCount(secondRequestContext, Times.Never(), Times.Never(), Times.Once());
+            VerifyQueryExecuteCallCount(secondRequestContext, Times.Never(), Times.AtMostOnce(), Times.AtMostOnce(), Times.Once());
             Assert.IsType<string>(error);
             Assert.NotEmpty((string)error);
-            Assert.Contains(Common.OwnerUri, queryService.ActiveQueries.Keys);
+            Assert.Equal(1, queryService.ActiveQueries.Count);
         }
 
         [Fact]
@@ -578,29 +735,36 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
             var queryParams = new QueryExecuteParams { OwnerUri = Common.OwnerUri, QuerySelection = Common.WholeDocument };
 
             // Note, we don't care about the results of the first request
-            var firstRequestContext = RequestContextMocks.SetupRequestContextMock<QueryExecuteResult, QueryExecuteCompleteParams>(null, QueryExecuteCompleteEvent.Type, null, null);
+            var firstRequestContext = RequestContextMocks.Create<QueryExecuteResult>(null);
-            await AwaitExecution(queryService, queryParams, firstRequestContext.Object);
+
+            queryService.HandleExecuteRequest(queryParams, firstRequestContext.Object).Wait();
 
             // ... And then I request another query after waiting for the first to complete
             QueryExecuteResult result = null;
             QueryExecuteCompleteParams complete = null;
-            var secondRequestContext =
+            QueryExecuteBatchCompleteParams batchComplete = null;
-                RequestContextMocks.SetupRequestContextMock<QueryExecuteResult, QueryExecuteCompleteParams>(qer => result = qer, QueryExecuteCompleteEvent.Type, (et, qecp) => complete = qecp, null);
+            var secondRequestContext = RequestContextMocks.Create<QueryExecuteResult>(qer => result = qer)
-            await AwaitExecution(queryService, queryParams, secondRequestContext.Object);
+                .AddEventHandling(QueryExecuteCompleteEvent.Type, (et, qecp) => complete = qecp)
+                .AddEventHandling(QueryExecuteBatchCompleteEvent.Type, (et, p) => batchComplete = p);
+            queryService.HandleExecuteRequest(queryParams, secondRequestContext.Object).Wait();
 
             // Then:
             // ... 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.Never());
+            VerifyQueryExecuteCallCount(secondRequestContext, Times.Once(), Times.Once(), Times.Once(), Times.Never());
             Assert.Null(result.Messages);
+
             Assert.False(complete.BatchSummaries.Any(b => b.HasError));
             Assert.Equal(1, queryService.ActiveQueries.Count);
+
+            Assert.NotNull(batchComplete);
+            Assert.False(batchComplete.BatchSummary.HasError);
+            Assert.Equal(complete.OwnerUri, batchComplete.OwnerUri);
         }
 
-        [Theory]
+        [Fact]
-        [InlineData(null)]
+        public async Task QueryExecuteMissingSelectionTest()
-        public async Task QueryExecuteMissingSelectionTest(SelectionData selection)
         {
 
             // Set up file for returning the query
@@ -613,19 +777,20 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
             // If:
             // ... I request to execute a query with a missing query string
             var queryService = await Common.GetPrimedExecutionService(Common.CreateMockFactory(null, false), true, workspaceService.Object);
-            var queryParams = new QueryExecuteParams { OwnerUri = Common.OwnerUri, QuerySelection = selection };
+            var queryParams = new QueryExecuteParams { OwnerUri = Common.OwnerUri, QuerySelection = null };
 
             object errorResult = null;
             var requestContext = RequestContextMocks.Create<QueryExecuteResult>(null)
                 .AddErrorHandling(error => errorResult = error);
             await queryService.HandleExecuteRequest(queryParams, requestContext.Object);
 
+
             // Then:
             // ... Am error should have been sent
             // ... No result should have been sent
             // ... No completion event should have been fired
             // ... An active query should not have been added
-            VerifyQueryExecuteCallCount(requestContext, Times.Never(), Times.Never(), Times.Once());
+            VerifyQueryExecuteCallCount(requestContext, Times.Never(), Times.Never(), Times.Never(), Times.Once());
             Assert.NotNull(errorResult);
             Assert.IsType<string>(errorResult);
             Assert.DoesNotContain(Common.OwnerUri, queryService.ActiveQueries.Keys);
@@ -651,19 +816,27 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
 
             QueryExecuteResult result = null;
             QueryExecuteCompleteParams complete = null;
-            var requestContext =
+            QueryExecuteBatchCompleteParams batchComplete = null;
-                RequestContextMocks.SetupRequestContextMock<QueryExecuteResult, QueryExecuteCompleteParams>(qer => result = qer, QueryExecuteCompleteEvent.Type, (et, qecp) => complete = qecp, null);
+            var requestContext = RequestContextMocks.Create<QueryExecuteResult>(qer => result = qer)
-            await AwaitExecution(queryService, queryParams, requestContext.Object);
+                .AddEventHandling(QueryExecuteCompleteEvent.Type, (et, qecp) => complete = qecp)
+                .AddEventHandling(QueryExecuteBatchCompleteEvent.Type, (et, p) => batchComplete = p);
+            queryService.HandleExecuteRequest(queryParams, requestContext.Object).Wait();
 
             // 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.Never());
+            VerifyQueryExecuteCallCount(requestContext, Times.Once(), Times.Once(), Times.Once(), Times.Never());
             Assert.Null(result.Messages);
+
             Assert.Equal(1, complete.BatchSummaries.Length);
             Assert.True(complete.BatchSummaries[0].HasError);
             Assert.NotEmpty(complete.BatchSummaries[0].Messages);
+
+            Assert.NotNull(batchComplete);
+            Assert.True(batchComplete.BatchSummary.HasError);
+            Assert.NotEmpty(batchComplete.BatchSummary.Messages);
+            Assert.Equal(complete.OwnerUri, batchComplete.OwnerUri);
         }
 
 #if USE_LIVE_CONNECTION
@@ -694,14 +867,18 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
         }
 #endif
 
-#endregion
+        #endregion
 
-        private static void VerifyQueryExecuteCallCount(Mock<RequestContext<QueryExecuteResult>> mock, Times sendResultCalls, Times sendEventCalls, Times sendErrorCalls)
+        private static void VerifyQueryExecuteCallCount(Mock<RequestContext<QueryExecuteResult>> mock, Times sendResultCalls, 
+            Times sendCompletionEventCalls, Times sendBatchCompletionEvent, Times sendErrorCalls)
         {
             mock.Verify(rc => rc.SendResult(It.IsAny<QueryExecuteResult>()), sendResultCalls);
             mock.Verify(rc => rc.SendEvent(
                 It.Is<EventType<QueryExecuteCompleteParams>>(m => m == QueryExecuteCompleteEvent.Type),
-                It.IsAny<QueryExecuteCompleteParams>()), sendEventCalls);
+                It.IsAny<QueryExecuteCompleteParams>()), sendCompletionEventCalls);
+            mock.Verify(rc => rc.SendEvent(
+                It.Is<EventType<QueryExecuteBatchCompleteParams>>(m => m == QueryExecuteBatchCompleteEvent.Type),
+                It.IsAny<QueryExecuteBatchCompleteParams>()), sendBatchCompletionEvent);
             mock.Verify(rc => rc.SendError(It.IsAny<object>()), sendErrorCalls);
         }
 

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

@@ -95,22 +95,25 @@ 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());
+            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
 
-        [Fact]
-        public void SubsetUnexecutedQueryTest()
-        {
-            // If I have a query that has *not* been executed
-            Query q = new Query(Common.StandardQuery, Common.CreateTestConnectionInfo(null, false), new QueryExecutionSettings(), Common.GetFileStreamFactory());
-
-            // ... And I ask for a subset with valid arguments
-            // Then:
-            // ... It should throw an exception
-            Assert.ThrowsAsync<InvalidOperationException>(() => q.GetSubset(0, 0, 0, 2)).Wait();
-        }
-
         [Theory]
         [InlineData(-1)]  // Invalid batch, too low
         [InlineData(2)]   // Invalid batch, too high

test/Microsoft.SqlTools.ServiceLayer.Test/Utility/RequestContextMocks.cs

@@ -50,11 +50,10 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.Utility
             this Mock<RequestContext<TResponse>> mock,
             Action<object> errorCallback)
         {
-
             // Setup the mock for SendError
             var sendErrorFlow = mock.Setup(rc => rc.SendError(It.IsAny<object>()))
                 .Returns(Task.FromResult(0));
-            if (mock != null && errorCallback != null)
+            if (errorCallback != null)
             {
                 sendErrorFlow.Callback(errorCallback);
             }