[azure-cosmos-db] Add ability to configure connection mode

vividus-framework · May 31, 2024 · 98120e3 · 98120e3
1 parent 108f14a
commit 98120e3
Show file tree

Hide file tree

Showing 4 changed files with 47 additions and 12 deletions.
diff --git a/docs/modules/plugins/pages/plugin-azure-cosmos-db.adoc b/docs/modules/plugins/pages/plugin-azure-cosmos-db.adoc
@@ -44,9 +44,11 @@ Where:
 
 . `account-key` - The key of the account configuration
 . `property-name` - The name of the container property. One of:
-.. `endpoint` - The endoint of the account.
+.. `endpoint` - The endpoint of the account.
 .. `key` - Either a master or readonly key used to perform authentication for Cosmos database service.
 If `key` is not set, <<_authentication,default authentication mechanism described below>> will be used.
+.. `connection-mode` - https://learn.microsoft.com/en-us/azure/cosmos-db/nosql/sdk-connection-modes [The mechanism],
+how the client connects to Azure Cosmos DB. Possible values: `DIRECT` or `GATEWAY`. If `connection-mode` is not set, `DIRECT` will be used.
 
 .Access configuration
 [source,properties]
@@ -71,6 +73,7 @@ azure.cosmos-db.database.commecrial.account-key=vividus
 # Accounts configuration contains connection details
 azure.cosmos-db.account.vividus.endpoint=https://vividus-cosmos.documents.azure.com:443/
 azure.cosmos-db.account.vividus.key=some key
+azure.cosmos-db.account.vividus.connection-mode=GATEWAY
 ----
 
 include::partial$azure-authentication.adoc[]

diff --git a/vividus-plugin-azure-cosmos-db/src/main/java/org/vividus/azure/cosmos/CosmosDbService.java b/vividus-plugin-azure-cosmos-db/src/main/java/org/vividus/azure/cosmos/CosmosDbService.java
@@ -16,6 +16,8 @@
 
 package org.vividus.azure.cosmos;
 
+import static com.azure.cosmos.ConnectionMode.GATEWAY;
+
 import java.util.stream.Collectors;
 
 import com.azure.core.credential.TokenCredential;
