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