diff --git a/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Batch.cs b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Batch.cs
index 99bf919e..a44c332c 100644
--- a/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Batch.cs
+++ b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Batch.cs
@@ -242,6 +242,14 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
             {
                 await BatchStart(this);
             }
+
+            // 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)
+            {
+                sqlConn.GetUnderlyingConnection().InfoMessage += ServerMessageHandler;
+            }
             
             try
             {
@@ -262,7 +270,6 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
             finally
             {
                 // Remove the message event handler from the connection
-                ReliableSqlConnection sqlConn = conn as ReliableSqlConnection;
                 if (sqlConn != null)
                 {
                     sqlConn.GetUnderlyingConnection().InfoMessage -= ServerMessageHandler;
@@ -438,9 +445,6 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
             DbCommand dbCommand;
             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 += ServerMessageHandler;
                 dbCommand = sqlConn.GetUnderlyingConnection().CreateCommand();
 
                 // Add a handler for when the command completes