Added Query and Scan limit features to save money on the new ondemand…

… feature

pom.xml

@@ -41,7 +41,7 @@
         <spring-data.version>2.1.2.RELEASE</spring-data.version>
 
         <hibernate-validator.version>6.0.9.Final</hibernate-validator.version>
-        <aws-java-sdk.version>1.11.443</aws-java-sdk.version>
+        <aws-java-sdk.version>1.11.515</aws-java-sdk.version>
         <junit.version>4.12</junit.version>
         <mockito.version>2.23.0</mockito.version>
         <cdi.version>1.2</cdi.version>

src/changes/changes.xml

@@ -27,6 +27,9 @@
       <action dev="derjust" issue="235" type="fix" date="2019-03-10">
         Upgrading to 5.1.0 results in the error Bean property 'dynamoDBMapperConfig' is not writable or has an invalid setter method
       </action>
+      <action dev="boostschicken" issue="226" type="add" date="2019-03-10">
+        @Query annotation to support query limiting
+      </action>
     </release>
     <release version="5.1.0" date="2019-01-28" description="Spring Boot 2.1 and Spring Data Lovelace-SR1 support">
       <action dev="boostschicken" type="add" date="2018-10-28">

src/main/java/org/socialsignin/spring/data/dynamodb/core/DynamoDBTemplate.java

@@ -51,6 +51,8 @@
 import java.util.function.Function;
 import java.util.stream.Collectors;
 
