Adding support for Select in QueryEnhancedRequest (#5356)

* Adding support for Select in QueryEnhancedRequest

* Adding support for Select in QueryEnhancedRequest

* Adding support for Select in QueryEnhancedRequest

* Delete services-custom/dynamodb-enhanced/dynamodb-local-metadata.json

* Fixing Checkstyle on QueryEnhancedRequest

Author: shetsa-amzn

.changes/next-release/feature-DynamoDBEnhancedClient-866b8ef.json

@@ -0,0 +1,6 @@
+{
+    "category": "DynamoDB Enhanced Client", 
+    "contributor": "shetsa-amzn", 
+    "type": "feature", 
+    "description": "Adding support for Select in QueryEnhancedRequest"
+}

services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/internal/operations/QueryOperation.java

@@ -37,7 +37,7 @@
 
 @SdkInternalApi
 public class QueryOperation<T> implements PaginatedTableOperation<T, QueryRequest, QueryResponse>,
-        PaginatedIndexOperation<T, QueryRequest, QueryResponse> {
+                                          PaginatedIndexOperation<T, QueryRequest, QueryResponse> {
 
     private final QueryEnhancedRequest request;
 
@@ -80,6 +80,7 @@ public QueryRequest generateRequest(TableSchema<T> tableSchema,
                                                         .expressionAttributeValues(expressionValues)
                                                         .expressionAttributeNames(expressionNames)
                                                         .scanIndexForward(this.request.scanIndexForward())
+                                                        .select(this.request.select())
                                                         .limit(this.request.limit())
                                                         .exclusiveStartKey(this.request.exclusiveStartKey())
                                                         .consistentRead(this.request.consistentRead())

services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/model/QueryEnhancedRequest.java

@@ -34,6 +34,7 @@
 import software.amazon.awssdk.services.dynamodb.model.QueryRequest;
 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;
 
 /**
@@ -52,6 +53,7 @@ public final class QueryEnhancedRequest {
     private final QueryConditional queryConditional;
     private final Map<String, AttributeValue> exclusiveStartKey;
     private final Boolean scanIndexForward;
+    private final Select select;
     private final Integer limit;
     private final Boolean consistentRead;
     private final Expression filterExpression;
@@ -62,13 +64,14 @@ private QueryEnhancedRequest(Builder builder) {
         this.queryConditional = builder.queryConditional;
         this.exclusiveStartKey = builder.exclusiveStartKey;
         this.scanIndexForward = builder.scanIndexForward;
+        this.select = builder.select;
         this.limit = builder.limit;
         this.consistentRead = builder.consistentRead;
         this.filterExpression = builder.filterExpression;
         this.returnConsumedCapacity = builder.returnConsumedCapacity;
         this.attributesToProject = builder.attributesToProject != null
-                ? Collections.unmodifiableList(builder.attributesToProject)
-                : null;
+                                   ? Collections.unmodifiableList(builder.attributesToProject)
+                                   : null;
     }
 
     /**
@@ -85,6 +88,7 @@ public Builder toBuilder() {
         return builder().queryConditional(queryConditional)
                         .exclusiveStartKey(exclusiveStartKey)
                         .scanIndexForward(scanIndexForward)
+                        .select(select)
                         .limit(limit)
                         .consistentRead(consistentRead)
                         .filterExpression(filterExpression)
@@ -114,6 +118,22 @@ public Boolean scanIndexForward() {
         return scanIndexForward;
     }
 
+    /**
+     * 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 value of limit set on this request object, or null if it doesn't exist.
      */
@@ -143,7 +163,8 @@ public Expression filterExpression() {
      */
     public List<String> attributesToProject() {
         return attributesToProject != null ? attributesToProject.stream()
-                .map(item -> String.join(".", item.elements())).collect(Collectors.toList()) : null;
+                                                                .map(item -> String.join(".", item.elements()))
+                                                                .collect(Collectors.toList()) : null;
     }
 
     /**
@@ -198,14 +219,20 @@ public boolean equals(Object o) {
             query.scanIndexForward != null) {
             return false;
         }
+
+        if (select != null ? ! select.equals(query.select) :
+            query.select != null) {
+            return false;
+        }
+
         if (limit != null ? ! limit.equals(query.limit) : query.limit != null) {
             return false;
         }
         if (consistentRead != null ? ! consistentRead.equals(query.consistentRead) : query.consistentRead != null) {
             return false;
         }
         if (attributesToProject != null
-                ? !attributesToProject.equals(query.attributesToProject) : query.attributesToProject != null) {
+            ? !attributesToProject.equals(query.attributesToProject) : query.attributesToProject != null) {
             return false;
         }
         if (returnConsumedCapacity != null
@@ -220,6 +247,7 @@ public int hashCode() {
         int result = queryConditional != null ? queryConditional.hashCode() : 0;
         result = 31 * result + (exclusiveStartKey != null ? exclusiveStartKey.hashCode() : 0);
         result = 31 * result + (scanIndexForward != null ? scanIndexForward.hashCode() : 0);
+        result = 31 * result + (select != null ? select.hashCode() : 0);
         result = 31 * result + (limit != null ? limit.hashCode() : 0);
         result = 31 * result + (consistentRead != null ? consistentRead.hashCode() : 0);
         result = 31 * result + (filterExpression != null ? filterExpression.hashCode() : 0);
@@ -238,6 +266,7 @@ public static final class Builder {
         private QueryConditional queryConditional;
         private Map<String, AttributeValue> exclusiveStartKey;
         private Boolean scanIndexForward;
+        private Select select;
         private Integer limit;
         private Boolean consistentRead;
         private Expression filterExpression;
@@ -271,6 +300,17 @@ public Builder scanIndexForward(Boolean scanIndexForward) {
             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;
+        }
+
         /**
          * The primary key of the first item that this operation will evaluate. By default, the operation will evaluate
          * the whole dataset. If used, normally this parameter is populated with the value that was returned for
@@ -353,8 +393,10 @@ public Builder attributesToProject(Collection<String> attributesToProject) {
                 this.attributesToProject.clear();
             }
             if (attributesToProject != null) {
-                addNestedAttributesToProject(new ArrayList<>(attributesToProject).stream()
-                        .map(NestedAttributeName::create).collect(Collectors.toList()));
+                addNestedAttributesToProject(new ArrayList<>(attributesToProject)
+                                                 .stream()
+                                                 .map(NestedAttributeName::create)
+                                                 .collect(Collectors.toList()));
             }
             return this;
         }
@@ -409,7 +451,7 @@ public Builder addAttributeToProject(String attributeToProject) {
         public Builder addNestedAttributesToProject(Collection<NestedAttributeName> nestedAttributeNames) {
             if (nestedAttributeNames != null) {
                 Validate.noNullElements(nestedAttributeNames,
-                        "nestedAttributeNames list must not contain null elements");
+                                        "nestedAttributeNames list must not contain null elements");
                 if (attributesToProject == null) {
                     this.attributesToProject = new ArrayList<>(nestedAttributeNames);
                 } else {

services-custom/dynamodb-enhanced/src/test/java/software/amazon/awssdk/enhanced/dynamodb/functionaltests/BasicQueryTest.java

@@ -56,6 +56,7 @@
 import software.amazon.awssdk.enhanced.dynamodb.model.QueryEnhancedRequest;
 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 BasicQueryTest extends LocalDynamoDbSyncTestBase {
 
@@ -574,4 +575,69 @@ public void queryAllRecordsDefaultSettings_withNestedProjectionNameEmptyNameMap(
 
                 });
     }
+
+    @Test
+    public void queryAllRecordsDefaultSettings_withSelect_specificAttributes() {
+        insertRecords();
+
+        Iterator<Page<Record>> results =
+            mappedTable.query(b -> b
+                .queryConditional(keyEqualTo(k -> k.partitionValue("id-value")))
+                .attributesToProject("value")
+                .select(Select.SPECIFIC_ATTRIBUTES)
+            ).iterator();
+
+        assertThat(results.hasNext(), is(true));
+        Page<Record> page = results.next();
+        assertThat(results.hasNext(), is(false));
+
+        assertThat(page.items().size(), is(RECORDS.size()));
+
+        Record firstRecord = page.items().get(0);
+        assertThat(firstRecord.id, is(nullValue()));
+        assertThat(firstRecord.sort, is(nullValue()));
+        assertThat(firstRecord.value, is(0));
+    }
+
+    @Test
+    public void queryAllRecordsDefaultSettings_withSelect_allAttributes() {
+        insertRecords();
+
+        Iterator<Page<Record>> results =
+            mappedTable.query(b -> b
+                .queryConditional(keyEqualTo(k -> k.partitionValue("id-value")))
+                .select(Select.ALL_ATTRIBUTES)
+            ).iterator();
+
+        assertThat(results.hasNext(), is(true));
+        Page<Record> page = results.next();
+        assertThat(results.hasNext(), is(false));
+
+        assertThat(page.items().size(), is(RECORDS.size()));
+
+        Record firstRecord = page.items().get(0);
+        assertThat(firstRecord.id, is("id-value"));
+        assertThat(firstRecord.sort, is(0));
+        assertThat(firstRecord.value, is(0));
+    }
+
+    @Test
+    public void queryAllRecordsDefaultSettings_withSelect_count() {
+        insertRecords();
+
+        Iterator<Page<Record>> results =
+            mappedTable.query(b -> b
+                .queryConditional(keyEqualTo(k -> k.partitionValue("id-value")))
+                .select(Select.COUNT)
+            ).iterator();
+
+        assertThat(results.hasNext(), is(true));
+        Page<Record> page = results.next();
+        assertThat(results.hasNext(), is(false));
+
+        assertThat(page.count(), is(RECORDS.size()));
+        assertThat(page.scannedCount(), is(RECORDS.size()));
+        assertThat(page.items(), is(empty()));
+        assertThat(page.lastEvaluatedKey(), is(nullValue()));
+    }
 }