From 53aa45a2bb77c84c606cc890665794f39cd9c231 Mon Sep 17 00:00:00 2001
From: Chris Eager <chris@signal.org>
Date: Tue, 14 Sep 2021 15:51:20 -0700
Subject: [PATCH] Use `queryPaginator` when deleting messages

---
 .../textsecuregcm/storage/MessagesDynamoDb.java               | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/storage/MessagesDynamoDb.java b/service/src/main/java/org/whispersystems/textsecuregcm/storage/MessagesDynamoDb.java
index 07e593198..771cbc94a 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/storage/MessagesDynamoDb.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/storage/MessagesDynamoDb.java
@@ -160,7 +160,7 @@ public class MessagesDynamoDb extends AbstractDynamoDbStore {
   @Nonnull
   private Optional<OutgoingMessageEntity> deleteItemsMatchingQueryAndReturnFirstOneActuallyDeleted(AttributeValue partitionKey, QueryRequest queryRequest) {
     Optional<OutgoingMessageEntity> result = Optional.empty();
-    for (Map<String, AttributeValue> item : db().query(queryRequest).items()) {
+    for (Map<String, AttributeValue> item : db().queryPaginator(queryRequest).items()) {
       final byte[] rangeKeyValue = item.get(KEY_SORT).b().asByteArray();
       DeleteItemRequest.Builder deleteItemRequest = DeleteItemRequest.builder()
           .tableName(tableName)
@@ -225,7 +225,7 @@ public class MessagesDynamoDb extends AbstractDynamoDbStore {
   }
 
   private void deleteRowsMatchingQuery(AttributeValue partitionKey, QueryRequest querySpec) {
-    writeInBatches(db().query(querySpec).items(), (itemBatch) -> deleteItems(partitionKey, itemBatch));
+    writeInBatches(db().queryPaginator(querySpec).items(), itemBatch -> deleteItems(partitionKey, itemBatch));
   }
 
   private void deleteItems(AttributeValue partitionKey, List<Map<String, AttributeValue>> items) {