hrynco-notification-service/HrynCo.NotificationService.DAL.EF/Core/EfUnitOfWork.cs

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

namespace HrynCo.NotificationService.DAL.EF.Core;

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

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

    public Task SaveChangesAsync(CancellationToken cancellationToken = default)
    {
        return _context.SaveChangesAsync(cancellationToken);
    }

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

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

    public ITransaction? GetCurrentTransaction()
    {
        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
        {
            if (ownsTransaction)
            {
                await tx.RollbackAsync();
            }

            throw;
        }
        finally
        {
            if (ownsTransaction)
            {
                await tx.DisposeAsync();
            }
        }
    }

    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
        {
            if (ownsTransaction)
            {
                await tx.RollbackAsync();
            }

            throw;
        }
        finally
        {
            if (ownsTransaction)
            {
                await tx.DisposeAsync();
            }
        }
    }
}