Skip to content

Commit

Permalink
Enable MongoDB replica set for the results cache and initialize it wh…
Browse files Browse the repository at this point in the history
…en starting the package.
  • Loading branch information
junhaoliao committed Dec 11, 2024
1 parent 6dd8fc1 commit fb4516e
Show file tree
Hide file tree
Showing 2 changed files with 22 additions and 1 deletion.
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
import sys

from pymongo import IndexModel, MongoClient
from pymongo.errors import OperationFailure

# Setup logging
# Create logger
Expand All @@ -15,6 +16,22 @@
logger.addHandler(logging_console_handler)


def initialize_replica_set(client, uri):
try:
result = client.admin.command("replSetGetStatus")
logger.info("Replica set already initialized: %s", result)
except OperationFailure as e:
logger.info("Initializing replica set")

# Explicit host specification is required, or the docker's ID would be used as the hostname.
config = {
"_id": "rs0",
"members": [{"_id": 0, "host": "localhost:27017"}],
}
client.admin.command("replSetInitiate", config)
logger.info("Replica set initialized successfully.")


def main(argv):
args_parser = argparse.ArgumentParser(description="Creates results cache indices for CLP.")
args_parser.add_argument("--uri", required=True, help="URI of the results cache.")
Expand All @@ -27,9 +44,11 @@ def main(argv):
stream_collection_name = parsed_args.stream_collection

try:
with MongoClient(results_cache_uri, directConnection=True) as results_cache_client:
initialize_replica_set(results_cache_client, results_cache_uri)

with MongoClient(results_cache_uri) as results_cache_client:
stream_collection = results_cache_client.get_default_database()[stream_collection_name]

file_split_id_index = IndexModel(["file_split_id"])
orig_file_id_index = IndexModel(["orig_file_id", "begin_msg_ix", "end_msg_ix"])
stream_collection.create_indexes([file_split_id_index, orig_file_id_index])
Expand Down
2 changes: 2 additions & 0 deletions components/package-template/src/etc/mongo/mongod.conf
Original file line number Diff line number Diff line change
@@ -1,5 +1,7 @@
net:
bindIp: 0.0.0.0
replication:
replSetName: "rs0"
systemLog:
destination: file
path: /var/log/mongodb/mongod.log
Expand Down

0 comments on commit fb4516e

Please sign in to comment.