From 93bf2af8bb5ce7e4245554496c534e9496de35d0 Mon Sep 17 00:00:00 2001 From: Mitchell Sternke Date: Fri, 2 Sep 2016 16:43:32 -0700 Subject: [PATCH] Fire connection changed event when USE statements are executed --- .../Connection/ConnectionService.cs | 48 +++++++++++++++++++ .../QueryExecution/Query.cs | 30 ++++++++++++ .../Connection/ConnectionServiceTests.cs | 35 ++++++++++++++ .../Utility/TestObjects.cs | 2 +- 4 files changed, 114 insertions(+), 1 deletion(-) diff --git a/src/Microsoft.SqlTools.ServiceLayer/Connection/ConnectionService.cs b/src/Microsoft.SqlTools.ServiceLayer/Connection/ConnectionService.cs index 2824ee38..4f80b94c 100644 --- a/src/Microsoft.SqlTools.ServiceLayer/Connection/ConnectionService.cs +++ b/src/Microsoft.SqlTools.ServiceLayer/Connection/ConnectionService.cs @@ -46,6 +46,16 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection private Dictionary ownerToConnectionMap = new Dictionary(); + /// + /// Service host object for sending/receiving requests/events. + /// Internal for testing purposes. + /// + internal IProtocolEndpoint ServiceHost + { + get; + set; + } + /// /// Default constructor is private since it's a singleton class /// @@ -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(); } + + /// + /// Change the database context of a connection. + /// + /// URI of the owner of the connection + /// Name of the database to change the connection to + 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()) + ); + } + } + } } } diff --git a/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Query.cs b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Query.cs index 73d91e5e..02bff4d3 100644 --- a/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Query.cs +++ b/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 } } + /// + /// "Error" code produced by SQL Server when the database context (name) for a connection changes. + /// + private const int DatabaseContextChangeErrorNumber = 5701; + + /// + /// Handler for database messages during query execution + /// + 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); + } + } + } + /// /// Retrieves a subset of the result sets /// diff --git a/test/Microsoft.SqlTools.ServiceLayer.Test/Connection/ConnectionServiceTests.cs b/test/Microsoft.SqlTools.ServiceLayer.Test/Connection/ConnectionServiceTests.cs index 29367d4c..a29720e8 100644 --- a/test/Microsoft.SqlTools.ServiceLayer.Test/Connection/ConnectionServiceTests.cs +++ b/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)); } + /// + /// Verify that a connection changed event is fired when the database context changes. + /// + [Fact] + public void ConnectionChangedEventIsFiredWhenDatabaseContextChanges() + { + var serviceHostMock = new Mock(); + + 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(ConnectionChangedNotification.Type, It.IsAny()), Times.Once()); + } + /// /// Verify that the SQL parser correctly detects errors in text /// diff --git a/test/Microsoft.SqlTools.ServiceLayer.Test/Utility/TestObjects.cs b/test/Microsoft.SqlTools.ServiceLayer.Test/Utility/TestObjects.cs index b1ee31bb..da079ea0 100644 --- a/test/Microsoft.SqlTools.ServiceLayer.Test/Utility/TestObjects.cs +++ b/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 } }