WIP for QueryExecution, mostly complete

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

@@ -0,0 +1,37 @@
+//
+// 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 Microsoft.SqlTools.ServiceLayer.Hosting.Protocol.Contracts;
+
+namespace Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts
+{
+    /// <summary>
+    /// Parameters for the query dispose request
+    /// </summary>
+    public class QueryDisposeParams
+    {
+        public string OwnerUri { get; set; }
+    }
+
+    /// <summary>
+    /// Parameters to return as the result of a query dispose request
+    /// </summary>
+    public class QueryDisposeResult
+    {
+        /// <summary>
+        /// Any error messages that occurred during disposing the result set. Optional, can be set
+        /// to null if there were no errors.
+        /// </summary>
+        public string Messages { get; set; }
+    }
+
+    public class QueryDisposeRequest
+    {
+        public static readonly
+            RequestType<QueryDisposeParams, QueryDisposeResult> Type =
+            RequestType<QueryDisposeParams, QueryDisposeResult>.Create("query/dispose");
+    }
+}

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

@@ -0,0 +1,34 @@
+using Microsoft.SqlTools.ServiceLayer.Hosting.Protocol.Contracts;
+
+namespace Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts
+{
+    public class QueryExecuteCompleteParams
+    {
+        /// <summary>
+        /// URI for the editor that owns the query
+        /// </summary>
+        public string OwnerUri { get; set; }
+
+        /// <summary>
+        /// Any messages that came back from the server during execution of the query
+        /// </summary>
+        public string[] Messages { get; set; }
+
+        /// <summary>
+        /// Whether or not the query was successful. True indicates errors, false indicates success
+        /// </summary>
+        public bool Error { get; set; }
+
+        /// <summary>
+        /// Summaries of the result sets that were returned with the query
+        /// </summary>
+        public ResultSetSummary[] ResultSetSummaries { get; set; }
+    }
+
+    public class QueryExecuteCompleteEvent
+    {
+        public static readonly 
+            EventType<QueryExecuteCompleteParams> Type =
+            EventType<QueryExecuteCompleteParams>.Create("query/complete");
+    }
+}

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

@@ -0,0 +1,44 @@
+//
+// 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 Microsoft.SqlTools.ServiceLayer.Hosting.Protocol.Contracts;
+
+namespace Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts
+{
+    /// <summary>
+    /// Parameters for the query execute request
+    /// </summary>
+    public class QueryExecuteParams
+    {
+        /// <summary>
+        /// The text of the query to execute
+        /// </summary>
+        public string QueryText { get; set; }
+
+        /// <summary>
+        /// URI for the editor that is asking for the query execute
+        /// </summary>
+        public string OwnerUri { get; set; }
+    }
+
+    /// <summary>
+    /// Parameters for the query execute result
+    /// </summary>
+    public class QueryExecuteResult
+    {
+        /// <summary>
+        /// Connection error messages. Optional, can be set to null to indicate no errors
+        /// </summary>
+        public string Messages { get; set; }
+    }
+
+    public class QueryExecuteRequest
+    {
+        public static readonly
+            RequestType<QueryExecuteParams, QueryExecuteResult> Type =
+            RequestType<QueryExecuteParams, QueryExecuteResult>.Create("query/execute");
+    }
+}

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

