tiku-backend.net/Tiku.Infrastructure/Persistence/TikuDbContext.cs

using System.Reflection;
using Microsoft.AspNetCore.DataProtection.EntityFrameworkCore;
using Microsoft.AspNetCore.Identity;
using Microsoft.AspNetCore.Identity.EntityFrameworkCore;
using Microsoft.EntityFrameworkCore;
using Tiku.Application.Security;
using Tiku.Domain.Common;
using Tiku.Domain.Identity;

namespace Tiku.Infrastructure.Persistence;

public sealed partial class TikuDbContext(
    DbContextOptions<TikuDbContext> options,
    ITenantContext tenantContext) : IdentityUserContext<User, Guid>(options), IDataProtectionKeyContext
{
    public TikuDbContext(DbContextOptions<TikuDbContext> options)
        : this(options, CreateToolingTenantContext())
    {
    }

    public Guid? CurrentTenantId => tenantContext.TenantId;
    public Guid CurrentTenantIdOrEmpty => tenantContext.TenantId ?? Guid.Empty;
    public bool IsTenantResolved => tenantContext.IsResolved;
    public bool IsSystemScope => tenantContext.IsSystem;
    public long SaveVersion { get; private set; }

    private static ITenantContext CreateToolingTenantContext()
    {
        var context = new TenantContext();
        context.InitializeSystem(null, "Direct DbContext construction for model tooling");
        return context;
    }

    protected override void OnModelCreating(ModelBuilder modelBuilder)
    {
        base.OnModelCreating(modelBuilder);
        modelBuilder.Entity<IdentityUserClaim<Guid>>().ToTable("user_claims");
        modelBuilder.Entity<IdentityUserLogin<Guid>>().ToTable("user_logins");
        modelBuilder.Entity<IdentityUserToken<Guid>>().ToTable("user_tokens");
        modelBuilder.HasPostgresExtension("citext");
        modelBuilder.HasPostgresExtension("ltree");
        modelBuilder.HasPostgresExtension("pg_trgm");
        modelBuilder.ApplyConfigurationsFromAssembly(typeof(TikuDbContext).Assembly);
        modelBuilder.Entity<DataProtectionKey>().ToTable("data_protection_keys");
        ApplyTenantQueryFilters(modelBuilder);
        ValidateTenantModel(modelBuilder);
        modelBuilder.UseSnakeCaseIdentifiers();
    }

    private void ApplyTenantQueryFilters(ModelBuilder modelBuilder)
    {
        var applyMethod = typeof(TikuDbContext)
            .GetMethod(nameof(ApplyTenantQueryFilter), BindingFlags.Instance | BindingFlags.NonPublic)!;

        foreach (var entityType in modelBuilder.Model.GetEntityTypes())
        {
            var tenantProperty = entityType.FindProperty("TenantId");
            if (tenantProperty?.ClrType != typeof(Guid) || entityType.BaseType is not null) continue;

            applyMethod.MakeGenericMethod(entityType.ClrType).Invoke(this, [modelBuilder]);
        }
    }

    private void ApplyTenantQueryFilter<TEntity>(ModelBuilder modelBuilder)
        where TEntity : class
    {
        modelBuilder.Entity<TEntity>().HasQueryFilter(entity =>
            IsSystemScope ||
            (IsTenantResolved &&
             EF.Property<Guid>(entity, "TenantId") == CurrentTenantIdOrEmpty));
    }

    private static void ValidateTenantModel(ModelBuilder modelBuilder)
    {
        var invalidUniqueIndexes = new List<string>();
        var invalidTenantForeignKeys = new List<string>();
        foreach (var entityType in modelBuilder.Model.GetEntityTypes())
        {
            var tenantProperty = entityType.FindProperty("TenantId");
            if (tenantProperty?.ClrType != typeof(Guid)) continue;

            if (!typeof(ITenantOwned).IsAssignableFrom(entityType.ClrType))
                throw new InvalidOperationException(
                    $"Tenant entity '{entityType.ClrType.Name}' must implement {nameof(ITenantOwned)}.");

            if (!entityType.GetDeclaredQueryFilters().Any())
                throw new InvalidOperationException(
                    $"Tenant entity '{entityType.ClrType.Name}' does not have a tenant query filter.");

            foreach (var index in entityType.GetDeclaredIndexes().Where(index => index.IsUnique))
                if (index.Properties.All(property => property.Name != "TenantId") &&
                    index.FindAnnotation("Tiku:GlobalUnique")?.Value is not true)
                    invalidUniqueIndexes.Add(
                        $"{entityType.ClrType.Name}({string.Join(",", index.Properties.Select(property => property.Name))})");


            foreach (var foreignKey in entityType.GetDeclaredForeignKeys().Where(foreignKey =>
                         typeof(ITenantOwned).IsAssignableFrom(foreignKey.PrincipalEntityType.ClrType)))
                if (foreignKey.PrincipalKey.Properties.All(property => property.Name != "TenantId"))
                    invalidTenantForeignKeys.Add(
                        $"{entityType.ClrType.Name}->{foreignKey.PrincipalEntityType.ClrType.Name}");
        }

        if (invalidUniqueIndexes.Count > 0)
            throw new InvalidOperationException(
                $"Unique indexes on tenant entities must include TenantId: {string.Join("; ", invalidUniqueIndexes)}");


        if (invalidTenantForeignKeys.Count > 0)
            throw new InvalidOperationException(
                $"Foreign keys between tenant entities must use a tenant-qualified principal key: {string.Join("; ", invalidTenantForeignKeys)}");
    }

    public override int SaveChanges(bool acceptAllChangesOnSuccess)
    {
        UpdateTimestamps();
        var result = base.SaveChanges(acceptAllChangesOnSuccess);
        SaveVersion++;
        return result;
    }

    public override async Task<int> SaveChangesAsync(
        bool acceptAllChangesOnSuccess,
        CancellationToken cancellationToken = default)
    {
        UpdateTimestamps();
        var result = await base.SaveChangesAsync(acceptAllChangesOnSuccess, cancellationToken);
        SaveVersion++;
        return result;
    }

    private void UpdateTimestamps()
    {
        var now = DateTimeOffset.UtcNow;

        foreach (var entry in ChangeTracker.Entries<User>().Where(entry => entry.State == EntityState.Modified))
            if (entry.Property(user => user.Status).IsModified ||
                entry.Property(user => user.PasswordHash).IsModified)
                entry.Entity.SecurityStamp = Guid.NewGuid().ToString("N");

        foreach (var entry in ChangeTracker.Entries<IHasTimestamps>())
        {
            if (entry.State == EntityState.Added) entry.Entity.CreatedAt = now;

            if (entry.State is EntityState.Added or EntityState.Modified) entry.Entity.UpdatedAt = now;
        }
    }
}