// Licensed to the .NET Foundation under one or more agreements.
// The .NET Foundation licenses this file to you under the MIT license.
 
using System.Data.Common;
using Microsoft.Extensions.Diagnostics.HealthChecks;
 
namespace HealthChecksSample;
 
public abstract class DbConnectionHealthCheck : IHealthCheck
{
    protected DbConnectionHealthCheck(string connectionString)
        : this(connectionString, testQuery: null)
    {
    }
 
    protected DbConnectionHealthCheck(string connectionString, string testQuery)
    {
        ConnectionString = connectionString ?? throw new ArgumentNullException(nameof(connectionString));
        TestQuery = testQuery;
    }
 
    protected string ConnectionString { get; }
 
    // This sample supports specifying a query to run as a boolean test of whether the database
    // is responding. It is important to choose a query that will return quickly or you risk
    // overloading the database.
    //
    // In most cases this is not necessary, but if you find it necessary, choose a simple query such as 'SELECT 1'.
    protected string TestQuery { get; }
 
    protected abstract DbConnection CreateConnection(string connectionString);
 
    public async Task<HealthCheckResult> CheckHealthAsync(HealthCheckContext context, CancellationToken cancellationToken = default(CancellationToken))
    {
        using (var connection = CreateConnection(ConnectionString))
        {
            try
            {
                await connection.OpenAsync(cancellationToken);
 
                if (TestQuery != null)
                {
                    var command = connection.CreateCommand();
                    command.CommandText = TestQuery;
 
                    await command.ExecuteNonQueryAsync(cancellationToken);
                }
            }
            catch (DbException ex)
            {
                return new HealthCheckResult(status: context.Registration.FailureStatus, exception: ex);
            }
        }
 
        return HealthCheckResult.Healthy();
    }
}