From 11ca99f419d113c6d0466fdc4aefb83697420d5c Mon Sep 17 00:00:00 2001 From: Benjamin Russell Date: Wed, 19 Oct 2016 11:10:13 -0700 Subject: [PATCH] Handling NOCOUNT being set (#96) After much thinking, this change brings the message behavior into line with SSMS, including if the NOCOUNT is set. This is a somewhat non-obvious solution, but the StatementCompleted event handler will only be fired if there is a record count to return. We'll add the message for number of records if the StatementCompleted event is fired, otherwise we won't add any messages while processing the resultsets of the batch. If any messages are returned from the server, we'll capture those. Then, if at the end of the batch, we haven't collected any messages from StatementCompleted events or server messages, then we'll add the "completed successfully" message. This matches behavior of SSMS that will only emit a "completed successfully" message if there were no other messages for the batch. * Solution to issue, some unit tests needed to be tweaked * Comments for the event handler --- .../QueryExecution/Batch.cs | 36 ++++++++++++++----- .../QueryExecution/ExecuteTests.cs | 11 ------ 2 files changed, 27 insertions(+), 20 deletions(-) diff --git a/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Batch.cs b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Batch.cs index 14a5e2c7..b104eb3f 100644 --- a/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Batch.cs +++ b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Batch.cs @@ -171,14 +171,17 @@ 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; if (sqlConn != null) { + // Register the message listener to *this instance* of the batch + // Note: This is being done to associate messages with batches sqlConn.GetUnderlyingConnection().InfoMessage += StoreDbMessage; command = sqlConn.GetUnderlyingConnection().CreateCommand(); + + // Add a handler for when the command completes + SqlCommand sqlCommand = (SqlCommand) command; + sqlCommand.StatementCompleted += StatementCompletedHandler; } else { @@ -204,10 +207,6 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution // Skip this result set if there aren't any rows (ie, UPDATE/DELETE/etc queries) if (!reader.HasRows && reader.FieldCount == 0) { - // Create a message with the number of affected rows -- IF the query affects rows - resultMessages.Add(new ResultMessage(reader.RecordsAffected >= 0 - ? SR.QueryServiceAffectedRows(reader.RecordsAffected) - : SR.QueryServiceCompletedSuccessfully)); continue; } @@ -220,9 +219,15 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution // Read until we hit the end of the result set await resultSet.ReadResultToEnd(cancellationToken).ConfigureAwait(false); - // Add a message for the number of rows the query returned - resultMessages.Add(new ResultMessage(SR.QueryServiceAffectedRows(resultSet.RowCount))); + } while (await reader.NextResultAsync(cancellationToken)); + + // If there were no messages, for whatever reason (NO COUNT set, messages + // were emitted, records returned), output a "successful" message + if (resultMessages.Count == 0) + { + resultMessages.Add(new ResultMessage(SR.QueryServiceCompletedSuccessfully)); + } } } } @@ -280,6 +285,19 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution #region Private Helpers + /// + /// Handler for when the StatementCompleted event is fired for this batch's command. This + /// will be executed ONLY when there is a rowcount to report. If this event is not fired + /// either NOCOUNT has been set or the command doesn't affect records. + /// + /// Sender of the event + /// Arguments for the event + private void StatementCompletedHandler(object sender, StatementCompletedEventArgs args) + { + // Add a message for the number of rows the query returned + resultMessages.Add(new ResultMessage(SR.QueryServiceAffectedRows(args.RecordCount))); + } + /// /// 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 diff --git a/test/Microsoft.SqlTools.ServiceLayer.Test/QueryExecution/ExecuteTests.cs b/test/Microsoft.SqlTools.ServiceLayer.Test/QueryExecution/ExecuteTests.cs index c5694fa1..72137474 100644 --- a/test/Microsoft.SqlTools.ServiceLayer.Test/QueryExecution/ExecuteTests.cs +++ b/test/Microsoft.SqlTools.ServiceLayer.Test/QueryExecution/ExecuteTests.cs @@ -73,9 +73,6 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution // ... There should be a message for how many rows were affected Assert.Equal(1, batch.ResultMessages.Count()); - Assert.Contains("1 ", batch.ResultMessages.First().Message); - // NOTE: 1 is expected because this test simulates a 'update' statement where 1 row was affected. - // The 1 in quotes is to make sure the 1 isn't part of a larger number } [Fact] @@ -107,7 +104,6 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution // ... There should be a message for how many rows were affected Assert.Equal(resultSets, batch.ResultMessages.Count()); - Assert.Contains(Common.StandardRows.ToString(), batch.ResultMessages.First().Message); } [Fact] @@ -149,13 +145,6 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution // ... Inside each result summary, there should be 5 column definitions Assert.Equal(Common.StandardColumns, rs.ColumnInfo.Length); } - - // ... There should be a message for how many rows were affected - Assert.Equal(resultSets, batch.ResultMessages.Count()); - foreach (var rsm in batch.ResultMessages) - { - Assert.Contains(Common.StandardRows.ToString(), rsm.Message); - } } [Fact]