Adding changes as requested for code review

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

@@ -3,7 +3,6 @@
 // Licensed under the MIT license. See LICENSE file in the project root for full license information.
 //
 
-using System;
 using Microsoft.SqlTools.ServiceLayer.Hosting.Protocol.Contracts;
 
 namespace Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts

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

@@ -1,7 +1,15 @@
-using Microsoft.SqlTools.ServiceLayer.Hosting.Protocol.Contracts;
+//
+// 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 with a query execution complete event
+    /// </summary>
     public class QueryExecuteCompleteParams
     {
         /// <summary>
@@ -17,7 +25,7 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts
         /// <summary>
         /// Whether or not the query was successful. True indicates errors, false indicates success
         /// </summary>
-        public bool Error { get; set; }
+        public bool HasError { get; set; }
 
         /// <summary>
         /// Summaries of the result sets that were returned with the query

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

@@ -3,7 +3,6 @@
 // Licensed under the MIT license. See LICENSE file in the project root for full license information.
 //
 
-using System;
 using Microsoft.SqlTools.ServiceLayer.Hosting.Protocol.Contracts;
 
 namespace Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts

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

@@ -3,7 +3,6 @@
 // Licensed under the MIT license. See LICENSE file in the project root for full license information.
 //
 
-using System;
 using Microsoft.SqlTools.ServiceLayer.Hosting.Protocol.Contracts;
 
 namespace Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts
@@ -42,7 +41,14 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts
     /// </summary>
     public class QueryExecuteSubsetResult
     {
+        /// <summary>
+        /// Subset request error messages. Optional, can be set to null to indicate no errors
+        /// </summary>
         public string Message { get; set; }
+
+        /// <summary>
+        /// The requested subset of results. Optional, can be set to null to indicate an error
+        /// </summary>
         public ResultSetSubset ResultSubset { get; set; }
     }
 

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

@@ -1,13 +1,24 @@
-using System;
+//
-using System.Collections.Generic;
+// Copyright (c) Microsoft. All rights reserved.
-using System.Linq;
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
-using System.Threading.Tasks;
+//
 
 namespace Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts
 {
+    /// <summary>
+    /// Class used to represent a subset of results from a query for transmission across JSON RPC
+    /// </summary>
     public class ResultSetSubset
     {
+        /// <summary>
+        /// The number of rows returned from result set, useful for determining if less rows were
+        /// returned than requested.
+        /// </summary>
         public int RowCount { get; set; }
+
+        /// <summary>
+        /// 2D array of the cell values requested from result set
+        /// </summary>
         public object[][] Rows { get; set; }
     }
 }

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

@@ -1,7 +1,15 @@
-using System.Data.Common;
+//
+// Copyright (c) Microsoft. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+//
+
+using System.Data.Common;
 
 namespace Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts
 {
+    /// <summary>
+    /// Represents a summary of information about a result without returning any cells of the results
+    /// </summary>
     public class ResultSetSummary
     {
         /// <summary>

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

@@ -1,4 +1,9 @@
-using System;
+//
+// 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;
@@ -10,18 +15,39 @@ using Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts;
 
 namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
 {
-    public class Query //: IDisposable
+    /// <summary>
+    /// Internal representation of an active query
+    /// </summary>
+    public class Query : IDisposable
     {
         #region Properties
 
-        public string QueryText { get; set; }
+        /// <summary>
-
+        /// Cancellation token source, used for cancelling async db actions
-        public ConnectionInfo EditorConnection { get; set; }
+        /// </summary>
-
         private readonly CancellationTokenSource cancellationSource;
 
+        /// <summary>
+        /// The connection info associated with the file editor owner URI, used to create a new
+        /// connection upon execution of the query
+        /// </summary>
+        public ConnectionInfo EditorConnection { get; set; }
+
+        public bool HasExecuted { get; set; }
+
+        /// <summary>
+        /// The text of the query to execute
+        /// </summary>
+        public string QueryText { get; set; }
+
+        /// <summary>
+        /// The result sets of the query execution
+        /// </summary>
         public List<ResultSet> ResultSets { get; set; }
 
+        /// <summary>
+        /// Property for generating a set result set summaries from the result sets
+        /// </summary>
         public ResultSetSummary[] ResultSummary
         {
             get
@@ -35,10 +61,13 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
             }
         }
 
-        public bool HasExecuted { get; set; }
-
         #endregion
 
+        /// <summary>
+        /// Constructor for a query
+        /// </summary>
+        /// <param name="queryText">The text of the query to execute</param>
+        /// <param name="connection">The information of the connection to use to execute the query</param>
         public Query(string queryText, ConnectionInfo connection)
         {
             // Sanity check for input
@@ -59,6 +88,9 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
             cancellationSource = new CancellationTokenSource();
         }
 
+        /// <summary>
+        /// Executes this query asynchronously and collects all result sets
+        /// </summary>
         public async Task Execute()
         {
             // Sanity check to make sure we haven't already run this query
@@ -67,11 +99,13 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
                 throw new InvalidOperationException("Query has already executed.");
             }
 
+            DbConnection conn = null;
+
             // Create a connection from the connection details
             try
             {
                 string connectionString = ConnectionService.BuildConnectionString(EditorConnection.ConnectionDetails);
-                using (DbConnection conn = EditorConnection.Factory.CreateSqlConnection(connectionString))
+                using (EditorConnection.Factory.CreateSqlConnection(connectionString))
                 {
                     await conn.OpenAsync(cancellationSource.Token);
 
@@ -112,6 +146,11 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
                     }
                 }
             }
+            catch (Exception)
+            {
+                // Dispose of the connection
+                conn?.Dispose();
+            }
             finally
             {
                 // Mark that we have executed
@@ -119,6 +158,13 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
             }
         }
 
+        /// <summary>
+        /// Retrieves a subset of the result sets
+        /// </summary>
+        /// <param name="resultSetIndex">The index for selecting the result set</param>
+        /// <param name="startRow">The starting row of the results</param>
+        /// <param name="rowCount">How many rows to retrieve</param>
+        /// <returns>A subset of results</returns>
         public ResultSetSubset GetSubset(int resultSetIndex, int startRow, int rowCount)
         {
             // Sanity check that the results are available
@@ -152,5 +198,37 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
                 RowCount = rows.Length
             };
         }
+
+        #region IDisposable Implementation
+
+        private bool disposed;
+
+        public void Dispose()
+        {
+            Dispose(true);
+            GC.SuppressFinalize(this);
+        }
+
+        protected virtual void Dispose(bool disposing)
+        {
+            if (disposed)
+            {
+                return;
+            }
+
+            if (disposing)
+            {
+                cancellationSource.Dispose();
+            }
+
+            disposed = true;
+        }
+
+        ~Query()
+        {
+            Dispose(false);
+        }
+
+        #endregion
     }
 }

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

