QrRapido/Services/HealthChecks/MongoDbHealthCheck.cs

using Microsoft.Extensions.Diagnostics.HealthChecks;
using MongoDB.Bson;
using MongoDB.Driver;
using QRRapidoApp.Data;
using System.Diagnostics;

namespace QRRapidoApp.Services.HealthChecks
{
    public class MongoDbHealthCheck : IHealthCheck
    {
        private readonly IServiceProvider _serviceProvider;
        private readonly IConfiguration _configuration;
        private readonly ILogger<MongoDbHealthCheck> _logger;
        
        private readonly int _timeoutSeconds;
        private readonly bool _includeDatabaseSize;
        private readonly bool _testQuery;

        public MongoDbHealthCheck(
            IServiceProvider serviceProvider, 
            IConfiguration configuration,
            ILogger<MongoDbHealthCheck> logger)
        {
            _serviceProvider = serviceProvider;
            _configuration = configuration;
            _logger = logger;
            
            _timeoutSeconds = configuration.GetValue<int>("HealthChecks:MongoDB:TimeoutSeconds", 5);
            _includeDatabaseSize = configuration.GetValue<bool>("HealthChecks:MongoDB:IncludeDatabaseSize", true);
            _testQuery = configuration.GetValue<bool>("HealthChecks:MongoDB:TestQuery", true);
        }

        public async Task<HealthCheckResult> CheckHealthAsync(HealthCheckContext context, CancellationToken cancellationToken = default)
        {
            var stopwatch = Stopwatch.StartNew();
            
            try
            {
                using var scope = _serviceProvider.CreateScope();
                var mongoContext = scope.ServiceProvider.GetService<MongoDbContext>();
                
                if (mongoContext?.Database == null)
                {
                    return HealthCheckResult.Degraded("MongoDB context not available - application running without database");
                }

                using var timeoutCts = new CancellationTokenSource(TimeSpan.FromSeconds(_timeoutSeconds));
                using var combinedCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, timeoutCts.Token);
                
                var data = new Dictionary<string, object>();
                
                // Test basic connectivity with ping
                var pingCommand = new BsonDocument("ping", 1);
                await mongoContext.Database.RunCommandAsync<BsonDocument>(pingCommand, cancellationToken: combinedCts.Token);
                
                var latencyMs = stopwatch.ElapsedMilliseconds;
                data["latency"] = $"{latencyMs}ms";
                data["status"] = latencyMs < 100 ? "fast" : latencyMs < 500 ? "normal" : "slow";
                
                // Test a simple query if enabled
                if (_testQuery)
                {
                    try
                    {
                        var testCollection = mongoContext.Database.GetCollection<BsonDocument>("Users");
                        var queryStopwatch = Stopwatch.StartNew();
                        await testCollection.CountDocumentsAsync(new BsonDocument(), cancellationToken: combinedCts.Token);
                        queryStopwatch.Stop();
                        
                        data["lastQuery"] = "successful";
                        data["queryLatencyMs"] = queryStopwatch.ElapsedMilliseconds;
                    }
                    catch (Exception queryEx)
                    {
                        _logger.LogWarning(queryEx, "MongoDB health check query failed");
                        data["lastQuery"] = "failed";
                        data["queryError"] = queryEx.Message;
                    }
                }
                
                // Get database size and basic stats if enabled
                if (_includeDatabaseSize)
                {
                    try
                    {
                        var dbStatsCommand = new BsonDocument("dbStats", 1);
                        var dbStats = await mongoContext.Database.RunCommandAsync<BsonDocument>(dbStatsCommand, cancellationToken: combinedCts.Token);
                        
                        var dataSize = dbStats.GetValue("dataSize", BsonValue.Create(0)).AsDouble;
                        var indexSize = dbStats.GetValue("indexSize", BsonValue.Create(0)).AsDouble;
                        var totalSizeMB = (dataSize + indexSize) / (1024 * 1024);
                        var documentCount = dbStats.GetValue("objects", BsonValue.Create(0)).ToInt64();
                        
                        data["databaseSizeMB"] = Math.Round(totalSizeMB, 1);
                        data["documentCount"] = documentCount;
                        data["indexSizeMB"] = Math.Round(indexSize / (1024 * 1024), 1);
                    }
                    catch (Exception statsEx)
                    {
                        _logger.LogWarning(statsEx, "Failed to get MongoDB database stats for health check");
                        data["databaseStatsError"] = statsEx.Message;
                    }
                }
                
                // Get MongoDB version
                try
                {
                    var serverStatus = await mongoContext.Database.RunCommandAsync<BsonDocument>(
                        new BsonDocument("serverStatus", 1), cancellationToken: combinedCts.Token);
                    data["version"] = serverStatus.GetValue("version", BsonValue.Create("unknown")).AsString;
                }
                catch (Exception versionEx)
                {
                    _logger.LogWarning(versionEx, "Failed to get MongoDB version for health check");
                    data["version"] = "unknown";
                }
                
                stopwatch.Stop();
                data["totalCheckTimeMs"] = stopwatch.ElapsedMilliseconds;
                
                // Determine health status based on performance
                if (latencyMs > 2000)
                {
                    return HealthCheckResult.Unhealthy($"MongoDB responding slowly ({latencyMs}ms)", data: data);
                }
                
                if (latencyMs > 1000)
                {
                    return HealthCheckResult.Degraded($"MongoDB performance degraded ({latencyMs}ms)", data: data);
                }
                
                return HealthCheckResult.Healthy($"MongoDB healthy ({latencyMs}ms)", data: data);
            }
            catch (OperationCanceledException)
            {
                return HealthCheckResult.Unhealthy($"MongoDB health check timed out after {_timeoutSeconds} seconds");
            }
            catch (Exception ex)
            {
                _logger.LogError(ex, "MongoDB health check failed");
                return HealthCheckResult.Unhealthy($"MongoDB health check failed: {ex.Message}");
            }
        }
    }
}