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]