Added retry policy for sleeping serverless error for SqlConnections (#2155)

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;

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)

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);
             }

src/Microsoft.SqlTools.ManagedBatchParser/ReliableConnection/ReliableSqlConnection.cs

@@ -60,9 +60,15 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection.ReliableConnection
         /// <param name="connectionString">The connection string used to open the SQL Azure database.</param>
         /// <param name="connectionRetryPolicy">The retry policy defining whether to retry a request if a connection fails to be established.</param>
         /// <param name="commandRetryPolicy">The retry policy defining whether to retry a request if a command fails to be executed.</param>
-        public ReliableSqlConnection(string connectionString, RetryPolicy connectionRetryPolicy, RetryPolicy commandRetryPolicy, string azureAccountToken)
+        /// <param name="retryProvider">Optional retry provider to handle errors in a special way</param>
+        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();
 

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
+    {
+        /// <summary>
+        /// Approved list of transient errors that require additional time to wait before connecting again.
+        /// </summary>
+        private static readonly HashSet<int> _retryableServerlessConnectivityError;
+        
+        /// <summary>
+        /// Max intervals between retries in seconds to wake up serverless instances.
+        /// </summary>
+        private const int _serverlessMaxIntervalTime = 30;
+        
+        /// <summary>
+        /// Maximum number of retries to wake up serverless instances.
+        /// </summary>
+        private const int _serverlessMaxRetries = 4;
+
+        static SqlRetryProviders()
+        {
+            _retryableServerlessConnectivityError = new HashSet<int>
+            {
+                //// 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, 
+            };
+        }
+
+        /// <summary>
+        /// Wait for SqlConnection to handle sleeping serverless instances (allows for them to wake up, otherwise it will result in errors).
+        /// </summary>
+        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;
+        }
+
+    }
+}

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))
                 {

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

@@ -52,7 +52,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection
         /// <summary>
         /// Properties used for creating/opening the SQL connection.
         /// </summary>
-        public ConnectionDetails ConnectionDetails { get; private set; }
+        public ConnectionDetails ConnectionDetails { get; set; }
 
         /// <summary>
         /// A map containing all connections to the database that are associated with 

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;
 
         /// <summary>
         /// 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<ConnectionCompleteParams?> 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;
+        }
+
         /// <summary>
         /// Tries to create and open a connection with the given ConnectParams.
         /// </summary>
@@ -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)

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
         /// <summary>
         /// Create a new SQL Connection object
         /// </summary>
-        DbConnection CreateSqlConnection(string connectionString, string azureAccountToken);
+        /// <param name="retryProvider">Optional retry provider to handle errors in a special way</param>
+        DbConnection CreateSqlConnection(string connectionString, string azureAccountToken, SqlRetryLogicBaseProvider retryProvider = null);
     }
 }

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
         /// <summary>
         /// Creates a new SqlConnection object
         /// </summary>
-        public DbConnection CreateSqlConnection(string connectionString, string azureAccountToken)
+        /// <param name="retryProvider">Optional retry provider to handle errors in a special way</param>
+        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);
         }
     }
 }

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;

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;

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

@@ -86,7 +86,7 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.Connection
                 });
 
             var mockFactory = new Mock<ISqlConnectionFactory>();
-            mockFactory.Setup(factory => factory.CreateSqlConnection(It.IsAny<string>(), It.IsAny<string>()))
+            mockFactory.Setup(factory => factory.CreateSqlConnection(It.IsAny<string>(), It.IsAny<string>(), It.IsAny<SqlRetryLogicBaseProvider>()))
                 .Returns(mockConnection.Object);
 
 
@@ -152,7 +152,7 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.Connection
                 .Returns(() => Task.Run(() => { }));
 
             var mockFactory = new Mock<ISqlConnectionFactory>();
-            mockFactory.SetupSequence(factory => factory.CreateSqlConnection(It.IsAny<string>(), It.IsAny<string>()))
+            mockFactory.SetupSequence(factory => factory.CreateSqlConnection(It.IsAny<string>(), It.IsAny<string>(), It.IsAny<SqlRetryLogicBaseProvider>()))
                 .Returns(mockConnection.Object)
                 .Returns(mockConnection2.Object);
 
@@ -215,7 +215,7 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.Connection
                 });
 
             var mockFactory = new Mock<ISqlConnectionFactory>();
