Merge pull request #16 from Microsoft/feature/queryExecuteServerMessages

Server Message Support for Query Execution

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

@@ -7,6 +7,7 @@ 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;
@@ -33,6 +34,14 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
         /// </summary>
         public ConnectionInfo EditorConnection { get; set; }
 
+        /// <summary>
+        /// Whether or not the query has an error
+        /// </summary>
+        public bool HasError { get; set; }
+
+        /// <summary>
+        /// Whether or not the query has completed executed, regardless of success or failure
+        /// </summary>
         public bool HasExecuted { get; set; }
 
         /// <summary>
@@ -40,6 +49,11 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
         /// </summary>
         public string QueryText { 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>
@@ -71,7 +85,7 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
         public Query(string queryText, ConnectionInfo connection)
         {
             // Sanity check for input
-            if (String.IsNullOrWhiteSpace(queryText))
+            if (String.IsNullOrEmpty(queryText))
             {
                 throw new ArgumentNullException(nameof(queryText), "Query text cannot be null");
             }
@@ -85,6 +99,7 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
             EditorConnection = connection;
             HasExecuted = false;
             ResultSets = new List<ResultSet>();
+            ResultMessages = new List<string>();
             cancellationSource = new CancellationTokenSource();
         }
 
@@ -107,6 +122,13 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
                 string connectionString = ConnectionService.BuildConnectionString(EditorConnection.ConnectionDetails);
                 using (conn = EditorConnection.Factory.CreateSqlConnection(connectionString))
                 {
+                    // If we have the message listener, bind to it
+                    SqlConnection sqlConn = conn as SqlConnection;
+                    if (sqlConn != null)
+                    {
+                        sqlConn.InfoMessage += StoreDbMessage;
+                    }
+
                     await conn.OpenAsync(cancellationSource.Token);
 
                     // Create a command that we'll use for executing the query
@@ -120,8 +142,13 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
                         {
                             do
                             {
-                                // TODO: This doesn't properly handle scenarios where the query is SELECT but does not have rows
+                                // Create a message with the number of affected rows
-                                if (!reader.HasRows)
+                                if (reader.RecordsAffected >= 0)
+                                {
+                                    ResultMessages.Add(String.Format("({0} row(s) affected)", reader.RecordsAffected));
+                                }
+
+                                if (!reader.HasRows && reader.FieldCount == 0)
                                 {
                                     continue;
                                 }
@@ -146,9 +173,15 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
                     }
                 }
             }
+            catch (DbException dbe)
+            {
+                HasError = true;
+                UnwrapDbException(dbe);
+                conn?.Dispose();
+            }
             catch (Exception)
             {
-                // Dispose of the connection
+                HasError = true;
                 conn?.Dispose();
                 throw;
             }
@@ -200,6 +233,48 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
             };
         }
 
+        /// <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);
+            }
+        }
+
         #region IDisposable Implementation
 
         private bool disposed;

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

@@ -5,7 +5,6 @@
 
 using System;
 using System.Collections.Concurrent;
-using System.Data.Common;
 using System.Threading.Tasks;
 using Microsoft.SqlTools.ServiceLayer.Connection;
 using Microsoft.SqlTools.ServiceLayer.Hosting;
@@ -242,31 +241,16 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
                 Messages = null
             });
 
-            try
+            // Wait for query execution and then send back the results
+            await Task.WhenAll(executeTask);
+            QueryExecuteCompleteParams eventParams = new QueryExecuteCompleteParams
             {
-                // Wait for query execution and then send back the results
+                HasError = query.HasError,
-                await Task.WhenAll(executeTask);
+                Messages = query.ResultMessages.ToArray(),
-                QueryExecuteCompleteParams eventParams = new QueryExecuteCompleteParams
+                OwnerUri = executeParams.OwnerUri,
-                {
+                ResultSetSummaries = query.ResultSummary
-                    HasError = false,
+            };
-                    Messages = new string[] { }, // TODO: Figure out how to get messages back from the server
+            await requestContext.SendEvent(QueryExecuteCompleteEvent.Type, eventParams);
-                    OwnerUri = executeParams.OwnerUri,
-                    ResultSetSummaries = query.ResultSummary
-                };
-                await requestContext.SendEvent(QueryExecuteCompleteEvent.Type, eventParams);
-            }
-            catch (DbException dbe)
-            {
-                // Dump the message to a complete event
-                QueryExecuteCompleteParams errorEvent = new QueryExecuteCompleteParams
-                {
-                    HasError = true,
-                    Messages = new[] {dbe.Message},
-                    OwnerUri = executeParams.OwnerUri,
-                    ResultSetSummaries = query.ResultSummary
-                };
-                await requestContext.SendEvent(QueryExecuteCompleteEvent.Type, errorEvent);
-            }
         }
 
         #endregion

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

@@ -2,6 +2,7 @@
 using System.Collections.Generic;
 using System.Data;
 using System.Data.Common;
+using System.Data.SqlClient;
 using System.Threading.Tasks;
 using Microsoft.SqlTools.ServiceLayer.Connection;
 using Microsoft.SqlTools.ServiceLayer.Connection.Contracts;
