DATACOUCH-187 - Allow counting queries and single projections.

This change allows simple projections to types long, boolean, String when an inline N1QL query uses an aggregation function like AVG, COUNT.

The projection's result must be a single row with a single key/value pair in the N1QL response. Only the value gets returned.

src/integration/java/org/springframework/data/couchbase/repository/N1qlCrudRepositoryTests.java

@@ -16,22 +16,21 @@
 
 package org.springframework.data.couchbase.repository;
 
-import static org.junit.Assert.assertFalse;
-import static org.junit.Assert.assertTrue;
-import static org.junit.Assert.fail;
+import static org.junit.Assert.*;
 
 import java.util.Date;
 import java.util.List;
 
 import com.couchbase.client.java.Bucket;
-import com.couchbase.client.java.error.DocumentDoesNotExistException;
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
 import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.dao.DataRetrievalFailureException;
 import org.springframework.data.couchbase.IntegrationTestApplicationConfig;
+import org.springframework.data.couchbase.core.CouchbaseQueryExecutionException;
 import org.springframework.data.couchbase.repository.config.RepositoryOperationsMapping;
 import org.springframework.data.couchbase.repository.support.CouchbaseRepositoryFactory;
 import org.springframework.data.couchbase.repository.support.IndexManager;
@@ -76,9 +75,9 @@ public class N1qlCrudRepositoryTests {
 
   @After
   public void cleanUp() {
-    try { itemRepository.delete(KEY_ITEM); } catch (DocumentDoesNotExistException e) {}
-    try { partyRepository.delete(KEY_PARTY); } catch (DocumentDoesNotExistException e) {}
-    try { partyRepository.delete(KEY_PARTY_KEYWORD); } catch (DocumentDoesNotExistException e) {}
+    try { itemRepository.delete(KEY_ITEM); } catch (DataRetrievalFailureException e) {}
+    try { partyRepository.delete(KEY_PARTY); } catch (DataRetrievalFailureException e) {}
+    try { partyRepository.delete(KEY_PARTY_KEYWORD); } catch (DataRetrievalFailureException e) {}
   }
 
   @Test
@@ -107,4 +106,56 @@ public class N1qlCrudRepositoryTests {
       }
     }
   }
+
+  @Test
+  public void shouldGenerateCountProjection() {
+    Party partyHasKeyword = new Party(KEY_PARTY_KEYWORD, "party", null, new Date(), 40, new Point(500, 500));
+    partyRepository.save(partyHasKeyword);
+    long countTotal = partyRepository.count();
+    long countCustom = partyRepository.countAllByDescriptionNotNull();
+    assertEquals(countTotal - 1, countCustom);
+  }
+
+  @Test
+  public void shouldCountWhenReturningLongAndUsingStringSelectFromSpEL() {
+    Party partyHasKeyword = new Party(KEY_PARTY_KEYWORD, "party", "second party", new Date(), 40, new Point(500, 500));
+    partyRepository.save(partyHasKeyword);
+
+    long countTotal = partyRepository.count();
+    long countCustom = partyRepository.countCustom();
+    assertEquals(countTotal, countCustom);
+  }
+
+  @Test
+  public void shouldCustomCountWhenReturningLongAndUsingStringWithoutSpEL() {
+    Party partyHasKeyword = new Party(KEY_PARTY_KEYWORD, "party", "second party", new Date(), 40, new Point(500, 500));
+    partyRepository.save(partyHasKeyword);
+
+    long countTotal = partyRepository.count();
+    long countCustom = partyRepository.countCustomPlusFive();
+    assertEquals(countTotal + 5, countCustom);
+  }
+
+  @Test(expected = CouchbaseQueryExecutionException.class)
+  public void shouldFailConversionWithStringReturnType() {
+    Party partyHasKeyword = new Party(KEY_PARTY_KEYWORD, "party", "desc is a N1QL keyword", new Date(), 40, new Point(500, 500));
+    partyRepository.save(partyHasKeyword);
+
+    String someString = partyRepository.findSomeString();
+  }
+
+  @Test
+  public void shouldDoNumericProjectionWithStringBasedQuery() {
+    Party partyHasKeyword = new Party(KEY_PARTY_KEYWORD, "party", "desc is a N1QL keyword", new Date(), 4000000, new Point(500, 500));
+    partyRepository.save(partyHasKeyword);
+
+    long max = partyRepository.findMaxAttendees();
+    assertEquals(4000000, max);
+  }
+
+  @Test
+  public void shouldDoBooleanProjectionWithStringBasedQuery() {
+    boolean someBoolean = partyRepository.justABoolean();
+    assertEquals(true, someBoolean);
+  }
 }

