From df4e06f8913bba00d702c9b66120b63bdbbbe1f3 Mon Sep 17 00:00:00 2001 From: Toshihiro Suzuki Date: Mon, 28 Mar 2022 09:21:14 +0900 Subject: [PATCH] Handle lastEvaluatedKey for query in DynamoDB (#534) --- .../storage/StorageIntegrationTestBase.java | 6 +- .../dynamo/DeleteStatementHandler.java | 24 +-- .../com/scalar/db/storage/dynamo/Dynamo.java | 32 +-- .../db/storage/dynamo/EmptyScanner.java | 29 +++ .../db/storage/dynamo/GetItemScanner.java | 73 +++++++ .../storage/dynamo/PutStatementHandler.java | 24 +-- .../db/storage/dynamo/QueryScanner.java | 100 +++++++++ .../scalar/db/storage/dynamo/ScannerImpl.java | 58 ----- .../dynamo/SelectStatementHandler.java | 59 +++--- .../db/storage/dynamo/GetItemScannerTest.java | 148 +++++++++++++ .../db/storage/dynamo/QueryScannerTest.java | 198 ++++++++++++++++++ .../dynamo/SelectStatementHandlerTest.java | 9 +- 12 files changed, 622 insertions(+), 138 deletions(-) create mode 100644 core/src/main/java/com/scalar/db/storage/dynamo/EmptyScanner.java create mode 100644 core/src/main/java/com/scalar/db/storage/dynamo/GetItemScanner.java create mode 100644 core/src/main/java/com/scalar/db/storage/dynamo/QueryScanner.java delete mode 100644 core/src/main/java/com/scalar/db/storage/dynamo/ScannerImpl.java create mode 100644 core/src/test/java/com/scalar/db/storage/dynamo/GetItemScannerTest.java create mode 100644 core/src/test/java/com/scalar/db/storage/dynamo/QueryScannerTest.java diff --git a/core/src/integration-test/java/com/scalar/db/storage/StorageIntegrationTestBase.java b/core/src/integration-test/java/com/scalar/db/storage/StorageIntegrationTestBase.java index 9be6f4b062..ab6a3b6c5a 100644 --- a/core/src/integration-test/java/com/scalar/db/storage/StorageIntegrationTestBase.java +++ b/core/src/integration-test/java/com/scalar/db/storage/StorageIntegrationTestBase.java @@ -54,6 +54,7 @@ public abstract class StorageIntegrationTestBase { private static final String COL_NAME3 = "c3"; private static final String COL_NAME4 = "c4"; private static final String COL_NAME5 = "c5"; + private static final String COL_NAME6 = "c6"; private static boolean initialized; private static DistributedStorage storage; @@ -96,6 +97,7 @@ private void createTable() throws ExecutionException { .addColumn(COL_NAME3, DataType.INT) .addColumn(COL_NAME4, DataType.INT) .addColumn(COL_NAME5, DataType.BOOLEAN) + .addColumn(COL_NAME6, DataType.BLOB) .addPartitionKey(COL_NAME1) .addClusteringKey(COL_NAME4) .addSecondaryIndex(COL_NAME3) @@ -1305,7 +1307,7 @@ public void scan_ScanLargeData_ShouldRetrieveExpectedValues() Key partitionKey = new Key(COL_NAME1, 1); for (int i = 0; i < 345; i++) { Key clusteringKey = new Key(COL_NAME4, i); - storage.put(new Put(partitionKey, clusteringKey)); + storage.put(new Put(partitionKey, clusteringKey).withValue(COL_NAME6, new byte[5000])); } Scan scan = new Scan(partitionKey); @@ -1329,7 +1331,7 @@ public void scan_ScanLargeDataWithOrdering_ShouldRetrieveExpectedValues() Key partitionKey = new Key(COL_NAME1, 1); for (int i = 0; i < 345; i++) { Key clusteringKey = new Key(COL_NAME4, i); - storage.put(new Put(partitionKey, clusteringKey)); + storage.put(new Put(partitionKey, clusteringKey).withValue(COL_NAME6, new byte[5000])); } Scan scan = new Scan(partitionKey).withOrdering(new Ordering(COL_NAME4, Order.ASC)); diff --git a/core/src/main/java/com/scalar/db/storage/dynamo/DeleteStatementHandler.java b/core/src/main/java/com/scalar/db/storage/dynamo/DeleteStatementHandler.java index b37fed2c8b..df332d5493 100644 --- a/core/src/main/java/com/scalar/db/storage/dynamo/DeleteStatementHandler.java +++ b/core/src/main/java/com/scalar/db/storage/dynamo/DeleteStatementHandler.java @@ -1,17 +1,15 @@ package com.scalar.db.storage.dynamo; +import static com.google.common.base.Preconditions.checkNotNull; + import com.scalar.db.api.Delete; import com.scalar.db.api.DeleteIfExists; -import com.scalar.db.api.Operation; import com.scalar.db.api.TableMetadata; import com.scalar.db.exception.storage.ExecutionException; import com.scalar.db.exception.storage.NoMutationException; import com.scalar.db.exception.storage.RetriableExecutionException; import com.scalar.db.util.TableMetadataManager; -import java.util.Collections; -import java.util.List; import java.util.Map; -import javax.annotation.Nonnull; import javax.annotation.concurrent.ThreadSafe; import software.amazon.awssdk.services.dynamodb.DynamoDbClient; import software.amazon.awssdk.services.dynamodb.model.AttributeValue; @@ -26,19 +24,17 @@ * @author Yuji Ito */ @ThreadSafe -public class DeleteStatementHandler extends StatementHandler { +public class DeleteStatementHandler { + private final DynamoDbClient client; + private final TableMetadataManager metadataManager; public DeleteStatementHandler(DynamoDbClient client, TableMetadataManager metadataManager) { - super(client, metadataManager); + this.client = checkNotNull(client); + this.metadataManager = checkNotNull(metadataManager); } - @Nonnull - @Override - public List> handle(Operation operation) throws ExecutionException { - checkArgument(operation, Delete.class); - Delete delete = (Delete) operation; - - TableMetadata tableMetadata = metadataManager.getTableMetadata(operation); + public void handle(Delete delete) throws ExecutionException { + TableMetadata tableMetadata = metadataManager.getTableMetadata(delete); try { delete(delete, tableMetadata); } catch (ConditionalCheckFailedException e) { @@ -48,8 +44,6 @@ public List> handle(Operation operation) throws Exec } catch (DynamoDbException e) { throw new ExecutionException(e.getMessage(), e); } - - return Collections.emptyList(); } private void delete(Delete delete, TableMetadata tableMetadata) { diff --git a/core/src/main/java/com/scalar/db/storage/dynamo/Dynamo.java b/core/src/main/java/com/scalar/db/storage/dynamo/Dynamo.java index 5bd0def4fc..8a955335f9 100644 --- a/core/src/main/java/com/scalar/db/storage/dynamo/Dynamo.java +++ b/core/src/main/java/com/scalar/db/storage/dynamo/Dynamo.java @@ -17,9 +17,9 @@ import com.scalar.db.storage.common.checker.OperationChecker; import com.scalar.db.util.ScalarDbUtils; import com.scalar.db.util.TableMetadataManager; +import java.io.IOException; import java.net.URI; import java.util.List; -import java.util.Map; import java.util.Optional; import javax.annotation.Nonnull; import javax.annotation.concurrent.ThreadSafe; @@ -30,7 +30,6 @@ import software.amazon.awssdk.regions.Region; import software.amazon.awssdk.services.dynamodb.DynamoDbClient; import software.amazon.awssdk.services.dynamodb.DynamoDbClientBuilder; -import software.amazon.awssdk.services.dynamodb.model.AttributeValue; /** * A storage implementation with DynamoDB for {@link DistributedStorage} @@ -82,16 +81,23 @@ public Optional get(Get get) throws ExecutionException { TableMetadata metadata = metadataManager.getTableMetadata(get); ScalarDbUtils.addProjectionsForKeys(get, metadata); - List> items = selectStatementHandler.handle(get); - if (items.size() > 1) { - throw new IllegalArgumentException("please use scan() for non-exact match selection"); - } - if (items.isEmpty() || items.get(0) == null) { - return Optional.empty(); + Scanner scanner = null; + try { + scanner = selectStatementHandler.handle(get); + Optional ret = scanner.one(); + if (scanner.one().isPresent()) { + throw new IllegalArgumentException("please use scan() for non-exact match selection"); + } + return ret; + } finally { + if (scanner != null) { + try { + scanner.close(); + } catch (IOException e) { + LOGGER.warn("failed to close the scanner", e); + } + } } - - return Optional.of( - new ResultInterpreter(get.getProjections(), metadata).interpret(items.get(0))); } @Override @@ -101,9 +107,7 @@ public Scanner scan(Scan scan) throws ExecutionException { TableMetadata metadata = metadataManager.getTableMetadata(scan); ScalarDbUtils.addProjectionsForKeys(scan, metadata); - List> items = selectStatementHandler.handle(scan); - - return new ScannerImpl(items, new ResultInterpreter(scan.getProjections(), metadata)); + return selectStatementHandler.handle(scan); } @Override diff --git a/core/src/main/java/com/scalar/db/storage/dynamo/EmptyScanner.java b/core/src/main/java/com/scalar/db/storage/dynamo/EmptyScanner.java new file mode 100644 index 0000000000..9d15432d71 --- /dev/null +++ b/core/src/main/java/com/scalar/db/storage/dynamo/EmptyScanner.java @@ -0,0 +1,29 @@ +package com.scalar.db.storage.dynamo; + +import com.scalar.db.api.Result; +import com.scalar.db.api.Scanner; +import java.util.Collections; +import java.util.Iterator; +import java.util.List; +import java.util.Optional; + +public class EmptyScanner implements Scanner { + + @Override + public Optional one() { + return Optional.empty(); + } + + @Override + public List all() { + return Collections.emptyList(); + } + + @Override + public void close() {} + + @Override + public Iterator iterator() { + return Collections.emptyIterator(); + } +} diff --git a/core/src/main/java/com/scalar/db/storage/dynamo/GetItemScanner.java b/core/src/main/java/com/scalar/db/storage/dynamo/GetItemScanner.java new file mode 100644 index 0000000000..6e231a8d62 --- /dev/null +++ b/core/src/main/java/com/scalar/db/storage/dynamo/GetItemScanner.java @@ -0,0 +1,73 @@ +package com.scalar.db.storage.dynamo; + +import com.scalar.db.api.Result; +import com.scalar.db.api.Scanner; +import com.scalar.db.storage.common.ScannerIterator; +import java.util.Collections; +import java.util.Iterator; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import javax.annotation.Nonnull; +import javax.annotation.concurrent.NotThreadSafe; +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; + +@NotThreadSafe +public final class GetItemScanner implements Scanner { + private final Map item; + private final ResultInterpreter resultInterpreter; + + private boolean hasNext; + private ScannerIterator scannerIterator; + + public GetItemScanner( + DynamoDbClient client, GetItemRequest request, ResultInterpreter resultInterpreter) { + GetItemResponse response = client.getItem(request); + if (response.hasItem()) { + item = response.item(); + hasNext = true; + } else { + item = null; + hasNext = false; + } + this.resultInterpreter = resultInterpreter; + } + + @Override + @Nonnull + public Optional one() { + if (!hasNext) { + return Optional.empty(); + } + + Result result = resultInterpreter.interpret(item); + hasNext = false; + return Optional.of(result); + } + + @Override + @Nonnull + public List all() { + if (!hasNext) { + return Collections.emptyList(); + } + Result result = resultInterpreter.interpret(item); + hasNext = false; + return Collections.singletonList(result); + } + + @Override + @Nonnull + public Iterator iterator() { + if (scannerIterator == null) { + scannerIterator = new ScannerIterator(this); + } + return scannerIterator; + } + + @Override + public void close() {} +} diff --git a/core/src/main/java/com/scalar/db/storage/dynamo/PutStatementHandler.java b/core/src/main/java/com/scalar/db/storage/dynamo/PutStatementHandler.java index fbbf6de56b..9ea122cbc1 100644 --- a/core/src/main/java/com/scalar/db/storage/dynamo/PutStatementHandler.java +++ b/core/src/main/java/com/scalar/db/storage/dynamo/PutStatementHandler.java @@ -1,6 +1,7 @@ package com.scalar.db.storage.dynamo; -import com.scalar.db.api.Operation; +import static com.google.common.base.Preconditions.checkNotNull; + import com.scalar.db.api.Put; import com.scalar.db.api.PutIfExists; import com.scalar.db.api.PutIfNotExists; @@ -9,10 +10,7 @@ import com.scalar.db.exception.storage.NoMutationException; import com.scalar.db.exception.storage.RetriableExecutionException; import com.scalar.db.util.TableMetadataManager; -import java.util.Collections; -import java.util.List; import java.util.Map; -import javax.annotation.Nonnull; import javax.annotation.concurrent.ThreadSafe; import software.amazon.awssdk.services.dynamodb.DynamoDbClient; import software.amazon.awssdk.services.dynamodb.model.AttributeValue; @@ -27,19 +25,17 @@ * @author Yuji Ito */ @ThreadSafe -public class PutStatementHandler extends StatementHandler { +public class PutStatementHandler { + private final DynamoDbClient client; + private final TableMetadataManager metadataManager; public PutStatementHandler(DynamoDbClient client, TableMetadataManager metadataManager) { - super(client, metadataManager); + this.client = checkNotNull(client); + this.metadataManager = checkNotNull(metadataManager); } - @Nonnull - @Override - public List> handle(Operation operation) throws ExecutionException { - checkArgument(operation, Put.class); - Put put = (Put) operation; - - TableMetadata tableMetadata = metadataManager.getTableMetadata(operation); + public void handle(Put put) throws ExecutionException { + TableMetadata tableMetadata = metadataManager.getTableMetadata(put); try { execute(put, tableMetadata); } catch (ConditionalCheckFailedException e) { @@ -49,8 +45,6 @@ public List> handle(Operation operation) throws Exec } catch (DynamoDbException e) { throw new ExecutionException(e.getMessage(), e); } - - return Collections.emptyList(); } private void execute(Put put, TableMetadata tableMetadata) { diff --git a/core/src/main/java/com/scalar/db/storage/dynamo/QueryScanner.java b/core/src/main/java/com/scalar/db/storage/dynamo/QueryScanner.java new file mode 100644 index 0000000000..1121ba55f6 --- /dev/null +++ b/core/src/main/java/com/scalar/db/storage/dynamo/QueryScanner.java @@ -0,0 +1,100 @@ +package com.scalar.db.storage.dynamo; + +import com.scalar.db.api.Result; +import com.scalar.db.api.Scanner; +import com.scalar.db.storage.common.ScannerIterator; +import java.util.ArrayList; +import java.util.Iterator; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import javax.annotation.Nonnull; +import javax.annotation.Nullable; +import software.amazon.awssdk.services.dynamodb.DynamoDbClient; +import software.amazon.awssdk.services.dynamodb.model.AttributeValue; +import software.amazon.awssdk.services.dynamodb.model.QueryRequest; +import software.amazon.awssdk.services.dynamodb.model.QueryResponse; + +public class QueryScanner implements Scanner { + + private final DynamoDbClient client; + private final QueryRequest request; + private final ResultInterpreter resultInterpreter; + + private Iterator> itemsIterator; + @Nullable private Map lastEvaluatedKey; + private int totalResultCount; + + private ScannerIterator scannerIterator; + + public QueryScanner( + DynamoDbClient client, QueryRequest request, ResultInterpreter resultInterpreter) { + this.client = client; + this.request = request; + this.resultInterpreter = resultInterpreter; + + query(request); + } + + @Override + @Nonnull + public Optional one() { + if (!hasNext()) { + return Optional.empty(); + } + + return Optional.of(resultInterpreter.interpret(itemsIterator.next())); + } + + private boolean hasNext() { + if (itemsIterator.hasNext()) { + return true; + } + if (lastEvaluatedKey != null) { + QueryRequest.Builder builder = request.toBuilder(); + builder.exclusiveStartKey(lastEvaluatedKey); + query(builder.build()); + return itemsIterator.hasNext(); + } + return false; + } + + private void query(QueryRequest request) { + QueryResponse response = client.query(request); + List> items = response.items(); + totalResultCount += items.size(); + itemsIterator = items.iterator(); + if ((request.limit() == null || totalResultCount < request.limit()) + && response.hasLastEvaluatedKey()) { + lastEvaluatedKey = response.lastEvaluatedKey(); + } else { + lastEvaluatedKey = null; + } + } + + @Override + @Nonnull + public List all() { + List ret = new ArrayList<>(); + while (true) { + Optional one = one(); + if (!one.isPresent()) { + break; + } + ret.add(one.get()); + } + return ret; + } + + @Override + @Nonnull + public Iterator iterator() { + if (scannerIterator == null) { + scannerIterator = new ScannerIterator(this); + } + return scannerIterator; + } + + @Override + public void close() {} +} diff --git a/core/src/main/java/com/scalar/db/storage/dynamo/ScannerImpl.java b/core/src/main/java/com/scalar/db/storage/dynamo/ScannerImpl.java deleted file mode 100644 index 6ac25fbcf2..0000000000 --- a/core/src/main/java/com/scalar/db/storage/dynamo/ScannerImpl.java +++ /dev/null @@ -1,58 +0,0 @@ -package com.scalar.db.storage.dynamo; - -import com.scalar.db.api.Result; -import com.scalar.db.api.Scanner; -import com.scalar.db.storage.common.ScannerIterator; -import java.util.ArrayList; -import java.util.Iterator; -import java.util.List; -import java.util.Map; -import java.util.Optional; -import javax.annotation.Nonnull; -import javax.annotation.concurrent.NotThreadSafe; -import software.amazon.awssdk.services.dynamodb.model.AttributeValue; - -@NotThreadSafe -public final class ScannerImpl implements Scanner { - private final List> items; - private final ResultInterpreter resultInterpreter; - - private ScannerIterator scannerIterator; - - public ScannerImpl(List> items, ResultInterpreter resultInterpreter) { - this.items = items; - this.resultInterpreter = resultInterpreter; - } - - @Override - @Nonnull - public Optional one() { - if (items.isEmpty()) { - return Optional.empty(); - } - Map item = items.remove(0); - - return Optional.of(resultInterpreter.interpret(item)); - } - - @Override - @Nonnull - public List all() { - List results = new ArrayList<>(); - items.forEach(i -> results.add(resultInterpreter.interpret(i))); - items.clear(); - return results; - } - - @Override - @Nonnull - public Iterator iterator() { - if (scannerIterator == null) { - scannerIterator = new ScannerIterator(this); - } - return scannerIterator; - } - - @Override - public void close() {} -} diff --git a/core/src/main/java/com/scalar/db/storage/dynamo/SelectStatementHandler.java b/core/src/main/java/com/scalar/db/storage/dynamo/SelectStatementHandler.java index 7b3ac649cd..a2e3fbce0d 100644 --- a/core/src/main/java/com/scalar/db/storage/dynamo/SelectStatementHandler.java +++ b/core/src/main/java/com/scalar/db/storage/dynamo/SelectStatementHandler.java @@ -1,13 +1,15 @@ package com.scalar.db.storage.dynamo; +import static com.google.common.base.Preconditions.checkNotNull; + import com.google.common.collect.ImmutableMap; import com.google.common.primitives.UnsignedBytes; import com.scalar.db.api.Consistency; import com.scalar.db.api.Get; -import com.scalar.db.api.Operation; import com.scalar.db.api.Scan; import com.scalar.db.api.Scan.Ordering; import com.scalar.db.api.Scan.Ordering.Order; +import com.scalar.db.api.Scanner; import com.scalar.db.api.Selection; import com.scalar.db.api.TableMetadata; import com.scalar.db.exception.storage.ExecutionException; @@ -19,7 +21,6 @@ import com.scalar.db.util.TableMetadataManager; import java.nio.ByteBuffer; import java.util.ArrayList; -import java.util.Collections; import java.util.HashMap; import java.util.List; import java.util.Map; @@ -32,9 +33,7 @@ import software.amazon.awssdk.services.dynamodb.model.DynamoDbException; import software.amazon.awssdk.services.dynamodb.model.DynamoDbRequest; import software.amazon.awssdk.services.dynamodb.model.GetItemRequest; -import software.amazon.awssdk.services.dynamodb.model.GetItemResponse; import software.amazon.awssdk.services.dynamodb.model.QueryRequest; -import software.amazon.awssdk.services.dynamodb.model.QueryResponse; /** * A handler class for select statement @@ -42,7 +41,9 @@ * @author Yuji Ito, Toshihiro Suzuki */ @ThreadSafe -public class SelectStatementHandler extends StatementHandler { +public class SelectStatementHandler { + private final DynamoDbClient client; + private final TableMetadataManager metadataManager; /** * Constructs a {@code SelectStatementHandler} with the specified {@link DynamoDbClient} and a new @@ -52,31 +53,30 @@ public class SelectStatementHandler extends StatementHandler { * @param metadataManager {@code TableMetadataManager} */ public SelectStatementHandler(DynamoDbClient client, TableMetadataManager metadataManager) { - super(client, metadataManager); + this.client = checkNotNull(client); + this.metadataManager = checkNotNull(metadataManager); } @Nonnull - @Override - public List> handle(Operation operation) throws ExecutionException { - checkArgument(operation, Get.class, Scan.class); - - TableMetadata tableMetadata = metadataManager.getTableMetadata(operation); + public Scanner handle(Selection selection) throws ExecutionException { + TableMetadata tableMetadata = metadataManager.getTableMetadata(selection); try { - if (ScalarDbUtils.isSecondaryIndexSpecified(operation, tableMetadata)) { - return executeQueryWithIndex((Selection) operation, tableMetadata); + if (ScalarDbUtils.isSecondaryIndexSpecified(selection, tableMetadata)) { + return executeQueryWithIndex(selection, tableMetadata); } - if (operation instanceof Get) { - return executeGet((Get) operation, tableMetadata); + if (selection instanceof Get) { + return executeGet((Get) selection, tableMetadata); } else { - return executeQuery((Scan) operation, tableMetadata); + return executeQuery((Scan) selection, tableMetadata); } + } catch (DynamoDbException e) { throw new ExecutionException(e.getMessage(), e); } } - private List> executeGet(Get get, TableMetadata tableMetadata) { + private Scanner executeGet(Get get, TableMetadata tableMetadata) { DynamoOperation dynamoOperation = new DynamoOperation(get, tableMetadata); GetItemRequest.Builder builder = @@ -92,16 +92,11 @@ private List> executeGet(Get get, TableMetadata tabl builder.consistentRead(true); } - GetItemResponse getItemResponse = client.getItem(builder.build()); - if (getItemResponse.hasItem()) { - return Collections.singletonList(getItemResponse.item()); - } else { - return Collections.emptyList(); - } + return new GetItemScanner( + client, builder.build(), new ResultInterpreter(get.getProjections(), tableMetadata)); } - private List> executeQueryWithIndex( - Selection selection, TableMetadata tableMetadata) { + private Scanner executeQueryWithIndex(Selection selection, TableMetadata tableMetadata) { DynamoOperation dynamoOperation = new DynamoOperation(selection, tableMetadata); Value keyValue = selection.getPartitionKey().get().get(0); String column = keyValue.getName(); @@ -130,17 +125,17 @@ private List> executeQueryWithIndex( } } - QueryResponse queryResponse = client.query(builder.build()); - return new ArrayList<>(queryResponse.items()); + return new QueryScanner( + client, builder.build(), new ResultInterpreter(selection.getProjections(), tableMetadata)); } - private List> executeQuery(Scan scan, TableMetadata tableMetadata) { + private Scanner executeQuery(Scan scan, TableMetadata tableMetadata) { DynamoOperation dynamoOperation = new DynamoOperation(scan, tableMetadata); QueryRequest.Builder builder = QueryRequest.builder().tableName(dynamoOperation.getTableName()); if (!setConditions(builder, scan, tableMetadata)) { - // if setConditions() fails, return empty list - return new ArrayList<>(); + // if setConditions() fails, return an empty scanner + return new EmptyScanner(); } if (!scan.getOrderings().isEmpty()) { @@ -163,8 +158,8 @@ private List> executeQuery(Scan scan, TableMetadata builder.consistentRead(true); } - QueryResponse queryResponse = client.query(builder.build()); - return new ArrayList<>(queryResponse.items()); + return new QueryScanner( + client, builder.build(), new ResultInterpreter(scan.getProjections(), tableMetadata)); } private void projectionExpression(DynamoDbRequest.Builder builder, Selection selection) { diff --git a/core/src/test/java/com/scalar/db/storage/dynamo/GetItemScannerTest.java b/core/src/test/java/com/scalar/db/storage/dynamo/GetItemScannerTest.java new file mode 100644 index 0000000000..a96ce33213 --- /dev/null +++ b/core/src/test/java/com/scalar/db/storage/dynamo/GetItemScannerTest.java @@ -0,0 +1,148 @@ +package com.scalar.db.storage.dynamo; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; +import static org.mockito.Mockito.when; + +import com.scalar.db.api.Result; +import java.util.Collections; +import java.util.Iterator; +import java.util.List; +import java.util.Map; +import java.util.NoSuchElementException; +import java.util.Optional; +import org.junit.Before; +import org.junit.Test; +import org.mockito.Mock; +import org.mockito.MockitoAnnotations; +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; + +public class GetItemScannerTest { + + @Mock private DynamoDbClient client; + @Mock private GetItemRequest request; + @Mock private ResultInterpreter resultInterpreter; + + @Mock private GetItemResponse response; + @Mock private Result result; + + @Before + public void setUp() throws Exception { + MockitoAnnotations.openMocks(this).close(); + + // Arrange + when(client.getItem(request)).thenReturn(response); + } + + @Test + public void one_WhenItemReturned_ShouldReturnResult() { + // Arrange + Map item = Collections.emptyMap(); + + when(response.hasItem()).thenReturn(true); + when(response.item()).thenReturn(item); + when(resultInterpreter.interpret(item)).thenReturn(result); + + GetItemScanner getItemScanner = new GetItemScanner(client, request, resultInterpreter); + + // Act + Optional actual1 = getItemScanner.one(); + Optional actual2 = getItemScanner.one(); + + // Assert + assertThat(actual1).isPresent(); + assertThat(actual1.get()).isEqualTo(result); + + assertThat(actual2).isNotPresent(); + } + + @Test + public void one_WhenNoItemReturned_ShouldReturnEmpty() { + // Arrange + when(response.hasItem()).thenReturn(false); + + GetItemScanner getItemScanner = new GetItemScanner(client, request, resultInterpreter); + + // Act + Optional actual = getItemScanner.one(); + + // Assert + assertThat(actual).isNotPresent(); + } + + @Test + public void all_WhenItemReturned_ShouldReturnResult() { + // Arrange + Map item = Collections.emptyMap(); + + when(response.hasItem()).thenReturn(true); + when(response.item()).thenReturn(item); + when(resultInterpreter.interpret(item)).thenReturn(result); + + GetItemScanner getItemScanner = new GetItemScanner(client, request, resultInterpreter); + + // Act + List results1 = getItemScanner.all(); + List results2 = getItemScanner.all(); + + // Assert + assertThat(results1.size()).isEqualTo(1); + assertThat(results1.get(0)).isEqualTo(result); + + assertThat(results2).isEmpty(); + } + + @Test + public void all_WhenNoItemReturned_ShouldReturnEmpty() { + // Arrange + when(response.hasItem()).thenReturn(false); + + GetItemScanner getItemScanner = new GetItemScanner(client, request, resultInterpreter); + + // Act + List results = getItemScanner.all(); + + // Assert + assertThat(results).isEmpty(); + } + + @Test + public void iterator_WhenItemReturned_ShouldReturnResult() { + // Arrange + Map item = Collections.emptyMap(); + + when(response.hasItem()).thenReturn(true); + when(response.item()).thenReturn(item); + when(resultInterpreter.interpret(item)).thenReturn(result); + + GetItemScanner getItemScanner = new GetItemScanner(client, request, resultInterpreter); + + // Act + Iterator iterator = getItemScanner.iterator(); + + // Assert + assertThat(iterator.hasNext()).isTrue(); + Result actual1 = iterator.next(); + assertThat(actual1).isEqualTo(result); + assertThat(iterator.hasNext()).isFalse(); + assertThatThrownBy(iterator::next).isInstanceOf(NoSuchElementException.class); + } + + @Test + public void iterator_WhenNoItemReturned_ShouldReturnEmpty() { + // Arrange + when(response.hasItem()).thenReturn(false); + + GetItemScanner getItemScanner = new GetItemScanner(client, request, resultInterpreter); + + // Act + Iterator iterator = getItemScanner.iterator(); + + // Assert + assertThat(iterator.hasNext()).isFalse(); + assertThatThrownBy(iterator::next).isInstanceOf(NoSuchElementException.class); + } +} diff --git a/core/src/test/java/com/scalar/db/storage/dynamo/QueryScannerTest.java b/core/src/test/java/com/scalar/db/storage/dynamo/QueryScannerTest.java new file mode 100644 index 0000000000..3774c9bb62 --- /dev/null +++ b/core/src/test/java/com/scalar/db/storage/dynamo/QueryScannerTest.java @@ -0,0 +1,198 @@ +package com.scalar.db.storage.dynamo; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import com.scalar.db.api.Result; +import java.util.Arrays; +import java.util.Collections; +import java.util.Iterator; +import java.util.List; +import java.util.Map; +import java.util.NoSuchElementException; +import java.util.Optional; +import org.junit.Before; +import org.junit.Test; +import org.mockito.Mock; +import org.mockito.MockitoAnnotations; +import software.amazon.awssdk.services.dynamodb.DynamoDbClient; +import software.amazon.awssdk.services.dynamodb.model.AttributeValue; +import software.amazon.awssdk.services.dynamodb.model.QueryRequest; +import software.amazon.awssdk.services.dynamodb.model.QueryResponse; + +public class QueryScannerTest { + + @Mock private DynamoDbClient client; + @Mock private QueryRequest request; + @Mock private ResultInterpreter resultInterpreter; + + @Mock private QueryRequest.Builder builder; + @Mock private QueryResponse response; + @Mock private Result result; + + @Before + public void setUp() throws Exception { + MockitoAnnotations.openMocks(this).close(); + + // Arrange + when(client.query(request)).thenReturn(response); + + when(request.limit()).thenReturn(null); + when(request.toBuilder()).thenReturn(builder); + when(builder.build()).thenReturn(request); + } + + @Test + public void one_ShouldReturnResult() { + // Arrange + Map item = Collections.emptyMap(); + List> items = Arrays.asList(item, item, item); + when(response.items()).thenReturn(items); + when(resultInterpreter.interpret(item)).thenReturn(result); + + QueryScanner queryScanner = new QueryScanner(client, request, resultInterpreter); + + // Act + Optional actual1 = queryScanner.one(); + Optional actual2 = queryScanner.one(); + Optional actual3 = queryScanner.one(); + Optional actual4 = queryScanner.one(); + + // Assert + assertThat(actual1).isPresent(); + assertThat(actual1.get()).isEqualTo(result); + assertThat(actual2).isPresent(); + assertThat(actual2.get()).isEqualTo(result); + assertThat(actual3).isPresent(); + assertThat(actual3.get()).isEqualTo(result); + assertThat(actual4).isNotPresent(); + + verify(resultInterpreter, times(3)).interpret(item); + } + + @Test + public void all_ShouldReturnResults() { + // Arrange + Map item = Collections.emptyMap(); + List> items = Arrays.asList(item, item, item); + when(response.items()).thenReturn(items); + when(resultInterpreter.interpret(item)).thenReturn(result); + + QueryScanner queryScanner = new QueryScanner(client, request, resultInterpreter); + + // Act + List results1 = queryScanner.all(); + List results2 = queryScanner.all(); + + // Assert + assertThat(results1.size()).isEqualTo(3); + assertThat(results1.get(0)).isEqualTo(result); + assertThat(results1.get(1)).isEqualTo(result); + assertThat(results1.get(2)).isEqualTo(result); + assertThat(results2).isEmpty(); + + verify(resultInterpreter, times(3)).interpret(item); + } + + @Test + public void iterator_ShouldReturnResults() { + // Arrange + Map item = Collections.emptyMap(); + List> items = Arrays.asList(item, item, item); + when(response.items()).thenReturn(items); + when(resultInterpreter.interpret(item)).thenReturn(result); + + QueryScanner queryScanner = new QueryScanner(client, request, resultInterpreter); + + // Act + Iterator iterator = queryScanner.iterator(); + + // Assert + assertThat(iterator.hasNext()).isTrue(); + assertThat(iterator.next()).isEqualTo(result); + assertThat(iterator.hasNext()).isTrue(); + assertThat(iterator.next()).isEqualTo(result); + assertThat(iterator.hasNext()).isTrue(); + assertThat(iterator.next()).isEqualTo(result); + assertThat(iterator.hasNext()).isFalse(); + assertThatThrownBy(iterator::next).isInstanceOf(NoSuchElementException.class); + + verify(resultInterpreter, times(3)).interpret(item); + } + + @Test + public void one_ResponseWithLastEvaluatedKey_ShouldReturnResults() { + // Arrange + Map item = Collections.emptyMap(); + List> items = Arrays.asList(item, item); + Map lastEvaluatedKey = Collections.emptyMap(); + + when(response.items()).thenReturn(items); + when(response.hasLastEvaluatedKey()).thenReturn(true).thenReturn(false); + when(response.lastEvaluatedKey()).thenReturn(lastEvaluatedKey); + when(resultInterpreter.interpret(item)).thenReturn(result); + + QueryScanner queryScanner = new QueryScanner(client, request, resultInterpreter); + + // Act + Optional actual1 = queryScanner.one(); + Optional actual2 = queryScanner.one(); + Optional actual3 = queryScanner.one(); + Optional actual4 = queryScanner.one(); + Optional actual5 = queryScanner.one(); + + // Assert + assertThat(actual1).isPresent(); + assertThat(actual1.get()).isEqualTo(result); + assertThat(actual2).isPresent(); + assertThat(actual2.get()).isEqualTo(result); + assertThat(actual3).isPresent(); + assertThat(actual3.get()).isEqualTo(result); + assertThat(actual4).isPresent(); + assertThat(actual4.get()).isEqualTo(result); + assertThat(actual5).isNotPresent(); + + verify(resultInterpreter, times(4)).interpret(item); + verify(builder).exclusiveStartKey(lastEvaluatedKey); + } + + @Test + public void one_RequestWithLimitAndResponseWithLastEvaluatedKey_ShouldReturnResults() { + // Arrange + Map item = Collections.emptyMap(); + List> items = Arrays.asList(item, item); + Map lastEvaluatedKey = Collections.emptyMap(); + + when(request.limit()).thenReturn(4); + when(response.items()).thenReturn(items); + when(response.hasLastEvaluatedKey()).thenReturn(true); + when(response.lastEvaluatedKey()).thenReturn(lastEvaluatedKey); + when(resultInterpreter.interpret(item)).thenReturn(result); + + QueryScanner queryScanner = new QueryScanner(client, request, resultInterpreter); + + // Act + Optional actual1 = queryScanner.one(); + Optional actual2 = queryScanner.one(); + Optional actual3 = queryScanner.one(); + Optional actual4 = queryScanner.one(); + Optional actual5 = queryScanner.one(); + + // Assert + assertThat(actual1).isPresent(); + assertThat(actual1.get()).isEqualTo(result); + assertThat(actual2).isPresent(); + assertThat(actual2.get()).isEqualTo(result); + assertThat(actual3).isPresent(); + assertThat(actual3.get()).isEqualTo(result); + assertThat(actual4).isPresent(); + assertThat(actual4.get()).isEqualTo(result); + assertThat(actual5).isNotPresent(); + + verify(resultInterpreter, times(4)).interpret(item); + verify(builder).exclusiveStartKey(lastEvaluatedKey); + } +} diff --git a/core/src/test/java/com/scalar/db/storage/dynamo/SelectStatementHandlerTest.java b/core/src/test/java/com/scalar/db/storage/dynamo/SelectStatementHandlerTest.java index a5f9cbf315..c9f3daaab7 100644 --- a/core/src/test/java/com/scalar/db/storage/dynamo/SelectStatementHandlerTest.java +++ b/core/src/test/java/com/scalar/db/storage/dynamo/SelectStatementHandlerTest.java @@ -12,8 +12,10 @@ import com.google.common.collect.ImmutableMap; import com.scalar.db.api.Get; import com.scalar.db.api.Operation; +import com.scalar.db.api.Result; import com.scalar.db.api.Scan; import com.scalar.db.api.Scan.Ordering.Order; +import com.scalar.db.api.Scanner; import com.scalar.db.api.TableMetadata; import com.scalar.db.exception.storage.ExecutionException; import com.scalar.db.io.Key; @@ -116,7 +118,7 @@ public void handle_GetOperationGiven_ShouldCallGetItem() { } @Test - public void handle_GetOperationNoItemReturned_ShouldReturnEmptyList() throws Exception { + public void handle_GetOperationNoItemReturned_ShouldReturnEmptyScanner() throws Exception { // Arrange when(client.getItem(any(GetItemRequest.class))).thenReturn(getResponse); when(getResponse.hasItem()).thenReturn(false); @@ -124,7 +126,10 @@ public void handle_GetOperationNoItemReturned_ShouldReturnEmptyList() throws Exc Get get = prepareGet(); // Act Assert - List> actual = handler.handle(get); + List actual; + try (Scanner scanner = handler.handle(get)) { + actual = scanner.all(); + } // Assert assertThat(actual).isEmpty();