Catch 401s that seem to come as throttle exception?

microsoft · Oct 25, 2023 · 9499150 · 9499150
1 parent d9619ab
commit 9499150
Showing 1 changed file with 18 additions and 5 deletions.
diff --git a/code/DeltaKustoIntegration/Kusto/KustoManagementGateway.cs b/code/DeltaKustoIntegration/Kusto/KustoManagementGateway.cs
@@ -21,9 +21,22 @@ namespace DeltaKustoIntegration.Kusto
     /// </summary>
     internal class KustoManagementGateway : IKustoManagementGateway
     {
-        private static readonly AsyncRetryPolicy _retryPolicy = Policy
-            .Handle<KustoException>(ex => !ex.IsPermanent)
-            .WaitAndRetryAsync(3, attempt => TimeSpan.FromSeconds(attempt));
+        private static readonly ResiliencePipeline _resiliencePipeline = new ResiliencePipelineBuilder()
+            .AddRetry(new RetryStrategyOptions
+            {
+                ShouldHandle = new PredicateBuilder().Handle<KustoException>(ex => !ex.IsPermanent),
+                MaxRetryAttempts = 3,
+                BackoffType = DelayBackoffType.Linear,
+                Delay = TimeSpan.FromSeconds(1)
+            })
+            .AddRetry(new RetryStrategyOptions
+            {
+                ShouldHandle = new PredicateBuilder().Handle<KustoException>(ex => ex.FailureCode==401),
+                MaxRetryAttempts = 1,
+                BackoffType = DelayBackoffType.Linear,
+                Delay = TimeSpan.FromSeconds(1)
+            })
+            .Build();
 
         private readonly Uri _clusterUri;
         private readonly string _database;
@@ -139,15 +152,15 @@ private async Task<IDataReader> ExecuteCommandAsync(
         {
             try
             {
-                return await _retryPolicy.ExecuteAsync(async () =>
+                return await _resiliencePipeline.ExecuteAsync(async (ct) =>
                 {
                     var reader = await _commandProvider.ExecuteControlCommandAsync(
                         _database,
                         commandScript,
                         _requestProperties);
 
                     return reader;
-                });
+                }, ct);
             }
             catch (Exception ex)
             {