@@ -62,7 +63,9 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
             // Setup the expected behavior
             if (throwOnRead)
             {
-                commandMockSetup.Throws(new Mock<DbException>().Object);
+                var mockException = new Mock<DbException>();
+                mockException.SetupGet(dbe => dbe.Message).Returns("Message");
+                commandMockSetup.Throws(mockException.Object);
             }
             else
             {

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

@@ -36,8 +36,9 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
             query.Execute().Wait();
 
             // Then:
-            // ... It should have executed
+            // ... It should have executed without error
             Assert.True(query.HasExecuted, "The query should have been marked executed.");
+            Assert.False(query.HasError);
             
             // ... The results should be empty
             Assert.Empty(query.ResultSets);
@@ -46,6 +47,9 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
             // ... The results should not be null
             Assert.NotNull(query.ResultSets);
             Assert.NotNull(query.ResultSummary);
+
+            // ... There should be a message for how many rows were affected
+            Assert.Equal(1, query.ResultMessages.Count);
         }
 
         [Fact]
@@ -61,8 +65,9 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
             query.Execute().Wait();
 
             // Then:
-            // ... It should have executed
+            // ... It should have executed without error
             Assert.True(query.HasExecuted, "The query should have been marked executed.");
+            Assert.False(query.HasError);
 
             // ... There should be exactly one result set
             Assert.Equal(resultSets, query.ResultSets.Count);
@@ -82,6 +87,9 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
 
             // ... Inside the result summary, there should be 5 rows
             Assert.Equal(rows, query.ResultSummary[0].RowCount);
+
+            // ... There should be a message for how many rows were affected
+            Assert.Equal(resultSets, query.ResultMessages.Count);
         }
 
         [Fact]
@@ -98,8 +106,9 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
             query.Execute().Wait();
 
             // Then:
-            // ... It should have executed
+            // ... It should have executed without error
             Assert.True(query.HasExecuted, "The query should have been marked executed.");
+            Assert.False(query.HasError);
 
             // ... There should be exactly two result sets
             Assert.Equal(resultSets, query.ResultSets.Count);
@@ -125,6 +134,9 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
                 // ... Inside each result summary, there should be 5 rows
                 Assert.Equal(rows, rs.RowCount);
             }
+
+            // ... There should be a message for how many rows were affected
+            Assert.Equal(resultSets, query.ResultMessages.Count);
         }
 
         [Fact]
@@ -134,10 +146,15 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
 
             // If I execute a query that is invalid
             Query query = new Query("Invalid query", ci);
+            query.Execute().Wait();
 
             // Then:
-            // ... It should throw an exception
+            // ... It should have executed with error
-            Exception e = Assert.Throws<AggregateException>(() => query.Execute().Wait());
+            Assert.True(query.HasExecuted);
+            Assert.True(query.HasError);
+            
+            // ... There should be plenty of messages for the eror
+            Assert.NotEmpty(query.ResultMessages);
         }
 
         [Fact]
@@ -150,8 +167,9 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
             query.Execute().Wait();
 
             // Then:
-            // ... It should have executed
+            // ... It should have executed without error
             Assert.True(query.HasExecuted, "The query should have been marked executed.");
+            Assert.False(query.HasError);
 
             // If I execute it again
             // Then:
@@ -160,7 +178,8 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
             Assert.Equal(1, ae.InnerExceptions.Count);
             Assert.IsType<InvalidOperationException>(ae.InnerExceptions[0]);
 
-            // ... The data should still be available
+            // ... The data should still be available without error
+            Assert.False(query.HasError);
             Assert.True(query.HasExecuted, "The query should still be marked executed.");
             Assert.NotEmpty(query.ResultSets);
             Assert.NotEmpty(query.ResultSummary);
@@ -208,12 +227,12 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
 
             // Then:
             // ... No Errors should have been sent
-            // ... A successful result should have been sent with no messages
+            // ... A successful result should have been sent with messages
             // ... A completion event should have been fired with empty results
             // ... There should be one active query
             VerifyQueryExecuteCallCount(requestContext, Times.Once(), Times.Once(), Times.Never());
             Assert.Null(result.Messages);
-            Assert.Empty(completeParams.Messages);
+            Assert.NotEmpty(completeParams.Messages);
             Assert.Empty(completeParams.ResultSetSummaries);
             Assert.False(completeParams.HasError);
             Assert.Equal(1, queryService.ActiveQueries.Count);
@@ -234,12 +253,12 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
 
             // Then:
             // ... No errors should have been sent
-            // ... A successful result should have been sent with no messages
+            // ... A successful result should have been sent with messages
             // ... A completion event should have been fired with one result
             // ... There should be one active query
             VerifyQueryExecuteCallCount(requestContext, Times.Once(), Times.Once(), Times.Never());
             Assert.Null(result.Messages);
-            Assert.Empty(completeParams.Messages);
+            Assert.NotEmpty(completeParams.Messages);
             Assert.NotEmpty(completeParams.ResultSetSummaries);
             Assert.False(completeParams.HasError);
             Assert.Equal(1, queryService.ActiveQueries.Count);
@@ -327,7 +346,6 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
 
         [Theory]
         [InlineData("")]
-        [InlineData("    ")]
         [InlineData(null)]
         public void QueryExecuteMissingQueryTest(string query)
         {