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/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}.