ShouldRetry

csharp-opensource · Nov 29, 2024 · b88c3e4 · b88c3e4
1 parent 3d1950f
commit b88c3e4
Show file tree

Hide file tree

Showing 3 changed files with 36 additions and 5 deletions.
diff --git a/CSharp.OpenSource.LinqToKql/Extensions/IQueryableExtension.cs b/CSharp.OpenSource.LinqToKql/Extensions/IQueryableExtension.cs
@@ -54,4 +54,11 @@ public static ILinqToKqlProvider<S> FromKQL<T, S>(this IQueryable<T> q, string t
         kql.TableOrKQL = tableOrKQL;
         return kql.Clone<S>(null);
     }
+
+    public static ILinqToKqlProvider<T> WithRetry<T>(this IQueryable<T> q, Func<Exception, Task<bool>> shouldRetry)
+    {
+        var kql = q.AsKQL();
+        kql.ShouldRetry = shouldRetry;
+        return kql;
+    }
 }
diff --git a/CSharp.OpenSource.LinqToKql/Provider/ILinqToKqlProvider.cs b/CSharp.OpenSource.LinqToKql/Provider/ILinqToKqlProvider.cs
@@ -3,11 +3,12 @@
 
 namespace CSharp.OpenSource.LinqToKql.Provider;
 
-public interface ILinqToKqlProvider<T> : IQueryable<T>, IQueryProvider, IOrderedQueryable<T>, IAsyncEnumerable<T>
+public interface ILinqToKqlProvider<T> : IQueryable<T>, IQueryProvider, IOrderedQueryable<T>, IAsyncEnumerable<T>, ICloneable
 {
     string? DefaultDbName { get; set; }
     string TableOrKQL { get; set; }
     LinqToKQLQueryTranslator Translator { get; }
+    Func<Exception, Task<bool>>? ShouldRetry { get; set; }
 
     LinqToKqlProvider<S> Clone<S>(Expression? expression = null);
     string TranslateToKQL(Expression? expression = null);

diff --git a/CSharp.OpenSource.LinqToKql/Provider/LinqToKqlProvider.cs b/CSharp.OpenSource.LinqToKql/Provider/LinqToKqlProvider.cs
@@ -14,20 +14,23 @@ public class LinqToKqlProvider<T> : ILinqToKqlProvider<T>
     public Expression Expression => _expression;
     public IQueryProvider Provider => this;
     protected ILinqToKqlProviderExecutor ProviderExecutor;
+    public Func<Exception, Task<bool>>? ShouldRetry { get; set; }
 
     public LinqToKqlProvider(
         string tableOrKQL,
         Expression? expression,
         ILinqToKqlProviderExecutor providerExecutor,
         LinqToKQLQueryTranslatorConfig? config = null,
-        string? defaultDbName = null)
+        string? defaultDbName = null,
+        Func<Exception, Task<bool>>? shouldRetry = null)
     {
         TableOrKQL = tableOrKQL;
         _expression = expression ?? Expression.Constant(this);
         ProviderExecutor = providerExecutor;
         config ??= new();
         Translator = new(config);
         DefaultDbName = defaultDbName;
+        ShouldRetry = shouldRetry;
     }
 
     public virtual object? Execute(Expression expression)
@@ -36,11 +39,23 @@ public LinqToKqlProvider(
     public virtual TResult Execute<TResult>(Expression expression)
         => ExecuteAsync<TResult>(expression).GetAwaiter().GetResult();
 
-    public virtual Task<TResult> ExecuteAsync<TResult>(Expression expression)
+    public virtual async Task<TResult> ExecuteAsync<TResult>(Expression expression)
     {
         if (ProviderExecutor == null) { throw new InvalidOperationException("ProviderExecutor is not set."); }
         var kql = TranslateToKQL(expression);
-        return ProviderExecutor.ExecuteAsync<TResult>(kql, DefaultDbName);
+        try
+        {
+            return await ProviderExecutor.ExecuteAsync<TResult>(kql, DefaultDbName);
+        }
+        catch (Exception ex) 
+        {
+            var shouldRetry = ShouldRetry == null ? false : await ShouldRetry(ex);
+            if (shouldRetry) 
+            {
+                return await ProviderExecutor.ExecuteAsync<TResult>(kql, DefaultDbName);
+            }
+            throw;
+        }
     }
 
     public virtual string TranslateToKQL(Expression? expression = null)
@@ -50,7 +65,14 @@ public virtual IQueryable<TElement> CreateQuery<TElement>(Expression expression)
         => Clone<TElement>(expression);
 
     public virtual LinqToKqlProvider<S> Clone<S>(Expression? expression = null)
-        => new LinqToKqlProvider<S>(TableOrKQL, expression, ProviderExecutor, Translator.Config, DefaultDbName);
+        => new LinqToKqlProvider<S>(
+                TableOrKQL,
+                expression,
+                ProviderExecutor,
+                Translator.Config,
+                DefaultDbName,
+                ShouldRetry
+           );
 
     public virtual async IAsyncEnumerator<T> GetAsyncEnumerator(CancellationToken cancellationToken = default)
     {
@@ -65,4 +87,5 @@ public virtual async IAsyncEnumerator<T> GetAsyncEnumerator(CancellationToken ca
     protected virtual IEnumerator<T> GetGenericEnumerator() => Provider.Execute<List<T>>(Expression).GetEnumerator();
     public virtual IEnumerator<T> GetEnumerator() => GetGenericEnumerator();
     IEnumerator IEnumerable.GetEnumerator() => GetGenericEnumerator();
+    public object Clone() => Clone<T>();
 }