Read database account properties using CosmosClient and CosmosAsyncClient

sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/models/CosmosDatabaseAccountResponseTest.java

@@ -0,0 +1,91 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package com.azure.cosmos.models;
+
+import com.azure.cosmos.CosmosAsyncClient;
+import com.azure.cosmos.CosmosClient;
+import com.azure.cosmos.CosmosClientBuilder;
+import com.azure.cosmos.rx.TestSuiteBase;
+import org.testng.annotations.AfterClass;
+import org.testng.annotations.BeforeClass;
+import org.testng.annotations.Factory;
+import org.testng.annotations.Test;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+public class CosmosDatabaseAccountResponseTest extends TestSuiteBase {
+    private CosmosClient client;
+    private CosmosAsyncClient asyncClient;
+
+    @Factory(dataProvider = "clientBuilders")
+    public CosmosDatabaseAccountResponseTest(CosmosClientBuilder clientBuilder) {
+        super(clientBuilder);
+    }
+
+    @BeforeClass(groups = {"fast"}, timeOut = SETUP_TIMEOUT)
+    public void before_CosmosDatabaseAccountTest() {
+        assertThat(this.client).isNull();
+        assertThat(this.asyncClient).isNull();
+
+        this.client = getClientBuilder().buildClient();
+        this.asyncClient = getClientBuilder().buildAsyncClient();
+    }
+
+    @AfterClass(groups = {"fast"}, timeOut = SHUTDOWN_TIMEOUT, alwaysRun = true)
+    public void afterClass() {
+        assertThat(this.client).isNotNull();
+        assertThat(this.asyncClient).isNotNull();
+
+        this.client.close();
+        this.asyncClient.close();
+    }
+
+    @Test(groups = {"fast"}, timeOut = TIMEOUT)
+    public void readCosmosDatabaseAccountWithClient() {
+        CosmosDatabaseAccountResponse latestDatabaseAccountResponse = client.readDatabaseAccount(false);
+
+        assertThat(latestDatabaseAccountResponse).isNotNull();
+        validateDatabaseAccountResponse(latestDatabaseAccountResponse);
+
+        CosmosDatabaseAccountResponse cachedResponse = client.readDatabaseAccount(true);
+
+        assertThat(cachedResponse).isNotNull();
+        validateDatabaseAccountResponse(cachedResponse);
+
+        validateEquality(latestDatabaseAccountResponse, cachedResponse);
+    }
+
+    @Test(groups = {"fast"}, timeOut = TIMEOUT)
+    public void readCosmosDatabaseAccountWithAsyncClient() {
+        CosmosDatabaseAccountResponse latestDatabaseAccountResponse = asyncClient.readDatabaseAccount(false).block();
+
+        assertThat(latestDatabaseAccountResponse).isNotNull();
+        validateDatabaseAccountResponse(latestDatabaseAccountResponse);
+
+        CosmosDatabaseAccountResponse cachedResponse = asyncClient.readDatabaseAccount(true).block();
+
+        assertThat(cachedResponse).isNotNull();
+        validateDatabaseAccountResponse(cachedResponse);
+
+        validateEquality(latestDatabaseAccountResponse, cachedResponse);
+    }
+
+    private void validateDatabaseAccountResponse(CosmosDatabaseAccountResponse response) {
+        assertThat(response).isNotNull();
+        assertThat(response.getId()).isNotNull();
+        assertThat(response.getId()).isNotEmpty();
+        assertThat(response.getReadRegions()).isNotEmpty();
+        assertThat(response.getWriteRegions()).isNotEmpty();
+        assertThat(response.isMultiWriteAccount()).isNotNull();
+        assertThat(response.getAccountLevelConsistency()).isNotNull();
+    }
+
+    private void validateEquality(CosmosDatabaseAccountResponse response1, CosmosDatabaseAccountResponse response2) {
+        assertThat(response1.getId()).isEqualTo(response2.getId());
+        assertThat(response1.getReadRegions()).isEqualTo(response2.getReadRegions());
+        assertThat(response1.getWriteRegions()).isEqualTo(response2.getWriteRegions());
+        assertThat(response1.isMultiWriteAccount()).isEqualTo(response2.isMultiWriteAccount());
+        assertThat(response1.getAccountLevelConsistency()).isEqualTo(response2.getAccountLevelConsistency());
+    }
+}

