Add keyword IN and NOT_IN, #144.

Signed-off-by: Pan Li <panli@microsoft.com>
microsoft · Aug 23, 2018 · 281fc09 · 281fc09
1 parent 9d1cfb5
commit 281fc09
Show file tree

Hide file tree

Showing 6 changed files with 100 additions and 8 deletions.
diff --git a/.../java/com/microsoft/azure/spring/data/cosmosdb/core/generator/AbstractQueryGenerator.java b/.../java/com/microsoft/azure/spring/data/cosmosdb/core/generator/AbstractQueryGenerator.java
@@ -48,8 +48,7 @@ private String generateBinaryQuery(@NonNull String left, @NonNull String right,
         return String.join(" ", left, type.getSqlKeyword(), right);
     }
 
-    private String generateINQuery(Criteria criteria) {
-        Assert.isTrue(CriteriaType.IN.equals(criteria.getType()), "Criteria type should be IN.");
+    private String generateInQuery(Criteria criteria) {
         Assert.isTrue(criteria.getSubjectValues().size() == 1, "Criteria should have only one subject value");
 
         if (!(criteria.getSubjectValues().get(0) instanceof Collection)) {
@@ -61,13 +60,19 @@ private String generateINQuery(Criteria criteria) {
 
         values.forEach(o -> {
             if (o instanceof Integer || o instanceof Long) {
-                inRangeValues.add(String.format("%d", (Long) o));
+                inRangeValues.add(String.format("%d", o));
+            } else if (o instanceof String) {
+                inRangeValues.add(String.format("'%s'", (String) o));
+            } else if (o instanceof Boolean) {
+                inRangeValues.add(String.format("%b", (Boolean) o));
             } else {
-                inRangeValues.add(String.format("'%s'", o.toString()));
+                throw new IllegalQueryException("IN keyword Range only support Number and String type.");
             }
         });
 
-        return String.format("r.%s IN (%s)", criteria.getSubject(), String.join(",", inRangeValues));
+        final String inRange = String.join(",", inRangeValues);
+
+        return String.format("r.%s %s (%s)", criteria.getSubject(), criteria.getType().getSqlKeyword(), inRange);
     }
 
     private String generateQueryBody(@NonNull Criteria criteria, @NonNull List<Pair<String, Object>> parameters) {
@@ -77,7 +82,8 @@ private String generateQueryBody(@NonNull Criteria criteria, @NonNull List<Pair<
             case ALL:
                 return "";
             case IN:
-                return this.generateINQuery(criteria);
+            case NOT_IN:
+                return this.generateInQuery(criteria);
             case IS_EQUAL:
             case BEFORE:
             case AFTER:

diff --git a/src/main/java/com/microsoft/azure/spring/data/cosmosdb/core/query/CriteriaType.java b/src/main/java/com/microsoft/azure/spring/data/cosmosdb/core/query/CriteriaType.java
@@ -24,6 +24,7 @@ public enum CriteriaType {
     BEFORE("<"),
     AFTER(">"),
     IN("IN"),
+    NOT_IN("NOT IN"),
     LESS_THAN("<"),
     LESS_THAN_EQUAL("<="),
     GREATER_THAN(">"),
@@ -42,6 +43,7 @@ public enum CriteriaType {
         map.put(Part.Type.BEFORE, CriteriaType.BEFORE);
         map.put(Part.Type.AFTER, CriteriaType.AFTER);
         map.put(Part.Type.IN, CriteriaType.IN);
+        map.put(Part.Type.NOT_IN, CriteriaType.NOT_IN);
         map.put(Part.Type.GREATER_THAN, CriteriaType.GREATER_THAN);
         map.put(Part.Type.GREATER_THAN_EQUAL, CriteriaType.GREATER_THAN_EQUAL);
         map.put(Part.Type.LESS_THAN, CriteriaType.LESS_THAN);
@@ -104,6 +106,8 @@ public static boolean isBinary(CriteriaType type) {
      */
     public static boolean isUnary(CriteriaType type) {
         switch (type) {
+            case IN:
+            case NOT_IN:
             case IS_EQUAL:
             case BEFORE:
             case AFTER:

diff --git a/src/main/java/com/microsoft/azure/spring/data/cosmosdb/core/query/DocumentQuery.java b/src/main/java/com/microsoft/azure/spring/data/cosmosdb/core/query/DocumentQuery.java
@@ -55,7 +55,12 @@ private boolean isCrossPartitionQuery(@NonNull String keyName) {
 
         final Optional<Criteria> criteria = this.getSubjectCriteria(this.criteria, keyName);
 
-        return criteria.map(criteria1 -> criteria1.getType() == CriteriaType.IS_EQUAL).orElse(true);
+        return criteria.map(criteria1 -> criteria1.getType() != CriteriaType.IS_EQUAL).orElse(true);
+    }
+
+    private boolean hasKeywordOr() {
+        // If there is OR keyword in DocumentQuery, the top node of Criteria must be OR type.
+        return this.criteria.getType() == CriteriaType.OR;
     }
 
     /**
@@ -75,7 +80,7 @@ public boolean isCrossPartitionQuery(@NonNull List<String> partitionKeys) {
             }
         }
 
-        return false;
+        return this.hasKeywordOr();
     }
 
     private Optional<Criteria> getSubjectCriteria(@NonNull Criteria criteria, @NonNull String keyName) {

diff --git a/src/test/java/com/microsoft/azure/spring/data/cosmosdb/core/query/CriteriaUnitTest.java b/src/test/java/com/microsoft/azure/spring/data/cosmosdb/core/query/CriteriaUnitTest.java
@@ -5,10 +5,13 @@
  */
 package com.microsoft.azure.spring.data.cosmosdb.core.query;
 
+import com.microsoft.azure.spring.data.cosmosdb.core.generator.FindQuerySpecGenerator;
+import com.microsoft.azure.spring.data.cosmosdb.exception.IllegalQueryException;
 import org.junit.Assert;
 import org.junit.Test;
 
 import java.util.Arrays;
+import java.util.Collections;
 import java.util.List;
 
 import static com.microsoft.azure.spring.data.cosmosdb.common.TestConstants.CRITERIA_KEY;
@@ -45,4 +48,22 @@ public void testBinaryCriteria() {
         Assert.assertEquals(criteria.getSubCriteria().get(0), leftCriteria);
         Assert.assertEquals(criteria.getSubCriteria().get(1), rightCriteria);
     }
+
+    @Test(expected = IllegalQueryException.class)
+    public void testInvalidInKeywordParameter() {
+        final List<Object> values = Collections.singletonList(CRITERIA_OBJECT);
+        final Criteria criteria = Criteria.getUnaryInstance(CriteriaType.IN, CRITERIA_KEY, values);
+        final DocumentQuery query = new DocumentQuery(criteria);
+
+        new FindQuerySpecGenerator().generate(query);
+    }
+
+    @Test(expected = IllegalQueryException.class)
+    public void testInvalidInKeywordType() {
+        final List<Object> values = Collections.singletonList(new IllegalQueryException(""));
+        final Criteria criteria = Criteria.getUnaryInstance(CriteriaType.IN, CRITERIA_KEY, values);
+        final DocumentQuery query = new DocumentQuery(criteria);
+
+        new FindQuerySpecGenerator().generate(query);
+    }
 }
diff --git a/.../com/microsoft/azure/spring/data/cosmosdb/repository/integration/ProjectRepositoryIT.java b/.../com/microsoft/azure/spring/data/cosmosdb/repository/integration/ProjectRepositoryIT.java
@@ -314,4 +314,56 @@ public void testFindByIn() {
 
         assertProjectListEquals(projects, Arrays.asList(PROJECT_0, PROJECT_4));
     }
+
+    @Test
+    public void testFindByInWithAnd() {
+        List<Project> projects = repository.findByCreatorInAndStarCountIn(Arrays.asList(CREATOR_0, CREATOR_1),
+                Arrays.asList(STAR_COUNT_2, STAR_COUNT_3));
+
+        Assert.assertTrue(projects.isEmpty());
+
+        projects = repository.findByCreatorInAndStarCountIn(Arrays.asList(CREATOR_0, CREATOR_1),
+                Arrays.asList(STAR_COUNT_0, STAR_COUNT_2));
+
+        assertProjectListEquals(projects, Arrays.asList(PROJECT_0, PROJECT_4));
+
+        projects = repository.findByCreatorInAndStarCountIn(Arrays.asList(CREATOR_0, CREATOR_1, CREATOR_2),
+                Arrays.asList(STAR_COUNT_0, STAR_COUNT_1, STAR_COUNT_2));
+
+        assertProjectListEquals(projects, Arrays.asList(PROJECT_0, PROJECT_1, PROJECT_2, PROJECT_4));
+    }
+
+    @Test
+    public void testFindByNotIn() {
+        List<Project> projects = repository.findByCreatorNotIn(
+                Arrays.asList(CREATOR_0, CREATOR_1, CREATOR_2, CREATOR_3));
+
+        Assert.assertTrue(projects.isEmpty());
+
+        projects = repository.findByCreatorNotIn(Arrays.asList(CREATOR_1, CREATOR_2));
+
+        assertProjectListEquals(projects, Arrays.asList(PROJECT_0, PROJECT_3, PROJECT_4));
+
+        projects = repository.findByCreatorNotIn(Arrays.asList(CREATOR_0, FAKE_CREATOR));
+
+        assertProjectListEquals(projects, Arrays.asList(PROJECT_1, PROJECT_2, PROJECT_3));
+    }
+
+    @Test
+    public void testFindByInWithNotIn() {
+        List<Project> projects = repository.findByCreatorInAndStarCountNotIn(Collections.singletonList(FAKE_CREATOR),
+                Arrays.asList(STAR_COUNT_2, STAR_COUNT_3));
+
+        Assert.assertTrue(projects.isEmpty());
+
+        projects = repository.findByCreatorInAndStarCountNotIn(Arrays.asList(CREATOR_0, CREATOR_1),
+                Arrays.asList(STAR_COUNT_0, STAR_COUNT_2));
+
+        assertProjectListEquals(projects, Collections.singletonList(PROJECT_1));
+
+        projects = repository.findByCreatorInAndStarCountNotIn(Arrays.asList(CREATOR_0, CREATOR_1, CREATOR_2),
+                Arrays.asList(STAR_COUNT_1, STAR_COUNT_2));
+
+        assertProjectListEquals(projects, Arrays.asList(PROJECT_0, PROJECT_4));
+    }
 }
diff --git a/...ava/com/microsoft/azure/spring/data/cosmosdb/repository/repository/ProjectRepository.java b/...ava/com/microsoft/azure/spring/data/cosmosdb/repository/repository/ProjectRepository.java
@@ -58,4 +58,8 @@ List<Project> findByNameOrCreatorAndForkCountOrStarCount(String name, String cre
     List<Project> findByCreatorIn(Collection<String> creators);
 
     List<Project> findByCreatorInAndStarCountIn(Collection<String> creators, Collection<Long> starCounts);
+
+    List<Project> findByCreatorNotIn(Collection<String> creators);
+
+    List<Project> findByCreatorInAndStarCountNotIn(Collection<String> creators, Collection<Long> starCounts);
 }