Add a tag->tag-set migration command

service/src/test/java/org/whispersystems/textsecuregcm/storage/IssuedReceiptsManagerTest.java

@@ -17,6 +17,8 @@ import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.CompletableFuture;
 import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+import java.util.stream.Stream;
 import org.assertj.core.api.Condition;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
@@ -26,11 +28,13 @@ import org.whispersystems.textsecuregcm.storage.DynamoDbExtensionSchema.Tables;
 import org.whispersystems.textsecuregcm.subscriptions.PaymentProvider;
 import org.whispersystems.textsecuregcm.util.AttributeValues;
 import org.whispersystems.textsecuregcm.util.TestRandomUtil;
+import reactor.core.scheduler.Schedulers;
 import software.amazon.awssdk.core.SdkBytes;
 import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
 import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
 import software.amazon.awssdk.services.dynamodb.model.GetItemRequest;
 import software.amazon.awssdk.services.dynamodb.model.GetItemResponse;
+import software.amazon.awssdk.services.dynamodb.model.UpdateItemRequest;
 
 class IssuedReceiptsManagerTest {
 
@@ -89,6 +93,74 @@ class IssuedReceiptsManagerTest {
     assertThat(future).succeedsWithin(Duration.ofSeconds(3));
   }
 
+  @Test
+  void testMigrateToTagSet() {
+    Instant now = Instant.ofEpochSecond(NOW_EPOCH_SECONDS);
+
+    issuedReceiptsManager
+        .recordIssuance("itemId", PaymentProvider.STRIPE, randomReceiptCredentialRequest(), now)
+        .join();
+    removeTagSet("itemId");
+
+    assertThat(getItem("itemId").item()).doesNotContainKey(IssuedReceiptsManager.KEY_ISSUED_RECEIPT_TAG_SET);
+
+    final IssuedReceiptsManager.IssuedReceipt issuedReceipt = issuedReceiptsManager
+        .receiptsWithoutTagSet(1, Schedulers.immediate())
+        .blockFirst();
+
+    issuedReceiptsManager.migrateToTagSet(issuedReceipt).join();
+
+    final Map<String, AttributeValue> item = getItem("itemId").item();
+    assertThat(item)
+        .containsKey(IssuedReceiptsManager.KEY_ISSUED_RECEIPT_TAG_SET)
+        .containsKey(IssuedReceiptsManager.KEY_ISSUED_RECEIPT_TAG);
+
+    final List<byte[]> tags = item
+        .get(IssuedReceiptsManager.KEY_ISSUED_RECEIPT_TAG_SET).bs()
+        .stream()
+        .map(SdkBytes::asByteArray)
+        .toList();
+    assertThat(tags).hasSize(1);
+
+    final byte[] tag = item.get(IssuedReceiptsManager.KEY_ISSUED_RECEIPT_TAG).b().asByteArray();
+    assertThat(tags).first().isEqualTo(tag);
+  }
+
+
+  @Test
+  void testReceiptsWithoutTagSet() {
+    Instant now = Instant.ofEpochSecond(NOW_EPOCH_SECONDS);
+
+    final int numItems = 100;
+    final List<String> expectedNoTagSet = IntStream.range(0, numItems)
+        .boxed()
+        .flatMap(i -> {
+          final String itemId = "item-%s".formatted(i);
+          issuedReceiptsManager.recordIssuance(itemId, PaymentProvider.STRIPE, randomReceiptCredentialRequest(), now).join();
+
+          if (i % 2 == 0) {
+            removeTagSet(itemId);
+            return Stream.of(itemId);
+          } else {
+            return Stream.empty();
+          }
+        }).toList();
+    final List<String> items = issuedReceiptsManager
+        .receiptsWithoutTagSet(1, Schedulers.immediate())
+        .map(IssuedReceiptsManager.IssuedReceipt::itemId)
+        .collectList().block();
+    assertThat(items).hasSize(numItems / 2);
+    assertThat(items).containsExactlyInAnyOrderElementsOf(expectedNoTagSet);
+  }
+
+  @Test
+  void testMigrateAfterRecordExpires() {
+    final IssuedReceiptsManager.IssuedReceipt issued = new IssuedReceiptsManager.IssuedReceipt("itemId",
+        TestRandomUtil.nextBytes(32));
+    // We should succeed but do nothing if the item is deleted by the time we try to migrate it
+    issuedReceiptsManager.migrateToTagSet(issued).join();
+    assertThat(getItem("itemId").hasItem()).isFalse();
+  }
 
   private GetItemResponse getItem(final String itemId) {
     final DynamoDbClient client = DYNAMO_DB_EXTENSION.getDynamoDbClient();
@@ -104,4 +176,15 @@ class IssuedReceiptsManagerTest {
     when(request.serialize()).thenReturn(bytes);
     return request;
   }
+
+  private void removeTagSet(final String itemId) {
+    final DynamoDbClient client = DYNAMO_DB_EXTENSION.getDynamoDbClient();
+    // Simulate an entry that was written before we wrote the tag set field
+    client.updateItem(UpdateItemRequest.builder()
+        .tableName(Tables.ISSUED_RECEIPTS.tableName())
+        .key(Map.of(IssuedReceiptsManager.KEY_PROCESSOR_ITEM_ID, AttributeValues.s(itemId)))
+        .updateExpression("REMOVE #tags")
+        .expressionAttributeNames(Map.of("#tags", IssuedReceiptsManager.KEY_ISSUED_RECEIPT_TAG_SET))
+        .build());
+  }
 }

service/src/test/java/org/whispersystems/textsecuregcm/workers/FinishPushNotificationExperimentCommandTest.java

@@ -81,6 +81,7 @@ class FinishPushNotificationExperimentCommandTest {
         null,
         null,
         null,
+        null,
         null);
 
     //noinspection unchecked

service/src/test/java/org/whispersystems/textsecuregcm/workers/NotifyIdleDevicesCommandTest.java

@@ -68,6 +68,7 @@ class NotifyIdleDevicesCommandTest {
           null,
           null,
           null,
+          null,
           null);
 
       this.idleDeviceNotificationScheduler = idleDeviceNotificationScheduler;

service/src/test/java/org/whispersystems/textsecuregcm/workers/StartPushNotificationExperimentCommandTest.java

@@ -70,6 +70,7 @@ class StartPushNotificationExperimentCommandTest {
           null,
           null,
           null,
+          null,
           null);
     }