sdk/cosmos/azure-cosmos/CHANGELOG.md

@@ -3,6 +3,7 @@
 ### 4.73.0-beta.1 (Unreleased)
 
 #### Features Added
+* Added an API `CosmosClient#readDatabaseAccount` to be able to fetch Cosmos DB account properties. - See [PR 45789](https://github.com/Azure/azure-sdk-for-java/pull/45789)
 
 #### Breaking Changes
 

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java

@@ -33,6 +33,7 @@
 import com.azure.cosmos.models.CosmosAuthorizationTokenResolver;
 import com.azure.cosmos.models.CosmosClientTelemetryConfig;
 import com.azure.cosmos.models.CosmosContainerIdentity;
+import com.azure.cosmos.models.CosmosDatabaseAccountResponse;
 import com.azure.cosmos.models.CosmosDatabaseProperties;
 import com.azure.cosmos.models.CosmosDatabaseRequestOptions;
 import com.azure.cosmos.models.CosmosDatabaseResponse;
@@ -729,6 +730,16 @@ private Mono<CosmosDatabaseResponse> createDatabaseInternal(Database database, C
                 requestOptions);
     }
 
+    /**
+     * Reads the Cosmos database account.
+     *
+     * @param shouldUseCache a boolean flag to determine whether to use the CosmosAsyncClient's-internal cache for reading the database account (setting shouldUseFlag to true can return stale data).
+     * @return the {@link CosmosDatabaseAccountResponse} with the read database account.
+     */
+    public Mono<CosmosDatabaseAccountResponse> readDatabaseAccount(boolean shouldUseCache) {
+        return this.asyncDocumentClient.readDatabaseAccount(shouldUseCache);
+    }
+
     ConsistencyLevel getEffectiveConsistencyLevel(
         OperationType operationType,
         ConsistencyLevel desiredConsistencyLevelOfOperation) {

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

@@ -6,6 +6,7 @@
 import com.azure.core.annotation.ServiceClient;
 import com.azure.cosmos.implementation.ImplementationBridgeHelpers;
 import com.azure.cosmos.models.CosmosContainerIdentity;
+import com.azure.cosmos.models.CosmosDatabaseAccountResponse;
 import com.azure.cosmos.models.CosmosDatabaseProperties;
 import com.azure.cosmos.models.CosmosDatabaseRequestOptions;
 import com.azure.cosmos.models.CosmosDatabaseResponse;
@@ -225,6 +226,19 @@ CosmosDatabaseResponse blockDatabaseResponse(Mono<CosmosDatabaseResponse> databa
         }
     }
 
+    CosmosDatabaseAccountResponse blockReadDatabaseAccount(Mono<CosmosDatabaseAccountResponse> databaseAccountMono) {
+        try {
+            return databaseAccountMono.block();
+        } catch (Exception ex) {
+            final Throwable throwable = Exceptions.unwrap(ex);
+            if (throwable instanceof CosmosException) {
+                throw (CosmosException) throwable;
+            } else {
+                throw Exceptions.propagate(ex);
+            }
+        }
+    }
+
     /**
      * Reads all Cosmos databases.
      * <!-- src_embed com.azure.cosmos.CosmosClient.readAllDatabases -->
@@ -311,6 +325,16 @@ public CosmosDatabase getDatabase(String id) {
         return new CosmosDatabase(id, this, asyncClientWrapper.getDatabase(id));
     }
 
+    /**
+     * Reads the Cosmos database account.
+     *
+     * @param shouldUseCache a boolean flag to determine whether to use the CosmosClient's-internal cache for reading the database account (setting shouldUseFlag to true can return stale data).
+     * @return the {@link CosmosDatabaseAccountResponse} with the read database account.
+     */
+    public CosmosDatabaseAccountResponse readDatabaseAccount(boolean shouldUseCache) {
+        return blockReadDatabaseAccount(this.asyncClientWrapper.readDatabaseAccount(shouldUseCache));
+    }
+
     CosmosAsyncClient asyncClient() {
         return this.asyncClientWrapper;
     }

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/AsyncDocumentClient.java

