Skip to content
This repository was archived by the owner on Jul 13, 2026. It is now read-only.

Commit e90f538

Browse files
TheovanKraayTheovanKraay
authored
Merge pull request #74 from TheovanKraay/bulk-sample-updates
Bulk sample updates
2 parents 59d21b9 + d63da2d commit e90f538

7 files changed

Lines changed: 494 additions & 110 deletions

File tree

‎pom.xml‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,11 @@
4747
<artifactId>azure-cosmos</artifactId>
4848
<version>LATEST</version>
4949
</dependency>
50+
<dependency>
51+
<groupId>com.azure</groupId>
52+
<artifactId>azure-identity</artifactId>
53+
<version>1.15.4</version>
54+
</dependency>
5055
<dependency>
5156
<groupId>org.apache.logging.log4j</groupId>
5257
<artifactId>log4j-api</artifactId>

‎src/main/java/com/azure/cosmos/examples/bulk/async/BulkWriter.java‎

Lines changed: 19 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -17,9 +17,11 @@
1717
import org.slf4j.Logger;
1818
import org.slf4j.LoggerFactory;
1919
import reactor.core.publisher.Flux;
20+
import reactor.core.publisher.Mono;
2021
import reactor.core.publisher.Sinks;
2122
import reactor.core.scheduler.Schedulers;
2223

24+
import java.time.Duration;
2325
import java.util.concurrent.Semaphore;
2426

