Fire connection changed event when USE statements are executed

src/Microsoft.SqlTools.ServiceLayer/Connection/ConnectionService.cs

@@ -46,6 +46,16 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection
            
         private Dictionary<string, ConnectionInfo> ownerToConnectionMap = new Dictionary<string, ConnectionInfo>();
 
+        /// <summary>
+        /// Service host object for sending/receiving requests/events.
+        /// Internal for testing purposes.
+        /// </summary>
+        internal IProtocolEndpoint ServiceHost
+        {
+            get;
+            set;
+        }
+
         /// <summary>
         /// Default constructor is private since it's a singleton class
         /// </summary>
@@ -251,6 +261,8 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection
 
         public void InitializeService(IProtocolEndpoint serviceHost)
         {
+            this.ServiceHost = serviceHost;
+
             // Register request and event handlers with the Service Host
             serviceHost.SetRequestHandler(ConnectionRequest.Type, HandleConnectRequest);
             serviceHost.SetRequestHandler(DisconnectRequest.Type, HandleDisconnectRequest);
@@ -480,5 +492,41 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection
 
             return connectionBuilder.ToString();
         }
+
+        /// <summary>
+        /// Change the database context of a connection.
+        /// </summary>
+        /// <param name="ownerUri">URI of the owner of the connection</param>
+        /// <param name="newDatabaseName">Name of the database to change the connection to</param>
+        public void ChangeConnectionDatabaseContext(string ownerUri, string newDatabaseName)
+        {
+            ConnectionInfo info;
+            if (TryFindConnection(ownerUri, out info))
+            {
+                try
+                {
+                    info.SqlConnection.ChangeDatabase(newDatabaseName);
+                    info.ConnectionDetails.DatabaseName = newDatabaseName;
+
+                    // Fire a connection changed event
+                    ConnectionChangedParams parameters = new ConnectionChangedParams();
+                    ConnectionSummary summary = (ConnectionSummary)(info.ConnectionDetails);
+                    parameters.Connection = summary.Clone();
+                    parameters.OwnerUri = ownerUri;
+                    ServiceHost.SendEvent(ConnectionChangedNotification.Type, parameters);
+                }
+                catch (Exception e)
+                {
+                    Logger.Write(
+                        LogLevel.Error, 
+                        string.Format(
+                            "Exception caught while trying to change database context to [{0}] for OwnerUri [{1}]. Exception:{2}", 
+                            newDatabaseName, 
+                            ownerUri, 
+                            e.ToString())
+                    );
+                }
+            }
+        }
     }
 }

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

@@ -5,6 +5,7 @@
 
 using System;
 using System.Data.Common;
+using System.Data.SqlClient;
 using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
@@ -145,6 +146,13 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
             {
                 await conn.OpenAsync();
 
+                if (conn.GetType() == typeof(SqlConnection))
+                {
+                    // Subscribe to database informational messages
+                    SqlConnection sqlConn = conn as SqlConnection;
+                    sqlConn.InfoMessage += OnInfoMessage;
+                }
+
                 // We need these to execute synchronously, otherwise the user will be very unhappy
                 foreach (Batch b in Batches)
                 {
@@ -153,6 +161,28 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
             }
         }
 
+        /// <summary>
+        /// "Error" code produced by SQL Server when the database context (name) for a connection changes.
+        /// </summary>
+        private const int DatabaseContextChangeErrorNumber = 5701;
+
+        /// <summary>
+        /// Handler for database messages during query execution
+        /// </summary>
+        private void OnInfoMessage(object sender, SqlInfoMessageEventArgs args)
+        {
+            SqlConnection conn = sender as SqlConnection;
+
+            foreach(SqlError error in args.Errors) 
+            {
+                // Did the database context change (error code 5701)?
+                if (error.Number == DatabaseContextChangeErrorNumber)
+                {
+                    ConnectionService.Instance.ChangeConnectionDatabaseContext(EditorConnection.OwnerUri, conn.Database);
+                }
+            }
+        }
+
         /// <summary>
         /// Retrieves a subset of the result sets
         /// </summary>