Allow user custom RequestOptions when configuration, #322. (#323)

* The config about PartitionKey and Throughput may be override. Signed-off-by: Pan Li <panli@microsoft.com>
2019-03-06 01:51:09 -06:00 · 2019-03-06 01:51:09 -06:00 · 0b264933f0
--- a/src/main/java/com/microsoft/azure/spring/data/cosmosdb/DocumentDbFactory.java
+++ b/src/main/java/com/microsoft/azure/spring/data/cosmosdb/DocumentDbFactory.java
@ -12,6 +12,7 @@ import com.microsoft.azure.spring.data.cosmosdb.common.MacAddress;
 import com.microsoft.azure.spring.data.cosmosdb.common.PropertyLoader;
 import com.microsoft.azure.spring.data.cosmosdb.common.TelemetrySender;
 import com.microsoft.azure.spring.data.cosmosdb.config.DocumentDBConfig;
+import lombok.Getter;
 import lombok.NonNull;
 import org.springframework.util.Assert;

@ -19,6 +20,7 @@ import javax.annotation.PostConstruct;

 public class DocumentDbFactory {

+    @Getter
    private final DocumentDBConfig config;

    private static final boolean IS_TELEMETRY_ALLOWED = PropertyLoader.isApplicationTelemetryAllowed();
--- a/src/main/java/com/microsoft/azure/spring/data/cosmosdb/common/CosmosdbUtils.java
+++ b/src/main/java/com/microsoft/azure/spring/data/cosmosdb/common/CosmosdbUtils.java
@ -0,0 +1,28 @@
+/**
+ * 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.common;
+
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.microsoft.azure.spring.data.cosmosdb.core.convert.ObjectMapperFactory;
+import com.microsoft.azure.spring.data.cosmosdb.exception.ConfigurationException;
+import lombok.NonNull;
+
+import java.io.IOException;
+
+public class CosmosdbUtils {
+
+    @SuppressWarnings("unchecked")
+    public static <T> T getCopyFrom(@NonNull T instance) {
+        final ObjectMapper mapper = ObjectMapperFactory.getObjectMapper();
+
+        try {
+            final String s = mapper.writeValueAsString(instance);
+            return (T) mapper.readValue(s, instance.getClass());
+        } catch (IOException e) {
+            throw new ConfigurationException("failed to get copy from " + instance.getClass().getName(), e);
+        }
+    }
+}
--- a/src/main/java/com/microsoft/azure/spring/data/cosmosdb/config/DocumentDBConfig.java
+++ b/src/main/java/com/microsoft/azure/spring/data/cosmosdb/config/DocumentDBConfig.java
@ -7,6 +7,7 @@ package com.microsoft.azure.spring.data.cosmosdb.config;

 import com.microsoft.azure.documentdb.ConnectionPolicy;
 import com.microsoft.azure.documentdb.ConsistencyLevel;
+import com.microsoft.azure.documentdb.RequestOptions;
 import com.microsoft.azure.spring.data.cosmosdb.exception.DocumentDBAccessException;
 import lombok.Builder;
 import lombok.Getter;
@ -27,13 +28,16 @@ public class DocumentDBConfig {

    private boolean allowTelemetry;

+    private RequestOptions requestOptions;
+
    public static DocumentDBConfigBuilder builder(String uri, String key, String database) {
        return defaultBuilder()
                .uri(uri)
                .key(key)
                .database(database)
                .connectionPolicy(ConnectionPolicy.GetDefault())
-                .consistencyLevel(ConsistencyLevel.Session);
+                .consistencyLevel(ConsistencyLevel.Session)
+                .requestOptions(new RequestOptions());
    }

    public static DocumentDBConfigBuilder builder(String connectionString, String database) {
--- a/src/main/java/com/microsoft/azure/spring/data/cosmosdb/core/DocumentDbTemplate.java
+++ b/src/main/java/com/microsoft/azure/spring/data/cosmosdb/core/DocumentDbTemplate.java
@ -9,6 +9,8 @@ package com.microsoft.azure.spring.data.cosmosdb.core;
 import com.microsoft.azure.documentdb.*;
 import com.microsoft.azure.documentdb.internal.HttpConstants;
 import com.microsoft.azure.spring.data.cosmosdb.DocumentDbFactory;
+import com.microsoft.azure.spring.data.cosmosdb.common.CosmosdbUtils;
+import com.microsoft.azure.spring.data.cosmosdb.config.DocumentDBConfig;
 import com.microsoft.azure.spring.data.cosmosdb.core.convert.MappingDocumentDbConverter;
 import com.microsoft.azure.spring.data.cosmosdb.core.generator.CountQueryGenerator;
 import com.microsoft.azure.spring.data.cosmosdb.core.generator.FindQuerySpecGenerator;
@ -33,7 +35,10 @@ import org.springframework.util.Assert;
 import org.springframework.util.StringUtils;

 import java.text.MessageFormat;
-import java.util.*;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.Iterator;
+import java.util.List;
 import java.util.stream.Collectors;

@Slf4j
@ -118,11 +123,8 @@ public class DocumentDbTemplate implements DocumentDbOperations, ApplicationCont
        assertValidId(id);

        try {
-            final RequestOptions options = new RequestOptions();
-
-            if (isIdFieldAsPartitionKey(domainClass)) {
-                options.setPartitionKey(new PartitionKey(id));
-            }
+            final PartitionKey partitionKey = isIdFieldAsPartitionKey(domainClass) ? new PartitionKey(id) : null;
+            final RequestOptions options = getRequestOptions(partitionKey, null);

            final String documentLink = getDocumentLink(this.databaseName, collectionName, id);
            final Resource document = getDocumentClient().readDocument(documentLink, options).getResource();
@ -314,8 +316,8 @@ public class DocumentDbTemplate implements DocumentDbOperations, ApplicationCont
        log.debug("execute deleteById in database {} collection {}", this.databaseName, collectionName);

        try {
-            getDocumentClient().deleteDocument(getDocumentLink(this.databaseName, collectionName, id.toString()),
-                    getRequestOptions(partitionKey, null));
+            final RequestOptions options = getRequestOptions(partitionKey, null);
+            getDocumentClient().deleteDocument(getDocumentLink(databaseName, collectionName, id.toString()), options);
        } catch (DocumentClientException ex) {
            throw new DocumentDBAccessException("deleteById exception", ex);
        }
@ -337,20 +339,23 @@ public class DocumentDbTemplate implements DocumentDbOperations, ApplicationCont
        return "/" + partitionKey;
    }

+    @NonNull
+    private DocumentDBConfig getDocumentDbConfig() {
+        return documentDbFactory.getConfig();
+    }
+
    private RequestOptions getRequestOptions(PartitionKey key, Integer requestUnit) {
-        if (key == null && requestUnit == null) {
-            return null;
-        }
+        final RequestOptions options = CosmosdbUtils.getCopyFrom(getDocumentDbConfig().getRequestOptions());

-        final RequestOptions requestOptions = new RequestOptions();
        if (key != null) {
-            requestOptions.setPartitionKey(key);
-        }
-        if (requestUnit != null) {
-            requestOptions.setOfferThroughput(requestUnit);
+            options.setPartitionKey(key);
        }

-        return requestOptions;
+        if (requestUnit != null) {
+            options.setOfferThroughput(requestUnit);
+        }
+
+        return options;
    }

    private <T> List<T> executeQuery(@NonNull SqlQuerySpec sqlQuerySpec, boolean isCrossPartition,
@ -419,14 +424,16 @@ public class DocumentDbTemplate implements DocumentDbOperations, ApplicationCont

    private void deleteDocument(@NonNull Document document, @NonNull List<String> partitionKeyNames) {
        try {
-            final RequestOptions options = new RequestOptions();
-
            Assert.isTrue(partitionKeyNames.size() <= 1, "Only one Partition is supported.");

+            PartitionKey partitionKey = null;
+
            if (!partitionKeyNames.isEmpty() && StringUtils.hasText(partitionKeyNames.get(0))) {
-                options.setPartitionKey(new PartitionKey(document.get(partitionKeyNames.get(0))));
+                partitionKey = new PartitionKey(document.get(partitionKeyNames.get(0)));
            }

+            final RequestOptions options = getRequestOptions(partitionKey, null);
+
            getDocumentClient().deleteDocument(document.getSelfLink(), options);
        } catch (DocumentClientException e) {
            throw new DocumentDBAccessException("Failed to delete document: " + document.getSelfLink(), e);
--- a/src/main/java/com/microsoft/azure/spring/data/cosmosdb/exception/ConfigurationException.java
+++ b/src/main/java/com/microsoft/azure/spring/data/cosmosdb/exception/ConfigurationException.java
@ -0,0 +1,19 @@
+/**
+ * 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.exception;
+
+import org.springframework.dao.DataAccessException;
+
+public class ConfigurationException extends DataAccessException {
+
+    public ConfigurationException(String msg) {
+        super(msg);
+    }
+
+    public ConfigurationException(String msg, Throwable cause) {
+        super(msg, cause);
+    }
+}
--- a/src/test/java/com/microsoft/azure/spring/data/cosmosdb/config/AbstractDocumentDbConfigurationUnitTest.java
+++ b/src/test/java/com/microsoft/azure/spring/data/cosmosdb/config/AbstractDocumentDbConfigurationUnitTest.java
@ -7,7 +7,9 @@
 package com.microsoft.azure.spring.data.cosmosdb.config;

 import com.fasterxml.jackson.databind.ObjectMapper;
+import com.microsoft.azure.documentdb.ConsistencyLevel;
 import com.microsoft.azure.documentdb.DocumentClient;
+import com.microsoft.azure.documentdb.RequestOptions;
 import com.microsoft.azure.spring.data.cosmosdb.Constants;
 import com.microsoft.azure.spring.data.cosmosdb.DocumentDbFactory;
 import com.microsoft.azure.spring.data.cosmosdb.common.TestConstants;
@ -53,6 +55,21 @@ public class AbstractDocumentDbConfigurationUnitTest {
        Assertions.assertThat(context.getBean(OBJECTMAPPER_BEAN_NAME)).isNotNull();
    }

+    @Test
+    public void testRequestOptionsConfigurable() {
+        final AbstractApplicationContext context = new AnnotationConfigApplicationContext(
+                RequestOptionsConfiguration.class);
+        final DocumentDbFactory factory = context.getBean(DocumentDbFactory.class);
+
+        Assertions.assertThat(factory).isNotNull();
+
+        final RequestOptions options = factory.getConfig().getRequestOptions();
+
+        Assertions.assertThat(options).isNotNull();
+        Assertions.assertThat(options.getConsistencyLevel()).isEqualTo(ConsistencyLevel.ConsistentPrefix);
+        Assertions.assertThat(options.getDisableRUPerMinuteUsage()).isTrue();
+        Assertions.assertThat(options.isScriptLoggingEnabled()).isTrue();
+    }

    @Configuration
    static class TestDocumentDbConfiguration extends AbstractDocumentDbConfiguration {
@ -77,4 +94,27 @@ public class AbstractDocumentDbConfigurationUnitTest {
            return new ObjectMapper();
        }
    }
+
+    @Configuration
+    static class RequestOptionsConfiguration extends AbstractDocumentDbConfiguration {
+
+        private RequestOptions getRequestOptions() {
+            final RequestOptions options = new RequestOptions();
+
+            options.setConsistencyLevel(ConsistencyLevel.ConsistentPrefix);
+            options.setDisableRUPerMinuteUsage(true);
+            options.setScriptLoggingEnabled(true);
+
+            return options;
+        }
+
+        @Override
+        public DocumentDBConfig getConfig() {
+            final RequestOptions options = getRequestOptions();
+            return DocumentDBConfig.builder("http://fake-uri", "fake-key", TestConstants.DB_NAME)
+                    .requestOptions(options)
+                    .build();
+        }
+
+    }
 }
--- a/src/test/java/com/microsoft/azure/spring/data/cosmosdb/repository/TestRepositoryConfig.java
+++ b/src/test/java/com/microsoft/azure/spring/data/cosmosdb/repository/TestRepositoryConfig.java
@ -5,6 +5,8 @@
 */
 package com.microsoft.azure.spring.data.cosmosdb.repository;

+import com.microsoft.azure.documentdb.ConsistencyLevel;
+import com.microsoft.azure.documentdb.RequestOptions;
 import com.microsoft.azure.spring.data.cosmosdb.common.TestConstants;
 import com.microsoft.azure.spring.data.cosmosdb.config.AbstractDocumentDbConfiguration;
 import com.microsoft.azure.spring.data.cosmosdb.config.DocumentDBConfig;
@ -30,13 +32,25 @@ public class TestRepositoryConfig extends AbstractDocumentDbConfiguration {
    @Value("${cosmosdb.database:}")
    private String database;

+    private RequestOptions getRequestOptions() {
+        final RequestOptions options = new RequestOptions();
+
+        options.setConsistencyLevel(ConsistencyLevel.ConsistentPrefix);
+        options.setDisableRUPerMinuteUsage(true);
+        options.setScriptLoggingEnabled(true);
+
+        return options;
+    }
+
    @Override
    public DocumentDBConfig getConfig() {
        final String dbName = StringUtils.hasText(this.database) ? this.database : TestConstants.DB_NAME;
+        final RequestOptions options = getRequestOptions();

        if (StringUtils.hasText(this.documentDbUri) && StringUtils.hasText(this.documentDbKey)) {
-            return DocumentDBConfig.builder(documentDbUri, documentDbKey, dbName).build();
+            return DocumentDBConfig.builder(documentDbUri, documentDbKey, dbName).requestOptions(options).build();
        }
-        return DocumentDBConfig.builder(connectionString, dbName).build();
+
+        return DocumentDBConfig.builder(connectionString, dbName).requestOptions(options).build();
    }
 }