src/integration/java/org/springframework/data/couchbase/repository/PartyRepository.java

@@ -3,11 +3,14 @@ package org.springframework.data.couchbase.repository;
 import java.util.Date;
 import java.util.List;
 
+import org.springframework.data.couchbase.core.query.Query;
 import org.springframework.data.couchbase.core.query.View;
+import org.springframework.data.couchbase.core.query.ViewIndexed;
 
 /**
  * @author Simon Baslé
  */
+@ViewIndexed(designDoc = "party", viewName = "all")
 public interface PartyRepository extends CouchbaseRepository<Party, String> {
 
   List<Party> findByAttendeesGreaterThanEqual(int minAttendees);
@@ -19,4 +22,20 @@ public interface PartyRepository extends CouchbaseRepository<Party, String> {
 
   List<Object> findAllByDescriptionNotNull();
 
+  long countAllByDescriptionNotNull();
+
+  @Query("SELECT MAX(attendees) FROM #{#n1ql.bucket} WHERE #{#n1ql.filter}")
+  long findMaxAttendees();
+
+  @Query("SELECT `desc` FROM #{#n1ql.bucket} WHERE #{#n1ql.filter}")
+  String findSomeString();
+
+  @Query("SELECT count(*) + 5 FROM #{#n1ql.bucket} WHERE #{#n1ql.filter}")
+  long countCustomPlusFive();
+
+  @Query("#{#n1ql.selectEntity} WHERE #{#n1ql.filter}")
+  long countCustom();
+
+  @Query("SELECT 1 = 1")
+  boolean justABoolean();
 }

src/main/asciidoc/repository.adoc

@@ -143,6 +143,10 @@ A few N1QL-specific values are provided through SpEL:
  - `#n1ql.bucket` will be replaced by the name of the bucket the entity is stored in, escaped in backticks.
  - `#n1ql.fields` will be replaced by the list of fields (eg. for a SELECT clause) necessary to reconstruct the entity.
 
+IMPORTANT: We recommend that you always use the `selectEntity` SpEL and a WHERE clause with a `filter` SpEL (since otherwise your query could be impacted by entities from other repositories).
+
+You can also do single projections in your N1QL queries (provided it selects only one field and returns only one result, usually an aggregation like `COUNT`, `AVG`, `MAX`...). Such projection would have a simple return type like `long`, `boolean` or `String`. This is *NOT* intended for projections to DTOs.
+
 Another example: "`#{#n1ql.selectEntity} WHERE #{#n1ql.filter} AND test = $1`", which is equivalent to
 `SELECT #{#n1ql.fields} FROM #{#n1ql.bucket} WHERE #{#n1ql.filter} AND test = $1`".
 

src/main/java/org/springframework/data/couchbase/repository/query/AbstractN1qlBasedQuery.java

@@ -16,9 +16,12 @@
 
 package org.springframework.data.couchbase.repository.query;
 
+import java.util.Collection;
 import java.util.List;
+import java.util.Map;
 
 import com.couchbase.client.java.document.json.JsonArray;
+import com.couchbase.client.java.error.QueryExecutionException;
 import com.couchbase.client.java.query.N1qlQuery;
 import com.couchbase.client.java.query.N1qlParams;
 import com.couchbase.client.java.query.Statement;
@@ -27,11 +30,14 @@ import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 import org.springframework.data.couchbase.core.CouchbaseOperations;
+import org.springframework.data.couchbase.core.CouchbaseQueryExecutionException;
 import org.springframework.data.domain.PageImpl;
 import org.springframework.data.domain.Pageable;
 import org.springframework.data.domain.SliceImpl;
+import org.springframework.data.repository.core.NamedQueries;
 import org.springframework.data.repository.query.ParameterAccessor;
 import org.springframework.data.repository.query.ParametersParameterAccessor;
+import org.springframework.data.repository.query.QueryCreationException;
 import org.springframework.data.repository.query.QueryMethod;
 import org.springframework.data.repository.query.RepositoryQuery;
 import org.springframework.data.util.StreamUtils;
@@ -53,8 +59,18 @@ public abstract class AbstractN1qlBasedQuery implements RepositoryQuery {
     this.couchbaseOperations = couchbaseOperations;
   }
 
+  /**
+   * The statement for a count() query. This must aggregate using count with the alias {@link CountFragment#COUNT_ALIAS}.
+   * @see CountFragment
+   */
   protected abstract Statement getCount(ParameterAccessor accessor, Object[] runtimeParameters);
 
+  /**
+   * @return true if the {@link #getCount(ParameterAccessor, Object[]) count statement} should also be used when
+   * the return type of the QueryMethod is a primitive type.
+   */
+  protected abstract boolean useGeneratedCountQuery();
+
   protected abstract Statement getStatement(ParameterAccessor accessor, Object[] runtimeParameters);
 
   protected abstract JsonArray getPlaceholderValues(ParameterAccessor accessor);
@@ -70,7 +86,6 @@ public abstract class AbstractN1qlBasedQuery implements RepositoryQuery {
         getCouchbaseOperations().getDefaultConsistency().n1qlConsistency());
 
     //prepare a count query
-    //TODO only do that when necessary (isPageQuery or isSliceQuery)
     Statement countStatement = getCount(accessor, parameters);
     N1qlQuery countQuery = buildQuery(countStatement, queryPlaceholderValues,
         getCouchbaseOperations().getDefaultConsistency().n1qlConsistency());
@@ -105,9 +120,18 @@ public abstract class AbstractN1qlBasedQuery implements RepositoryQuery {
       return executeCollection(query);
     } else if (queryMethod.isQueryForEntity()) {
       return executeEntity(query);
-    } else {
+    } else if (queryMethod.isStreamQuery()){
       return executeStream(query);
+    } else if (queryMethod.getReturnedObjectType().isPrimitive()
+        && useGeneratedCountQuery()) {
+      //attempt to execute the created COUNT query
+      return executeSingleProjection(countQuery);
+    } else {
+      //attempt a single projection on a simple type
+      // (ie, a single row with a single k->v entry where v is the desired value)
+      return executeSingleProjection(query);
     }
+    //more complex projections could be added in the future, like DTO direct mapping with a SELECT a,b,c FROM something
   }
 
   private void logIfNecessary(N1qlQuery query) {
@@ -157,6 +181,30 @@ public abstract class AbstractN1qlBasedQuery implements RepositoryQuery {
     return new SliceImpl(hasNext ? result.subList(0, pageSize) : result, pageable, hasNext);
   }
 
+  protected Object executeSingleProjection(N1qlQuery query) {
+    logIfNecessary(query);
+    //the structure of the response from N1QL gives us a JSON object even when selecting a single aggregation
+    List<Map> resultAsMap = couchbaseOperations.findByN1QLProjection(query, Map.class);
+
+    if (resultAsMap.size() != 1) {
+      throw new CouchbaseQueryExecutionException("Query returning a primitive type are expected to return " +
+          "exactly 1 result, got " + resultAsMap.size());
+    }
+
+    Map<String, Object> singleRow = (Map<String, Object>) resultAsMap.get(0);
+    if (singleRow.size() != 1) {
+      throw new CouchbaseQueryExecutionException("Query returning a simple type are expected to return " +
+          "a unique value, got " + singleRow.size());
+    }
+    Collection<Object> rowValues = singleRow.values();
+    if (rowValues.size() != 1) {
+      throw new CouchbaseQueryExecutionException("Query returning a simple type are expected to use a " +
+          "single aggregation/projection, got " + rowValues.size());
+    }
+
+    return rowValues.iterator().next();
+  }
+
   @Override
   public CouchbaseQueryMethod getQueryMethod() {
     return this.queryMethod;

src/main/java/org/springframework/data/couchbase/repository/query/CouchbaseQueryMethod.java

@@ -154,4 +154,9 @@ public class CouchbaseQueryMethod extends QueryMethod {
     String query = (String) AnnotationUtils.getValue(getN1qlAnnotation());
     return StringUtils.hasText(query) ? query : null;
   }
+
+  @Override
+  public String toString() {
+    return super.toString();
+  }
 }

src/main/java/org/springframework/data/couchbase/repository/query/PartTreeN1qlBasedQuery.java

@@ -91,4 +91,9 @@ public class PartTreeN1qlBasedQuery extends AbstractN1qlBasedQuery {
       return selectFromWhereOrderBy;
     }
   }
