diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java b/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java index d4a1f17dc..c56540f3d 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java +++ b/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java @@ -397,6 +397,7 @@ public class WhisperServerService extends Application { accountToCreate.setVersion(accountToCreate.getVersion() + 1); return (Void) null; @@ -478,7 +463,7 @@ public class Accounts extends AbstractDynamoDbStore { .transactItems(writeItems) .build(); - db().transactWriteItems(request); + dynamoDbClient.transactWriteItems(request); account.setVersion(account.getVersion() + 1); succeeded = true; @@ -603,7 +588,7 @@ public class Accounts extends AbstractDynamoDbStore { writeItems.add(UpdateAccountSpec.forAccount(accountsTableName, updatedAccount).transactItem()); - return asyncClient + return dynamoDbAsyncClient .transactWriteItems(TransactWriteItemsRequest.builder().transactItems(writeItems).build()) .thenRun(Util.NOOP) .exceptionally(ExceptionUtils.exceptionallyHandler(TransactionCanceledException.class, e -> { @@ -788,7 +773,7 @@ public class Accounts extends AbstractDynamoDbStore { holdToRemove.ifPresent(oldHold -> writeItems.add(releaseHoldIfAllowedTransactItem(updatedAccount.getUuid(), oldHold, now))); - return asyncClient.transactWriteItems(TransactWriteItemsRequest.builder().transactItems(writeItems).build()) + return dynamoDbAsyncClient.transactWriteItems(TransactWriteItemsRequest.builder().transactItems(writeItems).build()) .thenApply(ignored -> updatedAccount); }) .thenApply(updatedAccount -> { @@ -829,7 +814,7 @@ public class Accounts extends AbstractDynamoDbStore { // Otherwise, there's an existing link handle. If this is the result of an account being re-registered, we should // preserve the link handle. - return asyncClient.getItem(GetItemRequest.builder() + return dynamoDbAsyncClient.getItem(GetItemRequest.builder() .tableName(usernamesConstraintTableName) .key(Map.of(UsernameTable.KEY_USERNAME_HASH, AttributeValues.b(usernameHash))) .projectionExpression(UsernameTable.ATTR_RECLAIMABLE).build()) @@ -878,7 +863,7 @@ public class Accounts extends AbstractDynamoDbStore { // 2?: Adding that hold may have caused our account to exceed our maximum holds. Release an old hold holdToRemove.ifPresent(oldHold -> items.add(releaseHoldIfAllowedTransactItem(updatedAccount.getUuid(), oldHold, now))); - return asyncClient.transactWriteItems(TransactWriteItemsRequest.builder().transactItems(items).build()) + return dynamoDbAsyncClient.transactWriteItems(TransactWriteItemsRequest.builder().transactItems(items).build()) .thenAccept(ignored -> { account.setUsernameHash(null); account.setUsernameLinkDetails(null, null); @@ -1009,7 +994,7 @@ public class Accounts extends AbstractDynamoDbStore { .forAccount(accountsTableName, account) .updateItemRequest(); - return asyncClient.updateItem(updateItemRequest) + return dynamoDbAsyncClient.updateItem(updateItemRequest) .thenApply(response -> { account.setVersion(AttributeValues.getInt(response.attributes(), "V", account.getVersion() + 1)); return (Void) null; @@ -1059,7 +1044,7 @@ public class Accounts extends AbstractDynamoDbStore { writeItems.add(UpdateAccountSpec.forAccount(accountsTableName, account).transactItem()); writeItems.addAll(additionalWriteItems); - return asyncClient.transactWriteItems(TransactWriteItemsRequest.builder() + return dynamoDbAsyncClient.transactWriteItems(TransactWriteItemsRequest.builder() .transactItems(writeItems) .build()) .thenApply(response -> { @@ -1187,7 +1172,7 @@ public class Accounts extends AbstractDynamoDbStore { } public Optional findRecentlyDeletedAccountIdentifier(final UUID phoneNumberIdentifier) { - final GetItemResponse response = db().getItem(GetItemRequest.builder() + final GetItemResponse response = dynamoDbClient.getItem(GetItemRequest.builder() .tableName(deletedAccountsTableName) .consistentRead(true) .key(Map.of(DELETED_ACCOUNTS_KEY_ACCOUNT_PNI, AttributeValues.fromString(phoneNumberIdentifier.toString()))) @@ -1197,7 +1182,7 @@ public class Accounts extends AbstractDynamoDbStore { } public Optional findRecentlyDeletedPhoneNumberIdentifier(final UUID uuid) { - final QueryResponse response = db().query(QueryRequest.builder() + final QueryResponse response = dynamoDbClient.query(QueryRequest.builder() .tableName(deletedAccountsTableName) .indexName(DELETED_ACCOUNTS_UUID_TO_PNI_INDEX_NAME) .keyConditionExpression("#uuid = :uuid") @@ -1234,7 +1219,7 @@ public class Accounts extends AbstractDynamoDbStore { transactWriteItems.addAll(additionalWriteItems); - return asyncClient.transactWriteItems(TransactWriteItemsRequest.builder() + return dynamoDbAsyncClient.transactWriteItems(TransactWriteItemsRequest.builder() .transactItems(transactWriteItems) .build()) .thenRun(Util.NOOP); @@ -1251,7 +1236,7 @@ public class Accounts extends AbstractDynamoDbStore { return Flux.range(0, segments) .parallel() .runOn(scheduler) - .flatMap(segment -> asyncClient.scanPaginator(ScanRequest.builder() + .flatMap(segment -> dynamoDbAsyncClient.scanPaginator(ScanRequest.builder() .tableName(accountsTableName) .consistentRead(true) .segment(segment) @@ -1278,7 +1263,7 @@ public class Accounts extends AbstractDynamoDbStore { return Flux.range(0, segments) .parallel() .runOn(scheduler) - .flatMap(segment -> asyncClient.scanPaginator(ScanRequest.builder() + .flatMap(segment -> dynamoDbAsyncClient.scanPaginator(ScanRequest.builder() .tableName(deletedAccountsTableName) .consistentRead(true) .segment(segment) @@ -1346,7 +1331,7 @@ public class Accounts extends AbstractDynamoDbStore { @Nonnull private Optional> itemByKey(final String table, final String keyName, final AttributeValue keyValue) { - final GetItemResponse response = db().getItem(GetItemRequest.builder() + final GetItemResponse response = dynamoDbClient.getItem(GetItemRequest.builder() .tableName(table) .key(Map.of(keyName, keyValue)) .consistentRead(true) @@ -1356,7 +1341,7 @@ public class Accounts extends AbstractDynamoDbStore { @Nonnull private CompletableFuture>> itemByKeyAsync(final String table, final String keyName, final AttributeValue keyValue) { - return asyncClient.getItem(GetItemRequest.builder() + return dynamoDbAsyncClient.getItem(GetItemRequest.builder() .tableName(table) .key(Map.of(keyName, keyValue)) .consistentRead(true) @@ -1364,36 +1349,9 @@ public class Accounts extends AbstractDynamoDbStore { .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() - .tableName(table) - .indexName(indexName) - .keyConditionExpression("#gsiKey = :gsiValue") - .projectionExpression("#uuid") - .expressionAttributeNames(Map.of( - "#gsiKey", keyName, - "#uuid", KEY_ACCOUNT_UUID)) - .expressionAttributeValues(Map.of( - ":gsiValue", keyValue)) - .build()); - - if (response.count() == 0) { - return Optional.empty(); - } - - if (response.count() > 1) { - throw new IllegalStateException("More than one row located for GSI [%s], key-value pair [%s, %s]" - .formatted(indexName, keyName, keyValue)); - } - - final AttributeValue primaryKeyValue = response.items().get(0).get(KEY_ACCOUNT_UUID); - return itemByKey(table, KEY_ACCOUNT_UUID, primaryKeyValue); - } - @Nonnull private CompletableFuture>> itemByGsiKeyAsync(final String table, final String indexName, final String keyName, final AttributeValue keyValue) { - return asyncClient.query(QueryRequest.builder() + return dynamoDbAsyncClient.query(QueryRequest.builder() .tableName(table) .indexName(indexName) .keyConditionExpression("#gsiKey = :gsiValue") diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/workers/CommandDependencies.java b/service/src/main/java/org/whispersystems/textsecuregcm/workers/CommandDependencies.java index 10dcc0388..a05839ce3 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/workers/CommandDependencies.java +++ b/service/src/main/java/org/whispersystems/textsecuregcm/workers/CommandDependencies.java @@ -175,6 +175,7 @@ record CommandDependencies( new ClientPublicKeys(dynamoDbAsyncClient, configuration.getDynamoDbTables().getClientPublicKeys().getTableName()); Accounts accounts = new Accounts( + clock, dynamoDbClient, dynamoDbAsyncClient, configuration.getDynamoDbTables().getAccounts().getTableName(), diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountCreationDeletionIntegrationTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountCreationDeletionIntegrationTest.java index e8fe835e0..57ca4652a 100644 --- a/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountCreationDeletionIntegrationTest.java +++ b/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountCreationDeletionIntegrationTest.java @@ -102,6 +102,7 @@ public class AccountCreationDeletionIntegrationTest { DynamoDbExtensionSchema.Tables.CLIENT_PUBLIC_KEYS.tableName()); final Accounts accounts = new Accounts( + CLOCK, DYNAMO_DB_EXTENSION.getDynamoDbClient(), DYNAMO_DB_EXTENSION.getDynamoDbAsyncClient(), DynamoDbExtensionSchema.Tables.ACCOUNTS.tableName(), diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsManagerChangeNumberIntegrationTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsManagerChangeNumberIntegrationTest.java index ef1c2679f..20206f18c 100644 --- a/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsManagerChangeNumberIntegrationTest.java +++ b/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsManagerChangeNumberIntegrationTest.java @@ -93,6 +93,7 @@ class AccountsManagerChangeNumberIntegrationTest { DynamoDbExtensionSchema.Tables.CLIENT_PUBLIC_KEYS.tableName()); final Accounts accounts = new Accounts( + Clock.systemUTC(), DYNAMO_DB_EXTENSION.getDynamoDbClient(), DYNAMO_DB_EXTENSION.getDynamoDbAsyncClient(), Tables.ACCOUNTS.tableName(), diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsManagerConcurrentModificationIntegrationTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsManagerConcurrentModificationIntegrationTest.java index 3968d53ff..b363cdd38 100644 --- a/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsManagerConcurrentModificationIntegrationTest.java +++ b/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsManagerConcurrentModificationIntegrationTest.java @@ -91,6 +91,7 @@ class AccountsManagerConcurrentModificationIntegrationTest { when(dynamicConfigurationManager.getConfiguration()).thenReturn(new DynamicConfiguration()); accounts = new Accounts( + Clock.systemUTC(), DYNAMO_DB_EXTENSION.getDynamoDbClient(), DYNAMO_DB_EXTENSION.getDynamoDbAsyncClient(), Tables.ACCOUNTS.tableName(), diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsManagerUsernameIntegrationTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsManagerUsernameIntegrationTest.java index 0959cc130..523015337 100644 --- a/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsManagerUsernameIntegrationTest.java +++ b/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsManagerUsernameIntegrationTest.java @@ -102,6 +102,7 @@ class AccountsManagerUsernameIntegrationTest { ); accounts = Mockito.spy(new Accounts( + Clock.systemUTC(), DYNAMO_DB_EXTENSION.getDynamoDbClient(), DYNAMO_DB_EXTENSION.getDynamoDbAsyncClient(), Tables.ACCOUNTS.tableName(), 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 98ab1fdec..0ae814552 100644 --- a/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsTest.java +++ b/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsTest.java @@ -600,7 +600,9 @@ class AccountsTest { void testUpdateWithMockTransactionConflictException(boolean wrapException) { final DynamoDbAsyncClient dynamoDbAsyncClient = mock(DynamoDbAsyncClient.class); - accounts = new Accounts(mock(DynamoDbClient.class), + accounts = new Accounts( + clock, + mock(DynamoDbClient.class), dynamoDbAsyncClient, Tables.ACCOUNTS.tableName(), Tables.NUMBERS.tableName(), @@ -689,7 +691,9 @@ class AccountsTest { void testUpdateTransactionallyWithMockTransactionConflictException() { final DynamoDbAsyncClient dynamoDbAsyncClient = mock(DynamoDbAsyncClient.class); - accounts = new Accounts(mock(DynamoDbClient.class), + accounts = new Accounts( + clock, + mock(DynamoDbClient.class), dynamoDbAsyncClient, Tables.ACCOUNTS.tableName(), Tables.NUMBERS.tableName(), @@ -1081,7 +1085,9 @@ class AccountsTest { final Class expectedException) { final DynamoDbAsyncClient dbAsyncClient = mock(DynamoDbAsyncClient.class); - accounts = new Accounts(mock(DynamoDbClient.class), + accounts = new Accounts( + clock, + mock(DynamoDbClient.class), dbAsyncClient, Tables.ACCOUNTS.tableName(), Tables.NUMBERS.tableName(), @@ -1124,7 +1130,9 @@ class AccountsTest { final Class expectedException) { final DynamoDbAsyncClient dbAsyncClient = mock(DynamoDbAsyncClient.class); - accounts = new Accounts(mock(DynamoDbClient.class), + accounts = new Accounts( + clock, + mock(DynamoDbClient.class), dbAsyncClient, Tables.ACCOUNTS.tableName(), Tables.NUMBERS.tableName(), @@ -1225,7 +1233,9 @@ class AccountsTest { final Optional accountsCancellationString) { final DynamoDbAsyncClient dbAsyncClient = mock(DynamoDbAsyncClient.class); - accounts = new Accounts(mock(DynamoDbClient.class), + accounts = new Accounts( + clock, + mock(DynamoDbClient.class), dbAsyncClient, Tables.ACCOUNTS.tableName(), Tables.NUMBERS.tableName(), diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/storage/AddRemoveDeviceIntegrationTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/storage/AddRemoveDeviceIntegrationTest.java index aad025cb5..cf64e9904 100644 --- a/service/src/test/java/org/whispersystems/textsecuregcm/storage/AddRemoveDeviceIntegrationTest.java +++ b/service/src/test/java/org/whispersystems/textsecuregcm/storage/AddRemoveDeviceIntegrationTest.java @@ -101,6 +101,7 @@ public class AddRemoveDeviceIntegrationTest { DynamoDbExtensionSchema.Tables.CLIENT_PUBLIC_KEYS.tableName()); final Accounts accounts = new Accounts( + clock, DYNAMO_DB_EXTENSION.getDynamoDbClient(), DYNAMO_DB_EXTENSION.getDynamoDbAsyncClient(), DynamoDbExtensionSchema.Tables.ACCOUNTS.tableName(),