ChouBox/Utile/HuanMeng.DotNetCore/MultiTenant/MultiTenantDbContext.cs

using HuanMeng.DotNetCore.MultiTenant.Contract;

using Microsoft.EntityFrameworkCore;

using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Threading.Tasks;

namespace HuanMeng.DotNetCore.MultiTenant
{
    /// <summary>
    /// 基本多租户DbContext
    /// </summary>
    public class MultiTenantDbContext : DbContext, IMultiTenantDbContext
    {
        /// <summary>
        /// 租户信息
        /// </summary>
        public ITenantInfo? TenantInfo { get; set; }


        /// <summary>
        /// 构造函数
        /// </summary>
        /// <param name="tenantInfo"></param>
        public MultiTenantDbContext(ITenantInfo? tenantInfo)
        {
            this.TenantInfo = tenantInfo;
        }

        /// <summary>
        /// 构造函数
        /// </summary>
        /// <param name="tenantInfo"></param>
        /// <param name="options"></param>
        public MultiTenantDbContext(ITenantInfo? tenantInfo, DbContextOptions options)
        : base(options)
        {
            if (tenantInfo == null)
            {
                tenantInfo = new TenantInfo()
                {
                    TenantId = Guid.NewGuid(),
                    Identifier = "default",
                    Name = "default"
                };
            }
            this.TenantInfo = tenantInfo;
        }

        public void SetTenantInfo(ITenantInfo tenantInfo)
        {
            this.TenantInfo = tenantInfo;
        }

        public override int SaveChanges()
        {
            if (TenantInfo?.TenantId != null)
            {
                var entries = ChangeTracker.Entries()
                .Where(e => e.Entity is IMultiTenantEntity &&
                            (e.State == EntityState.Added || e.State == EntityState.Modified))
                .Select(e => e.Entity as IMultiTenantEntity)
                .ToList();

                foreach (var entity in entries)
                {
                    if (entity?.TenantId == null || entity?.TenantId == Guid.Empty)
                    {
                        entity.TenantId = TenantInfo.TenantId;
                    }
                }
            }
            return base.SaveChanges();
        }

        public override async Task<int> SaveChangesAsync(CancellationToken cancellationToken = default)
        {
            if (TenantInfo?.TenantId != null)
            {
                var entries = ChangeTracker.Entries()
                .Where(e => e.Entity is IMultiTenantEntity &&
                            (e.State == EntityState.Added || e.State == EntityState.Modified))
                .Select(e => e.Entity as IMultiTenantEntity)
                .ToList();

                foreach (var entity in entries)
                {
                    if (entity?.TenantId == null || entity?.TenantId == Guid.Empty)
                    {
                        entity.TenantId = TenantInfo.TenantId;
                    }
                }
            }
            return await base.SaveChangesAsync(cancellationToken);
        }

    }
}