Make the Host.Outbox capable of any Key type

Signed-off-by: Tomasz Maruszak <maruszaktomasz@gmail.com>
zarusz · Oct 18, 2024 · a9238c9 · a9238c9
1 parent 205d964
commit a9238c9
Showing 1 changed file with 17 additions and 11 deletions.
diff --git a/src/SlimMessageBus.Host.Outbox.Sql.DbContext/DbContextTransactionService.cs b/src/SlimMessageBus.Host.Outbox.Sql.DbContext/DbContextTransactionService.cs
@@ -4,25 +4,31 @@ public class DbContextTransactionService<TDbContext>(TDbContext dbContext, ISqlS
     : AbstractSqlTransactionService((SqlConnection)dbContext.Database.GetDbConnection())
     where TDbContext : Microsoft.EntityFrameworkCore.DbContext
 {
+    private IDbContextTransaction _currentTransaction;
+
     public TDbContext DbContext { get; } = dbContext;
 
-    public override SqlTransaction CurrentTransaction => (SqlTransaction)DbContext.Database.CurrentTransaction?.GetDbTransaction();
+    public override SqlTransaction CurrentTransaction => _currentTransaction?.GetDbTransaction() as SqlTransaction;
 
-    protected override Task OnBeginTransaction()
+    protected override async Task OnBeginTransaction()
     {
-        return DbContext.Database.BeginTransactionAsync(sqlSettings.TransactionIsolationLevel);
+        _currentTransaction = await DbContext.Database.BeginTransactionAsync(sqlSettings.TransactionIsolationLevel);
     }
 
-    protected override Task OnCompleteTransaction(bool transactionFailed)
+    protected override async Task OnCompleteTransaction(bool transactionFailed)
     {
-        if (transactionFailed)
-        {
-            DbContext.Database.RollbackTransaction();
-        }
-        else
+        if (_currentTransaction != null)
         {
-            DbContext.Database.CommitTransaction();
+            if (transactionFailed)
+            {
+                await _currentTransaction.RollbackAsync();
+            }
+            else
+            {
+                await _currentTransaction.CommitAsync();
+            }
+            _currentTransaction.Dispose();
+            _currentTransaction = null;
         }
-        return Task.CompletedTask;
     }
 }