From 81a21ef9e520817add4c3e77113e481faac8eb98 Mon Sep 17 00:00:00 2001
From: Ivan Maximov <sungam3r@yandex.ru>
Date: Sat, 22 Jul 2023 09:58:16 +0300
Subject: [PATCH] Add DynamoDBOptions.Limit and
 DynamoDBOptions.LastEvaluatedTableName

---
 .../DynamoDbHealthCheckBuilderExtensions.cs       |  2 +-
 src/HealthChecks.DynamoDb/DynamoDbHealthCheck.cs  | 15 ++++++++++++++-
 src/HealthChecks.DynamoDb/DynamoDbOptions.cs      | 10 ++++++++++
 .../HealthChecks.DynamoDb.approved.txt            |  2 ++
 4 files changed, 27 insertions(+), 2 deletions(-)

diff --git a/src/HealthChecks.DynamoDb/DependencyInjection/DynamoDbHealthCheckBuilderExtensions.cs b/src/HealthChecks.DynamoDb/DependencyInjection/DynamoDbHealthCheckBuilderExtensions.cs
index e49c4d5dfa..7b3cf30caa 100644
--- a/src/HealthChecks.DynamoDb/DependencyInjection/DynamoDbHealthCheckBuilderExtensions.cs
+++ b/src/HealthChecks.DynamoDb/DependencyInjection/DynamoDbHealthCheckBuilderExtensions.cs
@@ -36,7 +36,7 @@ public static IHealthChecksBuilder AddDynamoDb(
 
         return builder.Add(new HealthCheckRegistration(
             name ?? NAME,
-            sp => new DynamoDbHealthCheck(options),
+            _ => new DynamoDbHealthCheck(options),
             failureStatus,
             tags,
             timeout));
diff --git a/src/HealthChecks.DynamoDb/DynamoDbHealthCheck.cs b/src/HealthChecks.DynamoDb/DynamoDbHealthCheck.cs
index c710369eec..f6dd258d3f 100644
--- a/src/HealthChecks.DynamoDb/DynamoDbHealthCheck.cs
+++ b/src/HealthChecks.DynamoDb/DynamoDbHealthCheck.cs
@@ -1,13 +1,22 @@
+using System.Collections.Generic;
 using Amazon.DynamoDBv2;
+using Amazon.DynamoDBv2.Model;
 using Amazon.Runtime;
 using Microsoft.Extensions.Diagnostics.HealthChecks;
 
 namespace HealthChecks.DynamoDb;
 
+/// <summary>
+/// Health check for AWS DynamoDb database.
+/// </summary>
 public class DynamoDbHealthCheck : IHealthCheck
 {
     private readonly DynamoDBOptions _options;
 
+    /// <summary>
+    /// Creates health check for AWS DynamoDb database with the specified options.
+    /// </summary>
+    /// <param name="options"></param>
     public DynamoDbHealthCheck(DynamoDBOptions options)
     {
         _options = Guard.ThrowIfNull(options);
@@ -36,7 +45,11 @@ public async Task<HealthCheckResult> CheckHealthAsync(HealthCheckContext context
                 ? CreateClientWithCredentials(credentials)
                 : CreateClientWithoutCredentials();
 
-            _ = await client.ListTablesAsync(cancellationToken).ConfigureAwait(false);
+            var request = new ListTablesRequest { ExclusiveStartTableName = _options.LastEvaluatedTableName };
+            if (_options.Limit != null)
+                request.Limit = _options.Limit.Value;
+
+            var response = await client.ListTablesAsync(request, cancellationToken).ConfigureAwait(false);
 
             return HealthCheckResult.Healthy();
         }
diff --git a/src/HealthChecks.DynamoDb/DynamoDbOptions.cs b/src/HealthChecks.DynamoDb/DynamoDbOptions.cs
index 331d46d0cd..9770a6d02e 100644
--- a/src/HealthChecks.DynamoDb/DynamoDbOptions.cs
+++ b/src/HealthChecks.DynamoDb/DynamoDbOptions.cs
@@ -17,4 +17,14 @@ public class DynamoDBOptions
     public string? SecretKey { get; set; }
 
     public RegionEndpoint RegionEndpoint { get; set; } = null!;
+
+    /// <summary>
+    /// A maximum number of table names to read.
+    /// </summary>
+    public int? Limit { get; set; }
+
+    /// <summary>
+    /// The first table name to read.
+    /// </summary>
+    public string? LastEvaluatedTableName { get; set; }
 }
diff --git a/test/HealthChecks.DynamoDb.Tests/HealthChecks.DynamoDb.approved.txt b/test/HealthChecks.DynamoDb.Tests/HealthChecks.DynamoDb.approved.txt
index 0cc0758fab..38c1e4a4bd 100644
--- a/test/HealthChecks.DynamoDb.Tests/HealthChecks.DynamoDb.approved.txt
+++ b/test/HealthChecks.DynamoDb.Tests/HealthChecks.DynamoDb.approved.txt
@@ -7,6 +7,8 @@ namespace HealthChecks.DynamoDb
             "d")]
         public string? AccessKey { get; set; }
         public Amazon.Runtime.AWSCredentials? Credentials { get; set; }
+        public string? LastEvaluatedTableName { get; set; }
+        public int? Limit { get; set; }
         public Amazon.RegionEndpoint RegionEndpoint { get; set; }
         [System.Obsolete("Specify access key and secret as a BasicCredential to Credentials property instea" +
             "d")]