Cosmos Key Credential

sdk/cosmos/microsoft-azure-cosmos/src/main/java/com/azure/data/cosmos/CosmosClient.java

@@ -29,24 +29,27 @@ public class CosmosClient implements AutoCloseable {
     private final ConsistencyLevel desiredConsistencyLevel;
     private final List<Permission> permissions;
     private final TokenResolver tokenResolver;
+    private final CosmosKeyCredential cosmosKeyCredential;
 
 
      CosmosClient(CosmosClientBuilder builder) {
          this.configs = builder.configs();
          this.serviceEndpoint = builder.endpoint();
-        this.keyOrResourceToken = builder.key();
-        this.connectionPolicy = builder.connectionPolicy();
-        this.desiredConsistencyLevel = builder.consistencyLevel();
-        this.permissions = builder.permissions();
-        this.tokenResolver = builder.tokenResolver();
-        this.asyncDocumentClient = new AsyncDocumentClient.Builder()
-                .withServiceEndpoint(this.serviceEndpoint)
-                .withMasterKeyOrResourceToken(this.keyOrResourceToken)
-                .withConnectionPolicy(this.connectionPolicy)
-                .withConsistencyLevel(this.desiredConsistencyLevel)
-                .withConfigs(this.configs)
-                .withTokenResolver(this.tokenResolver)
-                .build();
+         this.keyOrResourceToken = builder.key();
+         this.connectionPolicy = builder.connectionPolicy();
+         this.desiredConsistencyLevel = builder.consistencyLevel();
+         this.permissions = builder.permissions();
+         this.tokenResolver = builder.tokenResolver();
+         this.cosmosKeyCredential = builder.cosmosKeyCredential();
+         this.asyncDocumentClient = new AsyncDocumentClient.Builder()
+             .withServiceEndpoint(this.serviceEndpoint)
+             .withMasterKeyOrResourceToken(this.keyOrResourceToken)
+             .withConnectionPolicy(this.connectionPolicy)
+             .withConsistencyLevel(this.desiredConsistencyLevel)
+             .withConfigs(this.configs)
+             .withTokenResolver(this.tokenResolver)
+             .withCosmosKeyCredential(this.cosmosKeyCredential)
+             .build();
     }
 
     AsyncDocumentClient getContextClient() {
@@ -121,10 +124,18 @@ TokenResolver getTokenResolver() {
         return tokenResolver;
     }
 
+    /**
+     * Gets the cosmos key credential
+     * @return cosmos key credential
+     */
+    CosmosKeyCredential getCosmosKeyCredential() {
+        return cosmosKeyCredential;
+    }
+
     /**
      * CREATE a Database if it does not already exist on the service
-     * 
-     * The {@link Mono} upon successful completion will contain a single cosmos database response with the 
+     *
+     * The {@link Mono} upon successful completion will contain a single cosmos database response with the
      * created or existing database.
      * @param databaseSettings CosmosDatabaseProperties
      * @return a {@link Mono} containing the cosmos database response with the created or existing database or
@@ -136,7 +147,7 @@ public Mono<CosmosDatabaseResponse> createDatabaseIfNotExists(CosmosDatabaseProp
 
     /**
      * CREATE a Database if it does not already exist on the service
-     * The {@link Mono} upon successful completion will contain a single cosmos database response with the 
+     * The {@link Mono} upon successful completion will contain a single cosmos database response with the
      * created or existing database.
      * @param id the id of the database
      * @return a {@link Mono} containing the cosmos database response with the created or existing database or
@@ -145,7 +156,7 @@ public Mono<CosmosDatabaseResponse> createDatabaseIfNotExists(CosmosDatabaseProp
     public Mono<CosmosDatabaseResponse> createDatabaseIfNotExists(String id) {
         return createDatabaseIfNotExistsInternal(getDatabase(id));
     }
-    
+
     private Mono<CosmosDatabaseResponse> createDatabaseIfNotExistsInternal(CosmosDatabase database){
         return database.read().onErrorResume(exception -> {
             if (exception instanceof CosmosClientException) {
@@ -160,12 +171,12 @@ private Mono<CosmosDatabaseResponse> createDatabaseIfNotExistsInternal(CosmosDat
 
     /**
      * Creates a database.
-     * 
-     * After subscription the operation will be performed. 
-     * The {@link Mono} upon successful completion will contain a single resource response with the 
+     *
+     * After subscription the operation will be performed.
+     * The {@link Mono} upon successful completion will contain a single resource response with the
      *      created database.
      * In case of failure the {@link Mono} will error.
-     * 
+     *
      * @param databaseSettings {@link CosmosDatabaseProperties}
      * @param options {@link CosmosDatabaseRequestOptions}
      * @return an {@link Mono} containing the single cosmos database response with the created database or an error.
@@ -183,12 +194,12 @@ public Mono<CosmosDatabaseResponse> createDatabase(CosmosDatabaseProperties data
 
     /**
      * Creates a database.
-     * 
-     * After subscription the operation will be performed. 
-     * The {@link Mono} upon successful completion will contain a single resource response with the 
+     *
+     * After subscription the operation will be performed.
+     * The {@link Mono} upon successful completion will contain a single resource response with the
      *      created database.
      * In case of failure the {@link Mono} will error.
-     * 
+     *
      * @param databaseSettings {@link CosmosDatabaseProperties}
      * @return an {@link Mono} containing the single cosmos database response with the created database or an error.
      */
@@ -198,12 +209,12 @@ public Mono<CosmosDatabaseResponse> createDatabase(CosmosDatabaseProperties data
 
     /**
      * Creates a database.
-     * 
-     * After subscription the operation will be performed. 
-     * The {@link Mono} upon successful completion will contain a single resource response with the 
+     *
+     * After subscription the operation will be performed.
+     * The {@link Mono} upon successful completion will contain a single resource response with the
      *      created database.
      * In case of failure the {@link Mono} will error.
-     * 
+     *
      * @param id id of the database
      * @return a {@link Mono} containing the single cosmos database response with the created database or an error.
      */
@@ -275,11 +286,11 @@ public Mono<CosmosDatabaseResponse> createDatabase(String id, int throughput) {
 
     /**
      * Reads all databases.
-     * 
-     * After subscription the operation will be performed. 
+     *
+     * After subscription the operation will be performed.
      * The {@link Flux} will contain one or several feed response of the read databases.
      * In case of failure the {@link Flux} will error.
-     * 
+     *
      * @param options {@link FeedOptions}
      * @return a {@link Flux} containing one or several feed response pages of read databases or an error.
      */
@@ -291,11 +302,11 @@ public Flux<FeedResponse<CosmosDatabaseProperties>> readAllDatabases(FeedOptions
 
     /**
      * Reads all databases.
-     * 
-     * After subscription the operation will be performed. 
+     *
+     * After subscription the operation will be performed.
      * The {@link Flux} will contain one or several feed response of the read databases.
      * In case of failure the {@link Flux} will error.
-     * 
+     *
      * @return a {@link Flux} containing one or several feed response pages of read databases or an error.
      */
     public Flux<FeedResponse<CosmosDatabaseProperties>> readAllDatabases() {

sdk/cosmos/microsoft-azure-cosmos/src/main/java/com/azure/data/cosmos/CosmosClientBuilder.java

@@ -4,6 +4,7 @@
 
 import com.azure.data.cosmos.internal.Configs;
 import com.azure.data.cosmos.internal.Permission;
+import org.apache.commons.lang3.StringUtils;
 
 import java.util.List;
 
@@ -33,6 +34,7 @@ public class CosmosClientBuilder {
     private ConsistencyLevel desiredConsistencyLevel;
     private List<Permission> permissions;
     private TokenResolver tokenResolver;
+    private CosmosKeyCredential cosmosKeyCredential;
 
     CosmosClientBuilder() {
     }
@@ -137,15 +139,15 @@ public CosmosClientBuilder permissions(List<Permission> permissions) {
     }
 
     /**
-     * Gets the (@link ConsistencyLevel) to be used
+     * Gets the {@link ConsistencyLevel} to be used
      * @return the consistency level
      */
     public ConsistencyLevel consistencyLevel() {
         return this.desiredConsistencyLevel;
     }
 
     /**
-     * Sets the (@link ConsistencyLevel) to be used
+     * Sets the {@link ConsistencyLevel} to be used
      * @param desiredConsistencyLevel {@link ConsistencyLevel}
      * @return current Builder
      */
@@ -163,7 +165,7 @@ public ConnectionPolicy connectionPolicy() {
     }
 
     /**
-     * Sets the (@link ConnectionPolicy) to be used
+     * Sets the {@link ConnectionPolicy} to be used
      * @param connectionPolicy {@link ConnectionPolicy}
      * @return current Builder
      */
@@ -172,6 +174,24 @@ public CosmosClientBuilder connectionPolicy(ConnectionPolicy connectionPolicy) {
         return this;
     }
 
+    /**
+     * Gets the {@link CosmosKeyCredential} to be used
+     * @return cosmosKeyCredential
+     */
+    public CosmosKeyCredential cosmosKeyCredential() {
+        return cosmosKeyCredential;
+    }
+
+    /**
+     * Sets the {@link CosmosKeyCredential} to be used
+     * @param cosmosKeyCredential {@link CosmosKeyCredential}
+     * @return current builder
+     */
+    public CosmosClientBuilder cosmosKeyCredential(CosmosKeyCredential cosmosKeyCredential) {
+        this.cosmosKeyCredential = cosmosKeyCredential;
+        return this;
+    }
+
     /**
      * Builds a cosmos configuration object with the provided properties
      * @return CosmosClient
@@ -180,8 +200,11 @@ public CosmosClient build() {
 
         ifThrowIllegalArgException(this.serviceEndpoint == null, "cannot build client without service endpoint");
         ifThrowIllegalArgException(
-            this.keyOrResourceToken == null && (permissions == null || permissions.isEmpty()) && this.tokenResolver == null,
-            "cannot build client without any one of key, resource token, permissions, and token resolver");
+            this.keyOrResourceToken == null && (permissions == null || permissions.isEmpty())
+                && this.tokenResolver == null && this.cosmosKeyCredential == null,
+            "cannot build client without any one of key, resource token, permissions, token resolver, and cosmos key credential");
+        ifThrowIllegalArgException(cosmosKeyCredential != null && StringUtils.isEmpty(cosmosKeyCredential.getMasterKey()),
+            "cannot build client without key credential");
 
         return new CosmosClient(this);
     }

sdk/cosmos/microsoft-azure-cosmos/src/main/java/com/azure/data/cosmos/CosmosKeyCredential.java

@@ -0,0 +1,25 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+package com.azure.data.cosmos;
+
+/**
+ * Cosmos Key Credential is used to store key credentials, in order to support dynamic key rotation.
+ * Singleton instance should be used to support multiple keys.
+ * SDK insures to use the updated key provided in the same singleton instance which was used when building {@link CosmosClient}
+ */
+public class CosmosKeyCredential {
+
+    private String masterKey;
+
+    public CosmosKeyCredential(String masterKey) {
+        this.masterKey = masterKey;
+    }
+
+    public String getMasterKey() {
+        return masterKey;
+    }
+
+    public void setMasterKey(String masterKey) {
+        this.masterKey = masterKey;
+    }
+}

sdk/cosmos/microsoft-azure-cosmos/src/main/java/com/azure/data/cosmos/internal/AsyncDocumentClient.java

@@ -5,10 +5,12 @@
 import com.azure.data.cosmos.ChangeFeedOptions;
 import com.azure.data.cosmos.ConnectionPolicy;
 import com.azure.data.cosmos.ConsistencyLevel;
+import com.azure.data.cosmos.CosmosKeyCredential;
 import com.azure.data.cosmos.FeedOptions;
 import com.azure.data.cosmos.FeedResponse;
 import com.azure.data.cosmos.SqlQuerySpec;
 import com.azure.data.cosmos.TokenResolver;
+import org.apache.commons.lang3.StringUtils;
 import reactor.core.publisher.Flux;
 
 import java.net.URI;
@@ -71,6 +73,7 @@ class Builder {
         String masterKeyOrResourceToken;
         URI serviceEndpoint;
         TokenResolver tokenResolver;
+        CosmosKeyCredential cosmosKeyCredential;
 
         public Builder withServiceEndpoint(String serviceEndpoint) {
             try {
@@ -134,23 +137,13 @@ public Builder withConnectionPolicy(ConnectionPolicy connectionPolicy) {
             return this;
         }
 
-        /**
-         * This method will accept tokenResolver which is rx function, it takes arguments<br>
-         * T1 requestVerb(STRING),<br>
-         * T2 resourceIdOrFullName(STRING),<br>
-         * T3 resourceType(com.azure.data.cosmos.internal.ResourceType),<br>
-         * T4 request headers(Map<STRING, STRING>)<br>
-         *<br>
-         * and return<br>
-         * R authenticationToken(STRING)<br>
-         *
-         * @param tokenResolver tokenResolver function for authentication.
-         * @return current Builder.
-         */
-        /*public Builder withTokenResolver(Func4<STRING, STRING, ResourceType, Map<STRING, STRING>, STRING> tokenResolver) {
-            this.tokenResolver = tokenResolver;
+        public Builder withCosmosKeyCredential(CosmosKeyCredential cosmosKeyCredential) {
+            if (cosmosKeyCredential != null && StringUtils.isEmpty(cosmosKeyCredential.getMasterKey())) {
+                throw new IllegalArgumentException("Cannot build client with empty key credential");
+            }
+            this.cosmosKeyCredential = cosmosKeyCredential;
             return this;
-        }*/
+        }
 
         /**
          * This method will accept functional interface TokenResolver which helps in generation authorization
@@ -173,16 +166,21 @@ public AsyncDocumentClient build() {
 
             ifThrowIllegalArgException(this.serviceEndpoint == null, "cannot build client without service endpoint");
             ifThrowIllegalArgException(
-                    this.masterKeyOrResourceToken == null && (permissionFeed == null || permissionFeed.isEmpty()) && this.tokenResolver == null,
-                    "cannot build client without any one of masterKey, resource token, permissionFeed and tokenResolver");
+                    this.masterKeyOrResourceToken == null && (permissionFeed == null || permissionFeed.isEmpty())
+                        && this.tokenResolver == null && this.cosmosKeyCredential == null,
+                    "cannot build client without any one of masterKey, " +
+                        "resource token, permissionFeed, tokenResolver and cosmos key credential");
+            ifThrowIllegalArgException(cosmosKeyCredential != null && StringUtils.isEmpty(cosmosKeyCredential.getMasterKey()),
+                "cannot build client without key credential");
 
             RxDocumentClientImpl client = new RxDocumentClientImpl(serviceEndpoint,
                                                                    masterKeyOrResourceToken,
                                                                    permissionFeed,
                                                                    connectionPolicy,
                                                                    desiredConsistencyLevel,
                                                                    configs,
-                                                                   tokenResolver);
+                                                                   tokenResolver,
+                                                                    cosmosKeyCredential);
             client.init();
             return client;
         }
@@ -242,6 +240,14 @@ public TokenResolver getTokenResolver() {
         public void setTokenResolver(TokenResolver tokenResolver) {
             this.tokenResolver = tokenResolver;
         }
+
+        public CosmosKeyCredential getCosmosKeyCredential() {
+            return cosmosKeyCredential;
+        }
+
+        public void setCosmosKeyCredential(CosmosKeyCredential cosmosKeyCredential) {
+            this.cosmosKeyCredential = cosmosKeyCredential;
+        }
     }
 
     /**