From 832634913798be14482479414236b69ce7f7ca78 Mon Sep 17 00:00:00 2001
From: Cody Robibero <cody@robibe.ro>
Date: Thu, 29 Dec 2022 19:40:39 -0700
Subject: [PATCH] Use custom database healthcheck (#8973)

---
 .../DbContextFactoryHealthCheck.cs            | 43 +++++++++++++++++++
 Jellyfin.Server/Startup.cs                    |  3 +-
 2 files changed, 45 insertions(+), 1 deletion(-)
 create mode 100644 Jellyfin.Server/HealthChecks/DbContextFactoryHealthCheck.cs

diff --git a/Jellyfin.Server/HealthChecks/DbContextFactoryHealthCheck.cs b/Jellyfin.Server/HealthChecks/DbContextFactoryHealthCheck.cs
new file mode 100644
index 0000000000..bf00dcd53f
--- /dev/null
+++ b/Jellyfin.Server/HealthChecks/DbContextFactoryHealthCheck.cs
@@ -0,0 +1,43 @@
+using System;
+using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.EntityFrameworkCore;
+using Microsoft.Extensions.Diagnostics.HealthChecks;
+
+namespace Jellyfin.Server.HealthChecks;
+
+/// <summary>
+/// Implementation of the <see cref="DbContextHealthCheck{TContext}"/> for a <see cref="IDbContextFactory{TContext}"/>.
+/// </summary>
+/// <typeparam name="TContext">The type of database context.</typeparam>
+public class DbContextFactoryHealthCheck<TContext> : IHealthCheck
+    where TContext : DbContext
+{
+    private readonly IDbContextFactory<TContext> _dbContextFactory;
+
+    /// <summary>
+    /// Initializes a new instance of the <see cref="DbContextFactoryHealthCheck{TContext}"/> class.
+    /// </summary>
+    /// <param name="contextFactory">Instance of the <see cref="IDbContextFactory{TContext}"/> interface.</param>
+    public DbContextFactoryHealthCheck(IDbContextFactory<TContext> contextFactory)
+    {
+        _dbContextFactory = contextFactory;
+    }
+
+    /// <inheritdoc />
+    public async Task<HealthCheckResult> CheckHealthAsync(HealthCheckContext context, CancellationToken cancellationToken = default)
+    {
+        ArgumentNullException.ThrowIfNull(context);
+
+        var dbContext = await _dbContextFactory.CreateDbContextAsync(cancellationToken).ConfigureAwait(false);
+        await using (dbContext.ConfigureAwait(false))
+        {
+            if (await dbContext.Database.CanConnectAsync(cancellationToken).ConfigureAwait(false))
+            {
+                return HealthCheckResult.Healthy();
+            }
+        }
+
+        return HealthCheckResult.Unhealthy();
+    }
+}
diff --git a/Jellyfin.Server/Startup.cs b/Jellyfin.Server/Startup.cs
index 49a57aa688..5d6a278c40 100644
--- a/Jellyfin.Server/Startup.cs
+++ b/Jellyfin.Server/Startup.cs
@@ -8,6 +8,7 @@ using System.Text;
 using Jellyfin.MediaEncoding.Hls.Extensions;
 using Jellyfin.Networking.Configuration;
 using Jellyfin.Server.Extensions;
+using Jellyfin.Server.HealthChecks;
 using Jellyfin.Server.Implementations;
 using Jellyfin.Server.Implementations.Extensions;
 using Jellyfin.Server.Infrastructure;
@@ -122,7 +123,7 @@ namespace Jellyfin.Server
                 .ConfigurePrimaryHttpMessageHandler(defaultHttpClientHandlerDelegate);
 
             services.AddHealthChecks()
-                .AddDbContextCheck<JellyfinDb>();
+                .AddCheck<DbContextFactoryHealthCheck<JellyfinDb>>(nameof(JellyfinDb));
 
             services.AddHlsPlaylistGenerator();
         }