Batch Paging Refinements (#7422)

src/GreenDonut/src/Core/CreateDataLoaderBranch.cs

@@ -0,0 +1,21 @@
+#if NET6_0_OR_GREATER
+namespace GreenDonut;
+
+/// <summary>
+/// Creates a branched DataLoader with a specific branch key.
+/// </summary>
+/// <typeparam name="TKey">
+/// The type of the DataLoader key.
+/// </typeparam>
+/// <typeparam name="TValue">
+/// The type of the DataLoader value.
+/// </typeparam>
+/// <typeparam name="TState">
+/// The custom state that is passed into the factory.
+/// </typeparam>
+public delegate IDataLoader CreateDataLoaderBranch<out TKey, TValue, in TState>(
+    string branchKey,
+    IDataLoader<TKey, TValue> dataLoader,
+    TState state)
+    where TKey : notnull;
+#endif

src/GreenDonut/src/Core/DataLoaderBase.cs

@@ -1,7 +1,4 @@
 using System.Collections.Immutable;
-#if NET8_0_OR_GREATER
-using GreenDonut.Projections;
-#endif
 using static GreenDonut.NoopDataLoaderDiagnosticEventListener;
 using static GreenDonut.Errors;
 
@@ -33,9 +30,9 @@ public abstract partial class DataLoaderBase<TKey, TValue>
     private readonly int _maxBatchSize;
     private readonly IDataLoaderDiagnosticEvents _diagnosticEvents;
     private readonly CancellationToken _ct;
-#if NET8_0_OR_GREATER
-    private ImmutableDictionary<string, ISelectionDataLoader<TKey, TValue>> _branches =
-        ImmutableDictionary<string, ISelectionDataLoader<TKey, TValue>>.Empty;
+#if NET6_0_OR_GREATER
+    private ImmutableDictionary<string, IDataLoader> _branches =
+        ImmutableDictionary<string, IDataLoader>.Empty;
 #endif
     private Batch<TKey>? _currentBatch;
 
@@ -80,14 +77,27 @@ protected DataLoaderBase(IBatchScheduler batchScheduler, DataLoaderOptions? opti
     public IImmutableDictionary<string, object?> ContextData { get; set; } =
         ImmutableDictionary<string, object?>.Empty;
 
-    private protected virtual bool AllowCachePropagation => true;
+    /// <summary>
+    /// Specifies if the values fetched by this DataLoader
+    /// are propagated through the cache.
+    /// </summary>
+    protected virtual bool AllowCachePropagation => true;
 
-    private protected virtual bool AllowBranching => true;
+    /// <summary>
+    /// Specifies if this DataLoader allows branching.
+    /// </summary>
+    protected virtual bool AllowBranching => true;
 
-    internal IBatchScheduler BatchScheduler
+    /// <summary>
+    /// Gets the batch scheduler of this DataLoader.
+    /// </summary>
+    protected internal IBatchScheduler BatchScheduler
         => _batchScheduler;
 
-    internal DataLoaderOptions Options
+    /// <summary>
+    /// Gets the options of this DataLoader.
+    /// </summary>
+    protected internal DataLoaderOptions Options
         => new()
         {
             MaxBatchSize = _maxBatchSize,
@@ -252,11 +262,24 @@ public void Set(TKey key, Task<TValue?> value)
             Cache.TryAdd(cacheKey, new Promise<TValue?>(value));
         }
     }
-#if NET8_0_OR_GREATER
+#if NET6_0_OR_GREATER
 
     /// <inheritdoc />
-    public ISelectionDataLoader<TKey, TValue> Branch(string key)
+    public IDataLoader Branch<TState>(
+        string key,
+        CreateDataLoaderBranch<TKey, TValue, TState> createBranch,
+        TState state)
     {
+        if (string.IsNullOrEmpty(key))
+        {
+            throw new ArgumentException("Value cannot be null or empty.", nameof(key));
+        }
+
+        if (createBranch == null)
+        {
+            throw new ArgumentNullException(nameof(createBranch));
+        }
+
         if(!AllowBranching)
         {
             throw new InvalidOperationException(
@@ -269,7 +292,7 @@ public ISelectionDataLoader<TKey, TValue> Branch(string key)
             {
                 if (!_branches.TryGetValue(key, out branch))
                 {
-                    var newBranch = new SelectionDataLoader<TKey, TValue>(this, key);
+                    var newBranch = createBranch(key, this, state);
                     _branches = _branches.Add(key, newBranch);
                     return newBranch;
                 }

src/GreenDonut/src/Core/DataLoaderExtensions.cs

@@ -1,5 +1,4 @@
 using System.Collections.Immutable;
-using System.Diagnostics.CodeAnalysis;
 
 namespace GreenDonut;
 
@@ -339,6 +338,110 @@ public static IDataLoader<TKey, TValue> TrySetState<TKey, TValue, TState>(
         return dataLoader;
     }
 
+    /// <summary>
+    /// Gets a state value from the <paramref name="dataLoader"/> or
+    /// creates a new one and stores it as state on the <paramref name="dataLoader"/>.
+    /// </summary>
+    /// <param name="dataLoader">
+    /// The data loader instance.
+    /// </param>
+    /// <param name="createValue">
+    /// A factory that creates the new state value.
+    /// </param>
+    /// <typeparam name="TKey">
+    /// The key type of the DataLoader.
+    /// </typeparam>
+    /// <typeparam name="TValue">
+    /// The value type of the DataLoader.
+    /// </typeparam>
+    /// <typeparam name="TState">
+    /// The state type.
+    /// </typeparam>
+    /// <returns>
+    /// Returns the state value.
+    /// </returns>
+    /// <exception cref="ArgumentNullException">
+    /// Throws if <paramref name="dataLoader"/> is <c>null</c>.
+    /// </exception>
+    public static TState GetOrSetState<TKey, TValue, TState>(
+        this IDataLoader<TKey, TValue> dataLoader,
+        Func<string, TState> createValue)
+        where TKey : notnull
+    {
+        if (dataLoader is null)
+        {
+            throw new ArgumentNullException(nameof(dataLoader));
+        }
+
+        var key = typeof(TState).FullName ?? typeof(TState).Name;
+
+        if(!dataLoader.ContextData.TryGetValue(key, out var internalValue))
+        {
+            internalValue = createValue(key);
+            dataLoader.ContextData = dataLoader.ContextData.SetItem(key, internalValue);
+        }
+
+        return (TState)internalValue!;
+    }
+
+    /// <summary>
+    /// Gets a state value from the <paramref name="dataLoader"/> or
+    /// creates a new one and stores it as state on the <paramref name="dataLoader"/>.
+    /// </summary>
+    /// <param name="dataLoader">
+    /// The data loader instance.
+    /// </param>
+    /// <param name="key">
+    /// The state key.
+    /// </param>
+    /// <param name="createValue">
+    /// A factory that creates the new state value.
+    /// </param>
+    /// <typeparam name="TKey">
+    /// The key type of the DataLoader.
+    /// </typeparam>
+    /// <typeparam name="TValue">
+    /// The value type of the DataLoader.
+    /// </typeparam>
+    /// <typeparam name="TState">
+    /// The state type.
+    /// </typeparam>
+    /// <returns>
+    /// Returns the state value.
+    /// </returns>
+    /// <exception cref="ArgumentNullException">
+    /// Throws if <paramref name="dataLoader"/> is <c>null</c>.
+    /// </exception>
+    /// <exception cref="ArgumentException">
+    /// Throws if <paramref name="key"/> is <c>null</c> or empty.
+    /// </exception>
+    public static TState GetOrSetState<TKey, TValue, TState>(
+        this IDataLoader<TKey, TValue> dataLoader,
+        string key,
+        Func<string, TState> createValue)
+        where TKey : notnull
+    {
+        if (dataLoader is null)
+        {
+            throw new ArgumentNullException(nameof(dataLoader));
+        }
+
+        if (string.IsNullOrEmpty(key))
+        {
+            throw new ArgumentException(
+                "The key must not be null or empty.",
+                nameof(key));
+        }
+
+        if(!dataLoader.ContextData.TryGetValue(key, out var internalValue))
+        {
+            internalValue = createValue(key);
+            dataLoader.ContextData = dataLoader.ContextData.SetItem(key, internalValue);
+        }
+
+        return (TState)internalValue!;
+    }
+
     /// <summary>
     /// Adds the value to a collection that is stored on the DataLoader state.
     /// </summary>

src/GreenDonut/src/Core/DataLoaderFetchContext.cs

@@ -1,6 +1,8 @@
 using System.Collections.Immutable;
-using System.Diagnostics.CodeAnalysis;
 #if NET8_0_OR_GREATER
+using System.Diagnostics.CodeAnalysis;
+#endif
+#if NET6_0_OR_GREATER
 using GreenDonut.Projections;
 #endif
 
@@ -21,6 +23,18 @@ public readonly struct DataLoaderFetchContext<TValue>(
 {
     public IImmutableDictionary<string, object?> ContextData { get; } = contextData;
 
+    /// <summary>
+    /// Gets a value from the DataLoader state snapshot.
+    /// </summary>
+    /// <typeparam name="TState">
+    /// The type of the state value.
+    /// </typeparam>
+    /// <returns>
+    /// Returns the state value if it exists.
+    /// </returns>
+    public TState? GetState<TState>()
+        => GetState<TState>(typeof(TState).FullName ?? typeof(TState).Name);
+
     /// <summary>
     /// Gets a value from the DataLoader state snapshot.
     /// </summary>
@@ -43,6 +57,21 @@ public readonly struct DataLoaderFetchContext<TValue>(
         return default;
     }
 
+    /// <summary>
+    /// Gets a required value from the DataLoader state snapshot.
+    /// </summary>
+    /// <typeparam name="TState">
+    /// The type of the state value.
+    /// </typeparam>
+    /// <returns>
+    /// Returns the state value if it exists.
+    /// </returns>
+    /// <exception cref="InvalidOperationException">
+    /// Throws an exception if the state value does not exist.
+    /// </exception>
+    public TState GetRequiredState<TState>()
+        => GetRequiredState<TState>(typeof(TState).FullName ?? typeof(TState).Name);
+
     /// <summary>
     /// Gets a required value from the DataLoader state snapshot.
     /// </summary>
@@ -69,6 +98,21 @@ public TState GetRequiredState<TState>(string key)
             $"The state `{key}` is not available on the DataLoader.");
     }
 
+    /// <summary>
+    /// Gets a value from the DataLoader state snapshot or returns a default value.
+    /// </summary>
+    /// <param name="defaultValue">
+    /// The default value to return if the state value does not exist.
+    /// </param>
+    /// <typeparam name="TState">
+    /// The type of the state value.
+    /// </typeparam>
+    /// <returns>
+    /// Returns the state value if it exists.
+    /// </returns>
+    public TState GetStateOrDefault<TState>(TState defaultValue)
+        => GetStateOrDefault(typeof(TState).FullName ?? typeof(TState).Name, defaultValue);
+
     /// <summary>
     /// Gets a value from the DataLoader state snapshot or returns a default value.
     /// </summary>
@@ -93,7 +137,7 @@ public TState GetStateOrDefault<TState>(string key, TState defaultValue)
 
         return defaultValue;
     }
-#if NET8_0_OR_GREATER
+#if NET6_0_OR_GREATER
 
     /// <summary>
     /// Gets the selector builder from the DataLoader state snapshot.
@@ -102,7 +146,9 @@ public TState GetStateOrDefault<TState>(string key, TState defaultValue)
     /// <returns>
     /// Returns the selector builder if it exists.
     /// </returns>
+#if NET8_0_OR_GREATER
     [Experimental(Experiments.Projections)]
+#endif
     public ISelectorBuilder GetSelector()
     {
         DefaultSelectorBuilder<TValue> context;

src/GreenDonut/src/Core/GreenDonut.csproj

@@ -4,9 +4,13 @@
     <PackageId>GreenDonut</PackageId>
     <AssemblyName>GreenDonut</AssemblyName>
     <RootNamespace>GreenDonut</RootNamespace>
-    <Description>Green Donut is a port of facebook's DataLoader utility, written in C# for .NET Core and .NET Framework.</Description>
+    <Description>GreenDonut is a port of facebook's DataLoader utility, written in C# for .NET Core and .NET Framework.</Description>
   </PropertyGroup>
 
+  <ItemGroup>
+    <InternalsVisibleTo Include="HotChocolate.Pagination.Batching" />
+  </ItemGroup>
+
   <ItemGroup>
     <PackageReference Include="Microsoft.Extensions.DependencyInjection.Abstractions" />
     <PackageReference Include="Microsoft.Extensions.ObjectPool" />

src/GreenDonut/src/Core/IDataLoader.cs

@@ -6,15 +6,18 @@
 namespace GreenDonut;
 
 /// <summary>
+/// <para>
 /// A <c>DataLoader</c> creates a public API for loading data from a
 /// particular data back-end with unique keys such as the `id` column of a
 /// SQL table or document name in a MongoDB database, given a batch loading
 /// function. -- facebook
-///
+/// </para>
+/// <para>
 /// Each <c>DataLoader</c> instance contains a unique memoized cache. Use
 /// caution when used in long-lived applications or those which serve many
 /// users with different access permissions and consider creating a new
 /// instance per web request. -- facebook
+/// </para>
 /// </summary>
 public interface IDataLoader
 {
@@ -162,18 +165,25 @@ public interface IDataLoader<in TKey, TValue>
     /// </exception>
     void Set(TKey key, Task<TValue?> value);
 
-#if NET8_0_OR_GREATER
+#if NET6_0_OR_GREATER
     /// <summary>
-    /// Branches the current <c>DataLoader</c> to allow for selections
-    /// to be applied to the data fetching.
+    /// Branches the current <c>DataLoader</c>.
     /// </summary>
     /// <param name="key">
     /// A unique key to identify the branch.
     /// </param>
+    /// <param name="createBranch">
+    /// Creates the branch of the current <c>DataLoader</c>.
+    /// </param>
+    /// <param name="state">
+    /// A custom state object that is passed to the branch factory.
+    /// </param>
     /// <returns>
-    /// A new <c>DataLoader</c> instance which allows for selections to be
-    /// applied to the data fetching.
+    /// A new <c>DataLoader</c> instance.
     /// </returns>
-    ISelectionDataLoader<TKey, TValue> Branch(string key);
+    IDataLoader Branch<TState>(
+        string key,
+        CreateDataLoaderBranch<TKey, TValue, TState> createBranch,
+        TState state);
 #endif
 }