FMWK-224 Add support for id projection query via Repository (#614)

aerospike · Aug 13, 2023 · f5eaebd · f5eaebd
1 parent 80138e6
commit f5eaebd
Show file tree

Hide file tree

Showing 8 changed files with 124 additions and 46 deletions.
diff --git a/src/main/java/org/springframework/data/aerospike/core/AerospikeTemplate.java b/src/main/java/org/springframework/data/aerospike/core/AerospikeTemplate.java
@@ -105,6 +105,11 @@ public AerospikeTemplate(IAerospikeClient client,
         this.indexRefresher = indexRefresher;
     }
 
+    @Override
+    public IAerospikeClient getAerospikeClient() {
+        return client;
+    }
+
     @Override
     public <T> void createIndex(Class<T> entityClass, String indexName,
                                 String binName, IndexType indexType) {
@@ -649,11 +654,6 @@ public <T> long count(Class<T> entityClass) {
         return count(setName);
     }
 
-    @Override
-    public IAerospikeClient getAerospikeClient() {
-        return client;
-    }
-
     @Override
     public long count(String setName) {
         Assert.notNull(setName, "Set for count must not be null!");

diff --git a/src/main/java/org/springframework/data/aerospike/query/Qualifier.java b/src/main/java/org/springframework/data/aerospike/query/Qualifier.java
@@ -39,7 +39,7 @@
  */
 public class Qualifier implements Map<String, Object>, Serializable {
 
-    protected static final String FIELD = "field";
+    public static final String FIELD = "field";
     protected static final String IGNORE_CASE = "ignoreCase";
     protected static final String VALUE1 = "value1";
     protected static final String VALUE2 = "value2";

diff --git a/...main/java/org/springframework/data/aerospike/repository/query/AerospikePartTreeQuery.java b/...main/java/org/springframework/data/aerospike/repository/query/AerospikePartTreeQuery.java
@@ -24,6 +24,7 @@
 import org.springframework.data.repository.query.QueryMethodEvaluationContextProvider;
 import org.springframework.data.repository.query.parser.AbstractQueryCreator;
 
+import java.util.Arrays;
 import java.util.List;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
@@ -52,11 +53,28 @@ public Object execute(Object[] parameters) {
 
         Class<?> targetClass = getTargetClass(accessor);
 
+        // "findById" has its own processing flow
+        if (isIdProjectionQuery(targetClass, parameters, query.getAerospikeCriteria())) {
+            Object accessorValue = accessor.getBindableValue(0);
+            Object result;
+            if (accessorValue == null) {
+                throw new IllegalStateException("Parameters accessor value is null while parameters quantity is > 0");
+            } else if (accessorValue.getClass().isArray()) {
+                result = aerospikeOperations.findByIds(Arrays.stream(((Object[]) accessorValue)).toList(),
+                    sourceClass, targetClass);
+            } else if (accessorValue instanceof Iterable<?>) {
+                result = aerospikeOperations.findByIds((Iterable<?>) accessorValue, sourceClass, targetClass);
+            } else {
+                result = aerospikeOperations.findById(accessorValue, sourceClass, targetClass);
+            }
+            return result;
+        }
+
         if (queryMethod.isPageQuery() || queryMethod.isSliceQuery()) {
             Stream<?> result = findByQuery(query, targetClass);
             List<?> results = result.toList();
             Pageable pageable = accessor.getPageable();
-            long numberOfAllResults = aerospikeOperations.count(query, queryMethod.getEntityInformation().getJavaType());
+            long numberOfAllResults = aerospikeOperations.count(query, sourceClass);
 
             if (queryMethod.isSliceQuery()) {
                 boolean hasNext = numberOfAllResults > pageable.getPageSize() * (pageable.getOffset() + 1);
@@ -75,25 +93,12 @@ public Object execute(Object[] parameters) {
         throw new UnsupportedOperationException("Query method " + queryMethod.getNamedQueryName() + " not supported.");
     }
 
-    private Class<?> getTargetClass(ParametersParameterAccessor accessor) {
-        // Dynamic projection
-        if (accessor.getParameters().hasDynamicProjection()) {
-            return accessor.findDynamicProjection();
-        }
-        // DTO projection
-        if (queryMethod.getReturnedObjectType() != queryMethod.getEntityInformation().getJavaType()) {
-            return queryMethod.getReturnedObjectType();
-        }
-        // No projection - target class will be the entity class.
-        return queryMethod.getEntityInformation().getJavaType();
-    }
-
     private Stream<?> findByQuery(Query query, Class<?> targetClass) {
         // Run query and map to different target class.
-        if (targetClass != queryMethod.getEntityInformation().getJavaType()) {
+        if (targetClass != sourceClass) {
             return aerospikeOperations.find(query, queryMethod.getEntityInformation().getJavaType(), targetClass);
         }
         // Run query and map to entity class type.
-        return aerospikeOperations.find(query, queryMethod.getEntityInformation().getJavaType());
+        return aerospikeOperations.find(query, sourceClass);
     }
 }
diff --git a/.../java/org/springframework/data/aerospike/repository/query/BaseAerospikePartTreeQuery.java b/.../java/org/springframework/data/aerospike/repository/query/BaseAerospikePartTreeQuery.java
@@ -28,6 +28,9 @@
 import org.springframework.util.ClassUtils;
 
 import java.lang.reflect.Constructor;
+import java.util.Objects;
+
+import static org.springframework.data.aerospike.query.Qualifier.FIELD;
 
 /**
  * @author Peter Milne
@@ -37,6 +40,7 @@
 public abstract class BaseAerospikePartTreeQuery implements RepositoryQuery {
 
     protected final QueryMethod queryMethod;
+    protected final Class<?> sourceClass;
     private final QueryMethodEvaluationContextProvider evaluationContextProvider;
     private final Class<? extends AbstractQueryCreator<?, ?>> queryCreator;
 
@@ -46,6 +50,7 @@ protected BaseAerospikePartTreeQuery(QueryMethod queryMethod,
         this.queryMethod = queryMethod;
         this.evaluationContextProvider = evalContextProvider;
         this.queryCreator = queryCreator;
+        this.sourceClass = queryMethod.getEntityInformation().getJavaType();
     }
 
     @Override
@@ -54,10 +59,10 @@ public QueryMethod getQueryMethod() {
     }
 
     protected Query prepareQuery(Object[] parameters, ParametersParameterAccessor accessor) {
-        PartTree tree = new PartTree(queryMethod.getName(), queryMethod.getEntityInformation().getJavaType());
+        PartTree tree = new PartTree(queryMethod.getName(), sourceClass);
         Query baseQuery = createQuery(accessor, tree);
 
-        AerospikeCriteria criteria = (AerospikeCriteria) baseQuery.getCriteria();
+        AerospikeCriteria criteria = baseQuery.getAerospikeCriteria();
         Query query = new Query(criteria);
 
         if (accessor.getPageable().isPaged()) {
@@ -80,18 +85,35 @@ protected Query prepareQuery(Object[] parameters, ParametersParameterAccessor ac
             query.setSort(baseQuery.getSort());
         }
 
-        if (query.getCriteria() instanceof SpelExpression) {
+        if (query.getCriteria() instanceof SpelExpression spelExpression) {
             EvaluationContext context = this.evaluationContextProvider.getEvaluationContext(queryMethod.getParameters(),
                 parameters);
-            ((SpelExpression) query.getCriteria()).setEvaluationContext(context);
+            spelExpression.setEvaluationContext(context);
         }
 
         return query;
     }
 
+    Class<?> getTargetClass(ParametersParameterAccessor accessor) {
+        // Dynamic projection
+        if (accessor.getParameters().hasDynamicProjection()) {
+            return accessor.findDynamicProjection();
+        }
+        // DTO projection
+        if (queryMethod.getReturnedObjectType() != queryMethod.getEntityInformation().getJavaType()) {
+            return queryMethod.getReturnedObjectType();
+        }
+        // No projection - target class will be the entity class.
+        return queryMethod.getEntityInformation().getJavaType();
+    }
+
     public Query createQuery(ParametersParameterAccessor accessor, PartTree tree) {
         Constructor<? extends AbstractQueryCreator<?, ?>> constructor = ClassUtils
             .getConstructorIfAvailable(queryCreator, PartTree.class, ParameterAccessor.class);
         return (Query) BeanUtils.instantiateClass(constructor, tree, accessor).createQuery();
     }
+
+    protected static boolean isIdProjectionQuery(Class<?> targetClass, Object[] params, AerospikeCriteria criteria) {
+        return targetClass != null && params != null && params.length > 0 && Objects.equals(criteria.get(FIELD), "id");
+    }
 }
diff --git a/src/main/java/org/springframework/data/aerospike/repository/query/Query.java b/src/main/java/org/springframework/data/aerospike/repository/query/Query.java
@@ -55,12 +55,19 @@ public Query(Sort sort) {
     }
 
     /**
-     * Get the criteria object.
+     * Get the CriteriaDefinition object.
      */
     public CriteriaDefinition getCriteria() {
         return criteria;
     }
 
+    /**
+     * Get the CriteriaDefinition object cast to AerospikeCriteria.
+     */
+    public AerospikeCriteria getAerospikeCriteria() {
+        return (AerospikeCriteria) criteria;
+    }
+
     /**
      * Get {@link Sort}.
      */

diff --git a/...a/org/springframework/data/aerospike/repository/query/ReactiveAerospikePartTreeQuery.java b/...a/org/springframework/data/aerospike/repository/query/ReactiveAerospikePartTreeQuery.java
@@ -22,6 +22,8 @@
 import org.springframework.data.repository.query.parser.AbstractQueryCreator;
 import reactor.core.publisher.Flux;
 
+import java.util.Arrays;
+
 /**
  * @author Igor Ermolenko
  */
@@ -42,20 +44,25 @@ public Object execute(Object[] parameters) {
         ParametersParameterAccessor accessor = new ParametersParameterAccessor(queryMethod.getParameters(), parameters);
         Query query = prepareQuery(parameters, accessor);
         Class<?> targetClass = getTargetClass(accessor);
-        return findByQuery(query, targetClass);
-    }
 
-    private Class<?> getTargetClass(ParametersParameterAccessor accessor) {
-        // Dynamic projection
-        if (accessor.getParameters().hasDynamicProjection()) {
-            return accessor.findDynamicProjection();
+        // "findById" has its own processing flow
+        if (isIdProjectionQuery(targetClass, parameters, query.getAerospikeCriteria())) {
+            Object accessorValue = accessor.getBindableValue(0);
+            Object result;
+            if (accessorValue == null) {
+                throw new IllegalStateException("Parameters accessor value is null while parameters quantity is > 0");
+            } else if (accessorValue.getClass().isArray()) {
+                result = aerospikeOperations.findByIds(Arrays.stream(((Object[]) accessorValue)).toList(),
+                    sourceClass, targetClass);
+            } else if (accessorValue instanceof Iterable<?>) {
+                result = aerospikeOperations.findByIds((Iterable<?>) accessorValue, sourceClass, targetClass);
+            } else {
+                result = aerospikeOperations.findById(accessorValue, sourceClass, targetClass);
+            }
+            return result;
         }
-        // DTO projection
-        if (queryMethod.getReturnedObjectType() != queryMethod.getEntityInformation().getJavaType()) {
-            return queryMethod.getReturnedObjectType();
-        }
-        // No projection - target class will be the entity class.
-        return queryMethod.getEntityInformation().getJavaType();
+
+        return findByQuery(query, targetClass);
     }
 
     private Flux<?> findByQuery(Query query, Class<?> targetClass) {

diff --git a/src/test/java/org/springframework/data/aerospike/repository/PersonRepositoryQueryTests.java b/src/test/java/org/springframework/data/aerospike/repository/PersonRepositoryQueryTests.java
@@ -1135,17 +1135,14 @@ public void findAllWithGivenIds() {
         List<Person> result = (List<Person>) repository.findAllById(List.of(dave.getId(), boyd.getId()));
 
         assertThat(result)
-            .hasSize(2)
-            .contains(dave)
-            .doesNotContain(oliver, carter, alicia);
+            .containsOnly(dave, boyd);
     }
 
     @Test
     public void findPersonsByLastName() {
         List<Person> result = repository.findByLastName("Beauford");
 
         assertThat(result)
-            .hasSize(1)
             .containsOnly(carter);
     }
 
@@ -1154,10 +1151,25 @@ public void findPersonsSomeFieldsByLastNameProjection() {
         List<PersonSomeFields> result = repository.findPersonSomeFieldsByLastName("Beauford");
 
         assertThat(result)
-            .hasSize(1)
             .containsOnly(carter.toPersonSomeFields());
     }
 
+    @Test
+    public void findPersonsSomeFieldsByIdProjection() {
+        List<PersonSomeFields> result = repository.findPersonSomeFieldsById(carter.getId());
+
+        assertThat(result)
+            .containsOnly(carter.toPersonSomeFields());
+    }
+
+    @Test
+    public void findPersonsSomeFieldsByIdsProjection() {
+        List<PersonSomeFields> result = repository.findPersonSomeFieldsById(carter.getId(), dave.getId(), boyd.getId());
+
+        assertThat(result)
+            .containsOnly(carter.toPersonSomeFields(), dave.toPersonSomeFields(), boyd.toPersonSomeFields());
+    }
+
     @Test
     public void findDynamicTypeByLastNameDynamicProjection() {
         List<PersonSomeFields> result = repository.findByLastName("Beauford", PersonSomeFields.class);
@@ -1167,6 +1179,23 @@ public void findDynamicTypeByLastNameDynamicProjection() {
             .containsOnly(carter.toPersonSomeFields());
     }
 
+    @Test
+    public void findDynamicTypeByIdDynamicProjection() {
+        List<PersonSomeFields> result = repository.findById(dave.getId(), PersonSomeFields.class);
+
+        assertThat(result)
+            .containsOnly(dave.toPersonSomeFields());
+    }
+
+    @Test
+    public void findDynamicTypeByIdsDynamicProjection() {
+        List<PersonSomeFields> result = repository.findById(List.of(boyd.getId(), dave.getId()),
+            PersonSomeFields.class);
+
+        assertThat(result)
+            .containsOnly(boyd.toPersonSomeFields(), dave.toPersonSomeFields());
+    }
+
     @Test
     public void findPersonsByFriendAge() {
         oliver.setFriend(alicia);

diff --git a/src/test/java/org/springframework/data/aerospike/sample/PersonRepository.java b/src/test/java/org/springframework/data/aerospike/sample/PersonRepository.java
@@ -42,9 +42,18 @@ public interface PersonRepository<P extends Person> extends AerospikeRepository<
     // DTO Projection
     List<PersonSomeFields> findPersonSomeFieldsByLastName(String lastName);
 
+    List<PersonSomeFields> findPersonSomeFieldsById(String id);
+
+    List<PersonSomeFields> findPersonSomeFieldsById(String... ids);
+
     // Dynamic Projection
     <T> List<T> findByLastName(String lastName, Class<T> type);
 
+    // Dynamic Projection
+    <T> List<T> findById(String id, Class<T> type);
+
+    <T> List<T> findById(Collection<String> ids, Class<T> type);
+
     Page<P> findByLastNameStartsWithOrderByAgeAsc(String prefix, Pageable pageable);
 
     List<P> findByLastNameEndsWith(String postfix);
@@ -67,8 +76,7 @@ public interface PersonRepository<P extends Person> extends AerospikeRepository<
 
     /**
      * Find all entities with firstName matching the given regex. POSIX Extended Regular Expression syntax is used to
-     * interpret the regex.
-     * The same as {@link #findByFirstNameLike(String)}
+     * interpret the regex. The same as {@link #findByFirstNameLike(String)}
      *
      * @param firstNameRegex Regex to find matching firstName
      */