2527
public class BulkWriter {
@@ -111,8 +113,17 @@ private void processResponseCode(
111113
"The operation for Item ID: [{}] Item PartitionKey Value: [{}] will be retried",
112114
itemOperation.getId(),
113115
itemOperation.getPartitionKeyValue());
114-
//re-scheduling
115-
scheduleWrites(itemOperation);
116+
long delay = itemResponse.getRetryAfterDuration() != null
117+
? itemResponse.getRetryAfterDuration().toMillis()
118+
: 0;
119+
120+
logger.info("Retrying after [{}] ms", delay);
121+
122+
Mono.delay(Duration.ofMillis(delay))
123+
// re-scheduling
124+
.doOnNext(t -> scheduleWrites(itemOperation))
125+
.subscribe();
126+
116127
} else {
117128
logger.info(
118129
"The operation for Item ID: [{}] Item PartitionKey Value: [{}] did not complete successfully " +
@@ -143,7 +154,11 @@ private void handleException(CosmosItemOperation itemOperation, Exception except
143154
}
144155

145156
private boolean shouldRetry(int statusCode) {
146-
return statusCode == HttpConstants.StatusCodes.REQUEST_TIMEOUT ||
147-
statusCode == HttpConstants.StatusCodes.TOO_MANY_REQUESTS;
157+
return statusCode == 408 ||
158+
statusCode == 429 ||
159+
statusCode == 503 ||
160+
statusCode == 500 ||
161+
statusCode == 449 ||
162+
statusCode == 410;
148163
}
149164
}

‎src/main/java/com/azure/cosmos/examples/bulk/async/SampleBulkHandleRetriesAsync.java‎

Lines changed: 191 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,191 @@
1+
// Copyright (c) Microsoft Corporation. All rights reserved.
2+
// Licensed under the MIT License.
3+
4+
package com.azure.cosmos.examples.bulk.async;
5+
6+
import com.azure.cosmos.ConsistencyLevel;
7+
import com.azure.cosmos.CosmosAsyncClient;
8+
import com.azure.cosmos.CosmosAsyncContainer;
9+
import com.azure.cosmos.CosmosAsyncDatabase;
10+
import com.azure.cosmos.CosmosClientBuilder;
11+
import com.azure.cosmos.GlobalThroughputControlConfig;
12+
import com.azure.cosmos.ThrottlingRetryOptions;
13+
import com.azure.cosmos.ThroughputControlGroupConfig;
14+
import com.azure.cosmos.ThroughputControlGroupConfigBuilder;
15+
import com.azure.cosmos.examples.common.AccountSettings;
16+
import com.azure.cosmos.examples.common.Families;
17+
import com.azure.cosmos.examples.common.Family;
18+
import com.azure.cosmos.models.CosmosBulkExecutionOptions;
19+
import com.azure.cosmos.models.CosmosBulkItemResponse;
20+
import com.azure.cosmos.models.CosmosBulkOperations;
21+
import com.azure.cosmos.models.CosmosContainerProperties;
22+
import com.azure.cosmos.models.CosmosContainerRequestOptions;
23+
import com.azure.cosmos.models.CosmosContainerResponse;
24+
import com.azure.cosmos.models.CosmosDatabaseResponse;
25+
import com.azure.cosmos.models.CosmosItemOperation;
26+
import com.azure.cosmos.models.CosmosItemRequestOptions;
27+
import com.azure.cosmos.models.CosmosPatchOperations;
28+
import com.azure.cosmos.models.PartitionKey;
29+
import com.azure.cosmos.models.ThroughputProperties;
30+
import com.azure.identity.DefaultAzureCredentialBuilder;
31+
import org.slf4j.Logger;
32+
import org.slf4j.LoggerFactory;
33+
import reactor.core.publisher.Flux;
34+
import reactor.core.publisher.Mono;
35+
36+
import java.time.Duration;
37+
import java.util.ArrayList;
38+
import java.util.List;
39+
import java.util.UUID;
40+
41+
public class SampleBulkHandleRetriesAsync {
42+
43+
private static final Logger logger = LoggerFactory.getLogger(SampleBulkHandleRetriesAsync.class);
44+
private final String databaseName = "AzureSampleFamilyDB";
45+
private final String containerName = "FamilyContainer";
46+
private CosmosAsyncClient client;
47+
private CosmosAsyncDatabase database;
48+
private CosmosAsyncContainer container;
49+
private int noOfItemsToCreate = 500;
50+
51+
public static void main(String[] args) {
52+
SampleBulkHandleRetriesAsync p = new SampleBulkHandleRetriesAsync();
53+
54+
try {
55+
logger.info("Starting ASYNC main");
56+
p.getStartedDemo();
57+
logger.info("Demo complete, please hold while resources are released");
58+
} catch (Exception e) {
59+
e.printStackTrace();
60+
logger.error(String.format("Cosmos getStarted failed with %s", e));
61+
} finally {
62+
logger.info("Closing the client");
63+
p.shutdown();
64+
}
65+
}
66+
67+
public void close() {
68+
client.close();
69+
}
70+
71+
private void getStartedDemo() {
72+
73+
logger.info("Using Azure Cosmos DB endpoint: " + AccountSettings.HOST);
74+
75+
/* Create an async client
76+
force rate limiting by reducing the max retry attempts
77+
to simulate throttling when under heavy load*/
78+
client = new CosmosClientBuilder()
79+
.endpoint(AccountSettings.HOST)
80+
.credential(new DefaultAzureCredentialBuilder().build())
81+
.throttlingRetryOptions(
82+
new ThrottlingRetryOptions()
83+
.setMaxRetryAttemptsOnThrottledRequests(1)
84+
.setMaxRetryWaitTime(Duration.ofSeconds(1)))
85+
.contentResponseOnWriteEnabled(true)
86+
.consistencyLevel(ConsistencyLevel.SESSION).buildAsyncClient();
87+
88+
89+
// Note: database must already exist (cannot be created via RBAC).
90+
database = client.getDatabase(databaseName);
91+
createContainerIfNotExists();
92+
93+
// Create 500 Family records
94+
ArrayList<Family> largeFamilies1 = new ArrayList<>();
95+
for (int i = 0; i < noOfItemsToCreate; i++) {
96+
Family family = new Family();
97+
family.setId(UUID.randomUUID().toString());
98+
family.setLastName("Family0");
99+
family.setRegistered(false);
100+
largeFamilies1.add(family);
101+
}
102+
logger.info("Ensure rate limiting with enough bulk upserts, retries handled by BulkWriter abstraction");
103+
largeBulkUpsertItemsWithBulkWriterAbstraction(largeFamilies1);
104+
}
105+
106+
private void createContainerIfNotExists() {
107+
logger.info("Create container " + containerName + " if not exists.");
108+
109+
// Create container if not exists
110+
// <CreateContainerIfNotExists>
111+
112+
CosmosContainerProperties containerProperties = new CosmosContainerProperties(
113+
containerName, "/lastName");
114+
ThroughputProperties throughputProperties = ThroughputProperties.createManualThroughput(400);
115+
Mono<CosmosContainerResponse> containerIfNotExists = database
116+
.createContainerIfNotExists(containerProperties, throughputProperties);
117+
118+
// Create container with 400 RU/s
119+
CosmosContainerResponse cosmosContainerResponse = containerIfNotExists.block();
120+
assert(cosmosContainerResponse != null);
121+
assert(cosmosContainerResponse.getProperties() != null);
122+
container = database.getContainer(cosmosContainerResponse.getProperties().getId());
123+
// </CreateContainerIfNotExists>
124+
125+
//Modify existing container
126+
containerProperties = cosmosContainerResponse.getProperties();
127+
Mono<CosmosContainerResponse> propertiesReplace =
128+
container.replace(containerProperties, new CosmosContainerRequestOptions());
129+
propertiesReplace.flatMap(containerResponse -> {
130+
logger.info(
131+
"setupContainer(): Container {}} in {} has been updated with it's new properties.",
132+
container.getId(),
133+
database.getId());
134+
return Mono.empty();
135+
}).onErrorResume((exception) -> {
136+
logger.error(
137+
"setupContainer(): Unable to update properties for container {} in database {}. e: {}",
138+
container.getId(),
139+
database.getId(),
140+
exception.getLocalizedMessage(),
141+
exception);
142+
return Mono.empty();
143+
}).block();
144+
145+
}
146+
147+
private void largeBulkUpsertItemsWithBulkWriterAbstraction(Iterable<Family> families) {
148+
List<CosmosItemOperation> cosmosItemOperations = new ArrayList<>();
149+
for (Family family : families) {
150+
cosmosItemOperations.add(CosmosBulkOperations.getUpsertItemOperation(family, new PartitionKey(family.getLastName())));
151+
}
152+
BulkWriter bulkWriter = new BulkWriter(container);
153+
for (CosmosItemOperation operation : cosmosItemOperations) {
154+
bulkWriter.scheduleWrites(operation);
155+
}
156+
bulkWriter.execute().subscribe();
157+
//get count of items in container
158+
try {
159+
logger.info("Waiting for bulk operations to complete...");
160+
Thread.sleep(20000); // Wait for the bulk operations to complete
161+
} catch (InterruptedException e) {
162+
throw new RuntimeException(e);
163+
}
164+
logger.info("Number of items to create was: " + noOfItemsToCreate);
165+
logger.info("Total items created after bulk load: " + container.readAllItems(new PartitionKey("Family0"), Family.class).count().block());
166+
}
167+
168+
169+
private void shutdown() {
170+
try {
171+
// To allow for the sequence to complete after subscribe() calls
172+
Thread.sleep(5000);
173+
//Clean shutdown
174+
logger.info("Deleting Cosmos DB resources");
175+
logger.info("-Deleting container...");
176+
if (container != null) container.delete().subscribe();
177+
logger.info("-Deleting database...");
178+
if (database != null) database.delete().subscribe();
179+
logger.info("-Closing the client...");
180+
} catch (InterruptedException err) {
181+
err.printStackTrace();
182+
} catch (Exception err) {
183+
logger.error("Deleting Cosmos DB resources failed, will still attempt to close the client. See stack " + "trace below.");
184+
err.printStackTrace();
185+
}
186+
client.close();
187+
logger.info("Done.");
188+
}
189+
190+
191+
}

‎src/main/java/com/azure/cosmos/examples/bulk/async/SampleBulkQuickStartAsync.java‎

Lines changed: 4 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
// <CosmosBulkOperationsImport>
1111
import com.azure.cosmos.models.*;
1212
// </CosmosBulkOperationsImport>
13+
import com.azure.identity.DefaultAzureCredentialBuilder;
1314
import org.slf4j.Logger;
1415
import org.slf4j.LoggerFactory;
1516
import reactor.core.publisher.Flux;
@@ -63,14 +64,15 @@ private void getStartedDemo() {
6364
// <CreateAsyncClient>
6465
client = new CosmosClientBuilder()
6566
.endpoint(AccountSettings.HOST)
66-
.key(AccountSettings.MASTER_KEY)
67+
.credential(new DefaultAzureCredentialBuilder().build())
6768
.preferredRegions(preferredRegions)
6869
.contentResponseOnWriteEnabled(true)
6970
.consistencyLevel(ConsistencyLevel.SESSION).buildAsyncClient();
7071

7172
// </CreateAsyncClient>
7273

73-
createDatabaseIfNotExists();
74+
// Note: database must already exist (cannot be created via RBAC).
75+
database = client.getDatabase(databaseName);
7476
createContainerIfNotExists();
7577

7678
// <AddDocsToStream>
@@ -131,20 +133,6 @@ private void getStartedDemo() {
131133
bulkCreateItemsWithBulkWriterAbstractionAndGlobalThroughputControl();
132134
}
133135

134-
private void createDatabaseIfNotExists() {
135-
logger.info("Create database " + databaseName + " if not exists.");
136-
137-
// Create database if not exists
138-
// <CreateDatabaseIfNotExists>
139-
Mono<CosmosDatabaseResponse> databaseIfNotExists = client.createDatabaseIfNotExists(databaseName);
140-
databaseIfNotExists.flatMap(databaseResponse -> {
141-
database = client.getDatabase(databaseResponse.getProperties().getId());
142-
logger.info("Checking database " + database.getId() + " completed!\n");
143-
return Mono.empty();
144-
}).block();
145-
// </CreateDatabaseIfNotExists>
146-
}
147-
148136
private void createContainerIfNotExists() {
149137
logger.info("Create container " + containerName + " if not exists.");
150138

@@ -349,8 +337,6 @@ private void shutdown() {
349337
logger.info("Deleting Cosmos DB resources");
350338
logger.info("-Deleting container...");
351339
if (container != null) container.delete().subscribe();
352-
logger.info("-Deleting database...");
353-
if (database != null) database.delete().subscribe();
354340
logger.info("-Closing the client...");
355341
} catch (InterruptedException err) {
356342
err.printStackTrace();

0 commit comments

Comments
 (0)