@@ -14,7 +14,10 @@ using Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts;
 
 namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
 {
-    public sealed class QueryExecutionService
+    /// <summary>
+    /// Service for executing queries
+    /// </summary>
+    public sealed class QueryExecutionService : IDisposable
     {
         #region Singleton Instance Implementation
 
@@ -39,24 +42,32 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
 
         #region Properties
 
-        private readonly Lazy<ConcurrentDictionary<string, Query>> queries =
+        /// <summary>
-            new Lazy<ConcurrentDictionary<string, Query>>(() => new ConcurrentDictionary<string, Query>());
+        /// The collection of active queries
-
+        /// </summary>
         internal ConcurrentDictionary<string, Query> ActiveQueries
         {
             get { return queries.Value; }
         }
 
+        /// <summary>
+        /// Instance of the connection service, used to get the connection info for a given owner URI
+        /// </summary>
         private ConnectionService ConnectionService { get; set; }
 
+        /// <summary>
+        /// Internal storage of active queries, lazily constructed as a threadsafe dictionary
+        /// </summary>
+        private readonly Lazy<ConcurrentDictionary<string, Query>> queries =
+            new Lazy<ConcurrentDictionary<string, Query>>(() => new ConcurrentDictionary<string, Query>());
+
         #endregion
 
-        #region Public Methods
-
         /// <summary>
-        /// 
+        /// Initializes the service with the service host, registers request handlers and shutdown
+        /// event handler.
         /// </summary>
-        /// <param name="serviceHost"></param>
+        /// <param name="serviceHost">The service host instance to register with</param>
         public void InitializeService(ServiceHost serviceHost)
         {
             // Register handlers for requests
@@ -64,11 +75,14 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
             serviceHost.SetRequestHandler(QueryExecuteSubsetRequest.Type, HandleResultSubsetRequest);
             serviceHost.SetRequestHandler(QueryDisposeRequest.Type, HandleDisposeRequest);
 
-            // Register handlers for events
+            // Register handler for shutdown event
+            serviceHost.RegisterShutdownTask((shutdownParams, requestContext) =>
+            {
+                Dispose();
+                return Task.FromResult(0);
+            });
         }
 
-        #endregion
-
         #region Request Handlers
 
         public async Task HandleExecuteRequest(QueryExecuteParams executeParams,
@@ -167,6 +181,8 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
 
         #endregion
 
+        #region Private Helpers
+
         private async Task<Query> CreateAndActivateNewQuery(QueryExecuteParams executeParams, RequestContext<QueryExecuteResult> requestContext)
         {
             try
@@ -232,7 +248,7 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
                 await Task.WhenAll(executeTask);
                 QueryExecuteCompleteParams eventParams = new QueryExecuteCompleteParams
                 {
-                    Error = false,
+                    HasError = false,
                     Messages = new string[] { }, // TODO: Figure out how to get messages back from the server
                     OwnerUri = executeParams.OwnerUri,
                     ResultSetSummaries = query.ResultSummary
@@ -244,7 +260,7 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
                 // Dump the message to a complete event
                 QueryExecuteCompleteParams errorEvent = new QueryExecuteCompleteParams
                 {
-                    Error = true,
+                    HasError = true,
                     Messages = new[] {dbe.Message},
                     OwnerUri = executeParams.OwnerUri,
                     ResultSetSummaries = query.ResultSummary
@@ -252,5 +268,42 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
                 await requestContext.SendEvent(QueryExecuteCompleteEvent.Type, errorEvent);
             }
         }
+
+        #endregion
+
+        #region IDisposable Implementation
+
+        private bool disposed;
+
+        public void Dispose()
+        {
+            Dispose(true);
+            GC.SuppressFinalize(this);
+        }
+
+        private void Dispose(bool disposing)
+        {
+            if (disposed)
+            {
+                return;
+            }
+
+            if (disposing)
+            {
+                foreach (var query in ActiveQueries)
+                {
+                    query.Value.Dispose();
+                }
+            }
+
+            disposed = true;
+        }
+
+        ~QueryExecutionService()
+        {
+            Dispose(false);
+        }
+
+        #endregion
     }
 }

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

@@ -1,4 +1,8 @@
-using System;
+//
+// Copyright (c) Microsoft. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+//
+
 using System.Collections.Generic;
 using System.Data.Common;
 

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

@@ -215,7 +215,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
             Assert.Null(result.Messages);
             Assert.Empty(completeParams.Messages);
             Assert.Empty(completeParams.ResultSetSummaries);
-            Assert.False(completeParams.Error);
+            Assert.False(completeParams.HasError);
             Assert.Equal(1, queryService.ActiveQueries.Count);
         }
 
@@ -241,7 +241,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
             Assert.Null(result.Messages);
             Assert.Empty(completeParams.Messages);
             Assert.NotEmpty(completeParams.ResultSetSummaries);
-            Assert.False(completeParams.Error);
+            Assert.False(completeParams.HasError);
             Assert.Equal(1, queryService.ActiveQueries.Count);
         }
 
@@ -321,7 +321,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
             // ... There should only be one active query
             VerifyQueryExecuteCallCount(secondRequestContext, Times.Once(), Times.Once(), Times.Never());
             Assert.Null(result.Messages);
-            Assert.False(complete.Error);
+            Assert.False(complete.HasError);
             Assert.Equal(1, queryService.ActiveQueries.Count);
         }
 
@@ -368,7 +368,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
             // ... A completion event should have been sent with error
             VerifyQueryExecuteCallCount(requestContext, Times.Once(), Times.Once(), Times.Never());
             Assert.Null(result.Messages);
-            Assert.True(complete.Error);
+            Assert.True(complete.HasError);
             Assert.NotEmpty(complete.Messages);
         }