+
+  @Override
+  protected boolean useGeneratedCountQuery() {
+    return false; //generated count query is just for Page/Slice, not projections
+  }
 }

src/main/java/org/springframework/data/couchbase/repository/query/StringN1qlBasedQuery.java

@@ -170,6 +170,11 @@ public class StringN1qlBasedQuery extends AbstractN1qlBasedQuery {
     return N1qlQuery.simple(parsedCountStatement).statement();
   }
 
+  @Override
+  protected boolean useGeneratedCountQuery() {
+    return this.originalStatement.contains(SPEL_SELECT_FROM_CLAUSE);
+  }
+
   /**
    * This class is exposed to SpEL parsing through the variable <code>#{@value StringN1qlBasedQuery#SPEL_PREFIX}</code>.
    * Use the attributes in your SpEL expressions: {@link #selectEntity}, {@link #fields}, {@link #bucket} and {@link #filter}.

src/test/java/org/springframework/data/couchbase/repository/query/AbstractN1qlBasedQueryTest.java

@@ -6,6 +6,7 @@ import static org.mockito.Mockito.*;
 
 import java.util.ArrayList;
 import java.util.List;
+import java.util.concurrent.CountDownLatch;
 
 import com.couchbase.client.java.document.json.JsonArray;
 import com.couchbase.client.java.document.json.JsonObject;
@@ -18,7 +19,9 @@ import com.couchbase.client.java.query.consistency.ScanConsistency;
 import org.junit.Test;
 import org.mockito.Mockito;
 
+import org.springframework.data.couchbase.core.mapping.event.User;
 import org.springframework.data.domain.Pageable;
+import org.springframework.data.repository.query.QueryCreationException;
 import org.springframework.data.repository.query.QueryMethod;
 
 public class AbstractN1qlBasedQueryTest {
@@ -90,6 +93,7 @@ public class AbstractN1qlBasedQueryTest {
     verify(mock, never()).executeStream(any(N1qlQuery.class));
     verify(mock, never()).executePaged(any(N1qlQuery.class), any(N1qlQuery.class), any(Pageable.class));
     verify(mock, never()).executeSliced(any(N1qlQuery.class), any(N1qlQuery.class), any(Pageable.class));
+    verify(mock, never()).executeSingleProjection(any(N1qlQuery.class));
   }
 
   @Test
@@ -109,6 +113,7 @@ public class AbstractN1qlBasedQueryTest {
     verify(mock, never()).executeStream(any(N1qlQuery.class));
     verify(mock, never()).executePaged(any(N1qlQuery.class), any(N1qlQuery.class), any(Pageable.class));
     verify(mock, never()).executeSliced(any(N1qlQuery.class), any(N1qlQuery.class), any(Pageable.class));
+    verify(mock, never()).executeSingleProjection(any(N1qlQuery.class));
   }
 
   @Test
@@ -128,6 +133,7 @@ public class AbstractN1qlBasedQueryTest {
     verify(mock).executeStream(any(N1qlQuery.class));
     verify(mock, never()).executePaged(any(N1qlQuery.class), any(N1qlQuery.class), any(Pageable.class));
     verify(mock, never()).executeSliced(any(N1qlQuery.class), any(N1qlQuery.class), any(Pageable.class));
+    verify(mock, never()).executeSingleProjection(any(N1qlQuery.class));
   }
 
   @Test
@@ -146,6 +152,7 @@ public class AbstractN1qlBasedQueryTest {
     verify(mock, never()).executeStream(any(N1qlQuery.class));
     verify(mock).executePaged(any(N1qlQuery.class), any(N1qlQuery.class), any(Pageable.class));
     verify(mock, never()).executeSliced(any(N1qlQuery.class), any(N1qlQuery.class), any(Pageable.class));
+    verify(mock, never()).executeSingleProjection(any(N1qlQuery.class));
   }
 
   @Test
@@ -165,10 +172,11 @@ public class AbstractN1qlBasedQueryTest {
     verify(mock, never()).executeStream(any(N1qlQuery.class));
     verify(mock, never()).executePaged(any(N1qlQuery.class), any(N1qlQuery.class), any(Pageable.class));
     verify(mock).executeSliced(any(N1qlQuery.class), any(N1qlQuery.class), any(Pageable.class));
+    verify(mock, never()).executeSingleProjection(any(N1qlQuery.class));
   }
 
   @Test
-  public void shouldThrowWhenUnsupportedType() throws NoSuchMethodException {
+  public void shouldThrowWhenModifyingType() {
     CouchbaseQueryMethod queryMethod = Mockito.mock(CouchbaseQueryMethod.class);
     N1qlQuery query = Mockito.mock(N1qlQuery.class);
     Pageable pageable = Mockito.mock(Pageable.class);
@@ -183,5 +191,48 @@ public class AbstractN1qlBasedQueryTest {
     verify(mock, never()).executeStream(any(N1qlQuery.class));
     verify(mock, never()).executePaged(any(N1qlQuery.class), any(N1qlQuery.class), any(Pageable.class));
     verify(mock, never()).executeSliced(any(N1qlQuery.class), any(N1qlQuery.class), any(Pageable.class));
+    verify(mock, never()).executeSingleProjection(any(N1qlQuery.class));
+  }
+
+  @Test
+  public void shouldExecuteSingleProjectionWhenRandomObjectReturnType() {
+    CouchbaseQueryMethod queryMethod = Mockito.mock(CouchbaseQueryMethod.class);
+    doReturn(CountDownLatch.class).when(queryMethod).getReturnedObjectType();
+
+    N1qlQuery query = Mockito.mock(N1qlQuery.class);
+    Pageable pageable = Mockito.mock(Pageable.class);
+    AbstractN1qlBasedQuery mock = mock(AbstractN1qlBasedQuery.class);
+    when(mock.executeDependingOnType(any(N1qlQuery.class), any(N1qlQuery.class), any(QueryMethod.class), any(Pageable.class),
+        anyBoolean(), anyBoolean(), anyBoolean()))
+        .thenCallRealMethod();
+
+    mock.executeDependingOnType(query, query, queryMethod, pageable, false, false, false);
+    verify(mock, never()).executeCollection(any(N1qlQuery.class));
+    verify(mock, never()).executeEntity(any(N1qlQuery.class));
+    verify(mock, never()).executeStream(any(N1qlQuery.class));
+    verify(mock, never()).executePaged(any(N1qlQuery.class), any(N1qlQuery.class), any(Pageable.class));
+    verify(mock, never()).executeSliced(any(N1qlQuery.class), any(N1qlQuery.class), any(Pageable.class));
+    verify(mock).executeSingleProjection(any(N1qlQuery.class));
+  }
+
+  @Test
+  public void shouldExecuteSingleProjectionWhenPrimitiveReturnType() {
+    CouchbaseQueryMethod queryMethod = Mockito.mock(CouchbaseQueryMethod.class);
+    doReturn(long.class).when(queryMethod).getReturnedObjectType();
+
+    N1qlQuery query = Mockito.mock(N1qlQuery.class);
+    Pageable pageable = Mockito.mock(Pageable.class);
+    AbstractN1qlBasedQuery mock = mock(AbstractN1qlBasedQuery.class);
+    when(mock.executeDependingOnType(any(N1qlQuery.class), any(N1qlQuery.class), any(QueryMethod.class), any(Pageable.class),
+        anyBoolean(), anyBoolean(), anyBoolean()))
+        .thenCallRealMethod();
+
+    mock.executeDependingOnType(query, query, queryMethod, pageable, false, false, false);
+    verify(mock, never()).executeCollection(any(N1qlQuery.class));
+    verify(mock, never()).executeEntity(any(N1qlQuery.class));
+    verify(mock, never()).executeStream(any(N1qlQuery.class));
+    verify(mock, never()).executePaged(any(N1qlQuery.class), any(N1qlQuery.class), any(Pageable.class));
+    verify(mock, never()).executeSliced(any(N1qlQuery.class), any(N1qlQuery.class), any(Pageable.class));
+    verify(mock).executeSingleProjection(any(N1qlQuery.class));
   }
 }