From ef02ea40be6e9d02f4853bb4d91d4ae28bfb9112 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Marek=20Smoli=C5=84ski?= <marek@goteleport.com>
Date: Wed, 28 Feb 2024 10:41:11 -0800
Subject: [PATCH] Fix DynamoDB BatchWriteItemInput items limit

---
 lib/backend/dynamo/dynamodbbk.go | 10 +++++++++-
 lib/backend/test/suite.go        | 10 ++++++++++
 2 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/lib/backend/dynamo/dynamodbbk.go b/lib/backend/dynamo/dynamodbbk.go
index a06773e9d6b36..f8d2aa3742b75 100644
--- a/lib/backend/dynamo/dynamodbbk.go
+++ b/lib/backend/dynamo/dynamodbbk.go
@@ -435,6 +435,14 @@ func (b *Backend) getAllRecords(ctx context.Context, startKey []byte, endKey []b
 	return nil, trace.BadParameter("backend entered endless loop")
 }
 
+const (
+	// batchOperationItemsLimit is the maximum number of items that can be put or deleted in a single batch operation.
+	// From https://docs.aws.amazon.com/amazondynamodb/latest/developerguide/Limits.html:
+	// A single call to BatchWriteItem can transmit up to 16MB of data over the network,
+	// consisting of up to 25 item put or delete operations.
+	batchOperationItemsLimit = 25
+)
+
 // DeleteRange deletes range of items with keys between startKey and endKey
 func (b *Backend) DeleteRange(ctx context.Context, startKey, endKey []byte) error {
 	if len(startKey) == 0 {
@@ -447,7 +455,7 @@ func (b *Backend) DeleteRange(ctx context.Context, startKey, endKey []byte) erro
 	// keep the very large limit, just in case if someone else
 	// keeps adding records
 	for i := 0; i < backend.DefaultRangeLimit/100; i++ {
-		result, err := b.getRecords(ctx, prependPrefix(startKey), prependPrefix(endKey), 100, nil)
+		result, err := b.getRecords(ctx, prependPrefix(startKey), prependPrefix(endKey), batchOperationItemsLimit, nil)
 		if err != nil {
 			return trace.Wrap(err)
 		}
diff --git a/lib/backend/test/suite.go b/lib/backend/test/suite.go
index 8de056200f810..b9265eb21b6e2 100644
--- a/lib/backend/test/suite.go
+++ b/lib/backend/test/suite.go
@@ -353,6 +353,16 @@ func testDeleteRange(t *testing.T, newBackend Constructor) {
 		require.NoError(t, err, "Failed creating value: %q => %q", item.Key, item.Value)
 	}
 
+	// Some Backends (e.g. DynamoDB) have a limit on the number of items that can
+	// be deleted in a single operation. This test is designed to be run with
+	// a backend that has a limit of 25 items per delete operation.
+	for i := 0; i < 100; i++ {
+		item := &backend.Item{Key: prefix(fmt.Sprintf("/prefix/c/cn%d", i)), Value: []byte(fmt.Sprintf("val cn%d", i))}
+		lease, err := uut.Create(ctx, *item)
+		require.NoError(t, err, "Failed creating value: %q => %q", item.Key, item.Value)
+		item.Revision = lease.Revision
+	}
+
 	err = uut.DeleteRange(ctx, prefix("/prefix/c"), backend.RangeEnd(prefix("/prefix/c")))
 	require.NoError(t, err)