diff --git a/docs/modules/plugins/pages/plugin-azure-cosmos-db.adoc b/docs/modules/plugins/pages/plugin-azure-cosmos-db.adoc index 6176cee940..ccc32b24b2 100644 --- 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. If `connection-mode` is not set, `DIRECT` will be used. Possible values: `DIRECT` or `GATEWAY`. .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 index f4eed0f7df..39f422e7fb 100644 --- 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,9 +16,13 @@ package org.vividus.azure.cosmos; +import static com.azure.cosmos.ConnectionMode.DIRECT; +import static com.azure.cosmos.ConnectionMode.GATEWAY; + import java.util.stream.Collectors; import com.azure.core.credential.TokenCredential; +import com.azure.cosmos.ConnectionMode; import com.azure.cosmos.CosmosClientBuilder; import com.azure.cosmos.CosmosContainer; import com.azure.cosmos.models.PartitionKey; @@ -44,7 +48,7 @@ public class CosmosDbService new CacheLoader() { @Override - public CosmosContainer load(CosmosDbContainer cosmosDbContainer) + public CosmosContainer load(CosmosDbContainer cosmosDbContainer) throws IllegalArgumentException { String databaseKey = cosmosDbContainer.getDbKey(); CosmosDbDatabase db = databases.get(databaseKey, @@ -66,6 +70,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/vividus-plugin-azure-cosmos-db/src/main/java/org/vividus/azure/cosmos/model/CosmosDbAccount.java b/vividus-plugin-azure-cosmos-db/src/main/java/org/vividus/azure/cosmos/model/CosmosDbAccount.java index b05dbce224..2ca0e05c9c 100644 --- a/vividus-plugin-azure-cosmos-db/src/main/java/org/vividus/azure/cosmos/model/CosmosDbAccount.java +++ b/vividus-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/vividus-plugin-azure-cosmos-db/src/test/java/org/vividus/azure/cosmos/CosmosDbServiceTests.java b/vividus-plugin-azure-cosmos-db/src/test/java/org/vividus/azure/cosmos/CosmosDbServiceTests.java index fa363823bd..27c0425d1b 100644 --- a/vividus-plugin-azure-cosmos-db/src/test/java/org/vividus/azure/cosmos/CosmosDbServiceTests.java +++ b/vividus-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 result = mock(CosmosPagedIterable.class); @@ -108,6 +101,7 @@ void shouldQueryDatabase() @Test void shouldDeleteItem() { + mockCosmosDBAccount(null, ConnectionMode.GATEWAY); testWithContainer((cosmosDbContainer, container) -> { @SuppressWarnings(UNCHECKED) CosmosItemResponse response = mock(CosmosItemResponse.class); @@ -120,6 +114,7 @@ void shouldDeleteItem() @Test void shouldInsertItem() { + mockCosmosDBAccount(KEY, ConnectionMode.DIRECT); testWithContainer((cosmosDbContainer, container) -> { @SuppressWarnings(UNCHECKED) CosmosItemResponse response = mock(CosmosItemResponse.class); @@ -132,6 +127,7 @@ void shouldInsertItem() @Test void shouldUpsertItem() { + mockCosmosDBAccount(null, null); testWithContainer((cosmosDbContainer, container) -> { @SuppressWarnings(UNCHECKED) CosmosItemResponse response = mock(CosmosItemResponse.class); @@ -144,6 +140,7 @@ void shouldUpsertItem() @Test void shouldReadItem() { + mockCosmosDBAccount(KEY, null); testWithContainer((cosmosDbContainer, container) -> { @SuppressWarnings(UNCHECKED) CosmosItemResponse 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 testToRun, String containerId) { CosmosDbContainer dbContainer = new CosmosDbContainer();