Skip to content

Commit

Permalink
Make Cosmos DB consistency level configurable (scalar-labs#1470)
Browse files Browse the repository at this point in the history
  • Loading branch information
brfrn169 authored Jan 29, 2024
1 parent e38ec16 commit ccb8cac
Show file tree
Hide file tree
Showing 7 changed files with 125 additions and 24 deletions.
10 changes: 1 addition & 9 deletions core/src/main/java/com/scalar/db/storage/cosmos/Cosmos.java
Original file line number Diff line number Diff line change
Expand Up @@ -2,9 +2,7 @@

import static com.google.common.base.Preconditions.checkArgument;

import com.azure.cosmos.ConsistencyLevel;
import com.azure.cosmos.CosmosClient;
import com.azure.cosmos.CosmosClientBuilder;
import com.google.common.annotations.VisibleForTesting;
import com.google.inject.Inject;
import com.scalar.db.api.Delete;
Expand Down Expand Up @@ -55,13 +53,7 @@ public Cosmos(DatabaseConfig databaseConfig) {

CosmosConfig config = new CosmosConfig(databaseConfig);

client =
new CosmosClientBuilder()
.endpoint(config.getEndpoint())
.key(config.getKey())
.directMode()
.consistencyLevel(ConsistencyLevel.STRONG)
.buildClient();
client = CosmosUtils.buildCosmosClient(config);

TableMetadataManager metadataManager =
new TableMetadataManager(
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -2,9 +2,7 @@

import static com.scalar.db.util.ScalarDbUtils.getFullTableName;

import com.azure.cosmos.ConsistencyLevel;
import com.azure.cosmos.CosmosClient;
import com.azure.cosmos.CosmosClientBuilder;
import com.azure.cosmos.CosmosContainer;
import com.azure.cosmos.CosmosDatabase;
import com.azure.cosmos.CosmosException;
Expand Down Expand Up @@ -72,13 +70,7 @@ public class CosmosAdmin implements DistributedStorageAdmin {
@Inject
public CosmosAdmin(DatabaseConfig databaseConfig) {
CosmosConfig config = new CosmosConfig(databaseConfig);
client =
new CosmosClientBuilder()
.endpoint(config.getEndpoint())
.key(config.getKey())
.directMode()
.consistencyLevel(ConsistencyLevel.STRONG)
.buildClient();
client = CosmosUtils.buildCosmosClient(config);
metadataDatabase = config.getMetadataDatabase();
}

Expand Down
11 changes: 11 additions & 0 deletions core/src/main/java/com/scalar/db/storage/cosmos/CosmosConfig.java
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,8 @@
import static com.scalar.db.config.ConfigUtils.getString;

import com.scalar.db.config.DatabaseConfig;
import java.util.Optional;
import javax.annotation.Nullable;
import javax.annotation.concurrent.Immutable;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
Expand All @@ -18,9 +20,12 @@ public class CosmosConfig {
@Deprecated
public static final String TABLE_METADATA_DATABASE = PREFIX + "table_metadata.database";

public static final String CONSISTENCY_LEVEL = PREFIX + "consistency_level";

private final String endpoint;
private final String key;
private final String metadataDatabase;
@Nullable private final String consistencyLevel;

public CosmosConfig(DatabaseConfig databaseConfig) {
String storage = databaseConfig.getStorage();
Expand Down Expand Up @@ -48,6 +53,8 @@ public CosmosConfig(DatabaseConfig databaseConfig) {
} else {
metadataDatabase = databaseConfig.getSystemNamespaceName();
}

consistencyLevel = getString(databaseConfig.getProperties(), CONSISTENCY_LEVEL, null);
}

// For the SpotBugs warning CT_CONSTRUCTOR_THROW
Expand All @@ -65,4 +72,8 @@ public String getKey() {
public String getMetadataDatabase() {
return metadataDatabase;
}

public Optional<String> getConsistencyLevel() {
return Optional.ofNullable(consistencyLevel);
}
}
33 changes: 33 additions & 0 deletions core/src/main/java/com/scalar/db/storage/cosmos/CosmosUtils.java
Original file line number Diff line number Diff line change
@@ -1,10 +1,43 @@
package com.scalar.db.storage.cosmos;

import com.azure.cosmos.ConsistencyLevel;
import com.azure.cosmos.CosmosClient;
import com.azure.cosmos.CosmosClientBuilder;
import com.google.common.annotations.VisibleForTesting;
import java.util.Locale;

public final class CosmosUtils {

private CosmosUtils() {}

public static String quoteKeyword(String keyword) {
return "[\"" + keyword + "\"]";
}

public static CosmosClient buildCosmosClient(CosmosConfig config) {
return new CosmosClientBuilder()
.endpoint(config.getEndpoint())
.key(config.getKey())
.directMode()
.consistencyLevel(getConsistencyLevel(config))
.buildClient();
}

@VisibleForTesting
static ConsistencyLevel getConsistencyLevel(CosmosConfig config) {
ConsistencyLevel consistencyLevel =
config
.getConsistencyLevel()
.map(c -> ConsistencyLevel.valueOf(c.toUpperCase(Locale.ROOT)))
.orElse(ConsistencyLevel.STRONG);

// Only STRONG and BOUNDED_STALENESS are supported
if (consistencyLevel != ConsistencyLevel.STRONG
&& consistencyLevel != ConsistencyLevel.BOUNDED_STALENESS) {
throw new IllegalArgumentException(
"The specified consistency level is not supported:" + consistencyLevel);
}

return consistencyLevel;
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@ public class CosmosConfigTest {
private static final String ANY_KEY = "any_key";
private static final String COSMOS_STORAGE = "cosmos";
private static final String ANY_TABLE_METADATA_DATABASE = "any_database";
private static final String ANY_CONSISTENCY_LEVEL = "any_consistency_level";

@Test
public void constructor_AllPropertiesGiven_ShouldLoadProperly() {
Expand All @@ -22,6 +23,7 @@ public void constructor_AllPropertiesGiven_ShouldLoadProperly() {
props.setProperty(DatabaseConfig.PASSWORD, ANY_KEY);
props.setProperty(DatabaseConfig.STORAGE, COSMOS_STORAGE);
props.setProperty(DatabaseConfig.SYSTEM_NAMESPACE_NAME, ANY_TABLE_METADATA_DATABASE);
props.setProperty(CosmosConfig.CONSISTENCY_LEVEL, ANY_CONSISTENCY_LEVEL);

// Act
CosmosConfig config = new CosmosConfig(new DatabaseConfig(props));
Expand All @@ -30,6 +32,7 @@ public void constructor_AllPropertiesGiven_ShouldLoadProperly() {
assertThat(config.getEndpoint()).isEqualTo(ANY_ENDPOINT);
assertThat(config.getKey()).isEqualTo(ANY_KEY);
assertThat(config.getMetadataDatabase()).isEqualTo(ANY_TABLE_METADATA_DATABASE);
assertThat(config.getConsistencyLevel()).hasValue(ANY_CONSISTENCY_LEVEL);
}

@Test
Expand All @@ -45,7 +48,7 @@ public void constructor_WithoutStorage_ShouldThrowIllegalArgumentException() {
}

@Test
public void constructor_WithoutTableMetadataDatabase_ShouldLoadProperly() {
public void constructor_WithoutSystemNamespaceNameAndConsistencyLevel_ShouldLoadProperly() {
// Arrange
Properties props = new Properties();
props.setProperty(DatabaseConfig.CONTACT_POINTS, ANY_ENDPOINT);
Expand All @@ -60,6 +63,7 @@ public void constructor_WithoutTableMetadataDatabase_ShouldLoadProperly() {
assertThat(config.getKey()).isEqualTo(ANY_KEY);
assertThat(config.getMetadataDatabase())
.isEqualTo(DatabaseConfig.DEFAULT_SYSTEM_NAMESPACE_NAME);
assertThat(config.getConsistencyLevel()).isEmpty();
}

@Test
Expand All @@ -78,6 +82,7 @@ public void constructor_WithTableMetadataDatabaseGiven_ShouldLoadProperly() {
assertThat(config.getEndpoint()).isEqualTo(ANY_ENDPOINT);
assertThat(config.getKey()).isEqualTo(ANY_KEY);
assertThat(config.getMetadataDatabase()).isEqualTo(ANY_TABLE_METADATA_DATABASE);
assertThat(config.getConsistencyLevel()).isEmpty();
}

@Test
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,67 @@
package com.scalar.db.storage.cosmos;

import static org.assertj.core.api.Assertions.assertThat;
import static org.mockito.Mockito.when;

import com.azure.cosmos.ConsistencyLevel;
import java.util.Optional;
import org.assertj.core.api.Assertions;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.mockito.Mock;
import org.mockito.MockitoAnnotations;

public class CosmosUtilsTest {

@Mock private CosmosConfig cosmosConfig;

@BeforeEach
public void setUp() throws Exception {
MockitoAnnotations.openMocks(this).close();
}

@Test
public void getConsistencyLevel_ShouldReturnStrongConsistency() {
// Arrange

// Act
ConsistencyLevel actual = CosmosUtils.getConsistencyLevel(cosmosConfig);

// Assert
assertThat(actual).isEqualTo(ConsistencyLevel.STRONG);
}

@Test
public void getConsistencyLevel_StrongGiven_ShouldReturnStrongConsistency() {
// Arrange
when(cosmosConfig.getConsistencyLevel()).thenReturn(Optional.of("STRONG"));

// Act
ConsistencyLevel actual = CosmosUtils.getConsistencyLevel(cosmosConfig);

// Assert
assertThat(actual).isEqualTo(ConsistencyLevel.STRONG);
}

@Test
public void getConsistencyLevel_BoundedStalenessGiven_ShouldReturnBoundedStalenessConsistency() {
// Arrange
when(cosmosConfig.getConsistencyLevel()).thenReturn(Optional.of("bounded_staleness"));

// Act
ConsistencyLevel actual = CosmosUtils.getConsistencyLevel(cosmosConfig);

// Assert
assertThat(actual).isEqualTo(ConsistencyLevel.BOUNDED_STALENESS);
}

@Test
public void getConsistencyLevel_InvalidConsistencyGiven_ShouldThrowIllegalArgumentException() {
// Arrange
when(cosmosConfig.getConsistencyLevel()).thenReturn(Optional.of("any"));

// Act Assert
Assertions.assertThatThrownBy(() -> CosmosUtils.getConsistencyLevel(cosmosConfig))
.isInstanceOf(IllegalArgumentException.class);
}
}
11 changes: 6 additions & 5 deletions docs/configurations.md
Original file line number Diff line number Diff line change
Expand Up @@ -80,11 +80,12 @@ The following configurations are available for Cassandra:

The following configurations are available for CosmosDB for NoSQL:

| Name | Description | Default |
|--------------------------------------|----------------------------------------------------------------------------------------------------------|------------|
| `scalar.db.storage` | `cosmos` must be specified. | - |
| `scalar.db.contact_points` | Azure Cosmos DB for NoSQL endpoint with which ScalarDB should communicate. | |
| `scalar.db.password` | Either a master or read-only key used to perform authentication for accessing Azure Cosmos DB for NoSQL. | |
| Name | Description | Default |
|--------------------------------------|----------------------------------------------------------------------------------------------------------|----------|
| `scalar.db.storage` | `cosmos` must be specified. | - |
| `scalar.db.contact_points` | Azure Cosmos DB for NoSQL endpoint with which ScalarDB should communicate. | |
| `scalar.db.password` | Either a master or read-only key used to perform authentication for accessing Azure Cosmos DB for NoSQL. | |
| `scalar.db.cosmos.consistency_level` | Consistency level used for Cosmos DB operations. `STRONG` or `BOUNDED_STALENESS` can be specified. | `STRONG` |

</div>
<div id="DynamoDB" class="tabcontent" markdown="1">
Expand Down

0 comments on commit ccb8cac

Please sign in to comment.