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]