stalwart-simplelogin-middleware/StalwartSimpleLoginMiddleware/Contexts/ApiKeyContext.cs

using Microsoft.EntityFrameworkCore;
using Npgsql;
using StalwartSimpleLoginMiddleware.Constants;
using StalwartSimpleLoginMiddleware.Entities;
using StalwartSimpleLoginMiddleware.Utilities;

namespace StalwartSimpleLoginMiddleware.Contexts;

public class ApiKeyContext : DbContext
{
    public DbSet<ApiKey> ApiKeys { get; set; }
    public DbSet<Member> Members { get; set; }

    public ApiKeyContext()
    {
    }

    public ApiKeyContext(DbContextOptions<ApiKeyContext> options) : base(options)
    {
    }

    protected override void OnConfiguring(DbContextOptionsBuilder optionsBuilder)
    {
        ConfigureOptions(optionsBuilder);
    }

    public static DbContextOptionsBuilder ConfigureOptions(DbContextOptionsBuilder optionsBuilder)
    {
        if (optionsBuilder.IsConfigured) return optionsBuilder;

        var connectionString = Environment.GetEnvironmentVariable(EnvironmentVariable.PostgresUrl);

        if (string.IsNullOrEmpty(connectionString))
        {
            // Default connection
            var defaultConnectionString = "Host=localhost;Port=5432;Database=postgres;";
            optionsBuilder.UseNpgsql(defaultConnectionString,
                nsqlOptions => nsqlOptions.MigrationsAssembly("StalwartSimpleLoginMiddleware"));
            return optionsBuilder;
        }

        var pgConnectionString = ConnectionHelper.GetPostgresConnectionString(connectionString);

        // Configure pooling
        var connectionStringBuilder = new NpgsqlConnectionStringBuilder(pgConnectionString)
        {
            MinPoolSize = 10
        };
        if (int.TryParse(Environment.GetEnvironmentVariable(EnvironmentVariable.PostgresMinPoolSize),
                out var minPoolSize))
        {
            connectionStringBuilder.MinPoolSize = minPoolSize;
        }

        if (int.TryParse(Environment.GetEnvironmentVariable(EnvironmentVariable.PostgresMaxPoolSize),
                out var maxPoolSize))
        {
            connectionStringBuilder.MaxPoolSize = maxPoolSize;
        }

        var pooledConnectionString = connectionStringBuilder.ToString();

        optionsBuilder.UseLazyLoadingProxies()
            .UseNpgsql(pooledConnectionString,
                nsqlOptions => nsqlOptions.MigrationsAssembly("StalwartSimpleLoginMiddleware"));
        return optionsBuilder;
    }

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

        modelBuilder.Entity<ApiKey>(entity =>
        {
            entity.HasKey(a => a.Key); // Primary key

            // Members navigation
            entity.HasMany(a => a.Members)
                .WithOne(m => m.ApiKey)
                .HasForeignKey(m => m.ApiKeyId)
                .IsRequired()
                .OnDelete(DeleteBehavior.Cascade);
        });

        // Configure Member entity
        modelBuilder.Entity<Member>(entity =>
        {
            // Composite key: ApiKeyId + Email
            entity.HasKey(m => new { m.ApiKeyId, m.Email });

            // Configure foreign key relationship with ApiKey
            entity.HasOne(m => m.ApiKey)
                .WithMany(a => a.Members)
                .HasForeignKey(m => m.ApiKeyId)
                .IsRequired()
                .OnDelete(DeleteBehavior.Cascade);
        });
    }

    // Method to check and apply migrations
    public async Task EnsureDatabaseMigrated(ILogger logger)
    {
        logger.LogInformation("Checking available migrations...");

        // Get all migrations and log them
        var migrations = Database.GetMigrations().ToList();
        logger.LogInformation($"Total migrations found: {migrations.Count}");

        foreach (var migration in migrations)
        {
            logger.LogInformation($"Migration: {migration}");
        }

        // Check applied migrations
        var pendingMigrations = (await Database.GetPendingMigrationsAsync()).ToArray();
        var appliedMigrations = (await Database.GetAppliedMigrationsAsync()).ToArray();

        logger.LogInformation($"Applied migrations: {string.Join(", ", appliedMigrations)}");

        if (pendingMigrations.Any())
        {
            logger.LogInformation($"Applying migrations: {string.Join(", ", pendingMigrations)}");
            await Database.MigrateAsync();
            logger.LogInformation("Database migrated.");

            if (appliedMigrations.Length == 0)
            {
                // Create Admin key on first migration
                var apiKey = ApiKeyHelper.GenerateKey();
                ApiKeys.Add(new ApiKey
                {
                    Key = apiKey,
                    OwnerEmail = "admin@domain.tld",
                    IsAdmin = true
                });
                await SaveChangesAsync();

                Console.WriteLine($"Generated Admin API Key: {apiKey}");
            }
        }
        else
        {
            logger.LogInformation("No pending migrations found.");
        }
    }
}