Wait unit after query executes to change database (#222)

Change logic for database change so OnInfoMessage does not get called until after the query execution is complete. This prevents multiple calls to open a single DataReader.

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

@@ -63,6 +63,11 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
         /// </summary>
         private IFileStreamFactory streamOutputFactory;
 
+        /// <summary>
+        /// Name of the new database if the database name was changed in the query
+        /// </summary>
+        private string newDatabaseName;
+
         /// <summary>
         /// ON keyword
         /// </summary>
@@ -426,6 +431,11 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
                     sqlConn.GetUnderlyingConnection().InfoMessage -= OnInfoMessage;
                 }
             }
+
+            if (newDatabaseName != null)
+            {
+                ConnectionService.Instance.ChangeConnectionDatabaseContext(editorConnection.OwnerUri, newDatabaseName);
+            }
         }
 
         /// <summary>
@@ -444,7 +454,7 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
                 // Did the database context change (error code 5701)?
                 if (error.Number == DatabaseContextChangeErrorNumber)
                 {
-                    ConnectionService.Instance.ChangeConnectionDatabaseContext(editorConnection.OwnerUri, conn.Database);
+                    newDatabaseName = conn.Database;
                 }
             }
         }