Factor out catching and transforming execution exceptions.

Gremlinq · Nov 16, 2023 · 6919631 · 6919631
1 parent 54d06bc
commit 6919631
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 30 deletions.
diff --git a/src/Core/Execution/GremlinQueryExecutor.cs b/src/Core/Execution/GremlinQueryExecutor.cs
@@ -128,36 +128,11 @@ public TransformExecutionExceptionGremlinQueryExecutor(IGremlinQueryExecutor bas
                 _exceptionTransformation = exceptionTransformation;
             }
 
-            public IAsyncEnumerable<T> Execute<T>(GremlinQueryExecutionContext context)
-            {
-                return Core(this, context);
-
-                static async IAsyncEnumerable<T> Core(TransformExecutionExceptionGremlinQueryExecutor @this, GremlinQueryExecutionContext context, [EnumeratorCancellation] CancellationToken ct = default)
-                {
-                    var enumerator = @this._baseExecutor
-                        .Execute<T>(context)
-                        .WithCancellation(ct)
-                        .GetAsyncEnumerator();
-
-                    await using (enumerator)
-                    {
-                        while (true)
-                        {
-                            try
-                            {
-                                if (!await enumerator.MoveNextAsync())
-                                    yield break;
-                            }
-                            catch (GremlinQueryExecutionException ex)
-                            {
-                                throw @this._exceptionTransformation(ex);
-                            }
-
-                            yield return enumerator.Current;
-                        }
-                    }
-                }
-            }
+            public IAsyncEnumerable<T> Execute<T>(GremlinQueryExecutionContext context) => _baseExecutor
+                .Execute<T>(context)
+                .Catch(ex => ex is GremlinQueryExecutionException executionException
+                    ? _exceptionTransformation(executionException)
+                    : ex);
         }
 
         private sealed class SerializingGremlinQueryExecutor : IGremlinQueryExecutor

diff --git a/src/Core/Extensions/AsyncEnumerable.cs b/src/Core/Extensions/AsyncEnumerable.cs
@@ -1,5 +1,7 @@
 using System.Runtime.CompilerServices;
 
+using ExRam.Gremlinq.Core.Execution;
+
 namespace ExRam.Gremlinq.Core
 {
     internal static class AsyncEnumerable
@@ -114,5 +116,31 @@ public static async ValueTask<TSource[]> ToArrayAsync<TSource>(this IAsyncEnumer
 
             return list.ToArray();
         }
+
+        public static IAsyncEnumerable<T> Catch<T>(this IAsyncEnumerable<T> source, Func<Exception, Exception> exceptionTransformation)
+        {
+            return Core(source, exceptionTransformation);
+
+            static async IAsyncEnumerable<T> Core(IAsyncEnumerable<T> source, Func<Exception, Exception> exceptionTransformation, [EnumeratorCancellation] CancellationToken ct = default)
+            {
+                await using (var enumerator = source.WithCancellation(ct).GetAsyncEnumerator())
+                {
+                    while (true)
+                    {
+                        try
+                        {
+                            if (!await enumerator.MoveNextAsync())
+                                yield break;
+                        }
+                        catch (Exception ex)
+                        {
+                            throw exceptionTransformation(ex);
+                        }
+
+                        yield return enumerator.Current;
+                    }
+                }
+            }
+        }
     }
 }