WIP update to support batch processing

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

@@ -0,0 +1,195 @@
+using System;
+using System.Collections.Generic;
+using System.Data;
+using System.Data.Common;
+using System.Data.SqlClient;
+using System.Linq;
+using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts;
+
+namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
+{
+    public class Batch
+    {
+        private const string RowsAffectedFormat = "({0} row(s) affected)";
+
+        #region Properties
+        /// <summary>
+        /// The text of batch that will be executed
+        /// </summary>
+        public string BatchText { get; set; }
+
+        /// <summary>
+        /// Whether or not the query has an error
+        /// </summary>
+        public bool HasError { get; set; }
+
+        public bool HasExecuted { get; set; }
+
+        /// <summary>
+        /// Messages that have come back from the server
+        /// </summary>
+        public List<string> ResultMessages { 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[] ResultSummaries
+        {
+            get
+            {
+                return ResultSets.Select((set, index) => new ResultSetSummary()
+                {
+                    ColumnInfo = set.Columns,
+                    Id = index,
+                    RowCount = set.Rows.Count
+                }).ToArray();
+            }
+        }
+
+        #endregion
+
+        public Batch(string batchText)
+        {
+            // Sanity check for input
+            if (string.IsNullOrEmpty(batchText))
+            {
+                throw new ArgumentNullException(nameof(batchText), "Query text cannot be null");
+            }
+
+            // Initialize the internal state
+            BatchText = batchText;
+            HasExecuted = false;
+            ResultSets = new List<ResultSet>();
+            ResultMessages = new List<string>();
+        }
+
+        public async Task Execute(DbConnection conn, CancellationToken cancellationToken)
+        {
+            try
+            {
+                // Register the message listener to *this instance* of the batch
+                // Note: This is being done to associate messages with batches
+                SqlConnection sqlConn = conn as SqlConnection;
+                if (sqlConn != null)
+                {
+                    sqlConn.InfoMessage += StoreDbMessage;
+                }
+
+                // Create a command that we'll use for executing the query
+                using (DbCommand command = conn.CreateCommand())
+                {
+                    command.CommandText = BatchText;
+                    command.CommandType = CommandType.Text;
+
+                    // Execute the command to get back a reader
+                    using (DbDataReader reader = await command.ExecuteReaderAsync(cancellationToken))
+                    {
+                        do
+                        {
+                            // Skip this result set if there aren't any rows
+                            if (!reader.HasRows && reader.FieldCount == 0)
+                            {
+                                // Create a message with the number of affected rows -- IF the query affects rows
+                                ResultMessages.Add(reader.RecordsAffected >= 0
+                                    ? string.Format(RowsAffectedFormat, reader.RecordsAffected)
+                                    : "Commad Executed Successfully");
+                                continue;
+                            }
+
+                            // Read until we hit the end of the result set
+                            ResultSet resultSet = new ResultSet();
+                            while (await reader.ReadAsync(cancellationToken))
+                            {
+                                resultSet.AddRow(reader);
+                            }
+
+                            // Read off the column schema information
+                            if (reader.CanGetColumnSchema())
+                            {
+                                resultSet.Columns = reader.GetColumnSchema().ToArray();
+                            }
+
+                            // Add the result set to the results of the query
+                            ResultSets.Add(resultSet);
+                        } while (await reader.NextResultAsync(cancellationToken));
+                    }
+                }
+            }
+            catch (DbException dbe)
+            {
+                HasError = true;
+                UnwrapDbException(dbe);
+            }
+            catch (Exception)
+            {
+                HasError = true;
+                throw;
+            }
+            finally
+            {
+                // Remove the message event handler from the connection
+                SqlConnection sqlConn = conn as SqlConnection;
+                if (sqlConn != null)
+                {
+                    sqlConn.InfoMessage -= StoreDbMessage;
+                }
+
+                // Mark that we have executed
+                HasExecuted = true;
+            }
+        }
+
+        #region Private Helpers
+
+        /// <summary>
+        /// Delegate handler for storing messages that are returned from the server
+        /// NOTE: Only messages that are below a certain severity will be returned via this
+        /// mechanism. Anything above that level will trigger an exception.
+        /// </summary>
+        /// <param name="sender">Object that fired the event</param>
+        /// <param name="args">Arguments from the event</param>
+        private void StoreDbMessage(object sender, SqlInfoMessageEventArgs args)
+        {
+            ResultMessages.Add(args.Message);
+        }
+
+        /// <summary>
+        /// Attempts to convert a <see cref="DbException"/> to a <see cref="SqlException"/> that
+        /// contains much more info about Sql Server errors. The exception is then unwrapped and
+        /// messages are formatted and stored in <see cref="ResultMessages"/>. If the exception
+        /// cannot be converted to SqlException, the message is written to the messages list.
+        /// </summary>
+        /// <param name="dbe">The exception to unwrap</param>
+        private void UnwrapDbException(DbException dbe)
+        {
+            SqlException se = dbe as SqlException;
+            if (se != null)
+            {
+                foreach (var error in se.Errors)
+                {
+                    SqlError sqlError = error as SqlError;
+                    if (sqlError != null)
+                    {
+                        string message = String.Format("Msg {0}, Level {1}, State {2}, Line {3}{4}{5}",
+                            sqlError.Number, sqlError.Class, sqlError.State, sqlError.LineNumber,
+                            Environment.NewLine, sqlError.Message);
+                        ResultMessages.Add(message);
+                    }
+                }
+            }
+            else
+            {
+                ResultMessages.Add(dbe.Message);
+            }
+        }
+
+        #endregion
+    }
+}