diff --git a/src/Microsoft.SqlTools.ManagedBatchParser/BatchParser/ConnectSqlCmdCommand.cs b/src/Microsoft.SqlTools.ManagedBatchParser/BatchParser/ConnectSqlCmdCommand.cs index 28e6ba47..3be8a80b 100644 --- a/src/Microsoft.SqlTools.ManagedBatchParser/BatchParser/ConnectSqlCmdCommand.cs +++ b/src/Microsoft.SqlTools.ManagedBatchParser/BatchParser/ConnectSqlCmdCommand.cs @@ -4,7 +4,6 @@ // using Microsoft.SqlServer.Management.Common; using System; -using System.Data; using System.Data.Common; using Microsoft.Data.SqlClient; @@ -64,13 +63,14 @@ namespace Microsoft.SqlTools.ServiceLayer.BatchParser return null; } - IDbConnection conn = null; + SqlConnection conn = null; try { string connString = ci.ConnectionString; connString += ";Pooling=false"; //turn off connection pooling (this is done in other tools so following the same pattern) conn = new SqlConnection(connString); + conn.RetryLogicProvider = Connection.ReliableConnection.SqlRetryProviders.ServerlessDBRetryProvider(); conn.Open(); return conn as DbConnection; diff --git a/src/Microsoft.SqlTools.ManagedBatchParser/BatchParser/ExecutionEngineCode/ExecutionEngine.cs b/src/Microsoft.SqlTools.ManagedBatchParser/BatchParser/ExecutionEngineCode/ExecutionEngine.cs index 350c8e04..1e68ee1b 100644 --- a/src/Microsoft.SqlTools.ManagedBatchParser/BatchParser/ExecutionEngineCode/ExecutionEngine.cs +++ b/src/Microsoft.SqlTools.ManagedBatchParser/BatchParser/ExecutionEngineCode/ExecutionEngine.cs @@ -922,6 +922,7 @@ namespace Microsoft.SqlTools.ServiceLayer.BatchParser.ExecutionEngineCode try { connection = new SqlConnection(connectionStringBuilder.ConnectionString); + connection.RetryLogicProvider = SqlRetryProviders.ServerlessDBRetryProvider(); connection.Open(); } catch (SqlException ex) diff --git a/src/Microsoft.SqlTools.ManagedBatchParser/ReliableConnection/ReliableConnectionHelper.cs b/src/Microsoft.SqlTools.ManagedBatchParser/ReliableConnection/ReliableConnectionHelper.cs index ea2bc879..e25c1f97 100644 --- a/src/Microsoft.SqlTools.ManagedBatchParser/ReliableConnection/ReliableConnectionHelper.cs +++ b/src/Microsoft.SqlTools.ManagedBatchParser/ReliableConnection/ReliableConnectionHelper.cs @@ -91,7 +91,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection.ReliableConnection connectionRetryPolicy = RetryPolicyFactory.CreateNoRetryPolicy(); } - ReliableSqlConnection connection = new ReliableSqlConnection(connectionString, connectionRetryPolicy, commandRetryPolicy, azureAccountToken); + ReliableSqlConnection connection = new ReliableSqlConnection(connectionString, connectionRetryPolicy, commandRetryPolicy, azureAccountToken, SqlRetryProviders.ServerlessDBRetryProvider()); try { @@ -426,6 +426,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection.ReliableConnection Validate.IsNotNull(nameof(readResult), readResult); using (var sqlConnection = new SqlConnection(connectionString)) { + sqlConnection.RetryLogicProvider = SqlRetryProviders.ServerlessDBRetryProvider(); sqlConnection.Open(); ExecuteReader(sqlConnection, commandText, readResult, initializeCommand, catchException); } diff --git a/src/Microsoft.SqlTools.ManagedBatchParser/ReliableConnection/ReliableSqlConnection.cs b/src/Microsoft.SqlTools.ManagedBatchParser/ReliableConnection/ReliableSqlConnection.cs index 9403fa5d..8801d348 100644 --- a/src/Microsoft.SqlTools.ManagedBatchParser/ReliableConnection/ReliableSqlConnection.cs +++ b/src/Microsoft.SqlTools.ManagedBatchParser/ReliableConnection/ReliableSqlConnection.cs @@ -60,9 +60,15 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection.ReliableConnection /// The connection string used to open the SQL Azure database. /// The retry policy defining whether to retry a request if a connection fails to be established. /// The retry policy defining whether to retry a request if a command fails to be executed. - public ReliableSqlConnection(string connectionString, RetryPolicy connectionRetryPolicy, RetryPolicy commandRetryPolicy, string azureAccountToken) + /// Optional retry provider to handle errors in a special way + public ReliableSqlConnection(string connectionString, RetryPolicy connectionRetryPolicy, RetryPolicy commandRetryPolicy, string azureAccountToken, SqlRetryLogicBaseProvider retryProvider = null) { _underlyingConnection = new SqlConnection(connectionString); + + if (retryProvider != null) { + _underlyingConnection.RetryLogicProvider = retryProvider; + } + _connectionRetryPolicy = connectionRetryPolicy ?? RetryPolicyFactory.CreateNoRetryPolicy(); _commandRetryPolicy = commandRetryPolicy ?? RetryPolicyFactory.CreateNoRetryPolicy(); diff --git a/src/Microsoft.SqlTools.ManagedBatchParser/ReliableConnection/SqlRetryProviders.cs b/src/Microsoft.SqlTools.ManagedBatchParser/ReliableConnection/SqlRetryProviders.cs new file mode 100644 index 00000000..b32a75e5 --- /dev/null +++ b/src/Microsoft.SqlTools.ManagedBatchParser/ReliableConnection/SqlRetryProviders.cs @@ -0,0 +1,66 @@ +// +// Copyright (c) Microsoft. All rights reserved. +// Licensed under the MIT license. See LICENSE file in the project root for full license information. +// + +using System; +using System.Collections.Generic; +using Microsoft.Data.SqlClient; +using Microsoft.SqlTools.BatchParser.Utility; + +namespace Microsoft.SqlTools.ServiceLayer.Connection.ReliableConnection +{ + public static class SqlRetryProviders + { + /// + /// Approved list of transient errors that require additional time to wait before connecting again. + /// + private static readonly HashSet _retryableServerlessConnectivityError; + + /// + /// Max intervals between retries in seconds to wake up serverless instances. + /// + private const int _serverlessMaxIntervalTime = 30; + + /// + /// Maximum number of retries to wake up serverless instances. + /// + private const int _serverlessMaxRetries = 4; + + static SqlRetryProviders() + { + _retryableServerlessConnectivityError = new HashSet + { + //// SQL Error Code: 40613 + //// Database XXXX on server YYYY is not currently available. Please retry the connection later. If the problem persists, contact customer + //// support, and provide them the session tracing ID of ZZZZZ. + 40613, + }; + } + + /// + /// Wait for SqlConnection to handle sleeping serverless instances (allows for them to wake up, otherwise it will result in errors). + /// + public static SqlRetryLogicBaseProvider ServerlessDBRetryProvider() + { + var serverlessRetryLogic = new SqlRetryLogicOption + { + NumberOfTries = _serverlessMaxRetries, + MaxTimeInterval = TimeSpan.FromSeconds(_serverlessMaxIntervalTime), + DeltaTime = TimeSpan.FromSeconds(1), + TransientErrors = _retryableServerlessConnectivityError + }; + + var provider = SqlConfigurableRetryFactory.CreateFixedRetryProvider(serverlessRetryLogic); + + provider.Retrying += (object s, SqlRetryingEventArgs e) => + { + Logger.Information($"attempt {e.RetryCount + 1} - current delay time:{e.Delay}"); + Logger.Information((e.Exceptions[e.Exceptions.Count - 1] is SqlException ex) ? $"{ex.Number}-{ex.Message}" : $"{e.Exceptions[e.Exceptions.Count - 1].Message}"); + }; + + return provider; + } + + } +} \ No newline at end of file diff --git a/src/Microsoft.SqlTools.ServiceLayer/Agent/Jobs/AgentNotebookHelper.cs b/src/Microsoft.SqlTools.ServiceLayer/Agent/Jobs/AgentNotebookHelper.cs index 91d8284c..b42d7c64 100644 --- a/src/Microsoft.SqlTools.ServiceLayer/Agent/Jobs/AgentNotebookHelper.cs +++ b/src/Microsoft.SqlTools.ServiceLayer/Agent/Jobs/AgentNotebookHelper.cs @@ -43,6 +43,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Agent } using (SqlConnection connection = new SqlConnection(ConnectionService.BuildConnectionString(connInfo.ConnectionDetails))) { + connection.RetryLogicProvider = Connection.ReliableConnection.SqlRetryProviders.ServerlessDBRetryProvider(); connection.Open(); using (SqlCommand sqlQueryCommand = new SqlCommand(sqlQuery, connection)) { diff --git a/src/Microsoft.SqlTools.ServiceLayer/Connection/ConnectionInfo.cs b/src/Microsoft.SqlTools.ServiceLayer/Connection/ConnectionInfo.cs index 6d902f4b..ea5cff08 100644 --- a/src/Microsoft.SqlTools.ServiceLayer/Connection/ConnectionInfo.cs +++ b/src/Microsoft.SqlTools.ServiceLayer/Connection/ConnectionInfo.cs @@ -52,7 +52,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection /// /// Properties used for creating/opening the SQL connection. /// - public ConnectionDetails ConnectionDetails { get; private set; } + public ConnectionDetails ConnectionDetails { get; set; } /// /// A map containing all connections to the database that are associated with diff --git a/src/Microsoft.SqlTools.ServiceLayer/Connection/ConnectionService.cs b/src/Microsoft.SqlTools.ServiceLayer/Connection/ConnectionService.cs index 5b0e3f1c..9029abf7 100644 --- a/src/Microsoft.SqlTools.ServiceLayer/Connection/ConnectionService.cs +++ b/src/Microsoft.SqlTools.ServiceLayer/Connection/ConnectionService.cs @@ -43,13 +43,14 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection public const int MaxTolerance = 2 * 60; // two minutes - standard tolerance across ADS for AAD tokens - public const int MaxServerlessReconnectTries = 5; // Max number of tries to wait for a serverless database to start up when its paused before giving up. - // SQL Error Code Constants // Referenced from: https://learn.microsoft.com/en-us/sql/relational-databases/errors-events/database-engine-events-and-errors?view=sql-server-ver16 private const int DoesNotMeetPWReqs = 18466; // Password does not meet complexity requirements. private const int PWCannotBeUsed = 18463; // Password cannot be used at this time. + // Default SQL constants (required to ensure connections such as serverless are able to wake up, connect, and retry properly). + private const int DefaultConnectTimeout = 30; + private const int DefaultCommandTimeout = 30; /// /// Singleton service instance @@ -385,7 +386,9 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection TrySetConnectionType(connectionParams); - connectionParams.Connection.ApplicationName = GetApplicationNameWithFeature(connectionParams.Connection.ApplicationName, connectionParams.Purpose); + // Fill in any details that are necessary (timeouts and application name) to ensure connection doesn't immediately disconnect if not specified (such as for serverless). + connectionParams.Connection = FillInDefaultDetailsForConnections(connectionParams.Connection, connectionParams.Purpose); + // If there is no ConnectionInfo in the map, create a new ConnectionInfo, // but wait until later when we are connected to add it to the map. ConnectionInfo connectionInfo; @@ -409,7 +412,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection } // Try to open a connection with the given ConnectParams - ConnectionCompleteParams? response = await this.TryOpenConnectionWithRetry(connectionInfo, connectionParams); + ConnectionCompleteParams? response = await this.TryOpenConnection(connectionInfo, connectionParams); if (response != null) { return response; @@ -432,34 +435,6 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection return completeParams; } - private async Task TryOpenConnectionWithRetry(ConnectionInfo connectionInfo, ConnectParams connectionParams) - { - int counter = 0; - ConnectionCompleteParams? response = null; - while (counter <= MaxServerlessReconnectTries) - { - // The OpenAsync function used in TryOpenConnection does not retry when a database is sleeping. - // SqlClient will be implemented at a later time, which will have automatic retries. - response = await TryOpenConnection(connectionInfo, connectionParams); - // If a serverless database is sleeping, it will return this error number and will need to be retried. - // See here for details: https://docs.microsoft.com/en-us/azure/azure-sql/database/serverless-tier-overview?view=azuresql#connectivity - if (response?.ErrorNumber == 40613) - { - counter++; - if (counter != MaxServerlessReconnectTries) - { - Logger.Information($"Database for connection {connectionInfo.OwnerUri} is paused, retrying connection. Attempt #{counter}"); - } - } - else - { - // Every other response, we can stop. - break; - } - } - return response; - } - private void TryCloseConnectionTemporaryConnection(ConnectParams connectionParams, ConnectionInfo connectionInfo) { try @@ -667,6 +642,25 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection return serverEdition; } + internal static ConnectionDetails FillInDefaultDetailsForConnections(ConnectionDetails inputConnectionDetails, string featureName) { + ConnectionDetails newConnectionDetails = inputConnectionDetails; + + if(string.IsNullOrWhiteSpace(newConnectionDetails.ApplicationName)) + { + newConnectionDetails.ApplicationName = ApplicationName; + } + else + { + newConnectionDetails.ApplicationName = GetApplicationNameWithFeature(newConnectionDetails.ApplicationName, featureName); + } + + newConnectionDetails.ConnectTimeout = Math.Max(DefaultConnectTimeout, newConnectionDetails.ConnectTimeout ?? 0); + + newConnectionDetails.CommandTimeout = Math.Max(DefaultCommandTimeout, newConnectionDetails.CommandTimeout ?? 0); + + return newConnectionDetails; + } + /// /// Tries to create and open a connection with the given ConnectParams. /// @@ -688,8 +682,8 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection // build the connection string from the input parameters string connectionString = BuildConnectionString(connectionInfo.ConnectionDetails); - // create a sql connection instance - connection = connectionInfo.Factory.CreateSqlConnection(connectionString, connectionInfo.ConnectionDetails.AzureAccountToken); + // create a sql connection instance (with enabled serverless retry logic to handle sleeping serverless databases) + connection = connectionInfo.Factory.CreateSqlConnection(connectionString, connectionInfo.ConnectionDetails.AzureAccountToken, SqlRetryProviders.ServerlessDBRetryProvider()); connectionInfo.AddConnection(connectionParams.Type, connection); // Add a cancellation token source so that the connection OpenAsync() can be cancelled @@ -1894,9 +1888,6 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection // allow pooling connections for language service feature to improve intellisense connection retention and performance. bool shouldForceDisablePooling = !EnableConnectionPooling && featureName != Constants.LanguageServiceFeature; - // increase the connection and command timeout to at least 30 seconds and and build connection string - connInfo.ConnectionDetails.ConnectTimeout = Math.Max(30, connInfo.ConnectionDetails.ConnectTimeout ?? 0); - connInfo.ConnectionDetails.CommandTimeout = Math.Max(30, connInfo.ConnectionDetails.CommandTimeout ?? 0); // enable PersistSecurityInfo to handle issues in SMO where the connection context is lost in reconnections connInfo.ConnectionDetails.PersistSecurityInfo = true; @@ -1905,7 +1896,9 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection { connInfo.ConnectionDetails.Pooling = false; } - connInfo.ConnectionDetails.ApplicationName = GetApplicationNameWithFeature(connInfo.ConnectionDetails.ApplicationName, featureName); + + // increase the connection and command timeout to at least 30 seconds and set application name. + connInfo.ConnectionDetails = FillInDefaultDetailsForConnections(connInfo.ConnectionDetails, featureName); // generate connection string string connectionString = ConnectionService.BuildConnectionString(connInfo.ConnectionDetails, shouldForceDisablePooling); @@ -1916,6 +1909,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection // open a dedicated binding server connection SqlConnection sqlConn = new SqlConnection(connectionString); + sqlConn.RetryLogicProvider = SqlRetryProviders.ServerlessDBRetryProvider(); // Fill in Azure authentication token if needed if (connInfo.ConnectionDetails.AzureAccountToken != null && connInfo.ConnectionDetails.AuthenticationType == AzureMFA) diff --git a/src/Microsoft.SqlTools.ServiceLayer/Connection/ISqlConnectionFactory.cs b/src/Microsoft.SqlTools.ServiceLayer/Connection/ISqlConnectionFactory.cs index abb38a42..b80b2d5d 100644 --- a/src/Microsoft.SqlTools.ServiceLayer/Connection/ISqlConnectionFactory.cs +++ b/src/Microsoft.SqlTools.ServiceLayer/Connection/ISqlConnectionFactory.cs @@ -6,6 +6,7 @@ #nullable disable using System.Data.Common; +using Microsoft.Data.SqlClient; namespace Microsoft.SqlTools.ServiceLayer.Connection { @@ -17,6 +18,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection /// /// Create a new SQL Connection object /// - DbConnection CreateSqlConnection(string connectionString, string azureAccountToken); + /// Optional retry provider to handle errors in a special way + DbConnection CreateSqlConnection(string connectionString, string azureAccountToken, SqlRetryLogicBaseProvider retryProvider = null); } } diff --git a/src/Microsoft.SqlTools.ServiceLayer/Connection/SqlConnectionFactory.cs b/src/Microsoft.SqlTools.ServiceLayer/Connection/SqlConnectionFactory.cs index 932fe5ea..9df1d034 100644 --- a/src/Microsoft.SqlTools.ServiceLayer/Connection/SqlConnectionFactory.cs +++ b/src/Microsoft.SqlTools.ServiceLayer/Connection/SqlConnectionFactory.cs @@ -7,6 +7,7 @@ using System.Data.Common; using Microsoft.SqlTools.ServiceLayer.Connection.ReliableConnection; +using Microsoft.Data.SqlClient; namespace Microsoft.SqlTools.ServiceLayer.Connection { @@ -20,11 +21,12 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection /// /// Creates a new SqlConnection object /// - public DbConnection CreateSqlConnection(string connectionString, string azureAccountToken) + /// Optional retry provider to handle errors in a special way + public DbConnection CreateSqlConnection(string connectionString, string azureAccountToken, SqlRetryLogicBaseProvider retryProvider = null) { RetryPolicy connectionRetryPolicy = RetryPolicyFactory.CreateDefaultConnectionRetryPolicy(); RetryPolicy commandRetryPolicy = RetryPolicyFactory.CreateDefaultConnectionRetryPolicy(); - return new ReliableSqlConnection(connectionString, connectionRetryPolicy, commandRetryPolicy, azureAccountToken); + return new ReliableSqlConnection(connectionString, connectionRetryPolicy, commandRetryPolicy, azureAccountToken, retryProvider); } } } diff --git a/src/Microsoft.SqlTools.ServiceLayer/Scripting/ScriptAsScriptingOperation.cs b/src/Microsoft.SqlTools.ServiceLayer/Scripting/ScriptAsScriptingOperation.cs index 7c45fba4..88585b28 100644 --- a/src/Microsoft.SqlTools.ServiceLayer/Scripting/ScriptAsScriptingOperation.cs +++ b/src/Microsoft.SqlTools.ServiceLayer/Scripting/ScriptAsScriptingOperation.cs @@ -47,6 +47,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Scripting public ScriptAsScriptingOperation(ScriptingParams parameters, string azureAccountToken) : base(parameters) { SqlConnection sqlConnection = new SqlConnection(this.Parameters.ConnectionString); + sqlConnection.RetryLogicProvider = Connection.ReliableConnection.SqlRetryProviders.ServerlessDBRetryProvider(); if (azureAccountToken != null) { sqlConnection.AccessToken = azureAccountToken; diff --git a/src/Microsoft.SqlTools.ServiceLayer/Scripting/SmoScriptingOperation.cs b/src/Microsoft.SqlTools.ServiceLayer/Scripting/SmoScriptingOperation.cs index a666ee3f..24abe388 100644 --- a/src/Microsoft.SqlTools.ServiceLayer/Scripting/SmoScriptingOperation.cs +++ b/src/Microsoft.SqlTools.ServiceLayer/Scripting/SmoScriptingOperation.cs @@ -78,6 +78,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Scripting string serverName = null; using (SqlConnection connection = new SqlConnection(connectionString)) { + connection.RetryLogicProvider = Connection.ReliableConnection.SqlRetryProviders.ServerlessDBRetryProvider(); if (azureAccessToken != null) { connection.AccessToken = azureAccessToken; diff --git a/test/Microsoft.SqlTools.ServiceLayer.UnitTests/Connection/ConnectionServiceTests.cs b/test/Microsoft.SqlTools.ServiceLayer.UnitTests/Connection/ConnectionServiceTests.cs index 33eb96ee..67590536 100644 --- a/test/Microsoft.SqlTools.ServiceLayer.UnitTests/Connection/ConnectionServiceTests.cs +++ b/test/Microsoft.SqlTools.ServiceLayer.UnitTests/Connection/ConnectionServiceTests.cs @@ -86,7 +86,7 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.Connection }); var mockFactory = new Mock(); - mockFactory.Setup(factory => factory.CreateSqlConnection(It.IsAny(), It.IsAny())) + mockFactory.Setup(factory => factory.CreateSqlConnection(It.IsAny(), It.IsAny(), It.IsAny())) .Returns(mockConnection.Object); @@ -152,7 +152,7 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.Connection .Returns(() => Task.Run(() => { })); var mockFactory = new Mock(); - mockFactory.SetupSequence(factory => factory.CreateSqlConnection(It.IsAny(), It.IsAny())) + mockFactory.SetupSequence(factory => factory.CreateSqlConnection(It.IsAny(), It.IsAny(), It.IsAny())) .Returns(mockConnection.Object) .Returns(mockConnection2.Object); @@ -215,7 +215,7 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.Connection }); var mockFactory = new Mock(); - mockFactory.Setup(factory => factory.CreateSqlConnection(It.IsAny(), It.IsAny())) + mockFactory.Setup(factory => factory.CreateSqlConnection(It.IsAny(), It.IsAny(), It.IsAny())) .Returns(mockConnection.Object); @@ -304,7 +304,7 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.Connection connectionMock.Setup(c => c.Database).Returns(expectedDbName); var mockFactory = new Mock(); - mockFactory.Setup(factory => factory.CreateSqlConnection(It.IsAny(), It.IsAny())) + mockFactory.Setup(factory => factory.CreateSqlConnection(It.IsAny(), It.IsAny(), It.IsAny())) .Returns(connectionMock.Object); var connectionService = new ConnectionService(mockFactory.Object); @@ -345,8 +345,8 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.Connection var dummySqlConnection = new TestSqlConnection(null); var mockFactory = new Mock(); - mockFactory.Setup(factory => factory.CreateSqlConnection(It.IsAny(), It.IsAny())) - .Returns((string connString, string azureAccountToken) => + mockFactory.Setup(factory => factory.CreateSqlConnection(It.IsAny(), It.IsAny(), It.IsAny())) + .Returns((string connString, string azureAccountToken, SqlRetryLogicBaseProvider retryProvider) => { dummySqlConnection.ConnectionString = connString; SqlConnectionStringBuilder scsb = new SqlConnectionStringBuilder(connString); @@ -1020,7 +1020,7 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.Connection { // Setup mock connection factory to inject query results var mockFactory = new Mock(); - mockFactory.Setup(factory => factory.CreateSqlConnection(It.IsAny(), It.IsAny())) + mockFactory.Setup(factory => factory.CreateSqlConnection(It.IsAny(), It.IsAny(), It.IsAny())) .Returns(CreateMockDbConnection(new[] { testdata })); var connectionService = new ConnectionService(mockFactory.Object); @@ -1654,8 +1654,8 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.Connection var connection = new TestSqlConnection(null); var mockFactory = new Mock(); - mockFactory.Setup(factory => factory.CreateSqlConnection(It.IsAny(), It.IsAny())) - .Returns((string connString, string azureAccountToken) => + mockFactory.Setup(factory => factory.CreateSqlConnection(It.IsAny(), It.IsAny(), It.IsAny())) + .Returns((string connString, string azureAccountToken, SqlRetryLogicBaseProvider retryProvider) => { connection.ConnectionString = connString; SqlConnectionStringBuilder scsb = new SqlConnectionStringBuilder(connString); @@ -1704,8 +1704,8 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.Connection var connection = mockConnection.Object; var mockFactory = new Mock(); - mockFactory.Setup(factory => factory.CreateSqlConnection(It.IsAny(), It.IsAny())) - .Returns((string connString, string azureAccountToken) => + mockFactory.Setup(factory => factory.CreateSqlConnection(It.IsAny(), It.IsAny(), It.IsAny())) + .Returns((string connString, string azureAccountToken, SqlRetryLogicBaseProvider retryProvider) => { connection.ConnectionString = connString; SqlConnectionStringBuilder scsb = new SqlConnectionStringBuilder(connString); @@ -1757,8 +1757,8 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.Connection var connection = mockConnection.Object; var mockFactory = new Mock(); - mockFactory.Setup(factory => factory.CreateSqlConnection(It.IsAny(), It.IsAny())) - .Returns((string connString, string azureAccountToken) => + mockFactory.Setup(factory => factory.CreateSqlConnection(It.IsAny(), It.IsAny(), It.IsAny())) + .Returns((string connString, string azureAccountToken, SqlRetryLogicBaseProvider retryProvider) => { connection.ConnectionString = connString; SqlConnectionStringBuilder scsb = new SqlConnectionStringBuilder(connString); @@ -1846,7 +1846,7 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.Connection { // Set up mock connection factory var mockFactory = new Mock(); - mockFactory.Setup(factory => factory.CreateSqlConnection(It.IsAny(), It.IsAny())) + mockFactory.Setup(factory => factory.CreateSqlConnection(It.IsAny(), It.IsAny(), It.IsAny())) .Returns(new TestSqlConnection(null)); var connectionService = new ConnectionService(mockFactory.Object); @@ -1865,7 +1865,7 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.Connection }); // Validate that the connection factory gets called with details NOT including an account token - mockFactory.Verify(factory => factory.CreateSqlConnection(It.IsAny(), It.Is(accountToken => accountToken == azureAccountToken)), Times.Once()); + mockFactory.Verify(factory => factory.CreateSqlConnection(It.IsAny(), It.Is(accountToken => accountToken == azureAccountToken), It.IsAny()), Times.Once()); } /// diff --git a/test/Microsoft.SqlTools.ServiceLayer.UnitTests/QueryExecution/Common.cs b/test/Microsoft.SqlTools.ServiceLayer.UnitTests/QueryExecution/Common.cs index ec6a2480..d2d9d523 100644 --- a/test/Microsoft.SqlTools.ServiceLayer.UnitTests/QueryExecution/Common.cs +++ b/test/Microsoft.SqlTools.ServiceLayer.UnitTests/QueryExecution/Common.cs @@ -20,6 +20,7 @@ using Microsoft.SqlTools.ServiceLayer.UnitTests.Utility; using Microsoft.SqlTools.ServiceLayer.Workspace; using Microsoft.SqlTools.ServiceLayer.Workspace.Contracts; using Microsoft.SqlTools.ServiceLayer.Test.Common; +using Microsoft.Data.SqlClient; using Moq; using Moq.Protected; using HostingProtocol = Microsoft.SqlTools.Hosting.Protocol; @@ -232,7 +233,7 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.QueryExecution private static ISqlConnectionFactory CreateMockFactory(TestResultSet[] data, bool throwOnExecute, bool throwOnRead) { var mockFactory = new Mock(); - mockFactory.Setup(factory => factory.CreateSqlConnection(It.IsAny(), It.IsAny())) + mockFactory.Setup(factory => factory.CreateSqlConnection(It.IsAny(), It.IsAny(), It.IsAny())) .Returns(() => CreateTestConnection(data, throwOnExecute, throwOnRead)); return mockFactory.Object; diff --git a/test/Microsoft.SqlTools.ServiceLayer.UnitTests/Utility/TestObjects.cs b/test/Microsoft.SqlTools.ServiceLayer.UnitTests/Utility/TestObjects.cs index ded6a4f1..822d4c28 100644 --- a/test/Microsoft.SqlTools.ServiceLayer.UnitTests/Utility/TestObjects.cs +++ b/test/Microsoft.SqlTools.ServiceLayer.UnitTests/Utility/TestObjects.cs @@ -14,6 +14,7 @@ using Microsoft.SqlTools.ServiceLayer.Connection; using Microsoft.SqlTools.ServiceLayer.Connection.Contracts; using Microsoft.SqlTools.ServiceLayer.LanguageServices; using Microsoft.SqlTools.ServiceLayer.Workspace.Contracts; +using Microsoft.Data.SqlClient; using Moq; namespace Microsoft.SqlTools.ServiceLayer.UnitTests.Utility @@ -353,7 +354,7 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.Utility /// public class TestSqlConnectionFactory : ISqlConnectionFactory { - public DbConnection CreateSqlConnection(string connectionString, string azureAccountToken) + public DbConnection CreateSqlConnection(string connectionString, string azureAccountToken, SqlRetryLogicBaseProvider retryProvider = null) { return new TestSqlConnection(null) {