hrynco-ef/HrynCo.DAL.EF/Core/EfUnitOfWork.cs

namespace HrynCo.DAL.EF.Core;

using HrynCo.DAL.Abstract;
using Microsoft.EntityFrameworkCore;
using Microsoft.EntityFrameworkCore.Storage;

public class EfUnitOfWork<TDbContext> : IUnitOfWork
    where TDbContext : DbContext
{
    private readonly TDbContext _context;
    private EfTransactionAdapter? _currentTransaction;

    protected EfUnitOfWork(TDbContext context)
    {
        _context = context;
    }

    public async Task<ITransaction> BeginTransactionAsync(CancellationToken cancellationToken = default)
    {
        if (_currentTransaction != null)
        {
            return _currentTransaction;
        }

        IDbContextTransaction transaction = await _context.Database.BeginTransactionAsync(cancellationToken);
        _currentTransaction = new EfTransactionAdapter(transaction);
        return _currentTransaction;
    }

    public async Task ExecuteInTransactionAsync(Func<Task> action)
    {
        ITransaction? existing = GetCurrentTransaction();
        bool ownsTransaction = existing is null;

        ITransaction tx = existing ?? await BeginTransactionAsync();

        try
        {
            await action();
            if (ownsTransaction)
            {
                await tx.CommitAsync();
            }
        }
        catch (Exception original)
        {
            if (ownsTransaction)
            {
                await RollbackPreservingFailureAsync(tx, original);
            }

            throw;
        }
        finally
        {
            if (ownsTransaction)
            {
                try
                {
                    await tx.DisposeAsync();
                }
                finally
                {
                    _currentTransaction = null;
                }
            }
        }
    }

    public async Task<TResult> ExecuteInTransactionAsync<TResult>(Func<Task<TResult>> action)
    {
        ITransaction? existing = GetCurrentTransaction();
        bool ownsTransaction = existing is null;

        ITransaction tx = existing ?? await BeginTransactionAsync();

        try
        {
            TResult result = await action();
            if (ownsTransaction)
            {
                await tx.CommitAsync();
            }

            return result;
        }
        catch (Exception original)
        {
            if (ownsTransaction)
            {
                await RollbackPreservingFailureAsync(tx, original);
            }

            throw;
        }
        finally
        {
            if (ownsTransaction)
            {
                try
                {
                    await tx.DisposeAsync();
                }
                finally
                {
                    _currentTransaction = null;
                }
            }
        }
    }

    private static async Task RollbackPreservingFailureAsync(ITransaction transaction, Exception original)
    {
        try
        {
            await transaction.RollbackAsync();
        }
        catch (Exception rollbackFailure)
        {
            throw new AggregateException("The transaction and its rollback both failed.", original, rollbackFailure);
        }
    }

    public ITransaction? GetCurrentTransaction()
    {
        return _currentTransaction;
    }
}