@@ -0,0 +1,55 @@
+//
+// 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 Microsoft.SqlTools.ServiceLayer.Hosting.Protocol.Contracts;
+
+namespace Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts
+{
+    /// <summary>
+    /// Parameters for a query result subset retrieval request
+    /// </summary>
+    public class QueryExecuteSubsetParams
+    {
+        /// <summary>
+        /// ID of the query to look up the results for
+        /// </summary>
+        public string OwnerId { get; set; }
+
+        /// <summary>
+        /// Index of the result set to get the results from
+        /// </summary>
+        public int ResultSetIndex { get; set; }
+
+        /// <summary>
+        /// Beginning index of the rows to return from the selected resultset. This index will be
+        /// included in the results.
+        /// </summary>
+        public int RowsStartIndex { get; set; }
+
+        /// <summary>
+        /// Number of rows to include in the result of this request. If the number of the rows 
+        /// exceeds the number of rows available after the start index, all available rows after
+        /// the start index will be returned.
+        /// </summary>
+        public int RowsCount { get; set; }
+    }
+
+    /// <summary>
+    /// Parameters for the result of a subset retrieval request
+    /// </summary>
+    public class QueryExecuteSubsetResult
+    {
+        public string Message { get; set; }
+        public ResultSetSubset ResultSubset { get; set; }
+    }
+
+    public class QueryExecuteSubsetRequest
+    {
+        public static readonly
+            RequestType<QueryExecuteSubsetParams, QueryExecuteSubsetResult> Type =
+            RequestType<QueryExecuteSubsetParams, QueryExecuteSubsetResult>.Create("query/subset");
+    }
+}

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

@@ -0,0 +1,13 @@
+using System;
+using System.Collections.Generic;
+using System.Linq;
+using System.Threading.Tasks;
+
+namespace Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts
+{
+    public class ResultSetSubset
+    {
+        public int RowCount { get; set; }
+        public object[][] Rows { get; set; }
+    }
+}

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

@@ -0,0 +1,22 @@
+using System.Data.Common;
+
+namespace Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts
+{
+    public class ResultSetSummary
+    {
+        /// <summary>
+        /// The ID of the result set within the query results
+        /// </summary>
+        public int Id { get; set; }
+
+        /// <summary>
+        /// The number of rows that was returned with the resultset
+        /// </summary>
+        public int RowCount { get; set; }
+
+        /// <summary>
+        /// Details about the columns that are provided as solutions
+        /// </summary>
+        public DbColumn[] ColumnInfo { get; set; }
+    }
+}

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

@@ -0,0 +1,144 @@
+using System;
+using System.Collections.Generic;
+using System.Data;
+using System.Data.Common;
+using System.Linq;
+using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.SqlTools.ServiceLayer.Connection;
+using Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts;
+
+namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
+{
+    public class Query //: IDisposable
+    {
+        #region Properties
+
+        public string QueryText { get; set; }
+
+        public ConnectionInfo EditorConnection { get; set; }
+
+        private readonly CancellationTokenSource cancellationSource;
+
+        public List<ResultSet> ResultSets { get; set; }
+
+        public ResultSetSummary[] ResultSummary
+        {
+            get
+            {
+                return ResultSets.Select((set, index) => new ResultSetSummary
+                {
+                    ColumnInfo = set.Columns,
+                    Id = index,
+                    RowCount = set.Rows.Count
+                }).ToArray();
+            }
+        }
+
+        public bool HasExecuted { get; set; }
+
+        #endregion
+
+        public Query(string queryText, ConnectionInfo connection)
+        {
+            // Sanity check for input
+            if (queryText == null)
+            {
+                throw new ArgumentNullException(nameof(queryText), "Query text cannot be null");
+            }
+            if (connection == null)
+            {
+                throw new ArgumentNullException(nameof(connection), "Connection cannot be null");
+            }
+
+            // Initialize the internal state
+            QueryText = queryText;
+            EditorConnection = connection;
+            HasExecuted = false;
+            ResultSets = new List<ResultSet>();
+            cancellationSource = new CancellationTokenSource();
+        }
+
+        public async Task Execute()
+        {
+            // Sanity check to make sure we haven't already run this query
+            if (HasExecuted)
+            {
+                throw new InvalidOperationException("Query has already executed.");
+            }
+
+            // Create a connection from the connection details
+            using (DbConnection conn = EditorConnection.Factory.CreateSqlConnection(EditorConnection.ConnectionDetails))
+            {
+                await conn.OpenAsync(cancellationSource.Token);
+
+                // Create a command that we'll use for executing the query
+                using (DbCommand command = conn.CreateCommand())
+                {
+                    command.CommandText = QueryText;
+                    command.CommandType = CommandType.Text;
+
+                    // Execute the command to get back a reader
+                    using (DbDataReader reader = await command.ExecuteReaderAsync(cancellationSource.Token))
+                    {
+                        do
+                        {
+                            // Create a new result set that we'll use to store all the data
+                            ResultSet resultSet = new ResultSet();
+                            if (reader.CanGetColumnSchema())
+                            {
+                                resultSet.Columns = reader.GetColumnSchema().ToArray();
+                            }
+
+                            // Read until we hit the end of the result set
+                            while (await reader.ReadAsync(cancellationSource.Token))
+                            {
+                                resultSet.AddRow(reader);
+                            }
+
+                            // Add the result set to the results of the query
+                            ResultSets.Add(resultSet);
+                        } while (await reader.NextResultAsync(cancellationSource.Token));
+                    }
+                }
+            }
+
+            // Mark that we have executed
+            HasExecuted = true;
+        }
+
+        public ResultSetSubset GetSubset(int resultSetIndex, int startRow, int rowCount)
+        {
+            // Sanity check that the results are available
+            if (!HasExecuted)
+            {
+                throw new InvalidOperationException("The query has not completed, yet.");
+            }
+
+            // Sanity check to make sure we have valid numbers
+            if (resultSetIndex < 0 || resultSetIndex >= ResultSets.Count)
+            {
+                throw new ArgumentOutOfRangeException(nameof(resultSetIndex), "Result set index cannot be less than 0" +
+                                                                             "or greater than the number of result sets");
+            }
+            ResultSet targetResultSet = ResultSets[resultSetIndex];
+            if (startRow < 0 || startRow >= targetResultSet.Rows.Count)
+            {
+                throw new ArgumentOutOfRangeException(nameof(startRow), "Start row cannot be less than 0 " +
+                                                                        "or greater than the number of rows in the resultset");
+            }
+            if (rowCount <= 0)
+            {
+                throw new ArgumentOutOfRangeException(nameof(rowCount), "Row count must be a positive integer");
+            }
+
+            // Retrieve the subset of the results as per the request
+            object[][] rows = targetResultSet.Rows.Skip(startRow).Take(rowCount).ToArray();
+            return new ResultSetSubset
+            {
+                Rows = rows,
+                RowCount = rows.Length
+            };
+        }
+    }
+}

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