+import static org.socialsignin.spring.data.dynamodb.repository.QueryConstants.QUERY_LIMIT_UNLIMITED;
+
 public class DynamoDBTemplate implements DynamoDBOperations, ApplicationContextAware {
 	private static final Logger LOGGER = LoggerFactory.getLogger(DynamoDBTemplate.class);
 	private final DynamoDBMapper dynamoDBMapper;
@@ -180,6 +182,15 @@ public List<FailedBatch> batchDelete(Iterable<?> entities) {
 	@Override
 	public <T> PaginatedQueryList<T> query(Class<T> clazz, QueryRequest queryRequest) {
 		QueryResult queryResult = amazonDynamoDB.query(queryRequest);
+
+		// If a limit is set, deactivate lazy loading of (matching) items after the
+		// limit
+		// via
+		// com.amazonaws.services.dynamodbv2.datamodeling.PaginatedQueryList.atEndOfResults()
+		if (queryRequest.getLimit() != null) {
+			queryResult.setLastEvaluatedKey(null);
+		}
+
 		return new PaginatedQueryList<T>(dynamoDBMapper, clazz, amazonDynamoDB, queryRequest, queryResult,
 				dynamoDBMapperConfig.getPaginationLoadingStrategy(), dynamoDBMapperConfig);
 	}

src/main/java/org/socialsignin/spring/data/dynamodb/repository/Query.java

@@ -21,6 +21,8 @@
 import java.lang.annotation.RetentionPolicy;
 import java.lang.annotation.Target;
 
+import static org.socialsignin.spring.data.dynamodb.repository.QueryConstants.QUERY_LIMIT_UNLIMITED;
+
 @Retention(RetentionPolicy.RUNTIME)
 @Target(ElementType.METHOD)
 @Documented
@@ -36,4 +38,13 @@
 	 *      Expressions</a>
 	 */
 	String fields() default "";
+
+	/**
+	 * An integer to limit the number of elements returned.
+	 *
+	 * @see <a href=
+	 *      "https://docs.aws.amazon.com/amazondynamodb/latest/developerguide/Expressions.ProjectionExpressions.html">Projection
+	 *      Expressions</a>
+	 */
+	int limit() default QUERY_LIMIT_UNLIMITED;
 }

src/main/java/org/socialsignin/spring/data/dynamodb/repository/QueryConstants.java

@@ -0,0 +1,25 @@
+/**
+ * Copyright © 2018 spring-data-dynamodb (https://github.com/derjust/spring-data-dynamodb)
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.socialsignin.spring.data.dynamodb.repository;
+
+public final class QueryConstants {
+
+	private QueryConstants() {
+	}
+
+	public static final int QUERY_LIMIT_UNLIMITED = Integer.MIN_VALUE;
+
+}

src/main/java/org/socialsignin/spring/data/dynamodb/repository/query/AbstractDynamoDBQueryCreator.java

@@ -46,21 +46,24 @@ public abstract class AbstractDynamoDBQueryCreator<T, ID, R>
 	protected final DynamoDBEntityInformation<T, ID> entityMetadata;
 	protected final DynamoDBOperations dynamoDBOperations;
 	protected final Optional<String> projection;
+	protected final Optional<Integer> limit;
 
 	public AbstractDynamoDBQueryCreator(PartTree tree, DynamoDBEntityInformation<T, ID> entityMetadata,
-			Optional<String> projection, DynamoDBOperations dynamoDBOperations) {
+			Optional<String> projection, Optional<Integer> limitResults, DynamoDBOperations dynamoDBOperations) {
 		super(tree);
 		this.entityMetadata = entityMetadata;
 		this.projection = projection;
+		this.limit = limitResults;
 		this.dynamoDBOperations = dynamoDBOperations;
 	}
 
 	public AbstractDynamoDBQueryCreator(PartTree tree, ParameterAccessor parameterAccessor,
 			DynamoDBEntityInformation<T, ID> entityMetadata, Optional<String> projection,
-			DynamoDBOperations dynamoDBOperations) {
+			Optional<Integer> limitResults, DynamoDBOperations dynamoDBOperations) {
 		super(tree, parameterAccessor);
 		this.entityMetadata = entityMetadata;
 		this.projection = projection;
+		this.limit = limitResults;
 		this.dynamoDBOperations = dynamoDBOperations;
 	}
 
@@ -89,7 +92,6 @@ protected DynamoDBQueryCriteria<T, ID> addCriteria(DynamoDBQueryCriteria<T, ID>
 		}
 
 		switch (part.getType()) {
-
 			case IN :
 				Object in = iterator.next();
 				Assert.notNull(in, "Creating conditions on null parameters not supported: please specify a value for '"

src/main/java/org/socialsignin/spring/data/dynamodb/repository/query/AbstractDynamoDBQueryCriteria.java

@@ -72,6 +72,7 @@ public abstract class AbstractDynamoDBQueryCriteria<T, ID> implements DynamoDBQu
 	protected String globalSecondaryIndexName;
 	protected Sort sort = Sort.unsorted();
 	protected Optional<String> projection = Optional.empty();
+	protected Optional<Integer> limit = Optional.empty();
 
 	public abstract boolean isApplicableForLoad();
 
@@ -132,6 +133,7 @@ protected QueryRequest buildQueryRequest(String tableName, String theIndexName,
 				queryRequest.setSelect(Select.ALL_PROJECTED_ATTRIBUTES);
 			}
 
+			limit.ifPresent(queryRequest::setLimit);
 			applySortIfSpecified(queryRequest, new ArrayList<>(new HashSet<>(allowedSortProperties)));
 		}
 		return queryRequest;
@@ -695,4 +697,10 @@ public DynamoDBQueryCriteria<T, ID> withProjection(Optional<String> projection)
 		this.projection = projection;
 		return this;
 	}
+
+	@Override
+	public DynamoDBQueryCriteria<T, ID> withLimit(Optional<Integer> limit) {
+		this.limit = limit;
+		return this;
+	}
 }

src/main/java/org/socialsignin/spring/data/dynamodb/repository/query/DynamoDBCountQueryCreator.java

@@ -32,13 +32,13 @@ public class DynamoDBCountQueryCreator<T, ID> extends AbstractDynamoDBQueryCreat
 
 	public DynamoDBCountQueryCreator(PartTree tree, DynamoDBEntityInformation<T, ID> entityMetadata,
 			DynamoDBOperations dynamoDBOperations, boolean pageQuery) {
-		super(tree, entityMetadata, Optional.empty(), dynamoDBOperations);
+		super(tree, entityMetadata, Optional.empty(), Optional.empty(), dynamoDBOperations);
 		this.pageQuery = pageQuery;
 	}
 
 	public DynamoDBCountQueryCreator(PartTree tree, ParameterAccessor parameterAccessor,
 			DynamoDBEntityInformation<T, ID> entityMetadata, DynamoDBOperations dynamoDBOperations, boolean pageQuery) {
-		super(tree, parameterAccessor, entityMetadata, Optional.empty(), dynamoDBOperations);
+		super(tree, parameterAccessor, entityMetadata, Optional.empty(), Optional.empty(), dynamoDBOperations);
 		this.pageQuery = pageQuery;
 
 	}

src/main/java/org/socialsignin/spring/data/dynamodb/repository/query/DynamoDBEntityWithHashAndRangeKeyCriteria.java

@@ -190,7 +190,7 @@ public DynamoDBQueryExpression<T> buildQueryExpression() {
 			queryExpression.setSelect(Select.SPECIFIC_ATTRIBUTES);
 			queryExpression.setProjectionExpression(projection.get());
 		}
-
+		limit.ifPresent(queryExpression::setLimit);
 		return queryExpression;
 	}
 

src/main/java/org/socialsignin/spring/data/dynamodb/repository/query/DynamoDBEntityWithHashKeyOnlyCriteria.java

@@ -116,6 +116,7 @@ public DynamoDBScanExpression buildScanExpression() {
 			scanExpression.setSelect(Select.SPECIFIC_ATTRIBUTES);
 			scanExpression.setProjectionExpression(projection.get());
 		}
+		limit.ifPresent(scanExpression::setLimit);
 		return scanExpression;
 	}
 

src/main/java/org/socialsignin/spring/data/dynamodb/repository/query/DynamoDBQueryCreator.java

@@ -29,9 +29,9 @@
 public class DynamoDBQueryCreator<T, ID> extends AbstractDynamoDBQueryCreator<T, ID, T> {
 
 	public DynamoDBQueryCreator(PartTree tree, ParameterAccessor parameterAccessor,
-			DynamoDBEntityInformation<T, ID> entityMetadata, Optional<String> projection,
+			DynamoDBEntityInformation<T, ID> entityMetadata, Optional<String> projection, Optional<Integer> limit,
 			DynamoDBOperations dynamoDBOperations) {
-		super(tree, parameterAccessor, entityMetadata, projection, dynamoDBOperations);
+		super(tree, parameterAccessor, entityMetadata, projection, limit, dynamoDBOperations);
 	}
 
 	@Override
@@ -41,7 +41,7 @@ protected Query<T> complete(@Nullable DynamoDBQueryCriteria<T, ID> criteria, Sor
 		} else {
 			criteria.withSort(sort);
 			criteria.withProjection(projection);
-
+			criteria.withLimit(limit);
 			return criteria.buildQuery(dynamoDBOperations);
 		}
 	}

src/main/java/org/socialsignin/spring/data/dynamodb/repository/query/DynamoDBQueryCriteria.java

@@ -43,6 +43,8 @@ DynamoDBQueryCriteria<T, ID> withSingleValueCriteria(String propertyName, Compar
 
 	DynamoDBQueryCriteria<T, ID> withProjection(Optional<String> projection);
 
+	DynamoDBQueryCriteria<T, ID> withLimit(Optional<Integer> limit);
+
 	Query<T> buildQuery(DynamoDBOperations dynamoDBOperations);
 
 	Query<Long> buildCountQuery(DynamoDBOperations dynamoDBOperations, boolean pageQuery);

src/main/java/org/socialsignin/spring/data/dynamodb/repository/query/DynamoDBQueryMethod.java

@@ -28,6 +28,8 @@
 import java.lang.reflect.Method;
 import java.util.Optional;
 
+import static org.socialsignin.spring.data.dynamodb.repository.QueryConstants.QUERY_LIMIT_UNLIMITED;
+
 /**
  * @author Michael Lavelle
  * @author Sebastian Just
@@ -38,6 +40,7 @@ public class DynamoDBQueryMethod<T, ID> extends QueryMethod {
 	private final boolean scanEnabledForRepository;
 	private final boolean scanCountEnabledForRepository;
 	private final Optional<String> projectionExpression;
+	private final Optional<Integer> limitResults;
 
 	public DynamoDBQueryMethod(Method method, RepositoryMetadata metadata, ProjectionFactory factory) {
 		super(method, metadata, factory);
@@ -54,8 +57,15 @@ public DynamoDBQueryMethod(Method method, RepositoryMetadata metadata, Projectio
 			} else {
 				this.projectionExpression = Optional.empty();
 			}
+			int limit = query.limit();
+			if (limit != QUERY_LIMIT_UNLIMITED) {
+				this.limitResults = Optional.of(query.limit());
+			} else {
+				this.limitResults = Optional.empty();
+			}
 		} else {
 			this.projectionExpression = Optional.empty();
+			this.limitResults = Optional.empty();
 		}
 	}
 
@@ -98,4 +108,7 @@ public Optional<String> getProjectionExpression() {
 		return this.projectionExpression;
 	}
 
+	public Optional<Integer> getLimitResults() {
+		return this.limitResults;
+	}
 }

src/main/java/org/socialsignin/spring/data/dynamodb/repository/query/PartTreeDynamoDBQuery.java

@@ -39,7 +39,7 @@ public PartTreeDynamoDBQuery(DynamoDBOperations dynamoDBOperations, DynamoDBQuer
 
 	protected DynamoDBQueryCreator<T, ID> createQueryCreator(ParametersParameterAccessor accessor) {
 		return new DynamoDBQueryCreator<>(tree, accessor, getQueryMethod().getEntityInformation(),
-				getQueryMethod().getProjectionExpression(), dynamoDBOperations);
+				getQueryMethod().getProjectionExpression(), getQueryMethod().getLimitResults(), dynamoDBOperations);
 	}
 
 	protected DynamoDBCountQueryCreator<T, ID> createCountQueryCreator(ParametersParameterAccessor accessor,

src/main/java/org/socialsignin/spring/data/dynamodb/repository/support/DynamoDBEntityInformation.java

@@ -50,4 +50,6 @@ default Object getRangeKey(ID id) {
 	}
 
 	Optional<String> getProjection();
+
+	Optional<Integer> getLimit();
 }

src/main/java/org/socialsignin/spring/data/dynamodb/repository/support/DynamoDBIdIsHashAndRangeKeyEntityInformationImpl.java

@@ -41,6 +41,7 @@ public class DynamoDBIdIsHashAndRangeKeyEntityInformationImpl<T, ID> extends Ref
 	private DynamoDBHashAndRangeKeyExtractingEntityMetadata<T, ID> metadata;
 	private HashAndRangeKeyExtractor<ID, ?> hashAndRangeKeyExtractor;
 	private Optional<String> projection = Optional.empty();
+	private Optional<Integer> limit = Optional.empty();
 
 	public DynamoDBIdIsHashAndRangeKeyEntityInformationImpl(Class<T> domainClass,
 			DynamoDBHashAndRangeKeyExtractingEntityMetadata<T, ID> metadata) {
@@ -54,6 +55,11 @@ public Optional<String> getProjection() {
 		return projection;
 	}
 
+	@Override
+	public Optional<Integer> getLimit() {
+		return limit;
+	}
+
 	@Override
 	public boolean isRangeKeyAware() {
 		return true;

src/main/java/org/socialsignin/spring/data/dynamodb/repository/support/DynamoDBIdIsHashKeyEntityInformationImpl.java

@@ -47,6 +47,7 @@ public class DynamoDBIdIsHashKeyEntityInformationImpl<T, ID> extends FieldAndGet
 	private DynamoDBHashKeyExtractingEntityMetadata<T> metadata;
 	private HashKeyExtractor<ID, ID> hashKeyExtractor;
 	private Optional<String> projection = Optional.empty();
+	private Optional<Integer> limit = Optional.empty();
 
 	public DynamoDBIdIsHashKeyEntityInformationImpl(Class<T> domainClass,
 			DynamoDBHashKeyExtractingEntityMetadata<T> metadata) {
@@ -60,6 +61,11 @@ public Optional<String> getProjection() {
 		return projection;
 	}
 
+	@Override
+	public Optional<Integer> getLimit() {
+		return limit;
+	}
+
 	@Override
 	public Object getHashKey(final ID id) {
 		Assert.isAssignable(getIdType(), id.getClass(),

src/test/java/org/socialsignin/spring/data/dynamodb/domain/sample/CRUDOperationsIT.java

@@ -83,6 +83,7 @@ public void testProjection() {
 		String postCode = "postCode";
 		String user1 = "projection1" + ThreadLocalRandom.current().nextLong();
 		String user2 = "projection2" + ThreadLocalRandom.current().nextLong();
+		String user3 = "projection2" + ThreadLocalRandom.current().nextLong();
 
 		User u1 = new User();
 		u1.setId("Id1" + ThreadLocalRandom.current().nextLong());
@@ -91,6 +92,13 @@ public void testProjection() {
 		u1.setPostCode(postCode);
 		u1.setNumberOfPlaylists(1);
 
+		User u3 = new User();
+		u3.setId("Id3" + ThreadLocalRandom.current().nextLong());
+		u3.setName(user3);
+		u3.setLeaveDate(Instant.now());
+		u3.setPostCode(postCode);
+		u3.setNumberOfPlaylists(1);
+
 		User u2 = new User();
 		u2.setId("Id2" + ThreadLocalRandom.current().nextLong());
 		u2.setName(user2);
@@ -100,11 +108,13 @@ public void testProjection() {
 
 		userRepository.save(u1);
 		userRepository.save(u2);
+		userRepository.save(u3);
 
 		List<User> actualList = new ArrayList<>();
 		userRepository.findAll().forEach(actualList::add);
 
 		List<User> projectedActuals = userRepository.findByPostCode(postCode);
+		// 2 matches but should be limited to 1 by @Query
 		assertEquals(1, projectedActuals.size());
 		User projectedActual = projectedActuals.get(0);
 		assertNull("Attribute not projected", projectedActual.getName());
@@ -113,11 +123,11 @@ public void testProjection() {
 		assertNull("Key not projected", projectedActual.getId());
 		assertNotNull("LeaveDate is projected", projectedActual.getLeaveDate());
 
-		List<User> fullActuals = userRepository.findByNameIn(Arrays.asList(user1, user2));
-		assertEquals(2, fullActuals.size());
+		List<User> fullActuals = userRepository.findByNameIn(Arrays.asList(user1, user2, user3));
+		assertEquals(3, fullActuals.size());
 		User fullActual = fullActuals.get(0);
-		assertThat(Arrays.asList(user1, user2), hasItems(fullActual.getName()));
-		assertThat(Arrays.asList(user1, user2), hasItems(fullActuals.get(1).getName()));
+		assertThat(Arrays.asList(user1, user2, user3), hasItems(fullActual.getName()));
+		assertThat(Arrays.asList(user1, user2, user3), hasItems(fullActuals.get(1).getName()));
 		assertNotNull(fullActual.getPostCode());
 		assertNotNull(fullActual.getNumberOfPlaylists());
 		assertNotNull(fullActual.getId());

src/test/java/org/socialsignin/spring/data/dynamodb/domain/sample/UserRepository.java

@@ -52,7 +52,7 @@ public interface UserRepository extends CrudRepository<User, String> {
 	@EnableScan
 	void deleteByIdAndName(String id, String name);
 
-	@Query(fields = "leaveDate")
+	@Query(fields = "leaveDate", limit = 1)
 	List<User> findByPostCode(String postCode);
 
 	@EnableScan

src/test/java/org/socialsignin/spring/data/dynamodb/repository/query/PartTreeDynamoDBQueryUnitTest.java

@@ -139,18 +139,6 @@ private <T, ID extends Serializable> void setupCommonMocksForThisRepositoryMetho
 			Mockito.when(mockEntityMetadata.isRangeKeyAware()).thenReturn(true);
 		}
 
-		try {
-			Field unwrappedReturnTypeField = mockDynamoDBQueryMethod.getClass() // org.socialsignin.spring.data.dynamodb.repository.query.DynamoDBQueryMethod
-					.getSuperclass() // org.springframework.data.repository.query.QueryMethod
-					.getDeclaredField("unwrappedReturnType");
-			unwrappedReturnTypeField.setAccessible(true); // It's final therefore unlocking the field
-			unwrappedReturnTypeField.set(mockDynamoDBQueryMethod, clazz);
-		} catch (Exception e) {
-			// There is little we can and want do if it fails - Aborting the whole test is
-			// fine
-			throw new RuntimeException(e);
-		}
-
 		Mockito.when(mockDynamoDBQueryMethod.getEntityType()).thenReturn(clazz);
 		Mockito.when(mockDynamoDBQueryMethod.getName()).thenReturn(repositoryMethodName);
 		Mockito.when(mockDynamoDBQueryMethod.getParameters()).thenReturn(mockParameters);