diff --git a/bigchaindb/backend/connection.py b/bigchaindb/backend/connection.py index cf6bece7..56b5cd82 100644 --- a/bigchaindb/backend/connection.py +++ b/bigchaindb/backend/connection.py @@ -16,7 +16,7 @@ logger = logging.getLogger(__name__) def connect(backend=None, host=None, port=None, name=None, max_tries=None, - connection_timeout=None, replicaset=None, ssl=False): + connection_timeout=None, replicaset=None, ssl=False, login=None, password=None): """Create a new connection to the database backend. All arguments default to the current configuration's values if not @@ -52,6 +52,8 @@ def connect(backend=None, host=None, port=None, name=None, max_tries=None, replicaset = replicaset or bigchaindb.config['database'].get('replicaset') ssl = bigchaindb.config['database'].get('ssl') if bigchaindb.config['database'].get('ssl') is not None \ else ssl + login = login or bigchaindb.config['database'].get('login') + password = password or bigchaindb.config['database'].get('password') try: module_name, _, class_name = BACKENDS[backend].rpartition('.') @@ -65,7 +67,7 @@ def connect(backend=None, host=None, port=None, name=None, max_tries=None, logger.debug('Connection: {}'.format(Class)) return Class(host=host, port=port, dbname=dbname, max_tries=max_tries, connection_timeout=connection_timeout, - replicaset=replicaset, ssl=ssl) + replicaset=replicaset, ssl=ssl, login=login, password=password) class Connection: diff --git a/bigchaindb/backend/mongodb/connection.py b/bigchaindb/backend/mongodb/connection.py index 274d64c1..9168190a 100644 --- a/bigchaindb/backend/mongodb/connection.py +++ b/bigchaindb/backend/mongodb/connection.py @@ -16,7 +16,7 @@ logger = logging.getLogger(__name__) class MongoDBConnection(Connection): - def __init__(self, replicaset=None, ssl=False, **kwargs): + def __init__(self, replicaset=None, ssl=False, login=None, password=None, **kwargs): """Create a new Connection instance. Args: @@ -30,6 +30,8 @@ class MongoDBConnection(Connection): self.replicaset = replicaset or bigchaindb.config['database']['replicaset'] self.ssl = bigchaindb.config['database'].get('ssl') if bigchaindb.config['database'].get('ssl') is not None \ else ssl + self.login = login or bigchaindb.config['database'].get('login') + self.password = password or bigchaindb.config['database'].get('password') @property def db(self): @@ -73,15 +75,20 @@ class MongoDBConnection(Connection): # we should only return a connection if the replica set is # initialized. initialize_replica_set will check if the # replica set is initialized else it will initialize it. - initialize_replica_set(self.host, self.port, self.connection_timeout, self.ssl) + initialize_replica_set(self.host, self.port, self.connection_timeout, self.dbname, self.ssl, self.login, self.password) # FYI: this might raise a `ServerSelectionTimeoutError`, # that is a subclass of `ConnectionFailure`. - return pymongo.MongoClient(self.host, - self.port, - replicaset=self.replicaset, - serverselectiontimeoutms=self.connection_timeout, - ssl=self.ssl) + client = pymongo.MongoClient(self.host, + self.port, + replicaset=self.replicaset, + serverselectiontimeoutms=self.connection_timeout, + ssl=self.ssl) + + if self.login is not None and self.password is not None: + client[self.dbname].authenticate(self.login, self.password) + + return client # `initialize_replica_set` might raise `ConnectionFailure` or `OperationFailure`. except (pymongo.errors.ConnectionFailure, @@ -89,7 +96,7 @@ class MongoDBConnection(Connection): raise ConnectionError() from exc -def initialize_replica_set(host, port, connection_timeout, ssl): +def initialize_replica_set(host, port, connection_timeout, dbname, ssl, login, password): """Initialize a replica set. If already initialized skip.""" # Setup a MongoDB connection @@ -100,6 +107,10 @@ def initialize_replica_set(host, port, connection_timeout, ssl): port=port, serverselectiontimeoutms=connection_timeout, ssl=ssl) + + if login is not None and password is not None: + conn[dbname].authenticate(login, password) + _check_replica_set(conn) host = '{}:{}'.format(bigchaindb.config['database']['host'], bigchaindb.config['database']['port'])