From b63db9cb39d52a68e0b109df79faee8a6251af7c Mon Sep 17 00:00:00 2001
From: Connor Quagliana <coquagli@microsoft.com>
Date: Fri, 3 Feb 2017 13:49:48 -0800
Subject: [PATCH] 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.
---
 .../QueryExecution/Query.cs                        | 14 ++++++++++++--
 1 file changed, 12 insertions(+), 2 deletions(-)

diff --git a/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Query.cs b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Query.cs
index f9d2ad10..bc4aa645 100644
--- a/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Query.cs
+++ b/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>
@@ -425,7 +430,12 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
                     // Subscribe to database informational messages
                     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;
                 }
             }
         }