@@ -27,6 +27,7 @@
 import com.azure.cosmos.models.CosmosChangeFeedRequestOptions;
 import com.azure.cosmos.models.CosmosClientTelemetryConfig;
 import com.azure.cosmos.models.CosmosContainerIdentity;
+import com.azure.cosmos.models.CosmosDatabaseAccountResponse;
 import com.azure.cosmos.models.CosmosItemIdentity;
 import com.azure.cosmos.models.CosmosPatchOperations;
 import com.azure.cosmos.models.CosmosQueryRequestOptions;
@@ -1646,6 +1647,8 @@ <T> Flux<FeedResponse<T>> readAllDocuments(
 
     ConsistencyLevel getDefaultConsistencyLevelOfAccount();
 
+    Mono<CosmosDatabaseAccountResponse> readDatabaseAccount(boolean shouldUseCache);
+
     /***
      * Configure fault injector provider.
      *

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/GlobalEndpointManager.java

@@ -273,7 +273,7 @@ private Mono<Void> refreshLocationPrivateAsync(DatabaseAccount databaseAccount)
                     Mono<DatabaseAccount> databaseAccountObs = getDatabaseAccountFromAnyLocationsAsync(
                             this.defaultEndpoint,
                             new ArrayList<>(this.getEffectivePreferredRegions()),
-                            this::getDatabaseAccountAsync);
+                        this::getDatabaseAccountAsync);
 
                     return databaseAccountObs.map(dbAccount -> {
                         this.databaseAccountWriteLock.lock();
@@ -340,7 +340,7 @@ private Mono<Void> startRefreshLocationTimerAsync(boolean initialization) {
 
                             logger.debug("startRefreshLocationTimerAsync() - Invoking refresh, I was registered on [{}]", now);
                             Mono<DatabaseAccount> databaseAccountObs = GlobalEndpointManager.getDatabaseAccountFromAnyLocationsAsync(this.defaultEndpoint, new ArrayList<>(this.getEffectivePreferredRegions()),
-                                    this::getDatabaseAccountAsync);
+                                this::getDatabaseAccountAsync);
 
                             return databaseAccountObs.flatMap(dbAccount -> {
                                 logger.info("db account retrieved {}", dbAccount);
@@ -360,7 +360,7 @@ public boolean hasThinClientReadLocations() {
         return this.hasThinClientReadLocations.get();
     }
 
-    private Mono<DatabaseAccount> getDatabaseAccountAsync(URI serviceEndpoint) {
+    public Mono<DatabaseAccount> getDatabaseAccountAsync(URI serviceEndpoint) {
         return this.owner.getDatabaseAccountFromEndpoint(serviceEndpoint)
             .doOnNext(databaseAccount -> {
                 if(databaseAccount != null) {
@@ -399,7 +399,7 @@ public ConnectionPolicy getConnectionPolicy() {
         return this.connectionPolicy;
     }
 
-    private List<String> getEffectivePreferredRegions() {
+    public List<String> getEffectivePreferredRegions() {
 
         if (this.connectionPolicy.getPreferredRegions() != null && !this.connectionPolicy.getPreferredRegions().isEmpty()) {
             return this.connectionPolicy.getPreferredRegions();
@@ -419,4 +419,8 @@ private List<String> getEffectivePreferredRegions() {
             this.databaseAccountReadLock.unlock();
         }
     }
+
+    public DatabaseAccount getDatabaseAccount() {
+        return this.latestDatabaseAccount;
+    }
 }

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ImplementationBridgeHelpers.java

@@ -57,6 +57,7 @@
 import com.azure.cosmos.models.CosmosClientTelemetryConfig;
 import com.azure.cosmos.models.CosmosContainerIdentity;
 import com.azure.cosmos.models.CosmosContainerProperties;
+import com.azure.cosmos.models.CosmosDatabaseAccountResponse;
 import com.azure.cosmos.models.CosmosItemIdentity;
 import com.azure.cosmos.models.CosmosItemRequestOptions;
 import com.azure.cosmos.models.CosmosItemResponse;
@@ -1901,4 +1902,43 @@ ReadConsistencyStrategy getEffectiveReadConsistencyStrategy(
                 ReadConsistencyStrategy clientLevelReadConsistencyStrategy);
         }
     }
+
+    public static final class CosmosDatabaseAccountResponseHelper {
+        private static final AtomicReference<CosmosDatabaseAccountResponseAccessor> accessor = new AtomicReference<>();
+        private static final AtomicBoolean cosmosDatabaseAccountResponseClassLoaded = new AtomicBoolean(false);
+
+        private CosmosDatabaseAccountResponseHelper() {}
+
+        public static void setCosmosDatabaseAccountResponseAccessor(final CosmosDatabaseAccountResponseAccessor newAccessor) {
+            if (!accessor.compareAndSet(null, newAccessor)) {
+                logger.debug("CosmosDatabaseAccountResponseAccessor already initialized!");
+            } else {
+                logger.debug("Setting CosmosDatabaseAccountResponseAccessor...");
+                cosmosDatabaseAccountResponseClassLoaded.set(true);
+            }
+        }
+
+        public static CosmosDatabaseAccountResponseAccessor getCosmosDatabaseAccountResponseAccessor() {
+            if (!cosmosDatabaseAccountResponseClassLoaded.get()) {
+                logger.debug("Initializing CosmosDatabaseAccountResponseAccessor...");
+                initializeAllAccessors();
+            }
+
+            CosmosDatabaseAccountResponseAccessor snapshot = accessor.get();
+            if (snapshot == null) {
+                logger.error("CosmosDatabaseAccountResponseAccessor is not initialized yet!");
+            }
+
+            return snapshot;
+        }
+
+        public interface CosmosDatabaseAccountResponseAccessor {
+            CosmosDatabaseAccountResponse build(
+                String id,
+                List<String> readableRegions,
+                List<String> writeableRegions,
+                boolean isMultiWriteAccount,
+                ConsistencyLevel consistencyLevel);
+        }
+    }
 }

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java

@@ -75,6 +75,7 @@
 import com.azure.cosmos.models.CosmosChangeFeedRequestOptions;
 import com.azure.cosmos.models.CosmosClientTelemetryConfig;
 import com.azure.cosmos.models.CosmosContainerIdentity;
+import com.azure.cosmos.models.CosmosDatabaseAccountResponse;
 import com.azure.cosmos.models.CosmosItemIdentity;
 import com.azure.cosmos.models.CosmosItemRequestOptions;
 import com.azure.cosmos.models.CosmosItemResponse;
@@ -125,7 +126,6 @@
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentLinkedQueue;
 import java.util.concurrent.ConcurrentMap;
-import java.util.concurrent.ThreadLocalRandom;
 import java.util.concurrent.TimeoutException;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
@@ -186,6 +186,9 @@ public class RxDocumentClientImpl implements AsyncDocumentClient, IAuthorization
     private static final ImplementationBridgeHelpers.ReadConsistencyStrategyHelper.ReadConsistencyStrategyAccessor readConsistencyStrategyAccessor =
         ImplementationBridgeHelpers.ReadConsistencyStrategyHelper.getReadConsistencyStrategyAccessor();
 
+    private static final ImplementationBridgeHelpers.CosmosDatabaseAccountResponseHelper.CosmosDatabaseAccountResponseAccessor databaseAccountAccessor =
+        ImplementationBridgeHelpers.CosmosDatabaseAccountResponseHelper.getCosmosDatabaseAccountResponseAccessor();
+
     private static final String tempMachineId = "uuid:" + UUIDs.nonBlockingRandomUUID();
     private static final AtomicInteger activeClientsCnt = new AtomicInteger(0);
     private static final Map<String, Integer> clientMap = new ConcurrentHashMap<>();
@@ -6274,6 +6277,60 @@ public ConsistencyLevel getDefaultConsistencyLevelOfAccount() {
         return this.gatewayConfigurationReader.getDefaultConsistencyLevel();
     }
 
+    @Override
+    public Mono<CosmosDatabaseAccountResponse> readDatabaseAccount(boolean shouldUseCachedAccountSnapshot) {
+
+        if (shouldUseCachedAccountSnapshot) {
+            return Mono.just(this.globalEndpointManager.getLatestDatabaseAccount())
+                .flatMap(databaseAccount -> {
+
+                    List<String> readableRegions
+                        = Utils.iterableToList(databaseAccount.getReadableLocations())
+                        .stream()
+                        .map(DatabaseAccountLocation::getName)
+                        .collect(Collectors.toList());
+
+                    List<String> writeableRegions
+                        = Utils.iterableToList(databaseAccount.getWritableLocations())
+                        .stream()
+                        .map(DatabaseAccountLocation::getName)
+                        .collect(Collectors.toList());
+
+                    return Mono.just(databaseAccountAccessor.build(
+                        databaseAccount.getId(),
+                        readableRegions,
+                        writeableRegions,
+                        databaseAccount.getEnableMultipleWriteLocations(),
+                        databaseAccount.getConsistencyPolicy().getDefaultConsistencyLevel()
+                    ));
+                });
+        } else {
+            return GlobalEndpointManager.getDatabaseAccountFromAnyLocationsAsync(this.serviceEndpoint, this.globalEndpointManager.getEffectivePreferredRegions(), this.globalEndpointManager::getDatabaseAccountAsync)
+                .flatMap(databaseAccount -> {
+
+                    List<String> readableRegions
+                        = Utils.iterableToList(databaseAccount.getReadableLocations())
+                        .stream()
+                        .map(DatabaseAccountLocation::getName)
+                        .collect(Collectors.toList());
+
+                    List<String> writeableRegions
+                        = Utils.iterableToList(databaseAccount.getWritableLocations())
+                        .stream()
+                        .map(DatabaseAccountLocation::getName)
+                        .collect(Collectors.toList());
+
+                    return Mono.just(databaseAccountAccessor.build(
+                        databaseAccount.getId(),
+                        readableRegions,
+                        writeableRegions,
+                        databaseAccount.getEnableMultipleWriteLocations(),
+                        databaseAccount.getConsistencyPolicy().getDefaultConsistencyLevel()
+                    ));
+                });
+        }
+    }
+
     /***
      * Configure fault injector provider.
      *

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Utils.java

@@ -37,10 +37,12 @@
 import java.time.ZoneId;
 import java.time.ZonedDateTime;
 import java.time.format.DateTimeFormatter;
+import java.util.ArrayList;
 import java.util.Base64;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.HashMap;
+import java.util.Iterator;
 import java.util.List;
 import java.util.Locale;
 import java.util.Map;
@@ -814,4 +816,14 @@ public static Duration min(Duration duration1, Duration duration2) {
             return duration1.compareTo(duration2) < 0 ? duration1 : duration2;
         }
     }
+
+    public static <T> List<T> iterableToList(Iterable<T> iterable) {
+        List<T> list = new ArrayList<>();
+
+        for (T item : iterable) {
+            list.add(item);
+        }
+
+        return list;
+    }
 }