CloudBuilder/CloudBuilder.Core/DatabaseAccessor/UnitOfWork.cs

using Microsoft.EntityFrameworkCore.Storage;
using Microsoft.EntityFrameworkCore;
using System;
using System.Collections.Generic;
using System.Data.Common;
using System.Data;
using System.Linq;
using System.Text;
using System.Threading.Tasks;
using CloudBuilder.Core.DatabaseAccessor.Entity;
using Microsoft.Extensions.DependencyInjection;

namespace CloudBuilder.Core.DatabaseAccessor
{
    public class UnitOfWork : IUnitOfWork
    {
        private readonly List<DbContext> _dbContexts = new List<DbContext>();
        private DbTransaction _transaction;
        private DbConnection _connection;
        private DbContext _dbContext;

        /// <summary>
        /// 初始化一个<see cref="UnitOfWork"/>类型的新实例
        /// </summary>
        public UnitOfWork([FromKeyedServices(MasterDbContext.MASTER_DB_CONTEXT)] DbContext context)
        {
            _dbContext = context;
            _dbContexts.Add(_dbContext);
            _connection = _dbContext.Database.GetDbConnection();
        }


        /// <summary>
        /// 获取 事务是否已提交
        /// </summary>
        public bool HasCommitted { get; private set; }

        /// <summary>
        /// 添加DbContext共享事务
        /// </summary>
        public void AddDbContext(DbContext dbContext)
        {
            _dbContexts.Add(dbContext);
        }

        public DbContext GetDbContext()
        {
            return _dbContext;
        }

        /// <summary>
        /// 对数据库连接开启事务
        /// </summary>
        public virtual void BeginOrUseTransaction()
        {
            if (_dbContexts.Count == 0)
            {
                return;
            }
            if (_transaction?.Connection == null)
            {
                if (_connection.State != ConnectionState.Open)
                {
                    _connection.Open();
                }

                _transaction = _connection.BeginTransaction();
            }

            foreach (DbContext context in _dbContexts)
            {
                if (context.Database.CurrentTransaction != null && context.Database.CurrentTransaction.GetDbTransaction() == _transaction)
                {
                    continue;
                }
                //if (context.IsRelationalTransaction())
                //{
                //共享事务仅限关系数据库
                context.Database.UseTransaction(_transaction);
                //}
                //else
                //{
                //    context.Database.BeginTransaction();
                //}
            }

            HasCommitted = false;
        }

        /// <summary>
        /// 对数据库连接开启事务
        /// </summary>
        /// <param name="cancellationToken">异步取消标记</param>
        /// <returns></returns>
        public virtual async Task BeginOrUseTransactionAsync(CancellationToken cancellationToken = default)
        {
            if (_dbContexts.Count == 0)
            {
                return;
            }
            if (_transaction?.Connection == null)
            {
                if (_connection.State != ConnectionState.Open)
                {
                    await _connection.OpenAsync(cancellationToken);
                }

                _transaction = _connection.BeginTransaction();
            }

            foreach (DbContext context in _dbContexts)
            {
                if (context.Database.CurrentTransaction != null && context.Database.CurrentTransaction.GetDbTransaction() == _transaction)
                {
                    continue;
                }
                //if (context.IsRelationalTransaction())
                //{
                context.Database.UseTransaction(_transaction);
                //}
                //else
                //{
                //    await context.Database.BeginTransactionAsync(cancellationToken);
                //}
            }

            HasCommitted = false;
        }

        /// <summary>
        /// 提交当前上下文的事务更改
        /// </summary>
        public virtual void Commit()
        {
            if (HasCommitted || _dbContexts.Count == 0 || _transaction == null)
            {
                return;
            }

            _transaction.Commit();
            foreach (DbContext context in _dbContexts)
            {
                //if (context.IsRelationalTransaction())
                //{
                context.Database.CurrentTransaction.Dispose();
                //关系型数据库共享事务
                continue;
                //}

                //context.Database.CommitTransaction();
            }

            HasCommitted = true;
        }

        /// <summary>
        /// 回滚所有事务
        /// </summary>
        public virtual void Rollback()
        {
            if (_transaction?.Connection != null)
            {
                _transaction.Rollback();
            }
            foreach (var context in _dbContexts)
            {
                //if (context.IsRelationalTransaction())
                //{
                CleanChanges(context);
                if (context.Database.CurrentTransaction != null)
                {
                    //context.Database.CurrentTransaction.Rollback();
                    context.Database.CurrentTransaction.Dispose();
                }
                continue;
                //}
                //context.Database.RollbackTransaction();
            }
            HasCommitted = true;
        }

        private static void CleanChanges(DbContext context)
        {
            var entries = context.ChangeTracker.Entries().ToArray();
            foreach (var entry in entries)
            {
                entry.State = EntityState.Detached;
            }
        }

        public void Dispose()
        {
            _transaction?.Dispose();
            foreach (DbContext context in _dbContexts)
            {
                context.Dispose();
            }
        }
    }
}