-            mockFactory.Setup(factory => factory.CreateSqlConnection(It.IsAny<string>(), It.IsAny<string>()))
+            mockFactory.Setup(factory => factory.CreateSqlConnection(It.IsAny<string>(), It.IsAny<string>(), It.IsAny<SqlRetryLogicBaseProvider>()))
                 .Returns(mockConnection.Object);
 
 
@@ -304,7 +304,7 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.Connection
             connectionMock.Setup(c => c.Database).Returns(expectedDbName);
 
             var mockFactory = new Mock<ISqlConnectionFactory>();
-            mockFactory.Setup(factory => factory.CreateSqlConnection(It.IsAny<string>(), It.IsAny<string>()))
+            mockFactory.Setup(factory => factory.CreateSqlConnection(It.IsAny<string>(), It.IsAny<string>(), It.IsAny<SqlRetryLogicBaseProvider>()))
                 .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<ISqlConnectionFactory>();
-            mockFactory.Setup(factory => factory.CreateSqlConnection(It.IsAny<string>(), It.IsAny<string>()))
-            .Returns((string connString, string azureAccountToken) =>
+            mockFactory.Setup(factory => factory.CreateSqlConnection(It.IsAny<string>(), It.IsAny<string>(), It.IsAny<SqlRetryLogicBaseProvider>()))
+            .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<ISqlConnectionFactory>();
-            mockFactory.Setup(factory => factory.CreateSqlConnection(It.IsAny<string>(), It.IsAny<string>()))
+            mockFactory.Setup(factory => factory.CreateSqlConnection(It.IsAny<string>(), It.IsAny<string>(), It.IsAny<SqlRetryLogicBaseProvider>()))
                 .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<ISqlConnectionFactory>();
-            mockFactory.Setup(factory => factory.CreateSqlConnection(It.IsAny<string>(), It.IsAny<string>()))
-            .Returns((string connString, string azureAccountToken) =>
+            mockFactory.Setup(factory => factory.CreateSqlConnection(It.IsAny<string>(), It.IsAny<string>(), It.IsAny<SqlRetryLogicBaseProvider>()))
+            .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<ISqlConnectionFactory>();
-            mockFactory.Setup(factory => factory.CreateSqlConnection(It.IsAny<string>(), It.IsAny<string>()))
-            .Returns((string connString, string azureAccountToken) =>
+            mockFactory.Setup(factory => factory.CreateSqlConnection(It.IsAny<string>(), It.IsAny<string>(), It.IsAny<SqlRetryLogicBaseProvider>()))
+            .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<ISqlConnectionFactory>();
-            mockFactory.Setup(factory => factory.CreateSqlConnection(It.IsAny<string>(), It.IsAny<string>()))
-            .Returns((string connString, string azureAccountToken) =>
+            mockFactory.Setup(factory => factory.CreateSqlConnection(It.IsAny<string>(), It.IsAny<string>(), It.IsAny<SqlRetryLogicBaseProvider>()))
+            .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<ISqlConnectionFactory>();
-            mockFactory.Setup(factory => factory.CreateSqlConnection(It.IsAny<string>(), It.IsAny<string>()))
+            mockFactory.Setup(factory => factory.CreateSqlConnection(It.IsAny<string>(), It.IsAny<string>(), It.IsAny<SqlRetryLogicBaseProvider>()))
                 .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<string>(), It.Is<string>(accountToken => accountToken == azureAccountToken)), Times.Once());
+            mockFactory.Verify(factory => factory.CreateSqlConnection(It.IsAny<string>(), It.Is<string>(accountToken => accountToken == azureAccountToken), It.IsAny<SqlRetryLogicBaseProvider>()), Times.Once());
         }
 
         /// <summary>

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<ISqlConnectionFactory>();
-            mockFactory.Setup(factory => factory.CreateSqlConnection(It.IsAny<string>(), It.IsAny<string>()))
+            mockFactory.Setup(factory => factory.CreateSqlConnection(It.IsAny<string>(), It.IsAny<string>(), It.IsAny<SqlRetryLogicBaseProvider>()))
                 .Returns(() => CreateTestConnection(data, throwOnExecute, throwOnRead));
 
             return mockFactory.Object;

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
     /// </summary>
     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)
             {