ColaFlow/colaflow-api/src/Modules/Identity/ColaFlow.Modules.Identity.Infrastructure/Persistence/IdentityDbContext.cs

using ColaFlow.Modules.Identity.Domain.Aggregates.Tenants;
using ColaFlow.Modules.Identity.Domain.Aggregates.Users;
using ColaFlow.Modules.Identity.Domain.Entities;
using ColaFlow.Modules.Identity.Infrastructure.Services;
using ColaFlow.Shared.Kernel.Common;
using MediatR;
using Microsoft.EntityFrameworkCore;

namespace ColaFlow.Modules.Identity.Infrastructure.Persistence;

public class IdentityDbContext(
    DbContextOptions<IdentityDbContext> options,
    ITenantContext tenantContext,
    IMediator mediator)
    : DbContext(options)
{
    public DbSet<Tenant> Tenants => Set<Tenant>();
    public DbSet<User> Users => Set<User>();
    public DbSet<RefreshToken> RefreshTokens => Set<RefreshToken>();
    public DbSet<UserTenantRole> UserTenantRoles => Set<UserTenantRole>();
    public DbSet<EmailVerificationToken> EmailVerificationTokens => Set<EmailVerificationToken>();

    protected override void OnModelCreating(ModelBuilder modelBuilder)
    {
        base.OnModelCreating(modelBuilder);

        // Apply all configurations from assembly
        modelBuilder.ApplyConfigurationsFromAssembly(typeof(IdentityDbContext).Assembly);

        // Configure Global Query Filters (automatic tenant data filtering)
        ConfigureGlobalQueryFilters(modelBuilder);
    }

    private void ConfigureGlobalQueryFilters(ModelBuilder modelBuilder)
    {
        // User entity global query filter
        // Automatically adds: WHERE tenant_id = @current_tenant_id
        modelBuilder.Entity<User>().HasQueryFilter(u =>
            !tenantContext.IsSet || u.TenantId == tenantContext.TenantId);

        // Tenant entity doesn't need filter (need to query all tenants)
    }

    /// <summary>
    /// Disable Query Filter (for admin operations)
    /// </summary>
    public IQueryable<T> WithoutTenantFilter<T>() where T : class
    {
        return Set<T>().IgnoreQueryFilters();
    }

    /// <summary>
    /// Override SaveChangesAsync to dispatch domain events before saving
    /// </summary>
    public override async Task<int> SaveChangesAsync(CancellationToken cancellationToken = default)
    {
        // Dispatch domain events BEFORE saving changes
        await DispatchDomainEventsAsync(cancellationToken);

        // Save changes to database
        return await base.SaveChangesAsync(cancellationToken);
    }

    /// <summary>
    /// Dispatch domain events to handlers via MediatR
    /// </summary>
    private async Task DispatchDomainEventsAsync(CancellationToken cancellationToken)
    {
        // Get all aggregate roots with pending domain events
        var domainEntities = ChangeTracker
            .Entries<AggregateRoot>()
            .Where(x => x.Entity.DomainEvents.Any())
            .Select(x => x.Entity)
            .ToList();

        // Collect all domain events
        var domainEvents = domainEntities
            .SelectMany(x => x.DomainEvents)
            .ToList();

        // Clear events from aggregates (prevent double-publishing)
        domainEntities.ForEach(entity => entity.ClearDomainEvents());

        // Publish each event via MediatR
        foreach (var domainEvent in domainEvents)
        {
            await mediator.Publish(domainEvent, cancellationToken);
        }
    }
}