Feat: UnitOfWorkRepository

Author: AnnaSasDev

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

@@ -20,5 +20,5 @@ public interface IUnitOfWork : IAsyncDisposable {
 
     ValueTask<TDbContext> GetDbContextAsync<TDbContext>(CancellationToken ct = default) where TDbContext : DbContext;
 
-    TRepo GetRepository<TRepo>() where TRepo : class, IUnitOfWorkRepository;
+    ValueTask<TRepo> GetRepositoryAsync<TRepo>(CancellationToken ct = default) where TRepo : class, IUnitOfWorkRepository;
 }

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

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

src/CodeOfChaos.Types.UnitOfWork/UnitOfWork.cs

@@ -84,11 +84,13 @@ public virtual async ValueTask<T> GetDbContextAsync<T>(CancellationToken ct = de
         return dbContext as T ?? throw new InvalidCastException($"Cannot cast DbContext of type '{dbContext.GetType()}' to '{typeof(T)}'");
     }
 
-    public virtual TRepo GetRepository<TRepo>() where TRepo : class, IUnitOfWorkRepository {
+    public virtual async ValueTask<TRepo> GetRepositoryAsync<TRepo>(CancellationToken ct = default) where TRepo : class, IUnitOfWorkRepository {
         if (AttachedRepositories.TryGetValue(typeof(TRepo), out IUnitOfWorkRepository? cachedRepo) && cachedRepo is TRepo castedCachedRepo) return castedCachedRepo;
 
         // Cache miss so we create a new instance
         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;
@@ -98,6 +100,10 @@ public virtual async ValueTask DisposeAsync() {
         if (_transaction != null) await TryRollbackTransactionAsync();
 
         if (!AttachedRepositories.IsEmpty) {
+            foreach (IUnitOfWorkRepository repository in AttachedRepositories.Values) {
+                if (repository is not UnitOfWorkRepository<TDbContext> castedRepo) continue;
+                castedRepo.Detach();
+            }
             AttachedRepositories.Clear();
         }
 

src/CodeOfChaos.Types.UnitOfWork/UnitOfWorkRepository.cs

@@ -0,0 +1,24 @@
+// ---------------------------------------------------------------------------------------------------------------------
+// Imports
+// ---------------------------------------------------------------------------------------------------------------------
+using Microsoft.EntityFrameworkCore;
+
+namespace CodeOfChaos.Types.UnitOfWork;
+
+// ---------------------------------------------------------------------------------------------------------------------
+// Code
+// ---------------------------------------------------------------------------------------------------------------------
+public abstract class UnitOfWorkRepository<TDbContext> : IUnitOfWorkRepository 
+    where TDbContext : DbContext
+{
+    private TDbContext? DbContext { get; set; }
+    
+    // -----------------------------------------------------------------------------------------------------------------
+    // Methods
+    // -----------------------------------------------------------------------------------------------------------------
+    internal async ValueTask AttachAsync(IUnitOfWork unitOfWork, CancellationToken ct = default) => DbContext = await unitOfWork.GetDbContextAsync<TDbContext>(ct);
+    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.");
+}

tests/Tests.CodeOfChaos.Types.UnitOfWork/UnitOfWorkTests.cs

@@ -200,7 +200,7 @@ public async Task GetRepository_ShouldRetrieveRepositoryFromServiceProvider() {
 
 
         // Act
-        var repository = _unitOfWork.GetRepository<IUnitOfWorkRepository>();
+        var repository = await _unitOfWork.GetRepositoryAsync<IUnitOfWorkRepository>();
 
         // Assert
         await Assert.That(repository).IsNotNull();