Adding support for Select in ScanEnhancedRequest (#5519)

aws · Sep 3, 2024 · 3f954e4 · 3f954e4
1 parent 71e129f
commit 3f954e4
Show file tree

Hide file tree

Showing 5 changed files with 199 additions and 0 deletions.
diff --git a/.changes/next-release/feature-DynamoDBEnhancedClient-da7de1f.json b/.changes/next-release/feature-DynamoDBEnhancedClient-da7de1f.json
@@ -0,0 +1,6 @@
+{
+    "category": "DynamoDB Enhanced Client", 
+    "contributor": "shetsa-amzn", 
+    "type": "feature", 
+    "description": "Adding support for Select in ScanEnhancedRequest"
+}
diff --git a/...main/java/software/amazon/awssdk/enhanced/dynamodb/internal/operations/ScanOperation.java b/...main/java/software/amazon/awssdk/enhanced/dynamodb/internal/operations/ScanOperation.java
@@ -82,6 +82,7 @@ public ScanRequest generateRequest(TableSchema<T> tableSchema,
             .exclusiveStartKey(this.request.exclusiveStartKey())
             .consistentRead(this.request.consistentRead())
             .returnConsumedCapacity(this.request.returnConsumedCapacity())
+            .select(this.request.select())
             .expressionAttributeValues(expressionValues)
             .expressionAttributeNames(expressionNames)
             .projectionExpression(projectionExpressionAsString);

diff --git a/...ced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/model/ScanEnhancedRequest.java b/...ced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/model/ScanEnhancedRequest.java
@@ -32,6 +32,7 @@
 import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
 import software.amazon.awssdk.services.dynamodb.model.ReturnConsumedCapacity;
 import software.amazon.awssdk.services.dynamodb.model.ScanRequest;
+import software.amazon.awssdk.services.dynamodb.model.Select;
 import software.amazon.awssdk.utils.Validate;
 
 /**
@@ -48,6 +49,7 @@ public final class ScanEnhancedRequest {
     private final Integer limit;
     private final Boolean consistentRead;
     private final Expression filterExpression;
+    private final Select select;
     private final List<NestedAttributeName> attributesToProject;
     private final Integer segment;
     private final Integer totalSegments;
@@ -60,6 +62,7 @@ private ScanEnhancedRequest(Builder builder) {
         this.totalSegments = builder.totalSegments;
         this.consistentRead = builder.consistentRead;
         this.filterExpression = builder.filterExpression;
+        this.select = builder.select;
         this.returnConsumedCapacity = builder.returnConsumedCapacity;
         this.attributesToProject = builder.attributesToProject != null
                 ? Collections.unmodifiableList(builder.attributesToProject)
@@ -83,6 +86,7 @@ public Builder toBuilder() {
                         .totalSegments(totalSegments)
                         .consistentRead(consistentRead)
                         .filterExpression(filterExpression)
+                        .select(select)
                         .returnConsumedCapacity(returnConsumedCapacity)
                         .addNestedAttributesToProject(attributesToProject);
     }
@@ -129,6 +133,24 @@ public Expression filterExpression() {
         return filterExpression;
     }
 
+    /**
+     * Returns the value of select, or null if it doesn't exist.
+     * @return
+     */
+    public Select select() {
+        return select;
+    }
+
+    /**
+     * Returns the value of select as a string, or null if it doesn't exist.
+     * @return
+     */
+    public String selectAsString() {
+        return String.valueOf(select);
+    }
+
+    /**
+
     /**
      * Returns the list of projected attributes on this request object, or an null if no projection is specified.
      * Nested attributes are represented using the '.' separator. Example : foo.bar is represented as "foo.bar" which is
@@ -204,6 +226,11 @@ public boolean equals(Object o) {
             ? !returnConsumedCapacity.equals(scan.returnConsumedCapacity) : scan.returnConsumedCapacity != null) {
             return false;
         }
+
+        if (select != null ? ! select.equals(scan.select) : scan.select != null) {
+            return false;
+        }
+
         return filterExpression != null ? filterExpression.equals(scan.filterExpression) : scan.filterExpression == null;
     }
 
@@ -215,6 +242,7 @@ public int hashCode() {
         result = 31 * result + (totalSegments != null ? totalSegments.hashCode() : 0);
         result = 31 * result + (consistentRead != null ? consistentRead.hashCode() : 0);
         result = 31 * result + (filterExpression != null ? filterExpression.hashCode() : 0);
+        result = 31 * result + (select != null ? select.hashCode() : 0);
         result = 31 * result + (attributesToProject != null ? attributesToProject.hashCode() : 0);
         result = 31 * result + (returnConsumedCapacity != null ? returnConsumedCapacity.hashCode() : 0);
         return result;
@@ -229,6 +257,7 @@ public static final class Builder {
         private Integer limit;
         private Boolean consistentRead;
         private Expression filterExpression;
+        private Select select;
         private List<NestedAttributeName> attributesToProject;
         private Integer segment;
         private Integer totalSegments;
@@ -335,6 +364,18 @@ public Builder filterExpression(Expression filterExpression) {
             return this;
         }
 
+        /**
+         * Determines the attributes to be returned in the result. See {@link Select} for examples and constraints.
+         * By default, all attributes are returned.
+         * @param select
+         * @return a builder of this type
+         */
+        public Builder select(Select select) {
+            this.select = select;
+            return this;
+        }
+
+
         /**
          * <p>
          * Sets a collection of the attribute names to be retrieved from the database. These attributes can include

diff --git a/...src/test/java/software/amazon/awssdk/enhanced/dynamodb/functionaltests/BasicScanTest.java b/...src/test/java/software/amazon/awssdk/enhanced/dynamodb/functionaltests/BasicScanTest.java
@@ -45,13 +45,15 @@
 import software.amazon.awssdk.enhanced.dynamodb.Expression;
 import software.amazon.awssdk.enhanced.dynamodb.NestedAttributeName;
 import software.amazon.awssdk.enhanced.dynamodb.TableSchema;
+import software.amazon.awssdk.enhanced.dynamodb.document.EnhancedDocument;
 import software.amazon.awssdk.enhanced.dynamodb.functionaltests.models.InnerAttributeRecord;
 import software.amazon.awssdk.enhanced.dynamodb.functionaltests.models.NestedTestRecord;
 import software.amazon.awssdk.enhanced.dynamodb.mapper.StaticTableSchema;
 import software.amazon.awssdk.enhanced.dynamodb.model.Page;
 import software.amazon.awssdk.enhanced.dynamodb.model.ScanEnhancedRequest;
 import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
 import software.amazon.awssdk.services.dynamodb.model.DeleteTableRequest;
+import software.amazon.awssdk.services.dynamodb.model.Select;
 
 public class BasicScanTest extends LocalDynamoDbSyncTestBase {
     private static class Record {
@@ -645,4 +647,98 @@ public void scanAllRecordsWithNestedProjectionNameEmptyNameMap() {
             Page<NestedTestRecord> next = resultsAttributeToProject.next();
         });
     }
+
+    @Test
+    public void scanAllRecordsWithSelect_specific_attr() {
+        insertRecords();
+        Map<String, AttributeValue> expressionValues = new HashMap<>();
+        expressionValues.put(":min_value", numberValue(3));
+        expressionValues.put(":max_value", numberValue(5));
+        Expression expression = Expression.builder()
+                                          .expression("#sort >= :min_value AND #sort <= :max_value")
+                                          .expressionValues(expressionValues)
+                                          .putExpressionName("#sort", "sort")
+                                          .build();
+
+        Iterator<Page<Record>> results =
+            mappedTable.scan(
+                ScanEnhancedRequest.builder()
+                                   .attributesToProject("sort")
+                                   .select(Select.SPECIFIC_ATTRIBUTES)
+                                   .filterExpression(expression)
+                                   .build()
+            ).iterator();
+
+        assertThat(results.hasNext(), is(true));
+        Page<Record> page = results.next();
+        assertThat(results.hasNext(), is(false));
+
+        assertThat(page.items(), hasSize(3));
+
+        Record record = page.items().get(0);
+
+        assertThat(record.id, is(nullValue()));
+        assertThat(record.sort, is(3));
+    }
+
+    @Test
+    public void scanAllRecordsWithSelect_All_Attr() {
+        insertRecords();
+        Map<String, AttributeValue> expressionValues = new HashMap<>();
+        expressionValues.put(":min_value", numberValue(3));
+        expressionValues.put(":max_value", numberValue(5));
+        Expression expression = Expression.builder()
+                                          .expression("#sort >= :min_value AND #sort <= :max_value")
+                                          .expressionValues(expressionValues)
+                                          .putExpressionName("#sort", "sort")
+                                          .build();
+
+        Iterator<Page<Record>> results =
+            mappedTable.scan(
+                ScanEnhancedRequest.builder()
+                                   .select(Select.ALL_ATTRIBUTES)
+                                   .filterExpression(expression)
+                                   .build()
+            ).iterator();
+
+        assertThat(results.hasNext(), is(true));
+        Page<Record> page = results.next();
+        assertThat(results.hasNext(), is(false));
+
+        assertThat(page.items(), hasSize(3));
+
+        Record record = page.items().get(0);
+
+        assertThat(record.id, is("id-value"));
+        assertThat(record.sort, is(3));
+    }
+
+    @Test
+    public void scanAllRecordsWithSelect_Count() {
+        insertRecords();
+        Map<String, AttributeValue> expressionValues = new HashMap<>();
+        expressionValues.put(":min_value", numberValue(3));
+        expressionValues.put(":max_value", numberValue(5));
+        Expression expression = Expression.builder()
+                                          .expression("#sort >= :min_value AND #sort <= :max_value")
+                                          .expressionValues(expressionValues)
+                                          .putExpressionName("#sort", "sort")
+                                          .build();
+
+        Iterator<Page<Record>> results =
+            mappedTable.scan(
+                ScanEnhancedRequest.builder()
+                                   .select(Select.COUNT)
+                                   .filterExpression(expression)
+                                   .build()
+            ).iterator();
+
+        assertThat(results.hasNext(), is(true));
+        Page<Record> page = results.next();
+        assertThat(results.hasNext(), is(false));
+
+        assertThat(page.count(), is(3));
+
+        assertThat(page.items().size(), is(0));
+    }
 }
diff --git a/...java/software/amazon/awssdk/enhanced/dynamodb/functionaltests/document/BasicScanTest.java b/...java/software/amazon/awssdk/enhanced/dynamodb/functionaltests/document/BasicScanTest.java
@@ -20,6 +20,7 @@
 import static org.hamcrest.Matchers.empty;
 import static org.hamcrest.Matchers.hasSize;
 import static org.hamcrest.Matchers.is;
+import static org.hamcrest.Matchers.isEmptyString;
 import static org.hamcrest.Matchers.nullValue;
 import static software.amazon.awssdk.enhanced.dynamodb.AttributeConverterProvider.defaultProvider;
 import static software.amazon.awssdk.enhanced.dynamodb.internal.AttributeValues.numberValue;
@@ -58,6 +59,7 @@
 import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
 import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
 import software.amazon.awssdk.services.dynamodb.model.DeleteTableRequest;
+import software.amazon.awssdk.services.dynamodb.model.Select;
 
 public class BasicScanTest extends LocalDynamoDbSyncTestBase {
     private DynamoDbClient lowLevelClient;
@@ -668,4 +670,57 @@ public void scanAllRecordsWithNestedProjectionNameEmptyNameMap() {
             Page<EnhancedDocument> next = resultsAttributeToProject.next();
         });
     }
+
+    @Test
+    public void scanAllRecordsDefaultSettings_select() {
+        insertDocuments();
+
+        Iterator<Page<EnhancedDocument>> results =
+            docMappedtable.scan(b -> b.select(Select.ALL_ATTRIBUTES)).iterator();
+
+        assertThat(results.hasNext(), is(true));
+        Page<EnhancedDocument> page = results.next();
+        assertThat(results.hasNext(), is(false));
+
+        assertThat(page.items().size(), is(DOCUMENTS.size()));
+
+        EnhancedDocument firstRecord = page.items().get(0);
+        assertThat(firstRecord.getString("id"), is("id-value"));
+        assertThat(firstRecord.getNumber("sort").intValue(), is(0));
+        assertThat(firstRecord.getNumber("value").intValue(), is(0));
+    }
+
+    @Test
+    public void scanAllRecordsDefaultSettings_select_specific_attr() {
+        insertDocuments();
+
+        Iterator<Page<EnhancedDocument>> results =
+            docMappedtable.scan(b -> b.attributesToProject("sort").select(Select.SPECIFIC_ATTRIBUTES)).iterator();
+
+        assertThat(results.hasNext(), is(true));
+        Page<EnhancedDocument> page = results.next();
+        assertThat(results.hasNext(), is(false));
+
+        assertThat(page.items().size(), is(DOCUMENTS.size()));
+
+        EnhancedDocument firstRecord = page.items().get(0);
+        assertThat(firstRecord.getString("id"), is(nullValue()));
+        assertThat(firstRecord.getNumber("sort").intValue(), is(0));
+    }
+
+
+    @Test
+    public void scanAllRecordsDefaultSettings_select_count() {
+        insertDocuments();
+
+        Iterator<Page<EnhancedDocument>> results =
+            docMappedtable.scan(b -> b.select(Select.COUNT)).iterator();
+
+        assertThat(results.hasNext(), is(true));
+        Page<EnhancedDocument> page = results.next();
+        assertThat(results.hasNext(), is(false));
+
+        assertThat(page.count(), is(DOCUMENTS.size()));
+        assertThat(page.items().size(), is(0));
+    }
 }