Feat: GetCachedDbSet

Author: AnnaSasDev

src/CodeOfChaos.Types.UnitOfWork.Contracts/IUnitOfWorkRepository.cs

@@ -5,4 +5,4 @@ namespace CodeOfChaos.Types.UnitOfWork;
 // ---------------------------------------------------------------------------------------------------------------------
 // Code
 // ---------------------------------------------------------------------------------------------------------------------
-public interface IUnitOfWorkRepository {}
+public interface IUnitOfWorkRepository;

src/CodeOfChaos.Types.UnitOfWork/UnitOfWork.cs

@@ -87,12 +87,17 @@ public virtual async ValueTask<TRepo> GetRepositoryAsync<TRepo>(CancellationToke
         if (AttachedRepositories.TryGetValue(typeof(TRepo), out IUnitOfWorkRepository? cachedRepo) && cachedRepo is TRepo castedCachedRepo) return castedCachedRepo;
 
         // Cache miss so we create a new instance
+        var repo = await CreateAndAttachRepositoryAsync<TRepo>(ct);
+        
+        AttachedRepositories.AddOrUpdate(typeof(TRepo), repo);
+        return repo;
+    }
+
+    public virtual async ValueTask<TRepo> CreateAndAttachRepositoryAsync<TRepo>(CancellationToken ct = default) where TRepo : class, IUnitOfWorkRepository {
         var repo = serviceScope.ServiceProvider.GetRequiredService<TRepo>();
         if (repo is not UnitOfWorkRepository<TDbContext> castedRepo) throw new InvalidCastException($"Cannot cast repository of type '{repo.GetType()}' to '{typeof(TRepo)}'");
 
         await castedRepo.AttachAsync(this, ct);
-
-        AttachedRepositories.AddOrUpdate(typeof(TRepo), repo);
         return repo;
     }
 

src/CodeOfChaos.Types.UnitOfWork/UnitOfWorkRepository.cs

@@ -2,6 +2,7 @@
 // Imports
 // ---------------------------------------------------------------------------------------------------------------------
 using Microsoft.EntityFrameworkCore;
+using System.Collections.Concurrent;
 
 namespace CodeOfChaos.Types.UnitOfWork;
 // ---------------------------------------------------------------------------------------------------------------------
@@ -10,6 +11,7 @@ namespace CodeOfChaos.Types.UnitOfWork;
 public abstract class UnitOfWorkRepository<TDbContext> : IUnitOfWorkRepository
     where TDbContext : DbContext {
     private TDbContext? DbContext { get; set; }
+    private ConcurrentDictionary<Type,object> DbSetCache {get;} = new();
 
     // -----------------------------------------------------------------------------------------------------------------
     // Methods
@@ -18,5 +20,11 @@ public abstract class UnitOfWorkRepository<TDbContext> : IUnitOfWorkRepository
     internal void Detach() => DbContext = null;// Remove the reference to the DbContext
 
     protected TDbContext GetDbContext() => DbContext ?? throw new InvalidOperationException("Repository is not attached to a UnitOfWork.");
-    protected DbSet<TModel> GetDbSet<TModel>() where TModel : class => DbContext?.Set<TModel>() ?? throw new InvalidOperationException("Repository is not attached to a UnitOfWork.");
+    protected DbSet<TModel> GetDbSet<TModel>() where TModel : class => GetDbContext().Set<TModel>();
+    protected DbSet<TModel> GetCachedDbSet<TModel>() where TModel : class 
+        => (DbSet<TModel>)DbSetCache.GetOrAdd(
+            typeof(TModel),
+            static (_, dbContext) => dbContext.Set<TModel>(), 
+            GetDbContext()
+        );
 }