From d6594a8a7097c06613b8102dc1fcfb06fc08eacc Mon Sep 17 00:00:00 2001
From: crobibero <cody@robibe.ro>
Date: Thu, 3 Sep 2020 09:26:22 -0600
Subject: [PATCH] Add db health check

---
 .../HealthChecks/JellyfinDbHealthCheck.cs     | 36 +++++++++++++++++++
 Jellyfin.Server/Jellyfin.Server.csproj        |  1 +
 Jellyfin.Server/Startup.cs                    |  8 ++++-
 3 files changed, 44 insertions(+), 1 deletion(-)
 create mode 100644 Jellyfin.Server/HealthChecks/JellyfinDbHealthCheck.cs

diff --git a/Jellyfin.Server/HealthChecks/JellyfinDbHealthCheck.cs b/Jellyfin.Server/HealthChecks/JellyfinDbHealthCheck.cs
new file mode 100644
index 000000000..aea684479
--- /dev/null
+++ b/Jellyfin.Server/HealthChecks/JellyfinDbHealthCheck.cs
@@ -0,0 +1,36 @@
+using System.Threading;
+using System.Threading.Tasks;
+using Jellyfin.Server.Implementations;
+using Microsoft.Extensions.Diagnostics.HealthChecks;
+
+namespace Jellyfin.Server.HealthChecks
+{
+    /// <summary>
+    /// Checks connectivity to the database.
+    /// </summary>
+    public class JellyfinDbHealthCheck : IHealthCheck
+    {
+        private readonly JellyfinDbProvider _dbProvider;
+
+        /// <summary>
+        /// Initializes a new instance of the <see cref="JellyfinDbHealthCheck"/> class.
+        /// </summary>
+        /// <param name="dbProvider">The jellyfin db provider.</param>
+        public JellyfinDbHealthCheck(JellyfinDbProvider dbProvider)
+        {
+            _dbProvider = dbProvider;
+        }
+
+        /// <inheritdoc />
+        public async Task<HealthCheckResult> CheckHealthAsync(HealthCheckContext context, CancellationToken cancellationToken = default)
+        {
+            await using var jellyfinDb = _dbProvider.CreateContext();
+            if (await jellyfinDb.Database.CanConnectAsync(cancellationToken).ConfigureAwait(false))
+            {
+                return HealthCheckResult.Healthy("Database connection successful.");
+            }
+
+            return HealthCheckResult.Unhealthy("Unable to connect to the database.");
+        }
+    }
+}
diff --git a/Jellyfin.Server/Jellyfin.Server.csproj b/Jellyfin.Server/Jellyfin.Server.csproj
index 7541707d9..6ea7e84de 100644
--- a/Jellyfin.Server/Jellyfin.Server.csproj
+++ b/Jellyfin.Server/Jellyfin.Server.csproj
@@ -43,6 +43,7 @@
     <PackageReference Include="CommandLineParser" Version="2.8.0" />
     <PackageReference Include="Microsoft.Extensions.Configuration.EnvironmentVariables" Version="3.1.6" />
     <PackageReference Include="Microsoft.Extensions.Configuration.Json" Version="3.1.6" />
+    <PackageReference Include="Microsoft.Extensions.Diagnostics.HealthChecks" Version="3.1.7" />
     <PackageReference Include="prometheus-net" Version="3.6.0" />
     <PackageReference Include="prometheus-net.AspNetCore" Version="3.6.0" />
     <PackageReference Include="Serilog.AspNetCore" Version="3.4.0" />
diff --git a/Jellyfin.Server/Startup.cs b/Jellyfin.Server/Startup.cs
index 995271aa3..08d59543a 100644
--- a/Jellyfin.Server/Startup.cs
+++ b/Jellyfin.Server/Startup.cs
@@ -1,11 +1,12 @@
 using System;
 using System.ComponentModel;
+using System.Data.OleDb;
 using System.Net.Http.Headers;
 using Jellyfin.Api.TypeConverters;
 using Jellyfin.Server.Extensions;
+using Jellyfin.Server.HealthChecks;
 using Jellyfin.Server.Middleware;
 using Jellyfin.Server.Models;
-using MediaBrowser.Common;
 using MediaBrowser.Common.Net;
 using MediaBrowser.Controller;
 using MediaBrowser.Controller.Configuration;
@@ -77,6 +78,9 @@ namespace Jellyfin.Server
                     c.DefaultRequestHeaders.UserAgent.Add(new ProductInfoHeaderValue($"({_serverApplicationHost.ApplicationUserAgentAddress})"));
                 })
                 .ConfigurePrimaryHttpMessageHandler(x => new DefaultHttpClientHandler());
+
+            services.AddHealthChecks()
+                .AddCheck<JellyfinDbHealthCheck>("JellyfinDb");
         }
 
         /// <summary>
@@ -132,6 +136,8 @@ namespace Jellyfin.Server
                 {
                     endpoints.MapMetrics(_serverConfigurationManager.Configuration.BaseUrl.TrimStart('/') + "/metrics");
                 }
+
+                endpoints.MapHealthChecks(_serverConfigurationManager.Configuration.BaseUrl.TrimStart('/') + "/health");
             });
 
             // Add type descriptor for legacy datetime parsing.