From 4a14c7c8f10577b0fbd31ca1ca1dde8e3d0ba2ab Mon Sep 17 00:00:00 2001
From: Rodolphe Marques <marques.rodolphe@gmail.com>
Date: Wed, 21 Dec 2016 15:53:25 +0100
Subject: [PATCH] Initialize replica set for mongodb when running bigchaindb
 init

---
 bigchaindb/backend/mongodb/schema.py | 70 ++++++++++++++++++++++++++++
 tests/backend/mongodb/test_schema.py | 35 ++++++++++++++
 2 files changed, 105 insertions(+)

diff --git a/bigchaindb/backend/mongodb/schema.py b/bigchaindb/backend/mongodb/schema.py
index d56c4e40..cc729b0f 100644
--- a/bigchaindb/backend/mongodb/schema.py
+++ b/bigchaindb/backend/mongodb/schema.py
@@ -1,8 +1,10 @@
 """Utils to initialize and drop the database."""
 
+import time
 import logging
 
 from pymongo import ASCENDING, DESCENDING
+from pymongo import errors
 
 from bigchaindb import backend
 from bigchaindb.common import exceptions
@@ -24,6 +26,9 @@ def create_database(conn, dbname):
     # TODO: read and write concerns can be declared here
     conn.conn.get_database(dbname)
 
+    # initialize the replica set
+    initialize_replica_set(conn)
+
 
 @register_schema(MongoDBConnection)
 def create_tables(conn, dbname):
@@ -85,3 +90,68 @@ def create_votes_secondary_index(conn, dbname):
                                              ('node_pubkey',
                                               ASCENDING)],
                                             name='block_and_voter')
+
+
+def initialize_replica_set(conn):
+    """Initialize a replica set. If already initialized skip."""
+    replica_set_name = _get_replica_set_name(conn)
+    config = {'_id': replica_set_name,
+              'members': [{'_id': 0, 'host': 'localhost:27017'}]}
+
+    try:
+        conn.conn.admin.command('replSetInitiate', config)
+    except errors.OperationFailure as exc_info:
+        if exc_info.details['codeName'] == 'AlreadyInitialized':
+            logger.info('Replica set already initialized')
+            return
+        else:
+            raise
+
+    _wait_for_replica_set_initialization(conn)
+    logger.info('Initialized replica set')
+
+
+def _get_replica_set_name(conn):
+    """Checks if the replSet option was enabled either through the command
+       line option or config file.
+
+       Note:
+           The setting we are looking for will have a different name depending
+           if it was set by the config file (`replSetName`) or by command
+           line arguments (`replSet`).
+
+       Returns:
+           The replica set name if enabled.
+
+        Raise:
+            :exc:`~ConfigurationError`: If mongod was not started with the
+            replSet option.
+    """
+    options = conn.conn.admin.command('getCmdLineOpts')
+    try:
+        repl_opts = options['parsed']['replication']
+        return repl_opts.get('replSetName', None) or repl_opts['replSet']
+    except KeyError:
+        raise exceptions.ConfigurationError('mongod was not started with'
+                                            ' the replSet option.')
+
+
+def _wait_for_replica_set_initialization(conn):
+    """Wait for a replica set to finish initialization.
+
+    If a replica set is being initialized for the first time it takes some
+    time. Nodes need to discover each other and an election needs to take
+    place. During this time the database is not writable so we need to wait
+    before continuing with the rest of the initialization
+    """
+
+    # I did not find a better way to do this for now.
+    # To check if the database is ready we will poll the mongodb logs until
+    # we find the line that says the database is ready
+    logger.info('Waiting for mongodb replica set initialization')
+    while True:
+        logs = conn.conn.admin.command('getLog', 'rs')['log']
+        for line in logs:
+            if 'database writes are now permitted' in line:
+                return
+        time.sleep(0.1)
diff --git a/tests/backend/mongodb/test_schema.py b/tests/backend/mongodb/test_schema.py
index 7cbaff64..257abe6d 100644
--- a/tests/backend/mongodb/test_schema.py
+++ b/tests/backend/mongodb/test_schema.py
@@ -1,4 +1,5 @@
 import pytest
+from unittest.mock import patch
 
 
 @pytest.mark.usefixtures('setup_database')
@@ -107,3 +108,37 @@ def test_drop():
 
     schema.drop_database(conn, dbname)
     assert dbname not in conn.conn.database_names()
+
+
+@pytest.mark.usefixtures('setup_database')
+def test_get_replica_set_name():
+    from pymongo.database import Database
+    from bigchaindb import backend
+    from bigchaindb.backend.mongodb.schema import _get_replica_set_name
+    from bigchaindb.common.exceptions import ConfigurationError
+
+    conn = backend.connect()
+
+    # no replSet option set
+    cmd_line_opts = {'argv': ['mongod', '--dbpath=/data'],
+                     'ok': 1.0,
+                     'parsed': {'storage': {'dbPath': '/data'}}}
+    with patch.object(Database, 'command', return_value=cmd_line_opts):
+        with pytest.raises(ConfigurationError):
+            _get_replica_set_name(conn)
+
+    # replSet option set through the command line
+    cmd_line_opts = {'argv': ['mongod', '--dbpath=/data', '--replSet=rs0'],
+                     'ok': 1.0,
+                     'parsed': {'replication': {'replSet': 'rs0'},
+                                'storage': {'dbPath': '/data'}}}
+    with patch.object(Database, 'command', return_value=cmd_line_opts):
+        assert _get_replica_set_name(conn) == 'rs0'
+
+    # replSet option set through the config file
+    cmd_line_opts = {'argv': ['mongod', '--dbpath=/data', '--replSet=rs0'],
+                     'ok': 1.0,
+                     'parsed': {'replication': {'replSetName': 'rs0'},
+                                'storage': {'dbPath': '/data'}}}
+    with patch.object(Database, 'command', return_value=cmd_line_opts):
+        assert _get_replica_set_name(conn) == 'rs0'