MirrorMirror
/
Signal-Server
mirror of https://github.com/signalapp/Signal-Server.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Add a generalized countItemsMatchingQuery method for DynamoDB stores.

This commit is contained in:
Jon Chambers 2021-02-04 17:28:44 -05:00 committed by Jon Chambers
parent 1dceee3fa0
commit e1f4deaacc
2 changed files with 23 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

22
service/src/main/java/org/whispersystems/textsecuregcm/storage/AbstractDynamoDbStore.java
View File

@ -7,7 +7,13 @@ package org.whispersystems.textsecuregcm.storage;
import com.amazonaws.services.dynamodbv2.document.BatchWriteItemOutcome;
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.Page;
import com.amazonaws.services.dynamodbv2.document.QueryOutcome;
import com.amazonaws.services.dynamodbv2.document.Table;
import com.amazonaws.services.dynamodbv2.document.TableWriteItems;
import com.amazonaws.services.dynamodbv2.document.spec.QuerySpec;
import io.micrometer.core.instrument.Counter;
import io.micrometer.core.instrument.Timer;
import org.slf4j.Logger;
@ -58,6 +64,22 @@ public class AbstractDynamoDbStore {
}
}
protected long countItemsMatchingQuery(final Table table, final QuerySpec querySpec) {
// This is very confusing, but does appear to be the intended behavior. See:
//
// - https://github.com/aws/aws-sdk-java/issues/693
// - https://github.com/aws/aws-sdk-java/issues/915
// - https://docs.aws.amazon.com/amazondynamodb/latest/developerguide/Query.html#Query.Count
long matchingItems = 0;
for (final Page<Item, QueryOutcome> page : table.query(querySpec).pages()) {
matchingItems += page.getLowLevelResult().getQueryResult().getCount();
}
return matchingItems;
}
static <T> void writeInBatches(final Iterable<T> items, final Consumer<List<T>> action) {
final List<T> batch = new ArrayList<>(DYNAMO_DB_MAX_BATCH_SIZE);

6
service/src/main/java/org/whispersystems/textsecuregcm/storage/KeysDynamoDb.java
View File

@ -128,11 +128,7 @@ public class KeysDynamoDb extends AbstractDynamoDbStore implements PreKeyStore {
.withSelect(Select.COUNT)
.withConsistentRead(false);
// This is very confusing, but does appear to be the intended behavior. See:
//
// - https://github.com/aws/aws-sdk-java/issues/693
// - https://github.com/aws/aws-sdk-java/issues/915
return table.query(querySpec).firstPage().getLowLevelResult().getQueryResult().getCount();
return (int)countItemsMatchingQuery(table, querySpec);
});
}