DocumentDb - Add option to check database and collection (#1985)

Xabaril · Aug 7, 2023 · 0f6818a · 0f6818a
1 parent 7492240
commit 0f6818a
Show file tree

Hide file tree

Showing 5 changed files with 49 additions and 9 deletions.
diff --git a/build/versions.props b/build/versions.props
@@ -17,7 +17,7 @@
     <HealthCheckConsul>7.0.0</HealthCheckConsul>
     <HealthCheckCosmosDb>7.0.0</HealthCheckCosmosDb>
     <HealthCheckDapr>7.0.0</HealthCheckDapr>
-    <HealthCheckDocumentDb>7.0.0</HealthCheckDocumentDb>
+    <HealthCheckDocumentDb>7.1.0</HealthCheckDocumentDb>
     <HealthCheckDynamoDb>7.0.0</HealthCheckDynamoDb>
     <HealthCheckElasticsearch>7.0.0</HealthCheckElasticsearch>
     <HealthCheckEventStoregRPC>7.0.0</HealthCheckEventStoregRPC>

diff --git a/src/HealthChecks.DocumentDb/DocumentDbHealthCheck.cs b/src/HealthChecks.DocumentDb/DocumentDbHealthCheck.cs
@@ -7,30 +7,40 @@ namespace HealthChecks.DocumentDb;
 public class DocumentDbHealthCheck : IHealthCheck
 {
     private static readonly ConcurrentDictionary<string, DocumentClient> _connections = new();
-    private readonly DocumentDbOptions _documentDbOptions = new();
+    private readonly DocumentDbOptions _options;
 
     public DocumentDbHealthCheck(DocumentDbOptions documentDbOptions)
     {
-        _documentDbOptions.UriEndpoint = Guard.ThrowIfNull(documentDbOptions.UriEndpoint);
-        _documentDbOptions.PrimaryKey = Guard.ThrowIfNull(documentDbOptions.PrimaryKey);
+        Guard.ThrowIfNull(documentDbOptions.UriEndpoint);
+        Guard.ThrowIfNull(documentDbOptions.PrimaryKey);
+
+        _options = documentDbOptions;
     }
 
     /// <inheritdoc />
     public async Task<HealthCheckResult> CheckHealthAsync(HealthCheckContext context, CancellationToken cancellationToken = default)
     {
         try
         {
-            if (!_connections.TryGetValue(_documentDbOptions.UriEndpoint, out var documentDbClient))
+            if (!_connections.TryGetValue(_options.UriEndpoint, out var documentDbClient))
             {
-                documentDbClient = new DocumentClient(new Uri(_documentDbOptions.UriEndpoint), _documentDbOptions.PrimaryKey);
+                documentDbClient = new DocumentClient(new Uri(_options.UriEndpoint), _options.PrimaryKey);
 
-                if (!_connections.TryAdd(_documentDbOptions.UriEndpoint, documentDbClient))
+                if (!_connections.TryAdd(_options.UriEndpoint, documentDbClient))
                 {
                     documentDbClient.Dispose();
-                    documentDbClient = _connections[_documentDbOptions.UriEndpoint];
+                    documentDbClient = _connections[_options.UriEndpoint];
                 }
             }
-            await documentDbClient.OpenAsync(cancellationToken).ConfigureAwait(false);
+
+            if (!string.IsNullOrWhiteSpace(_options.DatabaseName) && !string.IsNullOrWhiteSpace(_options.CollectionName))
+            {
+                await documentDbClient.ReadDocumentCollectionAsync(UriFactory.CreateDocumentCollectionUri(_options.DatabaseName, _options.CollectionName)).ConfigureAwait(false);
+            }
+            else
+            {
+                await documentDbClient.OpenAsync(cancellationToken).ConfigureAwait(false);
+            }
 
             return HealthCheckResult.Healthy();
         }

diff --git a/src/HealthChecks.DocumentDb/DocumentDbOptions.cs b/src/HealthChecks.DocumentDb/DocumentDbOptions.cs
@@ -8,4 +8,8 @@ public class DocumentDbOptions
     public string UriEndpoint { get; set; } = null!;
 
     public string PrimaryKey { get; set; } = null!;
+
+    public string? DatabaseName { get; set; }
+
+    public string? CollectionName { get; set; }
 }
diff --git a/test/HealthChecks.DocumentDb.Tests/DependencyInjection/RegistrationTests.cs b/test/HealthChecks.DocumentDb.Tests/DependencyInjection/RegistrationTests.cs
@@ -18,6 +18,7 @@ public void add_health_check_when_properly_configured()
         registration.Name.ShouldBe("documentdb");
         check.ShouldBeOfType<DocumentDbHealthCheck>();
     }
+
     [Fact]
     public void add_named_health_check_when_properly_configured()
     {
@@ -34,4 +35,27 @@ public void add_named_health_check_when_properly_configured()
         registration.Name.ShouldBe("my-documentdb-group");
         check.ShouldBeOfType<DocumentDbHealthCheck>();
     }
+
+    [Fact]
+    public void add_health_check_when_properly_configured_with_database_name_and_collection_name()
+    {
+        var services = new ServiceCollection();
+        services.AddHealthChecks()
+            .AddDocumentDb(_ =>
+            {
+                _.PrimaryKey = "key";
+                _.UriEndpoint = "endpoint";
+                _.DatabaseName = "database";
+                _.CollectionName = "collection";
+            });
+
+        using var serviceProvider = services.BuildServiceProvider();
+        var options = serviceProvider.GetRequiredService<IOptions<HealthCheckServiceOptions>>();
+
+        var registration = options.Value.Registrations.First();
+        var check = registration.Factory(serviceProvider);
+
+        registration.Name.ShouldBe("documentdb");
+        check.ShouldBeOfType<DocumentDbHealthCheck>();
+    }
 }
diff --git a/test/HealthChecks.DocumentDb.Tests/HealthChecks.DocumentDb.approved.txt b/test/HealthChecks.DocumentDb.Tests/HealthChecks.DocumentDb.approved.txt
@@ -8,6 +8,8 @@ namespace HealthChecks.DocumentDb
     public class DocumentDbOptions
     {
         public DocumentDbOptions() { }
+        public string? CollectionName { get; set; }
+        public string? DatabaseName { get; set; }
         public string PrimaryKey { get; set; }
         public string UriEndpoint { get; set; }
     }