@@ -0,0 +1,153 @@
+using System;
+using System.Collections.Concurrent;
+using System.Threading.Tasks;
+using Microsoft.SqlTools.ServiceLayer.Connection;
+using Microsoft.SqlTools.ServiceLayer.Hosting;
+using Microsoft.SqlTools.ServiceLayer.Hosting.Protocol;
+using Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts;
+
+namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
+{
+    public sealed class QueryExecutionService
+    {
+        #region Singleton Instance Implementation
+
+        private static readonly Lazy<QueryExecutionService> instance = new Lazy<QueryExecutionService>(() => new QueryExecutionService());
+
+        public static QueryExecutionService Instance
+        {
+            get { return instance.Value; }
+        }
+
+        private QueryExecutionService() { }
+
+        #endregion
+
+        #region Properties
+
+        private readonly Lazy<ConcurrentDictionary<string, Query>> queries =
+            new Lazy<ConcurrentDictionary<string, Query>>(() => new ConcurrentDictionary<string, Query>());
+
+        private ConcurrentDictionary<string, Query> ActiveQueries
+        {
+            get { return queries.Value; }
+        }
+
+        #endregion
+
+        #region Public Methods
+
+        /// <summary>
+        /// 
+        /// </summary>
+        /// <param name="serviceHost"></param>
+        public void InitializeService(ServiceHost serviceHost)
+        {
+            // Register handlers for requests
+            serviceHost.SetRequestHandler(QueryExecuteRequest.Type, HandleExecuteRequest);
+            serviceHost.SetRequestHandler(QueryExecuteSubsetRequest.Type, HandleResultSubsetRequest);
+            serviceHost.SetRequestHandler(QueryDisposeRequest.Type, HandleDisposeRequest);
+
+            // Register handlers for events
+        }
+
+        #endregion
+
+        #region Request Handlers
+
+        private async Task HandleExecuteRequest(QueryExecuteParams executeParams,
+            RequestContext<QueryExecuteResult> requestContext)
+        {
+            // Attempt to get the connection for the editor
+            ConnectionInfo connectionInfo;
+            if(!ConnectionService.Instance.TryFindConnection(executeParams.OwnerUri, out connectionInfo))
+            {
+                await requestContext.SendError("This editor is not connected to a database.");
+                return;
+            }
+
+            // If there is already an in-flight query, error out
+            Query newQuery = new Query(executeParams.QueryText, connectionInfo);
+            if (!ActiveQueries.TryAdd(executeParams.OwnerUri, newQuery))
+            {
+                await requestContext.SendError("A query is already in progress for this editor session." +
+                                               "Please cancel this query or wait for its completion.");
+                return;
+            }
+
+            // Launch the query and respond with successfully launching it
+            Task executeTask = newQuery.Execute();
+            await requestContext.SendResult(new QueryExecuteResult
+            {
+                Messages = null
+            });
+
+            // Wait for query execution and then send back the results
+            await Task.WhenAll(executeTask);
+            QueryExecuteCompleteParams eventParams = new QueryExecuteCompleteParams
+            {
+                Error = false,
+                Messages = new string[]{},  // TODO: Figure out how to get messages back from the server
+                OwnerUri = executeParams.OwnerUri,
+                ResultSetSummaries = newQuery.ResultSummary
+            };
+            await requestContext.SendEvent(QueryExecuteCompleteEvent.Type, eventParams);
+        }
+
+        private async Task HandleResultSubsetRequest(QueryExecuteSubsetParams subsetParams,
+            RequestContext<QueryExecuteSubsetResult> requestContext)
+        {
+            // Attempt to load the query
+            Query query;
+            if (!ActiveQueries.TryGetValue(subsetParams.OwnerId, out query))
+            {
+                var errorResult = new QueryExecuteSubsetResult
+                {
+                    Message = "The requested query does not exist."
+                };
+                await requestContext.SendResult(errorResult);
+                return;
+            }
+
+            try
+            {
+                // Retrieve the requested subset and return it
+                var result = new QueryExecuteSubsetResult
+                {
+                    Message = null,
+                    ResultSubset = query.GetSubset(
+                        subsetParams.ResultSetIndex, subsetParams.RowsStartIndex, subsetParams.RowsCount)
+                };
+                await requestContext.SendResult(result);
+            }
+            catch (Exception e)
+            {
+                await requestContext.SendResult(new QueryExecuteSubsetResult
+                {
+                    Message = e.Message
+                });
+            }
+        }
+
+        private async Task HandleDisposeRequest(QueryDisposeParams disposeParams,
+            RequestContext<QueryDisposeResult> requestContext)
+        {
+            // Attempt to remove the query for the owner uri
+            Query result;
+            if (!ActiveQueries.TryRemove(disposeParams.OwnerUri, out result))
+            {
+                await requestContext.SendError("Failed to dispose query, ID not found.");
+                return;
+            }
+
+            // Success
+            await requestContext.SendResult(new QueryDisposeResult
+            {
+                Messages = null
+            });
+        }
+
+        #endregion
+
+    }
+}

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

@@ -0,0 +1,33 @@
+using System;
+using System.Collections.Generic;
+using System.Data.Common;
+
+namespace Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts
+{
+    public class ResultSet
+    {
+        public DbColumn[] Columns { get; set; }
+
+        public List<object[]> Rows { get; private set; }
+
+        public ResultSet()
+        {
+            Rows = new List<object[]>();
+        }
+
+        /// <summary>
+        /// Add a row of data to the result set using a <see cref="DbDataReader"/> that has already
+        /// read in a row.
+        /// </summary>
+        /// <param name="reader">A <see cref="DbDataReader"/> that has already had a read performed</param>
+        public void AddRow(DbDataReader reader)
+        {
+            List<object> row = new List<object>();
+            for (int i = 0; i < reader.FieldCount; ++i)
+            {
+                row.Add(reader.GetValue(i));
+            }
+            Rows.Add(row.ToArray());
+        }
+    }
+}