Switch DynamoDB to AWSv2.

Switch from using com.amazonaws.services.dynamodbv2 to using
software.amazon.awssdk.services.dynamodb for all current DynamoDB uses.
This commit is contained in:
Graeme Connell
2021-05-24 16:43:56 -06:00
committed by gram-signal
parent cbd9681e3e
commit c545cff1b3
31 changed files with 1114 additions and 876 deletions

View File

@@ -11,26 +11,12 @@ import static org.mockito.ArgumentMatchers.any;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;
import com.amazonaws.services.dynamodbv2.AmazonDynamoDB;
import com.amazonaws.services.dynamodbv2.AmazonDynamoDBAsync;
import com.amazonaws.services.dynamodbv2.document.DynamoDB;
import com.amazonaws.services.dynamodbv2.document.Item;
import com.amazonaws.services.dynamodbv2.document.Page;
import com.amazonaws.services.dynamodbv2.document.ScanOutcome;
import com.amazonaws.services.dynamodbv2.document.Table;
import com.amazonaws.services.dynamodbv2.document.spec.GetItemSpec;
import com.amazonaws.services.dynamodbv2.document.spec.ScanSpec;
import com.amazonaws.services.dynamodbv2.model.AttributeDefinition;
import com.amazonaws.services.dynamodbv2.model.ConditionalCheckFailedException;
import com.amazonaws.services.dynamodbv2.model.CreateTableRequest;
import com.amazonaws.services.dynamodbv2.model.KeySchemaElement;
import com.amazonaws.services.dynamodbv2.model.KeyType;
import com.amazonaws.services.dynamodbv2.model.ScalarAttributeType;
import io.github.resilience4j.circuitbreaker.CallNotPermittedException;
import java.nio.charset.StandardCharsets;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashSet;
import java.util.Map;
import java.util.Optional;
import java.util.Random;
import java.util.Set;
@@ -46,7 +32,22 @@ import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.RegisterExtension;
import org.whispersystems.textsecuregcm.configuration.CircuitBreakerConfiguration;
import org.whispersystems.textsecuregcm.entities.SignedPreKey;
import org.whispersystems.textsecuregcm.util.UUIDUtil;
import org.whispersystems.textsecuregcm.util.AttributeValues;
import software.amazon.awssdk.services.dynamodb.DynamoDbAsyncClient;
import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
import software.amazon.awssdk.services.dynamodb.model.AttributeDefinition;
import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
import software.amazon.awssdk.services.dynamodb.model.ConditionalCheckFailedException;
import software.amazon.awssdk.services.dynamodb.model.CreateTableRequest;
import software.amazon.awssdk.services.dynamodb.model.GetItemRequest;
import software.amazon.awssdk.services.dynamodb.model.GetItemResponse;
import software.amazon.awssdk.services.dynamodb.model.KeySchemaElement;
import software.amazon.awssdk.services.dynamodb.model.KeyType;
import software.amazon.awssdk.services.dynamodb.model.ScalarAttributeType;
import software.amazon.awssdk.services.dynamodb.model.ScanRequest;
import software.amazon.awssdk.services.dynamodb.model.ScanResponse;
import software.amazon.awssdk.services.dynamodb.model.TransactWriteItemsRequest;
import software.amazon.awssdk.services.dynamodb.model.UpdateItemRequest;
class AccountsDynamoDbTest {
@@ -59,49 +60,75 @@ class AccountsDynamoDbTest {
static DynamoDbExtension dynamoDbExtension = DynamoDbExtension.builder()
.tableName(ACCOUNTS_TABLE_NAME)
.hashKey(AccountsDynamoDb.KEY_ACCOUNT_UUID)
.attributeDefinition(new AttributeDefinition(AccountsDynamoDb.KEY_ACCOUNT_UUID, ScalarAttributeType.B))
.attributeDefinition(AttributeDefinition.builder()
.attributeName(AccountsDynamoDb.KEY_ACCOUNT_UUID)
.attributeType(ScalarAttributeType.B)
.build())
.build();
private AccountsDynamoDb accountsDynamoDb;
private Table migrationDeletedAccountsTable;
private Table migrationRetryAccountsTable;
@BeforeEach
void setupAccountsDao() {
CreateTableRequest createNumbersTableRequest = CreateTableRequest.builder()
.tableName(NUMBERS_TABLE_NAME)
.keySchema(KeySchemaElement.builder()
.attributeName(AccountsDynamoDb.ATTR_ACCOUNT_E164)
.keyType(KeyType.HASH)
.build())
.attributeDefinitions(AttributeDefinition.builder()
.attributeName(AccountsDynamoDb.ATTR_ACCOUNT_E164)
.attributeType(ScalarAttributeType.S)
.build())
.provisionedThroughput(DynamoDbExtension.DEFAULT_PROVISIONED_THROUGHPUT)
.build();
CreateTableRequest createNumbersTableRequest = new CreateTableRequest()
.withTableName(NUMBERS_TABLE_NAME)
.withKeySchema(new KeySchemaElement(AccountsDynamoDb.ATTR_ACCOUNT_E164, KeyType.HASH))
.withAttributeDefinitions(new AttributeDefinition(AccountsDynamoDb.ATTR_ACCOUNT_E164, ScalarAttributeType.S))
.withProvisionedThroughput(DynamoDbExtension.DEFAULT_PROVISIONED_THROUGHPUT);
dynamoDbExtension.getDynamoDbClient().createTable(createNumbersTableRequest);
final Table numbersTable = dynamoDbExtension.getDynamoDB().createTable(createNumbersTableRequest);
final CreateTableRequest createMigrationDeletedAccountsTableRequest = CreateTableRequest.builder()
.tableName(MIGRATION_DELETED_ACCOUNTS_TABLE_NAME)
.keySchema(KeySchemaElement.builder()
.attributeName(MigrationDeletedAccounts.KEY_UUID)
.keyType(KeyType.HASH)
.build())
.attributeDefinitions(AttributeDefinition.builder()
.attributeName(MigrationDeletedAccounts.KEY_UUID)
.attributeType(ScalarAttributeType.B)
.build())
.provisionedThroughput(DynamoDbExtension.DEFAULT_PROVISIONED_THROUGHPUT)
.build();
final CreateTableRequest createMigrationDeletedAccountsTableRequest = new CreateTableRequest()
.withTableName(MIGRATION_DELETED_ACCOUNTS_TABLE_NAME)
.withKeySchema(new KeySchemaElement(MigrationDeletedAccounts.KEY_UUID, KeyType.HASH))
.withAttributeDefinitions(new AttributeDefinition(MigrationDeletedAccounts.KEY_UUID, ScalarAttributeType.B))
.withProvisionedThroughput(DynamoDbExtension.DEFAULT_PROVISIONED_THROUGHPUT);
dynamoDbExtension.getDynamoDbClient().createTable(createMigrationDeletedAccountsTableRequest);
migrationDeletedAccountsTable = dynamoDbExtension.getDynamoDB().createTable(createMigrationDeletedAccountsTableRequest);
MigrationDeletedAccounts migrationDeletedAccounts = new MigrationDeletedAccounts(
dynamoDbExtension.getDynamoDbClient(), MIGRATION_DELETED_ACCOUNTS_TABLE_NAME);
MigrationDeletedAccounts migrationDeletedAccounts = new MigrationDeletedAccounts(dynamoDbExtension.getDynamoDB(),
migrationDeletedAccountsTable.getTableName());
final CreateTableRequest createMigrationRetryAccountsTableRequest = CreateTableRequest.builder()
.tableName(MIGRATION_RETRY_ACCOUNTS_TABLE_NAME)
.keySchema(KeySchemaElement.builder()
.attributeName(MigrationRetryAccounts.KEY_UUID)
.keyType(KeyType.HASH)
.build())
.attributeDefinitions(AttributeDefinition.builder()
.attributeName(MigrationRetryAccounts.KEY_UUID)
.attributeType(ScalarAttributeType.B)
.build())
.provisionedThroughput(DynamoDbExtension.DEFAULT_PROVISIONED_THROUGHPUT)
.build();
final CreateTableRequest createMigrationRetryAccountsTableRequest = new CreateTableRequest()
.withTableName(MIGRATION_RETRY_ACCOUNTS_TABLE_NAME)
.withKeySchema(new KeySchemaElement(MigrationRetryAccounts.KEY_UUID, KeyType.HASH))
.withAttributeDefinitions(new AttributeDefinition(MigrationRetryAccounts.KEY_UUID, ScalarAttributeType.B))
.withProvisionedThroughput(DynamoDbExtension.DEFAULT_PROVISIONED_THROUGHPUT);
dynamoDbExtension.getDynamoDbClient().createTable(createMigrationRetryAccountsTableRequest);
migrationRetryAccountsTable = dynamoDbExtension.getDynamoDB().createTable(createMigrationRetryAccountsTableRequest);
MigrationRetryAccounts migrationRetryAccounts = new MigrationRetryAccounts((dynamoDbExtension.getDynamoDbClient()),
MIGRATION_RETRY_ACCOUNTS_TABLE_NAME);
MigrationRetryAccounts migrationRetryAccounts = new MigrationRetryAccounts((dynamoDbExtension.getDynamoDB()),
migrationRetryAccountsTable.getTableName());
this.accountsDynamoDb = new AccountsDynamoDb(dynamoDbExtension.getClient(), dynamoDbExtension.getAsyncClient(), new ThreadPoolExecutor(1, 1, 0, TimeUnit.SECONDS, new LinkedBlockingDeque<>()), dynamoDbExtension.getDynamoDB(), dynamoDbExtension.getTableName(), numbersTable.getTableName(),
migrationDeletedAccounts, migrationRetryAccounts);
this.accountsDynamoDb = new AccountsDynamoDb(
dynamoDbExtension.getDynamoDbClient(),
dynamoDbExtension.getDynamoDbAsyncClient(),
new ThreadPoolExecutor(1, 1, 0, TimeUnit.SECONDS, new LinkedBlockingDeque<>()),
dynamoDbExtension.getTableName(),
NUMBERS_TABLE_NAME,
migrationDeletedAccounts,
migrationRetryAccounts);
}
@Test
@@ -262,8 +289,12 @@ class AccountsDynamoDbTest {
verifyRecentlyDeletedAccountsTableItemCount(1);
assertThat(migrationDeletedAccountsTable
.getItem(MigrationDeletedAccounts.primaryKey(deletedAccount.getUuid()))).isNotNull();
Map<String, AttributeValue> primaryKey = MigrationDeletedAccounts.primaryKey(deletedAccount.getUuid());
assertThat(dynamoDbExtension.getDynamoDbClient().getItem(GetItemRequest.builder()
.tableName(MIGRATION_DELETED_ACCOUNTS_TABLE_NAME)
.key(Map.of(MigrationDeletedAccounts.KEY_UUID, primaryKey.get(MigrationDeletedAccounts.KEY_UUID)))
.build()))
.isNotNull();
accountsDynamoDb.deleteRecentlyDeletedUuids();
@@ -273,8 +304,10 @@ class AccountsDynamoDbTest {
private void verifyRecentlyDeletedAccountsTableItemCount(int expectedItemCount) {
int totalItems = 0;
for (Page<Item, ScanOutcome> page : migrationDeletedAccountsTable.scan(new ScanSpec()).pages()) {
for (Item ignored : page) {
for (ScanResponse page : dynamoDbExtension.getDynamoDbClient().scanPaginator(ScanRequest.builder()
.tableName(MIGRATION_DELETED_ACCOUNTS_TABLE_NAME)
.build())) {
for (Map<String, AttributeValue> item : page.items()) {
totalItems++;
}
}
@@ -306,16 +339,15 @@ class AccountsDynamoDbTest {
configuration.setRingBufferSizeInClosedState(2);
configuration.setFailureRateThreshold(50);
final AmazonDynamoDB client = mock(AmazonDynamoDB.class);
final DynamoDB dynamoDB = new DynamoDB(client);
final DynamoDbClient client = mock(DynamoDbClient.class);
when(client.transactWriteItems(any()))
when(client.transactWriteItems(any(TransactWriteItemsRequest.class)))
.thenThrow(RuntimeException.class);
when(client.updateItem(any()))
when(client.updateItem(any(UpdateItemRequest.class)))
.thenThrow(RuntimeException.class);
AccountsDynamoDb accounts = new AccountsDynamoDb(client, mock(AmazonDynamoDBAsync.class), mock(ThreadPoolExecutor.class), dynamoDB, ACCOUNTS_TABLE_NAME, NUMBERS_TABLE_NAME, mock(
AccountsDynamoDb accounts = new AccountsDynamoDb(client, mock(DynamoDbAsyncClient.class), mock(ThreadPoolExecutor.class), ACCOUNTS_TABLE_NAME, NUMBERS_TABLE_NAME, mock(
MigrationDeletedAccounts.class), mock(MigrationRetryAccounts.class));
Account account = generateAccount("+14151112222", UUID.randomUUID());
@@ -408,20 +440,22 @@ class AccountsDynamoDbTest {
}
private void verifyStoredState(String number, UUID uuid, Account expecting) {
final Table accounts = dynamoDbExtension.getDynamoDB().getTable(dynamoDbExtension.getTableName());
final DynamoDbClient db = dynamoDbExtension.getDynamoDbClient();
Item item = accounts.getItem(new GetItemSpec()
.withPrimaryKey(AccountsDynamoDb.KEY_ACCOUNT_UUID, UUIDUtil.toByteBuffer(uuid))
.withConsistentRead(true));
final GetItemResponse get = db.getItem(GetItemRequest.builder()
.tableName(dynamoDbExtension.getTableName())
.key(Map.of(AccountsDynamoDb.KEY_ACCOUNT_UUID, AttributeValues.fromUUID(uuid)))
.consistentRead(true)
.build());
if (item != null) {
String data = new String(item.getBinary(AccountsDynamoDb.ATTR_ACCOUNT_DATA), StandardCharsets.UTF_8);
if (get.hasItem()) {
String data = new String(get.item().get(AccountsDynamoDb.ATTR_ACCOUNT_DATA).b().asByteArray(), StandardCharsets.UTF_8);
assertThat(data).isNotEmpty();
assertThat(item.getNumber(AccountsDynamoDb.ATTR_MIGRATION_VERSION).intValue())
assertThat(AttributeValues.getInt(get.item(), AccountsDynamoDb.ATTR_MIGRATION_VERSION, -1))
.isEqualTo(expecting.getDynamoDbMigrationVersion());
Account result = AccountsDynamoDb.fromItem(item);
Account result = AccountsDynamoDb.fromItem(get.item());
verifyStoredState(number, uuid, result, expecting);
} else {
throw new AssertionError("No data");

View File

@@ -1,33 +1,35 @@
package org.whispersystems.textsecuregcm.storage;
import com.almworks.sqlite4java.SQLite;
import com.amazonaws.auth.AWSStaticCredentialsProvider;
import com.amazonaws.auth.BasicAWSCredentials;
import com.amazonaws.client.builder.AwsClientBuilder;
import com.amazonaws.services.dynamodbv2.AmazonDynamoDB;
import com.amazonaws.services.dynamodbv2.AmazonDynamoDBAsync;
import com.amazonaws.services.dynamodbv2.AmazonDynamoDBAsyncClientBuilder;
import com.amazonaws.services.dynamodbv2.AmazonDynamoDBClientBuilder;
import com.amazonaws.services.dynamodbv2.document.DynamoDB;
import com.amazonaws.services.dynamodbv2.local.main.ServerRunner;
import com.amazonaws.services.dynamodbv2.local.server.DynamoDBProxyServer;
import com.amazonaws.services.dynamodbv2.model.AttributeDefinition;
import com.amazonaws.services.dynamodbv2.model.CreateTableRequest;
import com.amazonaws.services.dynamodbv2.model.GlobalSecondaryIndex;
import com.amazonaws.services.dynamodbv2.model.KeySchemaElement;
import com.amazonaws.services.dynamodbv2.model.ProvisionedThroughput;
import java.net.ServerSocket;
import java.net.URI;
import java.util.ArrayList;
import java.util.List;
import org.junit.jupiter.api.extension.AfterEachCallback;
import org.junit.jupiter.api.extension.BeforeEachCallback;
import org.junit.jupiter.api.extension.ExtensionContext;
import software.amazon.awssdk.auth.credentials.AwsBasicCredentials;
import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider;
import software.amazon.awssdk.regions.Region;
import software.amazon.awssdk.services.dynamodb.DynamoDbAsyncClient;
import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
import software.amazon.awssdk.services.dynamodb.model.AttributeDefinition;
import software.amazon.awssdk.services.dynamodb.model.CreateTableRequest;
import software.amazon.awssdk.services.dynamodb.model.GlobalSecondaryIndex;
import software.amazon.awssdk.services.dynamodb.model.KeySchemaElement;
import software.amazon.awssdk.services.dynamodb.model.KeyType;
import software.amazon.awssdk.services.dynamodb.model.ProvisionedThroughput;
public class DynamoDbExtension implements BeforeEachCallback, AfterEachCallback {
static final String DEFAULT_TABLE_NAME = "test_table";
static final ProvisionedThroughput DEFAULT_PROVISIONED_THROUGHPUT = new ProvisionedThroughput(20L, 20L);
static final ProvisionedThroughput DEFAULT_PROVISIONED_THROUGHPUT = ProvisionedThroughput.builder()
.readCapacityUnits(20L)
.writeCapacityUnits(20L)
.build();
private DynamoDBProxyServer server;
private int port;
@@ -42,9 +44,8 @@ public class DynamoDbExtension implements BeforeEachCallback, AfterEachCallback
private final long readCapacityUnits;
private final long writeCapacityUnits;
private AmazonDynamoDB client;
private AmazonDynamoDBAsync asyncClient;
private DynamoDB dynamoDB;
private DynamoDbClient dynamoDB2;
private DynamoDbAsyncClient dynamoAsyncDB2;
private DynamoDbExtension(String tableName, String hashKey, String rangeKey, List<AttributeDefinition> attributeDefinitions, List<GlobalSecondaryIndex> globalSecondaryIndexes, long readCapacityUnits,
long writeCapacityUnits) {
@@ -87,26 +88,33 @@ public class DynamoDbExtension implements BeforeEachCallback, AfterEachCallback
KeySchemaElement[] keySchemaElements;
if (rangeKeyName == null) {
keySchemaElements = new KeySchemaElement[] {
new KeySchemaElement(hashKeyName, "HASH"),
KeySchemaElement.builder().attributeName(hashKeyName).keyType(KeyType.HASH).build(),
};
} else {
keySchemaElements = new KeySchemaElement[] {
new KeySchemaElement(hashKeyName, "HASH"),
new KeySchemaElement(rangeKeyName, "RANGE")
KeySchemaElement.builder().attributeName(hashKeyName).keyType(KeyType.HASH).build(),
KeySchemaElement.builder().attributeName(rangeKeyName).keyType(KeyType.RANGE).build(),
};
}
final CreateTableRequest createTableRequest = new CreateTableRequest()
.withTableName(tableName)
.withKeySchema(keySchemaElements)
.withAttributeDefinitions(attributeDefinitions.isEmpty() ? null : attributeDefinitions)
.withGlobalSecondaryIndexes(globalSecondaryIndexes.isEmpty() ? null : globalSecondaryIndexes)
.withProvisionedThroughput(new ProvisionedThroughput(readCapacityUnits, writeCapacityUnits));
final CreateTableRequest createTableRequest = CreateTableRequest.builder()
.tableName(tableName)
.keySchema(keySchemaElements)
.attributeDefinitions(attributeDefinitions.isEmpty() ? null : attributeDefinitions)
.globalSecondaryIndexes(globalSecondaryIndexes.isEmpty() ? null : globalSecondaryIndexes)
.provisionedThroughput(ProvisionedThroughput.builder()
.readCapacityUnits(readCapacityUnits)
.writeCapacityUnits(writeCapacityUnits)
.build())
.build();
getDynamoDB().createTable(createTableRequest);
getDynamoDbClient().createTable(createTableRequest);
}
private void startServer() throws Exception {
// Even though we're using AWS SDK v2, Dynamo's local implementation's canonical location
// is within v1 (https://github.com/aws/aws-sdk-java-v2/issues/982). This does support
// v2 clients, though.
SQLite.setLibraryPath("target/lib"); // if you see a library failed to load error, you need to run mvn test-compile at least once first
ServerSocket serverSocket = new ServerSocket(0);
serverSocket.setReuseAddress(false);
@@ -117,18 +125,18 @@ public class DynamoDbExtension implements BeforeEachCallback, AfterEachCallback
}
private void initializeClient() {
AmazonDynamoDBClientBuilder clientBuilder = AmazonDynamoDBClientBuilder.standard()
.withEndpointConfiguration(
new AwsClientBuilder.EndpointConfiguration("http://localhost:" + port, "local-test-region"))
.withCredentials(new AWSStaticCredentialsProvider(new BasicAWSCredentials("accessKey", "secretKey")));
client = clientBuilder.build();
asyncClient = AmazonDynamoDBAsyncClientBuilder.standard()
.withEndpointConfiguration(clientBuilder.getEndpoint())
.withCredentials(clientBuilder.getCredentials())
.build();
dynamoDB = new DynamoDB(client);
dynamoDB2 = DynamoDbClient.builder()
.endpointOverride(URI.create("http://localhost:" + port))
.region(Region.of("local-test-region"))
.credentialsProvider(StaticCredentialsProvider.create(
AwsBasicCredentials.create("accessKey", "secretKey")))
.build();
dynamoAsyncDB2 = DynamoDbAsyncClient.builder()
.endpointOverride(URI.create("http://localhost:" + port))
.region(Region.of("local-test-region"))
.credentialsProvider(StaticCredentialsProvider.create(
AwsBasicCredentials.create("accessKey", "secretKey")))
.build();
}
static class DynamoDbExtensionBuilder {
@@ -140,8 +148,8 @@ public class DynamoDbExtension implements BeforeEachCallback, AfterEachCallback
private List<AttributeDefinition> attributeDefinitions = new ArrayList<>();
private List<GlobalSecondaryIndex> globalSecondaryIndexes = new ArrayList<>();
private long readCapacityUnits = DEFAULT_PROVISIONED_THROUGHPUT.getReadCapacityUnits();
private long writeCapacityUnits = DEFAULT_PROVISIONED_THROUGHPUT.getWriteCapacityUnits();
private long readCapacityUnits = DEFAULT_PROVISIONED_THROUGHPUT.readCapacityUnits();
private long writeCapacityUnits = DEFAULT_PROVISIONED_THROUGHPUT.writeCapacityUnits();
private DynamoDbExtensionBuilder() {
@@ -178,16 +186,12 @@ public class DynamoDbExtension implements BeforeEachCallback, AfterEachCallback
}
}
public AmazonDynamoDB getClient() {
return client;
public DynamoDbClient getDynamoDbClient() {
return dynamoDB2;
}
public AmazonDynamoDBAsync getAsyncClient() {
return asyncClient;
}
public DynamoDB getDynamoDB() {
return dynamoDB;
public DynamoDbAsyncClient getDynamoDbAsyncClient() {
return dynamoAsyncDB2;
}
public String getTableName() {

View File

@@ -5,13 +5,13 @@
package org.whispersystems.textsecuregcm.storage;
import com.amazonaws.services.dynamodbv2.document.DynamoDB;
import com.amazonaws.services.dynamodbv2.model.AttributeDefinition;
import com.amazonaws.services.dynamodbv2.model.CreateTableRequest;
import com.amazonaws.services.dynamodbv2.model.KeySchemaElement;
import com.amazonaws.services.dynamodbv2.model.ProvisionedThroughput;
import com.amazonaws.services.dynamodbv2.model.ScalarAttributeType;
import org.whispersystems.textsecuregcm.tests.util.LocalDynamoDbRule;
import software.amazon.awssdk.services.dynamodb.model.AttributeDefinition;
import software.amazon.awssdk.services.dynamodb.model.CreateTableRequest;
import software.amazon.awssdk.services.dynamodb.model.KeySchemaElement;
import software.amazon.awssdk.services.dynamodb.model.KeyType;
import software.amazon.awssdk.services.dynamodb.model.ProvisionedThroughput;
import software.amazon.awssdk.services.dynamodb.model.ScalarAttributeType;
public class KeysDynamoDbRule extends LocalDynamoDbRule {
public static final String TABLE_NAME = "Signal_Keys_Test";
@@ -19,18 +19,22 @@ public class KeysDynamoDbRule extends LocalDynamoDbRule {
@Override
protected void before() throws Throwable {
super.before();
final DynamoDB dynamoDB = getDynamoDB();
final CreateTableRequest createTableRequest = new CreateTableRequest()
.withTableName(TABLE_NAME)
.withKeySchema(new KeySchemaElement(KeysDynamoDb.KEY_ACCOUNT_UUID, "HASH"),
new KeySchemaElement(KeysDynamoDb.KEY_DEVICE_ID_KEY_ID, "RANGE"))
.withAttributeDefinitions(new AttributeDefinition(KeysDynamoDb.KEY_ACCOUNT_UUID, ScalarAttributeType.B),
new AttributeDefinition(KeysDynamoDb.KEY_DEVICE_ID_KEY_ID, ScalarAttributeType.B))
.withProvisionedThroughput(new ProvisionedThroughput(20L, 20L));
dynamoDB.createTable(createTableRequest);
getDynamoDbClient().createTable(CreateTableRequest.builder()
.tableName(TABLE_NAME)
.keySchema(
KeySchemaElement.builder().attributeName(KeysDynamoDb.KEY_ACCOUNT_UUID).keyType(KeyType.HASH).build(),
KeySchemaElement.builder().attributeName(KeysDynamoDb.KEY_DEVICE_ID_KEY_ID).keyType(KeyType.RANGE)
.build())
.attributeDefinitions(AttributeDefinition.builder()
.attributeName(KeysDynamoDb.KEY_ACCOUNT_UUID)
.attributeType(ScalarAttributeType.B)
.build(),
AttributeDefinition.builder()
.attributeName(KeysDynamoDb.KEY_DEVICE_ID_KEY_ID)
.attributeType(ScalarAttributeType.B)
.build())
.provisionedThroughput(ProvisionedThroughput.builder().readCapacityUnits(20L).writeCapacityUnits(20L).build())
.build());
}
@Override

View File

@@ -9,6 +9,7 @@ import org.junit.Before;
import org.junit.ClassRule;
import org.junit.Test;
import org.whispersystems.textsecuregcm.entities.PreKey;
import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
import java.util.Collections;
import java.util.List;
@@ -17,6 +18,7 @@ import java.util.Optional;
import java.util.Set;
import java.util.UUID;
import static org.junit.Assert.assertArrayEquals;
import static org.junit.Assert.assertEquals;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;
@@ -34,7 +36,7 @@ public class KeysDynamoDbTest {
@Before
public void setup() {
keysDynamoDb = new KeysDynamoDb(dynamoDbRule.getDynamoDB(), KeysDynamoDbRule.TABLE_NAME);
keysDynamoDb = new KeysDynamoDb(dynamoDbRule.getDynamoDbClient(), KeysDynamoDbRule.TABLE_NAME);
account = mock(Account.class);
when(account.getNumber()).thenReturn(ACCOUNT_NUMBER);
@@ -133,4 +135,10 @@ public class KeysDynamoDbTest {
assertEquals(0, keysDynamoDb.getCount(account, DEVICE_ID));
assertEquals(1, keysDynamoDb.getCount(account, DEVICE_ID + 1));
}
@Test
public void testSortKeyPrefix() {
AttributeValue got = KeysDynamoDb.getSortKeyPrefix(123);
assertArrayEquals(new byte[]{0, 0, 0, 0, 0, 0, 0, 123}, got.b().asByteArray());
}
}

View File

@@ -9,12 +9,6 @@ import static org.junit.Assert.assertEquals;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;
import com.amazonaws.services.dynamodbv2.document.DynamoDB;
import com.amazonaws.services.dynamodbv2.document.Item;
import com.amazonaws.services.dynamodbv2.document.ItemCollection;
import com.amazonaws.services.dynamodbv2.document.ScanOutcome;
import com.amazonaws.services.dynamodbv2.document.Table;
import com.amazonaws.services.dynamodbv2.document.spec.ScanSpec;
import com.google.protobuf.ByteString;
import io.lettuce.core.cluster.SlotHash;
import java.nio.ByteBuffer;
@@ -22,6 +16,7 @@ import java.time.Duration;
import java.time.Instant;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.UUID;
import java.util.concurrent.ExecutorService;
@@ -38,6 +33,10 @@ import org.whispersystems.textsecuregcm.entities.MessageProtos;
import org.whispersystems.textsecuregcm.metrics.PushLatencyManager;
import org.whispersystems.textsecuregcm.redis.AbstractRedisClusterTest;
import org.whispersystems.textsecuregcm.tests.util.MessagesDynamoDbRule;
import org.whispersystems.textsecuregcm.util.AttributeValues;
import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
import software.amazon.awssdk.services.dynamodb.model.ScanRequest;
public class MessagePersisterIntegrationTest extends AbstractRedisClusterTest {
@@ -62,7 +61,7 @@ public class MessagePersisterIntegrationTest extends AbstractRedisClusterTest {
connection.sync().upstream().commands().configSet("notify-keyspace-events", "K$glz");
});
final MessagesDynamoDb messagesDynamoDb = new MessagesDynamoDb(messagesDynamoDbRule.getDynamoDB(), MessagesDynamoDbRule.TABLE_NAME, Duration.ofDays(7));
final MessagesDynamoDb messagesDynamoDb = new MessagesDynamoDb(messagesDynamoDbRule.getDynamoDbClient(), MessagesDynamoDbRule.TABLE_NAME, Duration.ofDays(7));
final AccountsManager accountsManager = mock(AccountsManager.class);
final DynamicConfigurationManager dynamicConfigurationManager = mock(DynamicConfigurationManager.class);
@@ -142,17 +141,16 @@ public class MessagePersisterIntegrationTest extends AbstractRedisClusterTest {
final List<MessageProtos.Envelope> persistedMessages = new ArrayList<>(messageCount);
DynamoDB dynamoDB = messagesDynamoDbRule.getDynamoDB();
Table table = dynamoDB.getTable(MessagesDynamoDbRule.TABLE_NAME);
final ItemCollection<ScanOutcome> scan = table.scan(new ScanSpec());
for (Item item : scan) {
persistedMessages.add(MessageProtos.Envelope.newBuilder()
.setServerGuid(convertBinaryToUuid(item.getBinary("U")).toString())
.setType(MessageProtos.Envelope.Type.valueOf(item.getInt("T")))
.setTimestamp(item.getLong("TS"))
.setServerTimestamp(extractServerTimestamp(item.getBinary("S")))
.setContent(ByteString.copyFrom(item.getBinary("C")))
.build());
DynamoDbClient dynamoDB = messagesDynamoDbRule.getDynamoDbClient();
for (Map<String, AttributeValue> item : dynamoDB
.scan(ScanRequest.builder().tableName(MessagesDynamoDbRule.TABLE_NAME).build()).items()) {
persistedMessages.add(MessageProtos.Envelope.newBuilder()
.setServerGuid(AttributeValues.getUUID(item, "U", null).toString())
.setType(MessageProtos.Envelope.Type.valueOf(AttributeValues.getInt(item, "T", -1)))
.setTimestamp(AttributeValues.getLong(item, "TS", -1))
.setServerTimestamp(extractServerTimestamp(AttributeValues.getByteArray(item, "S", null)))
.setContent(ByteString.copyFrom(AttributeValues.getByteArray(item, "C", null)))
.build());
}
assertEquals(expectedMessages, persistedMessages);

View File

@@ -4,10 +4,10 @@ import static org.junit.jupiter.api.Assertions.assertTrue;
import java.util.List;
import java.util.UUID;
import com.amazonaws.services.dynamodbv2.model.AttributeDefinition;
import com.amazonaws.services.dynamodbv2.model.ScalarAttributeType;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.RegisterExtension;
import software.amazon.awssdk.services.dynamodb.model.AttributeDefinition;
import software.amazon.awssdk.services.dynamodb.model.ScalarAttributeType;
class MigrationDeletedAccountsTest {
@@ -15,13 +15,16 @@ class MigrationDeletedAccountsTest {
static DynamoDbExtension dynamoDbExtension = DynamoDbExtension.builder()
.tableName("deleted_accounts_test")
.hashKey(MigrationDeletedAccounts.KEY_UUID)
.attributeDefinition(new AttributeDefinition(MigrationDeletedAccounts.KEY_UUID, ScalarAttributeType.B))
.attributeDefinition(AttributeDefinition.builder()
.attributeName(MigrationDeletedAccounts.KEY_UUID)
.attributeType(ScalarAttributeType.B)
.build())
.build();
@Test
void test() {
final MigrationDeletedAccounts migrationDeletedAccounts = new MigrationDeletedAccounts(dynamoDbExtension.getDynamoDB(),
final MigrationDeletedAccounts migrationDeletedAccounts = new MigrationDeletedAccounts(dynamoDbExtension.getDynamoDbClient(),
dynamoDbExtension.getTableName());
UUID firstUuid = UUID.randomUUID();

View File

@@ -2,12 +2,12 @@ package org.whispersystems.textsecuregcm.storage;
import static org.junit.jupiter.api.Assertions.assertTrue;
import com.amazonaws.services.dynamodbv2.model.AttributeDefinition;
import com.amazonaws.services.dynamodbv2.model.ScalarAttributeType;
import java.util.List;
import java.util.UUID;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.RegisterExtension;
import software.amazon.awssdk.services.dynamodb.model.AttributeDefinition;
import software.amazon.awssdk.services.dynamodb.model.ScalarAttributeType;
class MigrationRetryAccountsTest {
@@ -15,13 +15,16 @@ class MigrationRetryAccountsTest {
static DynamoDbExtension dynamoDbExtension = DynamoDbExtension.builder()
.tableName("account_migration_errors_test")
.hashKey(MigrationRetryAccounts.KEY_UUID)
.attributeDefinition(new AttributeDefinition(MigrationRetryAccounts.KEY_UUID, ScalarAttributeType.B))
.attributeDefinition(AttributeDefinition.builder()
.attributeName(MigrationRetryAccounts.KEY_UUID)
.attributeType(ScalarAttributeType.B)
.build())
.build();
@Test
void test() {
final MigrationRetryAccounts migrationRetryAccounts = new MigrationRetryAccounts(dynamoDbExtension.getDynamoDB(),
final MigrationRetryAccounts migrationRetryAccounts = new MigrationRetryAccounts(dynamoDbExtension.getDynamoDbClient(),
dynamoDbExtension.getTableName());
UUID firstUuid = UUID.randomUUID();

View File

@@ -9,8 +9,6 @@ import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertTrue;
import com.amazonaws.services.dynamodbv2.model.AttributeDefinition;
import com.amazonaws.services.dynamodbv2.model.ScalarAttributeType;
import java.time.Clock;
import java.time.Duration;
import java.time.Instant;
@@ -20,6 +18,8 @@ import java.util.UUID;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.RegisterExtension;
import software.amazon.awssdk.services.dynamodb.model.AttributeDefinition;
import software.amazon.awssdk.services.dynamodb.model.ScalarAttributeType;
class PushChallengeDynamoDbTest {
@@ -34,12 +34,15 @@ class PushChallengeDynamoDbTest {
static DynamoDbExtension dynamoDbExtension = DynamoDbExtension.builder()
.tableName(TABLE_NAME)
.hashKey(PushChallengeDynamoDb.KEY_ACCOUNT_UUID)
.attributeDefinition(new AttributeDefinition(PushChallengeDynamoDb.KEY_ACCOUNT_UUID, ScalarAttributeType.B))
.attributeDefinition(AttributeDefinition.builder()
.attributeName(PushChallengeDynamoDb.KEY_ACCOUNT_UUID)
.attributeType(ScalarAttributeType.B)
.build())
.build();
@BeforeEach
void setUp() {
this.pushChallengeDynamoDb = new PushChallengeDynamoDb(dynamoDbExtension.getDynamoDB(), TABLE_NAME, Clock.fixed(
this.pushChallengeDynamoDb = new PushChallengeDynamoDb(dynamoDbExtension.getDynamoDbClient(), TABLE_NAME, Clock.fixed(
Instant.ofEpochMilli(CURRENT_TIME_MILLIS), ZoneId.systemDefault()));
}

View File

@@ -4,13 +4,13 @@ import static org.junit.jupiter.api.Assertions.assertAll;
import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertTrue;
import com.amazonaws.services.dynamodbv2.model.AttributeDefinition;
import com.amazonaws.services.dynamodbv2.model.ScalarAttributeType;
import java.util.UUID;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.RegisterExtension;
import org.whispersystems.textsecuregcm.util.UUIDUtil;
import software.amazon.awssdk.services.dynamodb.model.AttributeDefinition;
import software.amazon.awssdk.services.dynamodb.model.ScalarAttributeType;
class ReportMessageDynamoDbTest {
@@ -22,13 +22,16 @@ class ReportMessageDynamoDbTest {
static DynamoDbExtension dynamoDbExtension = DynamoDbExtension.builder()
.tableName(TABLE_NAME)
.hashKey(ReportMessageDynamoDb.KEY_HASH)
.attributeDefinition(new AttributeDefinition(ReportMessageDynamoDb.KEY_HASH, ScalarAttributeType.B))
.attributeDefinition(AttributeDefinition.builder()
.attributeName(ReportMessageDynamoDb.KEY_HASH)
.attributeType(ScalarAttributeType.B)
.build())
.build();
@BeforeEach
void setUp() {
this.reportMessageDynamoDb = new ReportMessageDynamoDb(dynamoDbExtension.getDynamoDB(), TABLE_NAME);
this.reportMessageDynamoDb = new ReportMessageDynamoDb(dynamoDbExtension.getDynamoDbClient(), TABLE_NAME);
}
@Test

View File

@@ -20,7 +20,6 @@ import static org.mockito.Mockito.verifyNoMoreInteractions;
import static org.mockito.Mockito.verifyZeroInteractions;
import static org.mockito.Mockito.when;
import com.amazonaws.services.dynamodbv2.model.ConditionalCheckFailedException;
import io.lettuce.core.RedisException;
import io.lettuce.core.cluster.api.sync.RedisAdvancedClusterCommands;
import java.util.HashSet;
@@ -49,6 +48,7 @@ import org.whispersystems.textsecuregcm.storage.MessagesManager;
import org.whispersystems.textsecuregcm.storage.ProfilesManager;
import org.whispersystems.textsecuregcm.storage.UsernamesManager;
import org.whispersystems.textsecuregcm.tests.util.RedisClusterHelper;
import software.amazon.awssdk.services.dynamodb.model.ConditionalCheckFailedException;
class AccountsManagerTest {

View File

@@ -67,7 +67,7 @@ public class MessagesDynamoDbTest {
@Before
public void setup() {
messagesDynamoDb = new MessagesDynamoDb(dynamoDbRule.getDynamoDB(), MessagesDynamoDbRule.TABLE_NAME, Duration.ofDays(7));
messagesDynamoDb = new MessagesDynamoDb(dynamoDbRule.getDynamoDbClient(), MessagesDynamoDbRule.TABLE_NAME, Duration.ofDays(7));
}
@Test

View File

@@ -6,16 +6,16 @@
package org.whispersystems.textsecuregcm.tests.util;
import com.almworks.sqlite4java.SQLite;
import com.amazonaws.auth.AWSStaticCredentialsProvider;
import com.amazonaws.auth.BasicAWSCredentials;
import com.amazonaws.client.builder.AwsClientBuilder;
import com.amazonaws.services.dynamodbv2.AmazonDynamoDBClientBuilder;
import com.amazonaws.services.dynamodbv2.document.DynamoDB;
import com.amazonaws.services.dynamodbv2.local.main.ServerRunner;
import com.amazonaws.services.dynamodbv2.local.server.DynamoDBProxyServer;
import org.junit.rules.ExternalResource;
import software.amazon.awssdk.auth.credentials.AwsBasicCredentials;
import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider;
import software.amazon.awssdk.regions.Region;
import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
import java.net.ServerSocket;
import java.net.URI;
public class LocalDynamoDbRule extends ExternalResource {
private DynamoDBProxyServer server;
@@ -43,11 +43,12 @@ public class LocalDynamoDbRule extends ExternalResource {
super.after();
}
public DynamoDB getDynamoDB() {
AmazonDynamoDBClientBuilder clientBuilder =
AmazonDynamoDBClientBuilder.standard()
.withEndpointConfiguration(new AwsClientBuilder.EndpointConfiguration("http://localhost:" + port, "local-test-region"))
.withCredentials(new AWSStaticCredentialsProvider(new BasicAWSCredentials("accessKey", "secretKey")));
return new DynamoDB(clientBuilder.build());
public DynamoDbClient getDynamoDbClient() {
return DynamoDbClient.builder()
.endpointOverride(URI.create("http://localhost:" + port))
.region(Region.of("local-test-region"))
.credentialsProvider(StaticCredentialsProvider.create(
AwsBasicCredentials.create("accessKey", "secretKey")))
.build();
}
}

View File

@@ -5,15 +5,15 @@
package org.whispersystems.textsecuregcm.tests.util;
import com.amazonaws.services.dynamodbv2.document.DynamoDB;
import com.amazonaws.services.dynamodbv2.model.AttributeDefinition;
import com.amazonaws.services.dynamodbv2.model.CreateTableRequest;
import com.amazonaws.services.dynamodbv2.model.KeySchemaElement;
import com.amazonaws.services.dynamodbv2.model.LocalSecondaryIndex;
import com.amazonaws.services.dynamodbv2.model.Projection;
import com.amazonaws.services.dynamodbv2.model.ProjectionType;
import com.amazonaws.services.dynamodbv2.model.ProvisionedThroughput;
import com.amazonaws.services.dynamodbv2.model.ScalarAttributeType;
import software.amazon.awssdk.services.dynamodb.model.AttributeDefinition;
import software.amazon.awssdk.services.dynamodb.model.CreateTableRequest;
import software.amazon.awssdk.services.dynamodb.model.KeySchemaElement;
import software.amazon.awssdk.services.dynamodb.model.KeyType;
import software.amazon.awssdk.services.dynamodb.model.LocalSecondaryIndex;
import software.amazon.awssdk.services.dynamodb.model.Projection;
import software.amazon.awssdk.services.dynamodb.model.ProjectionType;
import software.amazon.awssdk.services.dynamodb.model.ProvisionedThroughput;
import software.amazon.awssdk.services.dynamodb.model.ScalarAttributeType;
public class MessagesDynamoDbRule extends LocalDynamoDbRule {
@@ -22,20 +22,21 @@ public class MessagesDynamoDbRule extends LocalDynamoDbRule {
@Override
protected void before() throws Throwable {
super.before();
DynamoDB dynamoDB = getDynamoDB();
CreateTableRequest createTableRequest = new CreateTableRequest()
.withTableName(TABLE_NAME)
.withKeySchema(new KeySchemaElement("H", "HASH"),
new KeySchemaElement("S", "RANGE"))
.withAttributeDefinitions(new AttributeDefinition("H", ScalarAttributeType.B),
new AttributeDefinition("S", ScalarAttributeType.B),
new AttributeDefinition("U", ScalarAttributeType.B))
.withProvisionedThroughput(new ProvisionedThroughput(20L, 20L))
.withLocalSecondaryIndexes(new LocalSecondaryIndex().withIndexName("Message_UUID_Index")
.withKeySchema(new KeySchemaElement("H", "HASH"),
new KeySchemaElement("U", "RANGE"))
.withProjection(new Projection().withProjectionType(ProjectionType.KEYS_ONLY)));
dynamoDB.createTable(createTableRequest);
getDynamoDbClient().createTable(CreateTableRequest.builder()
.tableName(TABLE_NAME)
.keySchema(KeySchemaElement.builder().attributeName("H").keyType(KeyType.HASH).build(),
KeySchemaElement.builder().attributeName("S").keyType(KeyType.RANGE).build())
.attributeDefinitions(
AttributeDefinition.builder().attributeName("H").attributeType(ScalarAttributeType.B).build(),
AttributeDefinition.builder().attributeName("S").attributeType(ScalarAttributeType.B).build(),
AttributeDefinition.builder().attributeName("U").attributeType(ScalarAttributeType.B).build())
.provisionedThroughput(ProvisionedThroughput.builder().readCapacityUnits(20L).writeCapacityUnits(20L).build())
.localSecondaryIndexes(LocalSecondaryIndex.builder().indexName("Message_UUID_Index")
.keySchema(KeySchemaElement.builder().attributeName("H").keyType(KeyType.HASH).build(),
KeySchemaElement.builder().attributeName("U").keyType(KeyType.RANGE).build())
.projection(Projection.builder().projectionType(ProjectionType.KEYS_ONLY).build())
.build())
.build());
}
@Override

View File

@@ -12,7 +12,6 @@ import java.io.InputStreamReader;
import java.net.Inet4Address;
import java.net.UnknownHostException;
import java.util.Optional;
import java.util.zip.GZIPInputStream;
import static org.junit.jupiter.api.Assertions.*;

View File

@@ -0,0 +1,60 @@
/*
* Copyright 2013-2021 Signal Messenger, LLC
* SPDX-License-Identifier: AGPL-3.0-only
*/
package org.whispersystems.textsecuregcm.util;
import org.junit.jupiter.api.Test;
import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
import java.nio.ByteBuffer;
import java.util.Map;
import java.util.UUID;
import static org.junit.jupiter.api.Assertions.*;
public class AttributeValuesTest {
@Test
void testUUIDRoundTrip() {
UUID orig = UUID.randomUUID();
AttributeValue av = AttributeValues.fromUUID(orig);
UUID returned = AttributeValues.getUUID(Map.of("foo", av), "foo", null);
assertEquals(orig, returned);
}
@Test
void testLongRoundTrip() {
long orig = 12345;
AttributeValue av = AttributeValues.fromLong(orig);
long returned = AttributeValues.getLong(Map.of("foo", av), "foo", -1);
assertEquals(orig, returned);
}
@Test
void testIntRoundTrip() {
int orig = 12345;
AttributeValue av = AttributeValues.fromInt(orig);
int returned = AttributeValues.getInt(Map.of("foo", av), "foo", -1);
assertEquals(orig, returned);
}
@Test
void testByteBuffer() {
byte[] bytes = {1, 2, 3};
ByteBuffer bb = ByteBuffer.wrap(bytes);
AttributeValue av = AttributeValues.fromByteBuffer(bb);
byte[] returned = av.b().asByteArray();
assertArrayEquals(bytes, returned);
returned = AttributeValues.getByteArray(Map.of("foo", av), "foo", null);
assertArrayEquals(bytes, returned);
}
@Test
void testByteBuffer2() {
final ByteBuffer byteBuffer = ByteBuffer.wrap(new byte[8]);
byteBuffer.putLong(123);
assertEquals(byteBuffer.remaining(), 0);
AttributeValue av = AttributeValues.fromByteBuffer(byteBuffer.flip());
assertArrayEquals(new byte[]{0, 0, 0, 0, 0, 0, 0, 123}, AttributeValues.getByteArray(Map.of("foo", av), "foo", null));
}
}

View File

@@ -79,7 +79,7 @@ public class WebSocketConnectionIntegrationTest extends AbstractRedisClusterTest
executorService = Executors.newSingleThreadExecutor();
messagesCache = new MessagesCache(getRedisCluster(), getRedisCluster(), executorService);
messagesDynamoDb = new MessagesDynamoDb(messagesDynamoDbRule.getDynamoDB(), MessagesDynamoDbRule.TABLE_NAME, Duration.ofDays(7));
messagesDynamoDb = new MessagesDynamoDb(messagesDynamoDbRule.getDynamoDbClient(), MessagesDynamoDbRule.TABLE_NAME, Duration.ofDays(7));
reportMessageManager = mock(ReportMessageManager.class);
account = mock(Account.class);
device = mock(Device.class);