Skip to content

Commit 34aa967

Browse files
evanchoolyevanchooly
authored
Add sharding support to MongoDBContainer (#6727)
1 parent f78c0c3 commit 34aa967

File tree

3 files changed

+148
-32
lines changed

3 files changed

+148
-32
lines changed

‎modules/mongodb/src/main/java/org/testcontainers/containers/MongoDBContainer.java‎

Lines changed: 60 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
import lombok.extern.slf4j.Slf4j;
77
import org.testcontainers.containers.wait.strategy.Wait;
88
import org.testcontainers.utility.DockerImageName;
9+
import org.testcontainers.utility.MountableFile;
910

1011
import java.io.IOException;
1112

@@ -29,6 +30,10 @@ public class MongoDBContainer extends GenericContainer<MongoDBContainer> {
2930

3031
private static final String MONGODB_DATABASE_NAME_DEFAULT = "test";
3132

33+
private static final String STARTER_SCRIPT = "/testcontainers_start.sh";
34+
35+
private boolean shardingEnabled;
36+
3237
/**
3338
* @deprecated use {@link MongoDBContainer(DockerImageName)} instead
3439
*/
@@ -46,8 +51,44 @@ public MongoDBContainer(final DockerImageName dockerImageName) {
4651
dockerImageName.assertCompatibleWith(DEFAULT_IMAGE_NAME);
4752

4853
withExposedPorts(MONGODB_INTERNAL_PORT);
49-
withCommand("--replSet", "docker-rs");
50-
waitingFor(Wait.forLogMessage("(?i).*waiting for connections.*", 1));
54+
}
55+
56+
@Override
57+
public void configure() {
58+
if (shardingEnabled) {
59+
withCreateContainerCmdModifier(cmd -> {
60+
cmd.withEntrypoint("sh");
61+
});
62+
withCommand("-c", "while [ ! -f " + STARTER_SCRIPT + " ]; do sleep 0.1; done; " + STARTER_SCRIPT);
63+
waitingFor(Wait.forLogMessage("(?i).*mongos ready.*", 1));
64+
} else {
65+
withCommand("--replSet", "docker-rs");
66+
waitingFor(Wait.forLogMessage("(?i).*waiting for connections.*", 1));
67+
}
68+
}
69+
70+
@Override
71+
protected void containerIsStarting(InspectContainerResponse containerInfo) {
72+
if (shardingEnabled) {
73+
copyFileToContainer(MountableFile.forClasspathResource("/sharding.sh", 0777), STARTER_SCRIPT);
74+
}
75+
}
76+
77+
/**
78+
* Enables sharding on the cluster
79+
*
80+
* @return this
81+
*/
82+
public MongoDBContainer withSharding() {
83+
this.shardingEnabled = true;
84+
return this;
85+
}
86+
87+
@Override
88+
protected void containerIsStarted(InspectContainerResponse containerInfo, boolean reused) {
89+
if (!shardingEnabled) {
90+
initReplicaSet(reused);
91+
}
5192
}
5293

5394
/**
@@ -81,15 +122,6 @@ public String getReplicaSetUrl(final String databaseName) {
81122
return getConnectionString() + "/" + databaseName;
82123
}
83124

84-
@Override
85-
protected void containerIsStarted(InspectContainerResponse containerInfo, boolean reused) {
86-
if (reused && isReplicationSetAlreadyInitialized()) {
87-
log.debug("Replica set already initialized.");
88-
} else {
89-
initReplicaSet();
90-
}
91-
}
92-
93125
private String[] buildMongoEvalCommand(final String command) {
94126
return new String[] {
95127
"sh",
@@ -133,20 +165,24 @@ private void checkMongoNodeExitCodeAfterWaiting(final Container.ExecResult execR
133165
}
134166

135167
@SneakyThrows(value = { IOException.class, InterruptedException.class })
136-
private void initReplicaSet() {
137-
log.debug("Initializing a single node node replica set...");
138-
final ExecResult execResultInitRs = execInContainer(buildMongoEvalCommand("rs.initiate();"));
139-
log.debug(execResultInitRs.getStdout());
140-
checkMongoNodeExitCode(execResultInitRs);
141-
142-
log.debug(
143-
"Awaiting for a single node replica set initialization up to {} attempts",
144-
AWAIT_INIT_REPLICA_SET_ATTEMPTS
145-
);
146-
final ExecResult execResultWaitForMaster = execInContainer(buildMongoEvalCommand(buildMongoWaitCommand()));
147-
log.debug(execResultWaitForMaster.getStdout());
168+
private void initReplicaSet(boolean reused) {
169+
if (reused && isReplicationSetAlreadyInitialized()) {
170+
log.debug("Replica set already initialized.");
171+
} else {
172+
log.debug("Initializing a single node node replica set...");
173+
final ExecResult execResultInitRs = execInContainer(buildMongoEvalCommand("rs.initiate();"));
174+
log.debug(execResultInitRs.getStdout());
175+
checkMongoNodeExitCode(execResultInitRs);
148176

149-
checkMongoNodeExitCodeAfterWaiting(execResultWaitForMaster);
177+
log.debug(
178+
"Awaiting for a single node replica set initialization up to {} attempts",
179+
AWAIT_INIT_REPLICA_SET_ATTEMPTS
180+
);
181+
final ExecResult execResultWaitForMaster = execInContainer(buildMongoEvalCommand(buildMongoWaitCommand()));
182+
log.debug(execResultWaitForMaster.getStdout());
183+
184+
checkMongoNodeExitCodeAfterWaiting(execResultWaitForMaster);
185+
}
150186
}
151187

152188
public static class ReplicaSetInitializationException extends RuntimeException {

‎modules/mongodb/src/main/resources/sharding.sh‎

Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,68 @@
1+
#!/bin/bash
2+
3+
CONFIGSVR=/tmp/mongod/configsvr
4+
SHARDSVR=/tmp/mongod/shardsvr
5+
6+
function retry() {
7+
COUNT=${COUNT:-0}
8+
if [ $COUNT == 5 ]
9+
then
10+
echo Failed $COUNT attempts
11+
exit 1
12+
fi
13+
14+
sleep $COUNT
15+
echo "Attempt #$[ $COUNT + 1 ] '$*' "
16+
eval $*
17+
if [ $? -ne 0 ]
18+
then
19+
COUNT=$[ $COUNT + 1 ]
20+
retry $*
21+
fi
22+
unset COUNT
23+
}
24+
25+
function initReplSet() {
26+
PORT=$1
27+
COUNT=${2:-1}
28+
29+
CMD="mongosh --quiet --port $PORT --eval \"if(db.adminCommand({replSetGetStatus: 1})['myState'] != 1) quit(900)\""
30+
eval $CMD
31+
retVal=$?
32+
if [ $retVal -ne 0 -a $COUNT -ne 5 ]
33+
then
34+
echo "Initiating replSet (attempt $COUNT)"
35+
mongosh --quiet --port $PORT --eval 'rs.initiate();'
36+
if [ $? -ne 0 ]
37+
then
38+
sleep $COUNT
39+
initReplSet $PORT $[ $COUNT + 1 ]
40+
fi
41+
fi
42+
unset COUNT
43+
}
44+
45+
rm -rf $CONFIGSVR $SHARDSVR
46+
mkdir -p $CONFIGSVR
47+
mkdir -p $SHARDSVR
48+
49+
echo "Starting configsvr"
50+
mongod --bind_ip_all --configsvr --port 27019 --replSet configsvr-rs --dbpath $CONFIGSVR --logpath /tmp/configsvr.log &
51+
echo "Initiating configsvr replSet"
52+
initReplSet 27019
53+
54+
echo "Starting shardsvr"
55+
mongod --bind_ip_all --shardsvr --port 27018 --replSet shardsvr-rs --dbpath $SHARDSVR --logpath /tmp/shardsvr.log &
56+
57+
echo "Initiating shardsvr replSet"
58+
initReplSet 27018
59+
60+
echo "Starting mongos"
61+
mongos --bind_ip_all --configdb configsvr-rs/localhost:27019 --logpath /tmp/mongos.log &
62+
63+
echo "Adding a shard"
64+
retry "mongosh --eval 'sh.addShard(\"shardsvr-rs/`hostname`:27018\");'"
65+
66+
echo "mongos ready"
67+
mongosh --quiet tctest --eval "db.testcollection.insertOne({});"
68+
sleep 36000

‎modules/mongodb/src/test/java/org/testcontainers/containers/MongoDBContainerTest.java‎

Lines changed: 20 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -35,11 +35,8 @@ public void shouldExecuteTransactions() {
3535
}
3636

3737
private void executeTx(MongoDBContainer mongoDBContainer) {
38-
final String mongoRsUrl = mongoDBContainer.getReplicaSetUrl();
39-
assertThat(mongoRsUrl).isNotNull();
40-
final String connectionString = mongoDBContainer.getConnectionString();
41-
final MongoClient mongoSyncClientBase = MongoClients.create(connectionString);
42-
final MongoClient mongoSyncClient = MongoClients.create(mongoRsUrl);
38+
final MongoClient mongoSyncClientBase = MongoClients.create(mongoDBContainer.getConnectionString());
39+
final MongoClient mongoSyncClient = MongoClients.create(mongoDBContainer.getReplicaSetUrl());
4340
mongoSyncClient
4441
.getDatabase("mydb1")
4542
.getCollection("foo")
@@ -103,10 +100,25 @@ public void shouldTestDatabaseName() {
103100
}
104101

105102
@Test
106-
public void supportsMongoDB_6() {
107-
try (final MongoDBContainer mongoDBContainer = new MongoDBContainer("mongo:6.0.1")) {
103+
public void shouldSupportSharding() {
104+
try (final MongoDBContainer mongoDBContainer = new MongoDBContainer("mongo:6").withSharding()) {
108105
mongoDBContainer.start();
109-
executeTx(mongoDBContainer);
106+
final MongoClient mongoClient = MongoClients.create(mongoDBContainer.getReplicaSetUrl());
107+
108+
mongoClient.getDatabase("mydb1").getCollection("foo").insertOne(new Document("abc", 0));
109+
110+
Document shards = mongoClient.getDatabase("config").getCollection("shards").find().first();
111+
assertThat(shards).isNotNull();
112+
assertThat(shards).isNotEmpty();
113+
assertThat(isReplicaSet(mongoClient));
110114
}
111115
}
116+
117+
private boolean isReplicaSet(MongoClient mongoClient) {
118+
return runIsMaster(mongoClient).get("setName") != null;
119+
}
120+
121+
private Document runIsMaster(MongoClient mongoClient) {
122+
return mongoClient.getDatabase("admin").runCommand(new Document("ismaster", 1));
123+
}
112124
}

0 commit comments

Comments
 (0)