From 8c6841d3c00931204eaba0e9058707629120c1da Mon Sep 17 00:00:00 2001
From: lqslan <lqslan@cn.ibm.com>
Date: Tue, 18 Nov 2014 15:38:51 +0800
Subject: [PATCH] Retry to connect database when DB2 or mongodb is restarted

The patch https://review.openstack.org/#/c/122387 works fine
with operations with get, record and update functions.
But exception would still occured with the operation of
db.collection.find() function.

This patch can give some benefit to tolerate DB restart
with find() function.
This patch also removes "test_mongo_find" test case since
it doesn't raise AutoReconnect exception at all.

Change-Id: Ia0474726960ce2b4b611fda0a1c304bb8ad96922
Closes-Bug: #1389985
---
 ceilometer/storage/mongo/utils.py             | 46 +++++++++++--------
 .../tests/storage/test_storage_scenarios.py   | 24 +++++-----
 2 files changed, 40 insertions(+), 30 deletions(-)

diff --git a/ceilometer/storage/mongo/utils.py b/ceilometer/storage/mongo/utils.py
index 40e1eb216b..cb6b6b9b22 100644
--- a/ceilometer/storage/mongo/utils.py
+++ b/ceilometer/storage/mongo/utils.py
@@ -253,13 +253,12 @@ class ConnectionPool(object):
         try:
             if cfg.CONF.database.mongodb_replica_set:
                 client = MongoProxy(
-                    Prefection(
-                        pymongo.MongoReplicaSetClient(
-                            url,
-                            replicaSet=cfg.CONF.database.mongodb_replica_set)))
+                    pymongo.MongoReplicaSetClient(
+                        url,
+                        replicaSet=cfg.CONF.database.mongodb_replica_set))
             else:
                 client = MongoProxy(
-                    Prefection(pymongo.MongoClient(url, safe=True)))
+                    pymongo.MongoClient(url, safe=True))
             return client
         except pymongo.errors.ConnectionFailure as e:
             LOG.warn(_('Unable to connect to the database server: '
@@ -439,6 +438,12 @@ class MongoProxy(object):
         """
         return MongoProxy(self.conn[item])
 
+    def find(self, *args, **kwargs):
+        # We need this modifying method to return a CursorProxy object so that
+        # we can handle the Cursor next function to catch the AutoReconnect
+        # exception.
+        return CursorProxy(self.conn.find(*args, **kwargs))
+
     def __getattr__(self, item):
         """Wrap MongoDB connection.
 
@@ -456,20 +461,25 @@ class MongoProxy(object):
         return self.conn(*args, **kwargs)
 
 
-class Prefection(pymongo.collection.Collection):
-    def __init__(self, conn):
-        self.conn = conn
+class CursorProxy(pymongo.cursor.Cursor):
+    def __init__(self, cursor):
+        self.cursor = cursor
 
-    def find(self, *args, **kwargs):
-        # We need this modifying method to check a connection for MongoDB
-        # in context of MongoProxy approach. Initially 'find' returns Cursor
-        # object and doesn't connect to db while Cursor is not used.
-        found = self.find(*args, **kwargs)
+    def __getitem__(self, item):
+        return self.cursor[item]
+
+    @safe_mongo_call
+    def next(self):
+        """Wrap Cursor next method.
+
+        This method will be executed before each Cursor next method call.
+        """
         try:
-            found[0]
-        except IndexError:
-                pass
-        return found
+            save_cursor = self.cursor.clone()
+            return self.cursor.next()
+        except pymongo.errors.AutoReconnect:
+            self.cursor = save_cursor
+            raise
 
     def __getattr__(self, item):
-        return getattr(self.conn, item)
+        return getattr(self.cursor, item)
diff --git a/ceilometer/tests/storage/test_storage_scenarios.py b/ceilometer/tests/storage/test_storage_scenarios.py
index 5e4482d904..f01eaa104b 100644
--- a/ceilometer/tests/storage/test_storage_scenarios.py
+++ b/ceilometer/tests/storage/test_storage_scenarios.py
@@ -3505,11 +3505,11 @@ class MongoAutoReconnectTest(DBTestBase,
         self.CONF.set_override('retry_interval', 0, group='database')
 
     def test_mongo_client(self):
-        if self.CONF.database.mongodb_replica_set:
-            self.assertIsInstance(self.conn.conn.conn.conn,
+        if cfg.CONF.database.mongodb_replica_set:
+            self.assertIsInstance(self.conn.conn.conn,
                                   pymongo.MongoReplicaSetClient)
         else:
-            self.assertIsInstance(self.conn.conn.conn.conn,
+            self.assertIsInstance(self.conn.conn.conn,
                                   pymongo.MongoClient)
 
     @staticmethod
@@ -3521,16 +3521,16 @@ class MongoAutoReconnectTest(DBTestBase,
                 return method(*args, **kwargs)
         return side_effect
 
-    def test_mongo_find(self):
+    def test_mongo_cursor_next(self):
+        expected_first_sample_timestamp = datetime.datetime(2012, 7, 2, 10, 39)
         raise_exc = [False, True]
-        method = self.conn.db.resource.find
-
-        with mock.patch('pymongo.collection.Collection.find',
-                        mock.Mock()) as mock_find:
-            mock_find.side_effect = self.create_side_effect(method, raise_exc)
-            mock_find.__name__ = 'find'
-            resources = list(self.conn.get_resources())
-            self.assertEqual(9, len(resources))
+        method = self.conn.db.resource.find().cursor.next
+        with mock.patch('pymongo.cursor.Cursor.next',
+                        mock.Mock()) as mock_next:
+            mock_next.side_effect = self.create_side_effect(method, raise_exc)
+            resource = self.conn.db.resource.find().next()
+            self.assertEqual(expected_first_sample_timestamp,
+                             resource['first_sample_timestamp'])
 
     def test_mongo_insert(self):
         raise_exc = [False, True]