PreExecute (#18)

csharp-opensource · Jan 27, 2025 · e537ab3 · e537ab3
1 parent 22cd3f3
commit e537ab3
Show file tree

Hide file tree

Showing 6 changed files with 40 additions and 6 deletions.
diff --git a/LinqToKql.Test/Translator/ExtenstionsTests.cs b/LinqToKql.Test/Translator/ExtenstionsTests.cs
@@ -0,0 +1,13 @@
+using CSharp.OpenSource.LinqToKql.Extensions;
+
+namespace CSharp.OpenSource.LinqToKql.Test.Translator;
+
+public class ExtenstionsTests : LinqToKQLQueryTranslatorBaseTest
+{
+    [Fact]
+    public Task WithPreExecute_TestAsync()
+        => AssertQueryAsync(
+            _q.WithPreExecute((p, kql) => $"{kql}{p.Translator.PipeWithIndentation}take 10").Select(x => x.Nested),
+            [_tableName, "project Nested" , "take 10"]
+        );
+}
diff --git a/LinqToKql.Test/Translator/LinqToKQLQueryTranslatorBaseTest.cs b/LinqToKql.Test/Translator/LinqToKQLQueryTranslatorBaseTest.cs
@@ -23,6 +23,10 @@ protected async Task AssertQueryAsync<T>(IQueryable<T> queryable, string[] expec
         );
         // Act
         var actual = kql.TranslateToKQL();
+        if (kql.PreExecute != null)
+        {
+            actual = kql.PreExecute(kql, actual);
+        }
 
         // Assert
         var partsExpected = expected.Split(config.PipeWithIndentation);

diff --git a/LinqToKql/Extensions/IQueryableExtension.cs b/LinqToKql/Extensions/IQueryableExtension.cs
@@ -72,6 +72,13 @@ public static ILinqToKqlProvider<T> WithRetry<T>(this IQueryable<T> q, Func<ILin
         return kql;
     }
 
+    public static ILinqToKqlProvider<T> WithPreExecute<T>(this IQueryable<T> q, Func<ILinqToKqlProvider, string, string> preExecute)
+    {
+        var kql = q.AsKQL();
+        kql.PreExecute = preExecute;
+        return kql;
+    }
+
     public static IQueryable<T> Or<T>(this IQueryable<T> q, List<Expression<Func<T, bool>>> predicates)
     {
         if (predicates == null || predicates.Count == 0) { return q; }

diff --git a/LinqToKql/Provider/ILinqToKqlProvider.cs b/LinqToKql/Provider/ILinqToKqlProvider.cs
@@ -14,6 +14,7 @@ public interface ILinqToKqlProvider
     LinqToKQLQueryTranslator Translator { get; set; }
     Func<ILinqToKqlProvider, Exception, Task<bool>>? ShouldRetry { get; set; }
     ILinqToKqlProviderExecutor ProviderExecutor { get; set; }
+    Func<ILinqToKqlProvider, string, string>? PreExecute { get; set; }
 
     LinqToKqlProvider<S> Clone<S>(Expression? expression = null, bool cloneExpressionOnNull = true);
     string TranslateToKQL(Expression? expression = null);

diff --git a/LinqToKql/Provider/KustoDbContext.cs b/LinqToKql/Provider/KustoDbContext.cs
@@ -16,7 +16,8 @@ public abstract class KustoDbContext : IKustoDbContext
 
     public LinqToKQLQueryTranslator Translator { get => dummyProvider.Translator; set => dummyProvider.Translator = value; }
 
-    public Func<ILinqToKqlProvider, Exception, Task<bool>>? ShouldRetry { get => throw new NotImplementedException(); set => throw new NotImplementedException(); }
+    public Func<ILinqToKqlProvider, Exception, Task<bool>>? ShouldRetry { get => dummyProvider.ShouldRetry; set => dummyProvider.ShouldRetry = value; }
+    public Func<ILinqToKqlProvider, string, string>? PreExecute { get => dummyProvider.PreExecute; set => dummyProvider.PreExecute = value; }
 
     protected virtual LinqToKQLQueryTranslatorConfig GetConfig() => new();
 

diff --git a/LinqToKql/Provider/LinqToKqlProvider.cs b/LinqToKql/Provider/LinqToKqlProvider.cs
@@ -15,14 +15,16 @@ public class LinqToKqlProvider<T> : ILinqToKqlProvider<T>
     public IQueryProvider Provider => this;
     public ILinqToKqlProviderExecutor ProviderExecutor { get; set; }
     public Func<ILinqToKqlProvider, Exception, Task<bool>>? ShouldRetry { get; set; }
+    public Func<ILinqToKqlProvider, string, string>? PreExecute { get; set; }
 
     public LinqToKqlProvider(
         string tableOrKQL,
         Expression? expression,
         ILinqToKqlProviderExecutor providerExecutor,
         LinqToKQLQueryTranslatorConfig? config = null,
         string? defaultDbName = null,
-        Func<ILinqToKqlProvider, Exception, Task<bool>>? shouldRetry = null)
+        Func<ILinqToKqlProvider, Exception, Task<bool>>? shouldRetry = null,
+        Func<ILinqToKqlProvider, string, string>? preExecute = null)
     {
         TableOrKQL = tableOrKQL;
         _expression = expression ?? Expression.Constant(this);
@@ -31,20 +33,25 @@ public LinqToKqlProvider(
         Translator = new(config);
         DefaultDbName = defaultDbName;
         ShouldRetry = shouldRetry;
+        PreExecute = preExecute;
     }
 
     public virtual object? Execute(Expression expression)
         => Execute<object>(expression);
 
     public virtual TResult Execute<TResult>(Expression expression)
-        => ExecuteAsync<TResult>(expression).GetAwaiter().GetResult();
+        => ExecuteAsync<TResult>(expression).GetAwaiter().GetResult()!;
 
-    public virtual async 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);
         try
         {
+            if (PreExecute != null)
+            {
+                kql = PreExecute(this, kql);
+            }
             return await ProviderExecutor.ExecuteAsync<TResult>(kql, DefaultDbName);
         }
         catch (Exception ex) 
@@ -83,14 +90,15 @@ public virtual LinqToKqlProvider<S> Clone<S>(Expression? expression = null, bool
             ProviderExecutor,
             Translator.Config,
             DefaultDbName,
-            ShouldRetry
+            ShouldRetry,
+            PreExecute
         );
     }
 
     public virtual async IAsyncEnumerator<T> GetAsyncEnumerator(CancellationToken cancellationToken = default)
     {
         var results = await ExecuteAsync<List<T>>(Expression);
-        foreach (var result in results)
+        foreach (var result in results!)
         {
             yield return result;
         }