diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/storage/Accounts.java b/service/src/main/java/org/whispersystems/textsecuregcm/storage/Accounts.java index a607392be..3dd61cf03 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/storage/Accounts.java +++ b/service/src/main/java/org/whispersystems/textsecuregcm/storage/Accounts.java @@ -23,6 +23,7 @@ import java.util.Map; import java.util.Objects; import java.util.Optional; import java.util.UUID; +import java.util.concurrent.CompletableFuture; import java.util.concurrent.CompletionException; import java.util.concurrent.CompletionStage; import java.util.concurrent.TimeUnit; @@ -631,12 +632,23 @@ public class Accounts extends AbstractDynamoDbStore { GET_BY_NUMBER_TIMER, phoneNumberConstraintTableName, ATTR_ACCOUNT_E164, AttributeValues.fromString(number)); } + @Nonnull + public CompletableFuture> getByE164Async(final String number) { + return getByIndirectLookupAsync( + GET_BY_NUMBER_TIMER, phoneNumberConstraintTableName, ATTR_ACCOUNT_E164, AttributeValues.fromString(number)); + } + @Nonnull public Optional getByPhoneNumberIdentifier(final UUID phoneNumberIdentifier) { return getByIndirectLookup( GET_BY_PNI_TIMER, phoneNumberIdentifierConstraintTableName, ATTR_PNI_UUID, AttributeValues.fromUUID(phoneNumberIdentifier)); } + @Nonnull + public CompletableFuture> getByPhoneNumberIdentifierAsync(final UUID phoneNumberIdentifier) { + return getByIndirectLookupAsync(GET_BY_PNI_TIMER, phoneNumberIdentifierConstraintTableName, ATTR_PNI_UUID, AttributeValues.fromUUID(phoneNumberIdentifier)); + } + @Nonnull public Optional getByUsernameHash(final byte[] usernameHash) { return getByIndirectLookup( @@ -662,6 +674,13 @@ public class Accounts extends AbstractDynamoDbStore { .map(Accounts::fromItem))); } + @Nonnull + public CompletableFuture> getByAccountIdentifierAsync(final UUID uuid) { + return record(GET_BY_UUID_TIMER, () -> itemByKeyAsync(accountsTableName, KEY_ACCOUNT_UUID, AttributeValues.fromUUID(uuid)) + .thenApply(maybeItem -> maybeItem.map(Accounts::fromItem))) + .toCompletableFuture(); + } + public void delete(final UUID uuid) { DELETE_TIMER.record(() -> getByAccountIdentifier(uuid).ifPresent(account -> { @@ -724,6 +743,16 @@ public class Accounts extends AbstractDynamoDbStore { return getByIndirectLookup(timer, tableName, keyName, keyValue, i -> true); } + @Nonnull + private CompletableFuture> getByIndirectLookupAsync( + final Timer timer, + final String tableName, + final String keyName, + final AttributeValue keyValue) { + + return getByIndirectLookupAsync(timer, tableName, keyName, keyValue, i -> true); + } + @Nonnull private Optional getByIndirectLookup( final Timer timer, @@ -739,6 +768,24 @@ public class Accounts extends AbstractDynamoDbStore { .map(Accounts::fromItem))); } + @Nonnull + private CompletableFuture> getByIndirectLookupAsync( + final Timer timer, + final String tableName, + final String keyName, + final AttributeValue keyValue, + final Predicate> predicate) { + + return record(timer, () -> itemByKeyAsync(tableName, keyName, keyValue) + .thenCompose(maybeItem -> maybeItem + .filter(predicate) + .map(item -> item.get(KEY_ACCOUNT_UUID)) + .map(uuid -> itemByKeyAsync(accountsTableName, KEY_ACCOUNT_UUID, uuid) + .thenApply(maybeAccountItem -> maybeAccountItem.map(Accounts::fromItem))) + .orElse(CompletableFuture.completedFuture(Optional.empty())))) + .toCompletableFuture(); + } + @Nonnull private Optional> itemByKey(final String table, final String keyName, final AttributeValue keyValue) { final GetItemResponse response = db().getItem(GetItemRequest.builder() @@ -749,6 +796,16 @@ public class Accounts extends AbstractDynamoDbStore { return Optional.ofNullable(response.item()).filter(m -> !m.isEmpty()); } + @Nonnull + private CompletableFuture>> itemByKeyAsync(final String table, final String keyName, final AttributeValue keyValue) { + return asyncClient.getItem(GetItemRequest.builder() + .tableName(table) + .key(Map.of(keyName, keyValue)) + .consistentRead(true) + .build()) + .thenApply(response -> Optional.ofNullable(response.item()).filter(item -> !item.isEmpty())); + } + @Nonnull private Optional> itemByGsiKey(final String table, final String indexName, final String keyName, final AttributeValue keyValue) { final QueryResponse response = db().query(QueryRequest.builder() diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsTest.java index c7297e0d3..83e396ab2 100644 --- a/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsTest.java +++ b/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsTest.java @@ -41,6 +41,7 @@ import java.util.stream.Collectors; import org.apache.commons.lang3.RandomUtils; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; import org.junit.jupiter.api.extension.RegisterExtension; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.ValueSource; @@ -73,6 +74,7 @@ import software.amazon.awssdk.services.dynamodb.model.TransactionCanceledExcepti import software.amazon.awssdk.services.dynamodb.model.TransactionConflictException; import software.amazon.awssdk.services.dynamodb.model.UpdateItemRequest; +@Timeout(value = 10, threadMode = Timeout.ThreadMode.SEPARATE_THREAD) class AccountsTest { private static final String BASE_64_URL_USERNAME_HASH_1 = "9p6Tip7BFefFOJzv4kv4GyXEYsBVfk_WbjNejdlOvQE"; @@ -573,6 +575,44 @@ class AccountsTest { assertThat(retrieved.isPresent()).isFalse(); } + @Test + void getByAccountIdentifierAsync() { + assertThat(accounts.getByAccountIdentifierAsync(UUID.randomUUID()).join()).isEmpty(); + + final Account account = + generateAccount("+14151112222", UUID.randomUUID(), UUID.randomUUID(), List.of(generateDevice(1))); + + accounts.create(account); + + assertThat(accounts.getByAccountIdentifierAsync(account.getUuid()).join()).isPresent(); + } + + @Test + void getByPhoneNumberIdentifierAsync() { + assertThat(accounts.getByPhoneNumberIdentifierAsync(UUID.randomUUID()).join()).isEmpty(); + + final Account account = + generateAccount("+14151112222", UUID.randomUUID(), UUID.randomUUID(), List.of(generateDevice(1))); + + accounts.create(account); + + assertThat(accounts.getByPhoneNumberIdentifierAsync(account.getPhoneNumberIdentifier()).join()).isPresent(); + } + + @Test + void getByE164Async() { + final String e164 = "+14151112222"; + + assertThat(accounts.getByE164Async(e164).join()).isEmpty(); + + final Account account = + generateAccount(e164, UUID.randomUUID(), UUID.randomUUID(), List.of(generateDevice(1))); + + accounts.create(account); + + assertThat(accounts.getByE164Async(e164).join()).isPresent(); + } + @Test void testCanonicallyDiscoverableSet() { Device device = generateDevice(1);