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.Aggregates.Invitations;
using ColaFlow.Modules.Identity.Domain.Entities;
using ColaFlow.Modules.Identity.Infrastructure.Services;
using ColaFlow.Shared.Kernel.Common;
using MediatR;
using Microsoft.EntityFrameworkCore;
using Microsoft.Extensions.Hosting;
using Microsoft.Extensions.Logging;
using System.Diagnostics;

namespace ColaFlow.Modules.Identity.Infrastructure.Persistence;

public class IdentityDbContext(
    DbContextOptions<IdentityDbContext> options,
    ITenantContext tenantContext,
    IMediator mediator,
    IHostEnvironment environment,
    ILogger<IdentityDbContext> logger)
    : 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>();
    public DbSet<PasswordResetToken> PasswordResetTokens => Set<PasswordResetToken>();
    public DbSet<Invitation> Invitations => Set<Invitation>();
    public DbSet<EmailRateLimit> EmailRateLimits => Set<EmailRateLimit>();

    protected override void OnConfiguring(DbContextOptionsBuilder optionsBuilder)
    {
        base.OnConfiguring(optionsBuilder);

        // Enable query logging in development for performance analysis
        if (environment.IsDevelopment())
        {
            optionsBuilder
                .EnableSensitiveDataLogging()
                .LogTo(Console.WriteLine, LogLevel.Information)
                .EnableDetailedErrors();
        }
    }

    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)
    {
        var stopwatch = Stopwatch.StartNew();

        // Dispatch domain events BEFORE saving changes
        await DispatchDomainEventsAsync(cancellationToken);

        // Save changes to database
        var result = await base.SaveChangesAsync(cancellationToken);

        stopwatch.Stop();

        // Log slow database operations (> 1 second)
        if (stopwatch.ElapsedMilliseconds > 1000)
        {
            logger.LogWarning(
                "Slow database operation detected: SaveChangesAsync took {ElapsedMs}ms",
                stopwatch.ElapsedMilliseconds);
        }

        return result;
    }

    /// <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);
        }
    }
}