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>

test/Microsoft.SqlTools.ServiceLayer.Test/Connection/ConnectionServiceTests.cs

@@ -11,6 +11,7 @@ using System.Reflection;
 using System.Threading.Tasks;
 using Microsoft.SqlTools.ServiceLayer.Connection;
 using Microsoft.SqlTools.ServiceLayer.Connection.Contracts;
+using Microsoft.SqlTools.ServiceLayer.Hosting.Protocol;
 using Microsoft.SqlTools.ServiceLayer.Test.Utility;
 using Microsoft.SqlTools.Test.Utility;
 using Moq;
@@ -287,6 +288,40 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.Connection
             Assert.True(connectionString.Contains(connectionStringMarker));
         }
 
+        /// <summary>
+        /// Verify that a connection changed event is fired when the database context changes.
+        /// </summary>
+        [Fact]
+        public void ConnectionChangedEventIsFiredWhenDatabaseContextChanges()
+        {
+            var serviceHostMock = new Mock<IProtocolEndpoint>();
+
+            var connectionService = TestObjects.GetTestConnectionService();
+            connectionService.ServiceHost = serviceHostMock.Object;
+
+            // Set up an initial connection
+            string ownerUri = "file://my/sample/file.sql";
+            var connectionResult =
+                connectionService
+                .Connect(new ConnectParams()
+                {
+                    OwnerUri = ownerUri,
+                    Connection = TestObjects.GetTestConnectionDetails()
+                });
+
+            // verify that a valid connection id was returned
+            Assert.NotEmpty(connectionResult.ConnectionId);
+
+            ConnectionInfo info;
+            Assert.True(connectionService.TryFindConnection(ownerUri, out info));
+
+            // Tell the connection manager that the database change ocurred
+            connectionService.ChangeConnectionDatabaseContext(ownerUri, "myOtherDb");
+
+            // Verify that the connection changed event was fired
+            serviceHostMock.Verify(x => x.SendEvent<ConnectionChangedParams>(ConnectionChangedNotification.Type, It.IsAny<ConnectionChangedParams>()), Times.Once());
+        }
+
         /// <summary>
         /// Verify that the SQL parser correctly detects errors in text
         /// </summary>

test/Microsoft.SqlTools.ServiceLayer.Test/Utility/TestObjects.cs

@@ -194,7 +194,7 @@ namespace Microsoft.SqlTools.Test.Utility
 
         public override void ChangeDatabase(string databaseName)
         {
-            throw new NotImplementedException();
+            // No Op
         }
     }