Skip to content

Commit a7e960e

Browse files
committed
Added implementation for handling projection expression as string value.
1 parent b034c13 commit a7e960e

File tree

9 files changed

+183
-51
lines changed

9 files changed

+183
-51
lines changed

services-custom/dynamodb-enhanced/src/it/java/software/amazon/awssdk/enhanced/dynamodb/ScanQueryIntegrationTest.java

Lines changed: 46 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -43,7 +43,6 @@
4343
import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
4444
import software.amazon.awssdk.services.dynamodb.model.ConsumedCapacity;
4545
import software.amazon.awssdk.services.dynamodb.model.ReturnConsumedCapacity;
46-
import software.amazon.awssdk.services.dynamodb.model.Select;
4746

4847
public class ScanQueryIntegrationTest extends DynamoDbEnhancedIntegrationTestBase {
4948

@@ -183,46 +182,60 @@ public void query_withReturnConsumedCapacityAndDifferentReadConsistency_checksCo
183182
assertThat(strongConsumedCapacity.capacityUnits(), is(greaterThanOrEqualTo(eventualConsumedCapacity.capacityUnits())));
184183
}
185184

186-
private Map<String, AttributeValue> getKeyMap(int sort) {
187-
Map<String, AttributeValue> result = new HashMap<>();
188-
result.put("id", stringValue(RECORDS.get(sort).getId()));
189-
result.put("sort", numberValue(RECORDS.get(sort).getSort()));
190-
return Collections.unmodifiableMap(result);
191-
}
192-
193185
@Test
194-
public void query_withStringSelect_returnsSpecifiedAttributes() {
186+
public void query_withStringProjectionExpression_checksProjectedAttributes() {
195187
insertRecords();
196188

197-
QueryEnhancedRequest request = QueryEnhancedRequest.builder()
198-
.queryConditional(sortBetween(k -> k.partitionValue("id-value").sortValue(2),
199-
k -> k.partitionValue("id-value").sortValue(6)))
200-
.select("ALL_ATTRIBUTES")
201-
.build();
202-
203-
Iterator<Page<Record>> results = mappedTable.query(request).iterator();
204-
205-
while (results.hasNext()) {
206-
Page<Record> page = results.next();
207-
for (Record record : page.items()) {
208-
assertThat(record.getId(), is(notNullValue()));
209-
assertThat(record.getSort(), is(notNullValue()));
210-
assertThat(record.getValue(), is(notNullValue()));
211-
assertThat(record.getStringAttribute(), is(notNullValue()));
212-
}
213-
}
189+
String projectionExpression = "id, sort";
190+
Iterator<Page<Record>> results =
191+
mappedTable.query(QueryEnhancedRequest.builder()
192+
.queryConditional(sortBetween(k -> k.partitionValue("id-value").sortValue(2),
193+
k -> k.partitionValue("id-value").sortValue(6)))
194+
.returnStringProjectionExpression(projectionExpression)
195+
.limit(3)
196+
.build())
197+
.iterator();
198+
199+
Page<Record> page1 = results.next();
200+
assertThat(results.hasNext(), is(true));
201+
Page<Record> page2 = results.next();
202+
assertThat(results.hasNext(), is(false));
203+
204+
assertThat(page1.items().get(0).getId(), is(notNullValue()));
205+
assertThat(page1.items().get(0).getSort(), is(notNullValue()));
206+
207+
assertThat(page2.items().get(0).getId(), is(notNullValue()));
208+
assertThat(page2.items().get(0).getSort(), is(notNullValue()));
214209
}
215210

216211
@Test
217-
public void query_withInvalidStringSelect_returnsUnknown() {
212+
public void scan_withStringProjectionExpression_checksProjectedAttributes() {
218213
insertRecords();
219214

220-
QueryEnhancedRequest request = QueryEnhancedRequest.builder()
221-
.queryConditional(sortBetween(k -> k.partitionValue("id-value").sortValue(2),
222-
k -> k.partitionValue("id-value").sortValue(6)))
223-
.select("INVALID_SELECT")
224-
.build();
215+
String projectionExpression = "id, sort";
216+
Iterator<Page<Record>> results =
217+
mappedTable.scan(ScanEnhancedRequest.builder()
218+
.returnStringProjectionExpression(projectionExpression)
219+
.limit(5)
220+
.build())
221+
.iterator();
222+
223+
Page<Record> page1 = results.next();
224+
assertThat(results.hasNext(), is(true));
225+
Page<Record> page2 = results.next();
226+
assertThat(results.hasNext(), is(false));
225227

226-
assertThat(request.select(), is(equalTo(Select.UNKNOWN_TO_SDK_VERSION)));
228+
assertThat(page1.items().get(0).getId(), is(notNullValue()));
229+
assertThat(page1.items().get(0).getSort(), is(notNullValue()));
230+
231+
assertThat(page2.items().get(0).getId(), is(notNullValue()));
232+
assertThat(page2.items().get(0).getSort(), is(notNullValue()));
233+
}
234+
235+
private Map<String, AttributeValue> getKeyMap(int sort) {
236+
Map<String, AttributeValue> result = new HashMap<>();
237+
result.put("id", stringValue(RECORDS.get(sort).getId()));
238+
result.put("sort", numberValue(RECORDS.get(sort).getSort()));
239+
return Collections.unmodifiableMap(result);
227240
}
228241
}

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -67,8 +67,8 @@ public QueryRequest generateRequest(TableSchema<T> tableSchema,
6767
expressionNames = Expression.joinNames(expressionNames, this.request.filterExpression().expressionNames());
6868
}
6969

70-
String projectionExpressionAsString = null;
71-
if (this.request.attributesToProject() != null) {
70+
String projectionExpressionAsString = this.request.stringProjectionExpression();
71+
if (projectionExpressionAsString == null && this.request.attributesToProject() != null) {
7272
ProjectionExpression attributesToProject = ProjectionExpression.create(this.request.nestedAttributesToProject());
7373
projectionExpressionAsString = attributesToProject.projectionExpressionAsString().orElse(null);
7474
expressionNames = Expression.joinNames(expressionNames, attributesToProject.expressionAttributeNames());

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

Lines changed: 26 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -59,6 +59,7 @@ public final class QueryEnhancedRequest {
5959
private final Expression filterExpression;
6060
private final List<NestedAttributeName> attributesToProject;
6161
private final String returnConsumedCapacity;
62+
private final String stringProjectionExpression;
6263

6364
private QueryEnhancedRequest(Builder builder) {
6465
this.queryConditional = builder.queryConditional;
@@ -69,6 +70,7 @@ private QueryEnhancedRequest(Builder builder) {
6970
this.consistentRead = builder.consistentRead;
7071
this.filterExpression = builder.filterExpression;
7172
this.returnConsumedCapacity = builder.returnConsumedCapacity;
73+
this.stringProjectionExpression = builder.stringProjectionExpression;
7274
this.attributesToProject = builder.attributesToProject != null
7375
? Collections.unmodifiableList(builder.attributesToProject)
7476
: null;
@@ -176,6 +178,13 @@ public List<NestedAttributeName> nestedAttributesToProject() {
176178
return attributesToProject;
177179
}
178180

181+
/**
182+
* Returns the String projection expression for the query operation.
183+
*/
184+
public String stringProjectionExpression() {
185+
return stringProjectionExpression;
186+
}
187+
179188

180189
/**
181190
* Whether to return the capacity consumed by this operation.
@@ -239,6 +248,10 @@ public boolean equals(Object o) {
239248
? !returnConsumedCapacity.equals(query.returnConsumedCapacity) : query.returnConsumedCapacity != null) {
240249
return false;
241250
}
251+
if (stringProjectionExpression != null
252+
? !stringProjectionExpression.equals(query.stringProjectionExpression) : query.stringProjectionExpression != null) {
253+
return false;
254+
}
242255
return filterExpression != null ? filterExpression.equals(query.filterExpression) : query.filterExpression == null;
243256
}
244257

@@ -253,6 +266,7 @@ public int hashCode() {
253266
result = 31 * result + (filterExpression != null ? filterExpression.hashCode() : 0);
254267
result = 31 * result + (attributesToProject != null ? attributesToProject.hashCode() : 0);
255268
result = 31 * result + (returnConsumedCapacity != null ? returnConsumedCapacity.hashCode() : 0);
269+
result = 31 * result + (stringProjectionExpression != null ? stringProjectionExpression.hashCode() : 0);
256270
return result;
257271
}
258272

@@ -272,6 +286,7 @@ public static final class Builder {
272286
private Expression filterExpression;
273287
private List<NestedAttributeName> attributesToProject;
274288
private String returnConsumedCapacity;
289+
private String stringProjectionExpression;
275290

276291
private Builder() {
277292
}
@@ -311,17 +326,6 @@ public Builder select(Select select) {
311326
return this;
312327
}
313328

314-
/**
315-
* Determines the attributes to be returned in the result. See {@link Select} for examples and constraints.
316-
* If not found, returns {@link Select#UNKNOWN_TO_SDK_VERSION}.
317-
* @param select the selection criteria as a string
318-
* @return a builder of this type
319-
*/
320-
public Builder select(String select) {
321-
this.select = Select.fromValue(select);
322-
return this;
323-
}
324-
325329
/**
326330
* The primary key of the first item that this operation will evaluate. By default, the operation will evaluate
327331
* the whole dataset. If used, normally this parameter is populated with the value that was returned for
@@ -525,6 +529,17 @@ public Builder returnConsumedCapacity(String returnConsumedCapacity) {
525529
return this;
526530
}
527531

532+
/**
533+
* Sets the projection expression for the query operation.
534+
*
535+
* @param stringProjectionExpression the projection expression as a string
536+
* @return a builder of this type
537+
*/
538+
public Builder returnStringProjectionExpression(String stringProjectionExpression) {
539+
this.stringProjectionExpression = stringProjectionExpression;
540+
return this;
541+
}
542+
528543
public QueryEnhancedRequest build() {
529544
return new QueryEnhancedRequest(this);
530545
}

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

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,7 @@ public final class ScanEnhancedRequest {
5454
private final Integer segment;
5555
private final Integer totalSegments;
5656
private final String returnConsumedCapacity;
57+
private String stringProjectionExpression;
5758

5859
private ScanEnhancedRequest(Builder builder) {
5960
this.exclusiveStartKey = builder.exclusiveStartKey;
@@ -245,6 +246,7 @@ public int hashCode() {
245246
result = 31 * result + (select != null ? select.hashCode() : 0);
246247
result = 31 * result + (attributesToProject != null ? attributesToProject.hashCode() : 0);
247248
result = 31 * result + (returnConsumedCapacity != null ? returnConsumedCapacity.hashCode() : 0);
249+
result = 31 * result + (stringProjectionExpression != null ? stringProjectionExpression.hashCode() : 0);
248250
return result;
249251
}
250252

@@ -262,6 +264,7 @@ public static final class Builder {
262264
private Integer segment;
263265
private Integer totalSegments;
264266
private String returnConsumedCapacity;
267+
private String stringProjectionExpression;
265268

266269
private Builder() {
267270
}
@@ -520,6 +523,16 @@ public Builder returnConsumedCapacity(String returnConsumedCapacity) {
520523
return this;
521524
}
522525

526+
/**
527+
* Whether to return the stringProjectionExpression for this operation.
528+
* @param stringProjectionExpression
529+
* @return a builder of this type
530+
*/
531+
public Builder returnStringProjectionExpression(String stringProjectionExpression) {
532+
this.stringProjectionExpression = stringProjectionExpression;
533+
return this;
534+
}
535+
523536
public ScanEnhancedRequest build() {
524537
return new ScanEnhancedRequest(this);
525538
}

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

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -295,4 +295,27 @@ public void queryExclusiveStartKey() {
295295
assertThat(page.items(), is(KEYS_ONLY_RECORDS.subList(8, 10)));
296296
assertThat(page.lastEvaluatedKey(), is(nullValue()));
297297
}
298+
299+
@Test
300+
public void queryWithStringProjectionExpression() {
301+
insertRecords();
302+
303+
String projectionExpression = "id, sort";
304+
QueryEnhancedRequest request = QueryEnhancedRequest.builder()
305+
.queryConditional(keyEqualTo(k -> k.partitionValue("gsi-id-value")))
306+
.returnStringProjectionExpression(projectionExpression)
307+
.build();
308+
309+
SdkPublisher<Page<Record>> publisher = keysOnlyMappedIndex.query(request);
310+
311+
List<Page<Record>> results = drainPublisher(publisher, 1);
312+
Page<Record> page = results.get(0);
313+
314+
assertThat(page.items().size(), is(KEYS_ONLY_RECORDS.size()));
315+
316+
Record firstRecord = page.items().get(0);
317+
assertThat(firstRecord.getId(), is("id-value"));
318+
assertThat(firstRecord.getSort(), is(0));
319+
assertThat(firstRecord.getValue(), is(nullValue())); // Ensure non-projected attributes are not present
320+
}
298321
}

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

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -741,4 +741,26 @@ public void scanAllRecordsWithSelect_Count() {
741741

742742
assertThat(page.items().size(), is(0));
743743
}
744+
745+
@Test
746+
public void scanWithStringProjectionExpression() {
747+
insertRecords();
748+
749+
String projectionExpression = "id, sort";
750+
ScanEnhancedRequest request = ScanEnhancedRequest.builder()
751+
.returnStringProjectionExpression(projectionExpression)
752+
.build();
753+
754+
Iterator<Page<Record>> results = mappedTable.scan(request).iterator();
755+
756+
assertThat(results.hasNext(), is(true));
757+
Page<Record> page = results.next();
758+
assertThat(results.hasNext(), is(false));
759+
760+
assertThat(page.items().size(), is(RECORDS.size()));
761+
762+
Record firstRecord = page.items().get(0);
763+
assertThat(firstRecord.getId(), is("id-value"));
764+
assertThat(firstRecord.getSort(), is(0));
765+
}
744766
}

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

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -676,4 +676,27 @@ public void queryAllRecordsDefaultSettings_withNestedProjectionNameEmptyNameMap(
676676

677677
});
678678
}
679+
680+
@Test
681+
public void queryWithStringProjectionExpression() {
682+
insertDocuments();
683+
684+
String projectionExpression = "id, sort";
685+
QueryEnhancedRequest request = QueryEnhancedRequest.builder()
686+
.queryConditional(keyEqualTo(k -> k.partitionValue("id-value")))
687+
.returnStringProjectionExpression(projectionExpression)
688+
.build();
689+
690+
Iterator<Page<EnhancedDocument>> results = docMappedtable.query(request).iterator();
691+
692+
assertThat(results.hasNext(), is(true));
693+
Page<EnhancedDocument> page = results.next();
694+
assertThat(results.hasNext(), is(false));
695+
696+
assertThat(page.items().size(), is(DOCUMENTS.size()));
697+
698+
EnhancedDocument firstRecord = page.items().get(0);
699+
assertThat(firstRecord.getString("id"), is("id-value"));
700+
assertThat(firstRecord.getNumber("sort").intValue(), is(0));
701+
}
679702
}

services-custom/dynamodb-enhanced/src/test/java/software/amazon/awssdk/enhanced/dynamodb/internal/operations/QueryOperationTest.java

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -434,6 +434,21 @@ public void queryItem_withExtension_correctlyTransformsItem() {
434434
.items(attributeMap).build()));
435435
}
436436

437+
@Test
438+
public void generateRequest_withStringProjectionExpression() {
439+
String projectionExpression = "id, sort";
440+
QueryOperation<FakeItem> queryToTest =
441+
QueryOperation.create(QueryEnhancedRequest.builder()
442+
.queryConditional(keyEqualTo(k -> k.partitionValue(keyItem.getId())))
443+
.returnStringProjectionExpression(projectionExpression)
444+
.build());
445+
QueryRequest queryRequest = queryToTest.generateRequest(FakeItem.getTableSchema(),
446+
PRIMARY_CONTEXT,
447+
null);
448+
449+
assertThat(queryRequest.projectionExpression(), is(projectionExpression));
450+
}
451+
437452
private static QueryResponse generateFakeQueryResults(List<Map<String, AttributeValue>> queryItemMapsPage) {
438453
return QueryResponse.builder().items(queryItemMapsPage).build();
439454
}

services-custom/dynamodb-enhanced/src/test/java/software/amazon/awssdk/enhanced/dynamodb/model/QueryEnhancedRequestTest.java

Lines changed: 13 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,6 @@
2323
import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
2424

2525
import java.util.*;
26-
import software.amazon.awssdk.services.dynamodb.model.Select;
2726

2827
import static java.util.Collections.singletonMap;
2928
import static org.assertj.core.api.Assertions.assertThatThrownBy;
@@ -184,12 +183,21 @@ public void toBuilder() {
184183
assertThat(copiedObject, is(builtObject));
185184
}
186185

187-
@org.junit.jupiter.api.Test
188-
public void testSelectWithStringInput() {
186+
@Test
187+
public void builder_withStringProjectionExpression() {
188+
String projectionExpression = "id, sort";
189189
QueryEnhancedRequest request = QueryEnhancedRequest.builder()
190-
.select("ALL_ATTRIBUTES")
190+
.returnStringProjectionExpression(projectionExpression)
191191
.build();
192192

193-
assertThat(Select.ALL_ATTRIBUTES, is(request.select()));
193+
assertThat(request.stringProjectionExpression(), is(projectionExpression));
194+
}
195+
196+
@Test
197+
public void builder_withoutStringProjectionExpression() {
198+
QueryEnhancedRequest request = QueryEnhancedRequest.builder().build();
199+
200+
assertThat(request.stringProjectionExpression(), is(nullValue()));
194201
}
202+
195203
}

0 commit comments

Comments
 (0)