Add conn/command retry to simple OE (#2356)

* add conn and command retry to simple OE

* add a enableRetry parameter

* docs

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

@@ -83,6 +83,29 @@ public ReliableSqlConnection(string connectionString, RetryPolicy connectionRetr
             }
         }
 
+        /// <summary>
+        /// Initialize a new instance of the ReliableSqlConnection class with a given connection
+        /// </summary>
+        /// <param name="connection">The connection used to open the SQL DB</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>
+        /// <param name="retryProvider">Optional retry provider to handle errors in a special way</param>
+        public ReliableSqlConnection(SqlConnection connection, RetryPolicy connectionRetryPolicy, RetryPolicy commandRetryPolicy, SqlRetryLogicBaseProvider retryProvider = null)
+        {
+            _underlyingConnection = connection;
+
+            if (retryProvider != null) {
+                _underlyingConnection.RetryLogicProvider = retryProvider;
+            }
+
+            _connectionRetryPolicy = connectionRetryPolicy ?? RetryPolicyFactory.CreateNoRetryPolicy();
+            _commandRetryPolicy = commandRetryPolicy ?? RetryPolicyFactory.CreateNoRetryPolicy();
+
+            _underlyingConnection.StateChange += OnConnectionStateChange;
+            _connectionRetryPolicy.RetryOccurred += RetryConnectionCallback;
+            _commandRetryPolicy.RetryOccurred += RetryCommandCallback;
+        }
+
         /// <summary>
         /// Performs application-defined tasks associated with freeing, releasing, or
         ///  resetting managed and unmanaged resources.

src/Microsoft.SqlTools.SqlCore/SimpleObjectExplorer/ObjectExplorer.cs

@@ -8,14 +8,15 @@
 using System.Linq;
 using System.Threading.Tasks;
 using Microsoft.Data.SqlClient;
+using Microsoft.SqlTools.ServiceLayer.Connection.ReliableConnection;
 
 namespace Microsoft.SqlTools.SqlCore.SimpleObjectExplorer
 {
     public static class ObjectExplorer
     {
-        public static async Task<TreeNode> GetObjectExplorerModel(SqlConnection connection)
+        public static async Task<TreeNode> GetObjectExplorerModel(SqlConnection connection, bool enableRetry = true)
         {
-            ObjectMetadata[] metdata = await FetchObjectExplorerMetadataTable(connection);
+            ObjectMetadata[] metdata = await FetchObjectExplorerMetadataTable(connection, enableRetry);
             TreeNode root = new DatabaseNode(null, new ObjectMetadata() { Name = connection.Database, Type = "Database", DisplayName = connection.Database });
             // Load all the children
             Stack<TreeNode> stack = new Stack<TreeNode>();
@@ -39,13 +40,20 @@ public static async Task<TreeNode> GetObjectExplorerModel(SqlConnection connecti
             return root;
         }
 
-        private static async Task<ObjectMetadata[]> FetchObjectExplorerMetadataTable(SqlConnection connection)
+        private static async Task<ObjectMetadata[]> FetchObjectExplorerMetadataTable(SqlConnection connection, bool enableRetry)
         {
             string[] metadataQueries = ObjectExplorerModelQueries.Queries.Values.ToArray();
             string combinedQuery = string.Join(Environment.NewLine + "UNION ALL" + Environment.NewLine, metadataQueries);
-            using (SqlCommand command = new SqlCommand(combinedQuery, connection))
+            ReliableSqlConnection reliableSqlConnection = new ReliableSqlConnection(connection, 
+                enableRetry ? RetryPolicyFactory.CreateDefaultDataConnectionRetryPolicy() : RetryPolicyFactory.CreateNoRetryPolicy(), 
+                RetryPolicyFactory.CreateDefaultSchemaCommandRetryPolicy(useRetry: enableRetry));
+
+            // ReliableSqlConnection only opens the underlying SqlConnection if it is not already open
+            await reliableSqlConnection.OpenAsync();
+            using (ReliableSqlConnection.ReliableSqlCommand command = new ReliableSqlConnection.ReliableSqlCommand(reliableSqlConnection))
             {
-                using (SqlDataReader reader = await command.ExecuteReaderAsync())
+                command.CommandText = combinedQuery;
+                using (System.Data.Common.DbDataReader reader = await command.ExecuteReaderAsync())
                 {
                     List<ObjectMetadata> metadata = new List<ObjectMetadata>();
                     while (reader.Read())