@@ -44,7 +46,7 @@ public class CosmosDbService
                 new CacheLoader<CosmosDbContainer, CosmosContainer>()
                 {
                     @Override
-                    public CosmosContainer load(CosmosDbContainer cosmosDbContainer)
+                    public CosmosContainer load(CosmosDbContainer cosmosDbContainer) throws IllegalArgumentException
                     {
                         String databaseKey = cosmosDbContainer.getDbKey();
                         CosmosDbDatabase db = databases.get(databaseKey,
@@ -66,6 +68,16 @@ public CosmosContainer load(CosmosDbContainer cosmosDbContainer)
                         {
                             builder.key(key);
                         }
+
+                        if (GATEWAY.equals(account.getConnectionMode()))
+                        {
+                            builder.gatewayMode();
+                        }
+                        else
+                        {
+                            builder.directMode();
+                        }
+
                         return builder.buildClient()
                                 .getDatabase(db.getId())
                                 .getContainer(cosmosDbContainer.getId());

diff --git a/...-plugin-azure-cosmos-db/src/main/java/org/vividus/azure/cosmos/model/CosmosDbAccount.java b/...-plugin-azure-cosmos-db/src/main/java/org/vividus/azure/cosmos/model/CosmosDbAccount.java
@@ -1,5 +1,5 @@
 /*
- * Copyright 2019-2021 the original author or authors.
+ * Copyright 2019-2024 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -16,10 +16,13 @@
 
 package org.vividus.azure.cosmos.model;
 
+import com.azure.cosmos.ConnectionMode;
+
 public class CosmosDbAccount
 {
     private String key;
     private String endpoint;
+    private ConnectionMode connectionMode;
 
     public String getKey()
     {
@@ -40,4 +43,14 @@ public void setEndpoint(String endpoint)
     {
         this.endpoint = endpoint;
     }
+
+    public ConnectionMode getConnectionMode()
+    {
+        return connectionMode;
+    }
+
+    public void setConnectionMode(ConnectionMode connectionMode)
+    {
+        this.connectionMode = connectionMode;
+    }
 }
diff --git a/...s-plugin-azure-cosmos-db/src/test/java/org/vividus/azure/cosmos/CosmosDbServiceTests.java b/...s-plugin-azure-cosmos-db/src/test/java/org/vividus/azure/cosmos/CosmosDbServiceTests.java
@@ -27,6 +27,7 @@
 import java.util.stream.Stream;
 
 import com.azure.core.credential.TokenCredential;
+import com.azure.cosmos.ConnectionMode;
 import com.azure.cosmos.CosmosClient;
 import com.azure.cosmos.CosmosClientBuilder;
 import com.azure.cosmos.CosmosContainer;
@@ -80,21 +81,13 @@ void beforeEach()
         database.setAccountKey(ACCOUNT_KEY);
         database.setId(DB);
         when(databases.get(DB_KEY, "Database configuration not found for the key: %s", DB_KEY)).thenReturn(database);
-        CosmosDbAccount cosmosDbAccount = new CosmosDbAccount();
-        cosmosDbAccount.setEndpoint(ENDPOINT);
-        cosmosDbAccount.setKey(KEY);
-        when(accounts.get(ACCOUNT_KEY, ACCOUNT_MSG, ACCOUNT_KEY))
-            .thenReturn(cosmosDbAccount);
         cosmosDbService = new CosmosDbService(jsonUtils, accounts, databases, tokenCredential);
     }
 
     @Test
     void shouldQueryDatabase()
     {
-        CosmosDbAccount cosmosDbAccount = new CosmosDbAccount();
-        cosmosDbAccount.setEndpoint(ENDPOINT);
-        when(accounts.get(ACCOUNT_KEY, ACCOUNT_MSG, ACCOUNT_KEY))
-                .thenReturn(cosmosDbAccount);
+        mockCosmosDBAccount(KEY, ConnectionMode.GATEWAY);
         testWithContainer((cosmosDbContainer, container) -> {
             @SuppressWarnings(UNCHECKED)
             CosmosPagedIterable<JsonNode> result = mock(CosmosPagedIterable.class);
@@ -108,6 +101,7 @@ void shouldQueryDatabase()
     @Test
     void shouldDeleteItem()
     {
+        mockCosmosDBAccount(null, ConnectionMode.GATEWAY);
         testWithContainer((cosmosDbContainer, container) -> {
             @SuppressWarnings(UNCHECKED)
             CosmosItemResponse<Object> response = mock(CosmosItemResponse.class);
@@ -120,6 +114,7 @@ void shouldDeleteItem()
     @Test
     void shouldInsertItem()
     {
+        mockCosmosDBAccount(KEY, ConnectionMode.DIRECT);
         testWithContainer((cosmosDbContainer, container) -> {
             @SuppressWarnings(UNCHECKED)
             CosmosItemResponse<JsonNode> response = mock(CosmosItemResponse.class);
@@ -132,6 +127,7 @@ void shouldInsertItem()
     @Test
     void shouldUpsertItem()
     {
+        mockCosmosDBAccount(null, null);
         testWithContainer((cosmosDbContainer, container) -> {
             @SuppressWarnings(UNCHECKED)
             CosmosItemResponse<JsonNode> response = mock(CosmosItemResponse.class);
@@ -144,6 +140,7 @@ void shouldUpsertItem()
     @Test
     void shouldReadItem()
     {
+        mockCosmosDBAccount(KEY, null);
         testWithContainer((cosmosDbContainer, container) -> {
             @SuppressWarnings(UNCHECKED)
             CosmosItemResponse<JsonNode> response = mock(CosmosItemResponse.class);
@@ -154,6 +151,16 @@ void shouldReadItem()
         }, "http://azure.com/read");
     }
 
+    private void mockCosmosDBAccount(String key, ConnectionMode connectionMode)
+    {
+        CosmosDbAccount cosmosDbAccount = new CosmosDbAccount();
+        cosmosDbAccount.setEndpoint(ENDPOINT);
+        cosmosDbAccount.setKey(key);
+        cosmosDbAccount.setConnectionMode(connectionMode);
+        when(accounts.get(ACCOUNT_KEY, ACCOUNT_MSG, ACCOUNT_KEY))
+                .thenReturn(cosmosDbAccount);
+    }
+
     private void testWithContainer(BiConsumer<CosmosDbContainer, CosmosContainer> testToRun, String containerId)
     {
         CosmosDbContainer dbContainer = new CosmosDbContainer();