Add find and delete on documentQuery, #144. (#151)

* Add find and delete on documentQuery, #144.

Signed-off-by: Pan Li <panli@microsoft.com>

src/main/java/com/microsoft/azure/spring/data/cosmosdb/core/DocumentDbOperations.java

@@ -9,7 +9,7 @@
 import com.microsoft.azure.documentdb.DocumentCollection;
 import com.microsoft.azure.documentdb.PartitionKey;
 import com.microsoft.azure.spring.data.cosmosdb.core.convert.MappingDocumentDbConverter;
-import com.microsoft.azure.spring.data.cosmosdb.core.query.Query;
+import com.microsoft.azure.spring.data.cosmosdb.core.query.DocumentQuery;
 import com.microsoft.azure.spring.data.cosmosdb.repository.support.DocumentDbEntityInformation;
 
 import java.util.List;
@@ -32,7 +32,7 @@ <T> T findById(String collectionName,
                    Object id,
                    Class<T> entityClass);
 
-    <T> List<T> find(Query query,
+    <T> List<T> find(DocumentQuery query,
                      Class<T> entityClass,
                      String collectionName);
 
@@ -54,7 +54,7 @@ <T> void deleteById(String collectionName,
 
     void deleteAll(String collectionName);
 
-    <T> List<T> delete(Query query, Class<T> entityClass, String collectionName);
+    <T> List<T> delete(DocumentQuery query, Class<T> entityClass, String collectionName);
 
     MappingDocumentDbConverter getConverter();
 }

src/main/java/com/microsoft/azure/spring/data/cosmosdb/core/DocumentDbTemplate.java

@@ -10,6 +10,10 @@
 import com.microsoft.azure.documentdb.internal.HttpConstants;
 import com.microsoft.azure.spring.data.cosmosdb.DocumentDbFactory;
 import com.microsoft.azure.spring.data.cosmosdb.core.convert.MappingDocumentDbConverter;
+import com.microsoft.azure.spring.data.cosmosdb.core.criteria.Criteria;
+import com.microsoft.azure.spring.data.cosmosdb.core.generator.FindQuerySpecGenerator;
+import com.microsoft.azure.spring.data.cosmosdb.core.generator.QuerySpecGenerator;
+import com.microsoft.azure.spring.data.cosmosdb.core.query.DocumentQuery;
 import com.microsoft.azure.spring.data.cosmosdb.core.query.Query;
 import com.microsoft.azure.spring.data.cosmosdb.exception.DatabaseCreationException;
 import com.microsoft.azure.spring.data.cosmosdb.exception.DocumentDBAccessException;
@@ -29,6 +33,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import java.util.stream.Collectors;
 
 public class DocumentDbTemplate implements DocumentDbOperations, ApplicationContextAware {
     private static final Logger LOGGER = LoggerFactory.getLogger(DocumentDbTemplate.class);
@@ -281,7 +286,7 @@ private Database createDatabaseIfNotExists(String dbName) {
     }
 
     private DocumentCollection createCollection(@NonNull String dbName, String partitionKeyFieldName,
-                                               @NonNull DocumentDbEntityInformation information) {
+                                                @NonNull DocumentDbEntityInformation information) {
         DocumentCollection collection = new DocumentCollection();
         final String collectionName = information.getCollectionName();
         final IndexingPolicy policy = information.getIndexingPolicy();
@@ -395,6 +400,57 @@ private RequestOptions getRequestOptions(PartitionKey key, Integer requestUnit)
         return requestOptions;
     }
 
+    public <T> List<T> find(@NonNull DocumentQuery query, Class<T> domainClass, @NonNull String collectionName) {
+        Assert.hasText(collectionName, "collection should not be null, empty or only whitespaces");
+
+        final FeedOptions feedOptions = new FeedOptions();
+        final DocumentCollection collection = getDocCollection(collectionName);
+        final Optional<String> partitionKeyName = getPartitionKeyField(domainClass);
+        final QuerySpecGenerator generator = new FindQuerySpecGenerator(domainClass);
+        final Optional<Criteria> partitionCriteria = query.getSubjectCriteria(partitionKeyName.orElse(""));
+
+        feedOptions.setEnableCrossPartitionQuery(!partitionCriteria.isPresent());
+
+        final SqlQuerySpec sqlQuerySpec = generator.generate(query);
+        final List<Document> result = documentDbFactory.getDocumentClient()
+                .queryDocuments(collection.getSelfLink(), sqlQuerySpec, feedOptions).getQueryIterable().toList();
+
+        return result.stream().map(r -> getConverter().read(domainClass, r)).collect(Collectors.toList());
+    }
+
+    @Override
+    public <T> List<T> delete(@NonNull DocumentQuery query, Class<T> domainClass, @NonNull String collectionName) {
+        Assert.hasText(collectionName, "collection should not be null, empty or only whitespaces");
+
+        final FeedOptions feedOptions = new FeedOptions();
+        final DocumentCollection collection = getDocCollection(collectionName);
+        final Optional<String> partitionKeyName = getPartitionKeyField(domainClass);
+        final Optional<Criteria> partitionCriteria = query.getSubjectCriteria(partitionKeyName.orElse(""));
+        final QuerySpecGenerator generator = new FindQuerySpecGenerator(domainClass);
+
+        feedOptions.setEnableCrossPartitionQuery(!partitionCriteria.isPresent());
+
+        final SqlQuerySpec sqlQuerySpec = generator.generate(query);
+        final List<Document> results = documentDbFactory.getDocumentClient()
+                .queryDocuments(collection.getSelfLink(), sqlQuerySpec, feedOptions).getQueryIterable().toList();
+        final RequestOptions options = new RequestOptions();
+
+        partitionCriteria.ifPresent(c -> options.setPartitionKey(new PartitionKey(c.getSubValues().get(0))));
+
+        final List<T> deletedResult = new ArrayList<>();
+
+        for (final Document d : results) {
+            try {
+                documentDbFactory.getDocumentClient().deleteDocument(d.getSelfLink(), options);
+                deletedResult.add(getConverter().read(domainClass, d));
+            } catch (DocumentClientException e) {
+                throw new DocumentDBAccessException(String.format("Failed to delete document %s", d.getSelfLink()), e);
+            }
+        }
+
+        return deletedResult;
+    }
+
     public <T> List<T> find(Query query, Class<T> domainClass, String collectionName) {
         Assert.notNull(query, "query should not be null");
         Assert.notNull(domainClass, "domainClass should not be null");
@@ -482,7 +538,6 @@ public MappingDocumentDbConverter getConverter() {
         return this.mappingDocumentDbConverter;
     }
 
-    @Override
     public <T> List<T> delete(Query query, Class<T> entityClass, String collectionName) {
         Assert.notNull(query, "query should not be null");
         Assert.notNull(entityClass, "entityClass should not be null");

src/main/java/com/microsoft/azure/spring/data/cosmosdb/core/criteria/CriteriaType.java

@@ -5,6 +5,13 @@
  */
 package com.microsoft.azure.spring.data.cosmosdb.core.criteria;
 
+import lombok.Getter;
+import org.springframework.data.repository.query.parser.Part;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+
 import static com.microsoft.azure.spring.data.cosmosdb.Constants.SQL_KEYWORD_AND;
 import static com.microsoft.azure.spring.data.cosmosdb.Constants.SQL_KEYWORD_OR;
 
@@ -13,6 +20,17 @@ public enum CriteriaType {
     OR,
     AND;
 
+    @Getter
+    private static final Map<Part.Type, CriteriaType> criteriaMap;
+
+    static {
+        final Map<Part.Type, CriteriaType> map = new HashMap<>();
+
+        map.put(Part.Type.SIMPLE_PROPERTY, CriteriaType.IS_EQUAL);
+
+        criteriaMap = Collections.unmodifiableMap(map);
+    }
+
     public static String toSqlKeyword(CriteriaType type) {
         switch (type) {
             case AND:
@@ -23,4 +41,5 @@ public static String toSqlKeyword(CriteriaType type) {
                 throw new UnsupportedOperationException("Unsupported criteria type.");
         }
     }
+
 }

src/main/java/com/microsoft/azure/spring/data/cosmosdb/core/generator/AbstractQueryGenerator.java

@@ -0,0 +1,92 @@
+/**
+ * Copyright (c) Microsoft Corporation. All rights reserved.
+ * Licensed under the MIT License. See LICENSE in the project root for
+ * license information.
+ */
+package com.microsoft.azure.spring.data.cosmosdb.core.generator;
+
+import com.microsoft.azure.spring.data.cosmosdb.core.criteria.Criteria;
+import com.microsoft.azure.spring.data.cosmosdb.core.criteria.CriteriaType;
+import com.microsoft.azure.spring.data.cosmosdb.core.query.DocumentQuery;
+import com.microsoft.azure.spring.data.cosmosdb.repository.support.DocumentDbEntityInformation;
+import org.javatuples.Pair;
+import org.springframework.lang.NonNull;
+import org.springframework.util.Assert;
+
+import java.util.ArrayList;
+import java.util.List;
+
+import static com.microsoft.azure.spring.data.cosmosdb.Constants.ID_PROPERTY_NAME;
+
+public abstract class AbstractQueryGenerator {
+
+    protected final DocumentDbEntityInformation information;
+
+    @SuppressWarnings("unchecked")
+    protected <T> AbstractQueryGenerator(@NonNull Class<T> domainClass) {
+        this.information = new DocumentDbEntityInformation(domainClass);
+    }
+
+    private String getCriteriaSubject(@NonNull Criteria criteria) {
+        String subject = criteria.getSubject();
+
+        if (subject.equals(information.getIdField().getName())) {
+            subject = ID_PROPERTY_NAME;
+        }
+
+        return subject;
+    }
+
+    private String generateIsEqual(@NonNull Criteria criteria, @NonNull List<Pair<String, Object>> parameters) {
+        final String subject = this.getCriteriaSubject(criteria);
+
+        Assert.isTrue(criteria.getSubValues().size() == 1, "IS_EQUAL should have only one subject value");
+
+        parameters.add(Pair.with(subject, criteria.getSubValues().get(0)));
+
+        return String.format("r.%s=@%s", subject, subject);
+    }
+
+    private String generateBinaryQuery(@NonNull String left, @NonNull String right, CriteriaType type) {
+        Assert.isTrue(Criteria.isBinaryOperation(type), "Criteria type should be binary operation");
+
+        final String keyword = CriteriaType.toSqlKeyword(type);
+
+        return String.join(" ", left, keyword, right);
+    }
+
+    private String generateQueryBodyDfs(@NonNull Criteria criteria, @NonNull List<Pair<String, Object>> parameters) {
+        final CriteriaType type = criteria.getType();
+
+        switch (type) {
+            case IS_EQUAL:
+                return this.generateIsEqual(criteria, parameters);
+            case AND:
+            case OR:
+                Assert.isTrue(criteria.getSubCriteria().size() == 2, "criteria should have two SubCriteria");
+
+                final String left = generateQueryBodyDfs(criteria.getSubCriteria().get(0), parameters);
+                final String right = generateQueryBodyDfs(criteria.getSubCriteria().get(1), parameters);
+
+                return generateBinaryQuery(left, right, type);
+            default:
+                throw new UnsupportedOperationException("unsupported Criteria type" + type);
+        }
+    }
+
+    /**
+     * Generate a query body for interface QuerySpecGenerator.
+     * The query body compose of Sql query String and its' parameters.
+     * The parameters organized as a list of Pair, for each pair compose parameter name and value.
+     *
+     * @param query the representation for query method.
+     * @return A pair tuple compose of Sql query.
+     */
+    @NonNull
+    protected Pair<String, List<Pair<String, Object>>> generateQueryBody(@NonNull DocumentQuery query) {
+        final List<Pair<String, Object>> parameters = new ArrayList<>();
+        final String queryString = this.generateQueryBodyDfs(query.getCriteria(), parameters);
+
+        return Pair.with(queryString, parameters);
+    }
+}

src/main/java/com/microsoft/azure/spring/data/cosmosdb/core/generator/FindQuerySpecGenerator.java

@@ -0,0 +1,40 @@
+/**
+ * Copyright (c) Microsoft Corporation. All rights reserved.
+ * Licensed under the MIT License. See LICENSE in the project root for
+ * license information.
+ */
+package com.microsoft.azure.spring.data.cosmosdb.core.generator;
+
+import com.microsoft.azure.documentdb.SqlParameter;
+import com.microsoft.azure.documentdb.SqlParameterCollection;
+import com.microsoft.azure.documentdb.SqlQuerySpec;
+import com.microsoft.azure.spring.data.cosmosdb.core.query.DocumentQuery;
+import org.javatuples.Pair;
+import org.springframework.lang.NonNull;
+
+import java.util.List;
+import java.util.stream.Collectors;
+
+import static com.microsoft.azure.spring.data.cosmosdb.core.convert.MappingDocumentDbConverter.toDocumentDBValue;
+
+public class FindQuerySpecGenerator extends AbstractQueryGenerator implements QuerySpecGenerator {
+
+    public <T> FindQuerySpecGenerator(@NonNull Class<T> domainClass) {
+        super(domainClass);
+    }
+
+    @Override
+    public SqlQuerySpec generate(@NonNull DocumentQuery query) {
+        final Pair<String, List<Pair<String, Object>>> queryBody = super.generateQueryBody(query);
+        final String queryHeader = "SELECT * FROM ROOT r WHERE";
+        final String queryString = queryHeader + " " + queryBody.getValue0();
+        final List<Pair<String, Object>> parameters = queryBody.getValue1();
+        final SqlParameterCollection sqlParameters = new SqlParameterCollection();
+
+        sqlParameters.addAll(parameters.stream()
+                .map(p -> new SqlParameter("@" + p.getValue0(), toDocumentDBValue(p.getValue1())))
+                .collect(Collectors.toList()));
+
+        return new SqlQuerySpec(queryString, sqlParameters);
+    }
+}

src/main/java/com/microsoft/azure/spring/data/cosmosdb/repository/query/AbstractDocumentDbQuery.java

@@ -6,7 +6,7 @@
 package com.microsoft.azure.spring.data.cosmosdb.repository.query;
 
 import com.microsoft.azure.spring.data.cosmosdb.core.DocumentDbOperations;
-import com.microsoft.azure.spring.data.cosmosdb.core.query.Query;
+import com.microsoft.azure.spring.data.cosmosdb.core.query.DocumentQuery;
 import org.springframework.data.repository.query.RepositoryQuery;
 import org.springframework.data.repository.query.ResultProcessor;
 
@@ -22,17 +22,17 @@ public AbstractDocumentDbQuery(DocumentDbQueryMethod method, DocumentDbOperation
 
     public Object execute(Object[] parameters) {
         final DocumentDbParameterAccessor accessor = new DocumentDbParameterParameterAccessor(method, parameters);
-        final Query query = createQuery(accessor);
+        final DocumentQuery query = createQuery(accessor);
 
         final ResultProcessor processor = method.getResultProcessor().withDynamicProjection(accessor);
         final String collection = ((DocumentDbEntityMetadata) method.getEntityInformation()).getCollectionName();
 
-        final DocumentDbQueryExecution execution = getExecution(query, accessor);
+        final DocumentDbQueryExecution execution = getExecution();
         return execution.execute(query, processor.getReturnedType().getDomainType(), collection);
     }
 
 
-    private DocumentDbQueryExecution getExecution(Query query, DocumentDbParameterAccessor accessor) {
+    private DocumentDbQueryExecution getExecution() {
         if (isDeleteQuery()) {
             return new DocumentDbQueryExecution.DeleteExecution(operations);
         } else {
@@ -44,7 +44,7 @@ public DocumentDbQueryMethod getQueryMethod() {
         return method;
     }
 
-    protected abstract Query createQuery(DocumentDbParameterAccessor accessor);
+    protected abstract DocumentQuery createQuery(DocumentDbParameterAccessor accessor);
 
     protected abstract boolean isDeleteQuery();