From 1557645e946bb554a0877a267b0837308db9b590 Mon Sep 17 00:00:00 2001
From: vrde <agranzot@gmail.com>
Date: Tue, 31 Jan 2017 02:07:36 +0100
Subject: [PATCH] Better exception handling

---
 bigchaindb/backend/exceptions.py         |  3 +++
 bigchaindb/backend/mongodb/changefeed.py |  6 ++++--
 bigchaindb/backend/mongodb/connection.py | 18 ++++++++++--------
 bigchaindb/commands/bigchain.py          |  7 ++++++-
 4 files changed, 23 insertions(+), 11 deletions(-)

diff --git a/bigchaindb/backend/exceptions.py b/bigchaindb/backend/exceptions.py
index 1d55bc52..1f7c57a7 100644
--- a/bigchaindb/backend/exceptions.py
+++ b/bigchaindb/backend/exceptions.py
@@ -1,5 +1,8 @@
 from bigchaindb.exceptions import BigchainDBError
 
 
+class ConnectionError(BigchainDBError):
+    """Exception raised when the connection to the DataBase fails."""
+
 class DatabaseOpFailedError(BigchainDBError):
     """Exception for database operation errors."""
diff --git a/bigchaindb/backend/mongodb/changefeed.py b/bigchaindb/backend/mongodb/changefeed.py
index c54bd5da..ecf36bcc 100644
--- a/bigchaindb/backend/mongodb/changefeed.py
+++ b/bigchaindb/backend/mongodb/changefeed.py
@@ -8,7 +8,8 @@ from bigchaindb import backend
 from bigchaindb.backend.changefeed import ChangeFeed
 from bigchaindb.backend.utils import module_dispatch_registrar
 from bigchaindb.backend.mongodb.connection import MongoDBConnection
-
+from bigchaindb.backend.exceptions import (DatabaseOpFailedError,
+                                           ConnectionError)
 
 logger = logging.getLogger(__name__)
 register_changefeed = module_dispatch_registrar(backend.changefeed)
@@ -31,7 +32,8 @@ class MongoDBChangeFeed(ChangeFeed):
                 break
             except (errors.ConnectionFailure, errors.OperationFailure,
                     errors.AutoReconnect,
-                    errors.ServerSelectionTimeoutError) as exc:
+                    errors.ServerSelectionTimeoutError,
+                    DatabaseOpFailedError, ConnectionError) as exc:
                 logger.exception(exc)
                 time.sleep(1)
 
diff --git a/bigchaindb/backend/mongodb/connection.py b/bigchaindb/backend/mongodb/connection.py
index c641d726..61cc430d 100644
--- a/bigchaindb/backend/mongodb/connection.py
+++ b/bigchaindb/backend/mongodb/connection.py
@@ -8,6 +8,7 @@ from pymongo import errors
 import bigchaindb
 from bigchaindb.utils import Lazy
 from bigchaindb.common import exceptions
+from bigchaindb.backend import exceptions as backend_exceptions
 from bigchaindb.backend.connection import Connection
 
 logger = logging.getLogger(__name__)
@@ -51,7 +52,7 @@ class MongoDBConnection(Connection):
     @property
     def conn(self):
         if self.connection is None:
-            self.connection = self._connect()
+            self._connect()
         return self.connection
 
     @property
@@ -77,23 +78,24 @@ class MongoDBConnection(Connection):
             try:
                 # FYI: this might raise a `ServerSelectionTimeoutError`,
                 # that is a subclass of `ConnectionFailure`.
-                connection = MongoClient(self.host,
-                                         self.port,
-                                         replicaset=self.replicaset,
-                                         serverselectiontimeoutms=self.connection_timeout)
+                self.connection = MongoClient(self.host,
+                                              self.port,
+                                              replicaset=self.replicaset,
+                                              serverselectiontimeoutms=self.connection_timeout)
 
                 # 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)
-                return connection
             except (errors.ConnectionFailure, errors.AutoReconnect) as exc:
                 logger.warning('Attempt %s/%s. Connection to %s:%s failed after %sms.',
                                attempt, self.max_tries if self.max_tries != 0 else '∞',
                                self.host, self.port, self.connection_timeout)
                 if attempt == self.max_tries:
-                    logger.exception('Cannot connect to the Database. Giving up.')
-                    raise errors.ConnectionFailure() from exc
+                    logger.critical('Cannot connect to the Database. Giving up.')
+                    raise backend_exceptions.ConnectionError() from exc
+            else:
+                break
 
 
 def collection(name):
diff --git a/bigchaindb/commands/bigchain.py b/bigchaindb/commands/bigchain.py
index 6661e902..4e18543d 100644
--- a/bigchaindb/commands/bigchain.py
+++ b/bigchaindb/commands/bigchain.py
@@ -23,7 +23,8 @@ from bigchaindb.utils import ProcessGroup
 from bigchaindb import backend
 from bigchaindb.backend import schema
 from bigchaindb.backend.admin import set_replicas, set_shards
-from bigchaindb.backend.exceptions import DatabaseOpFailedError
+from bigchaindb.backend.exceptions import (DatabaseOpFailedError,
+                                           ConnectionError)
 from bigchaindb.commands import utils
 from bigchaindb import processes
 
@@ -157,6 +158,8 @@ def run_init(args):
     except DatabaseAlreadyExists:
         print('The database already exists.', file=sys.stderr)
         print('If you wish to re-initialize it, first drop it.', file=sys.stderr)
+    except ConnectionError:
+        print('Cannot connect to the database.', file=sys.stderr)
 
 
 def run_drop(args):
@@ -201,6 +204,8 @@ def run_start(args):
         _run_init()
     except DatabaseAlreadyExists:
         pass
+    except ConnectionError:
+        print('Cannot connect to the database.', file=sys.stderr)
     except KeypairNotFoundException:
         sys.exit("Can't start BigchainDB, no keypair found. "
                  'Did you run `bigchaindb configure`?')