Merge pull request arngarden#23 from quantopian/master

Upstreaming multiple changes Quantopian has been using internally

Author: arngarden

.gitignore

@@ -19,6 +19,7 @@ var/
 *.egg-info/
 .installed.cfg
 *.egg
+.eggs/
 
 # Installer logs
 pip-log.txt

MANIFEST.in

@@ -1 +1 @@
-include *.txt *.md
+include *.txt *.md *.rst

README.md

@@ -0,0 +1,24 @@
+MongoDBProxy
+============
+
+MongoDBProxy is used to create a proxy around a MongoDB-connection in order to
+automatically handle AutoReconnect-exceptions.  You use MongoDBProxy in the
+same way you would an ordinary MongoDB-connection but don't need to worry about
+handling AutoReconnects by yourself.
+
+Usage::
+
+    >>> import pymongo
+    >>> import mongo_proxy
+    >>> safe_conn = mongo_proxy.MongoProxy(pymongo.MongoReplicaSetClient(replicaSet='blog_rs')
+    >>> safe_conn.blogs.posts.insert(post)  # Automatically handles AutoReconnect.
+
+**See here for more details:**
+`<http://www.arngarden.com/2013/04/29/handling-mongodb-autoreconnect-exceptions-in-python-using-a-proxy/>`_
+
+**Contributors**:
+
+- Jonathan Kamens (`<https://github.com/jikamens>`_)
+- Michael Cetrulo (`<https://github.com/git2samus>`_)
+- Richard Frank (`<https://github.com/richafrank>`_)
+- David Lindquist (`<https://github.com/dlindquist>`_)

README.rst

@@ -1,2 +1,8 @@
 from .mongodb_proxy import MongoProxy
 from .durable_cursor import DurableCursor, MongoReconnectFailure
+
+__all__ = [
+    'MongoProxy',
+    'DurableCursor',
+    'MongoReconnectFailure',
+]

mongo_proxy/__init__.py

@@ -15,15 +15,26 @@
 """
 
 import logging
-import sys
 import time
-from pymongo.errors import AutoReconnect, OperationFailure
+from pymongo.errors import (
+    AutoReconnect,
+    CursorNotFound,
+    ExecutionTimeout,
+    OperationFailure,
+    WTimeoutError,
+)
 
 # How long we are willing to attempt to reconnect when the replicaset
 # fails over.  We double the delay between each attempt.
 MAX_RECONNECT_TIME = 60
 MAX_SLEEP = 5
 RECONNECT_INITIAL_DELAY = 1
+RETRYABLE_OPERATION_FAILURE_CLASSES = (
+    AutoReconnect,  # AutoReconnect is raised when the primary node fails
+    CursorNotFound,
+    ExecutionTimeout,
+    WTimeoutError,
+)
 
 
 class MongoReconnectFailure(Exception):
@@ -136,46 +147,41 @@ def reload_cursor(self):
     def alive(self):
         return self.tailable and self.cursor.alive
 
-    def next(self):
-        try:
-            next_record = self.cursor.next()
-
-        # OperationFailure is raised when an operation fails inside
-        # the remote DB.  This most commonly occurs when our cursor
-        # has been inactive for 10 minutes or more.
-        except OperationFailure as exc:
-            self.logger.info("""
-Attempting to handle cursor timeout.
-OperationFailure exception catches a lot of failure cases.
-The current exception is actually:
-{exc}
-TODO: Inspect the exc name and only reload the cursor on timeouts.
+    def __next__(self):
+        next_record = self._with_retry(get_next=True, f=lambda: next(self.cursor))
+        # Increment count before returning so we know how many records
+        # to skip if a failure occurs later.
+        self.counter += 1
+        return next_record
 
-The query spec that timed out was:
-{spec}
-""".strip().format(exc=exc, spec=self.spec))
+    next = __next__
 
+    def _with_retry(self, get_next, f, *args, **kwargs):
+        try:
+            next_record = f(*args, **kwargs)
+        except RETRYABLE_OPERATION_FAILURE_CLASSES as exc:
+            self.logger.info(
+                "Got {!r}; attempting recovery. The query spec was: {}"
+                .format(exc, self.spec)
+            )
             # Try to reload the cursor and continue where we left off
-            self.reload_cursor()
-            next_record = self.cursor.next()
-            self.logger.info("Cursor reload after timeout successful.")
-
-        # AutoReconnect is raised when the primary node fails and we
-        # attempt to reconnect to the replica set.
-        except AutoReconnect:
-            self.logger.info("Got AutoReconnect; attempting recovery",
-                             exc_info=sys.exc_info())
-
-            # Try for up to self.max_reconnect_time to reconnect to
-            # the replicaset before giving up.  If the reconnect is
-            # successful, we return success == True along with the
-            # next record to return. Otherwise we return (False,
-            # None).
-            next_record = self.try_reconnect()
+            next_record = self.try_reconnect(get_next=get_next)
+            self.logger.info("Cursor reload after {!r} successful."
+                             .format(exc))
+
+        except OperationFailure as exc:
+            # No special subclass for this:
+            if 'interrupted at shutdown' in str(exc.args[0]):
+                self.logger.info(
+                    "Got {!r}; attempting recovery. The query spec was: {}"
+                    .format(exc, self.spec)
+                )
+                next_record = self.try_reconnect(get_next=get_next)
+                self.logger.info("Cursor reload after {!r} successful."
+                                 .format(exc))
+            else:
+                raise
 
-        # Increment count before returning so we know how many records
-        # to skip if a failure occurs later.
-        self.counter += 1
         return next_record
 
     def try_reconnect(self, get_next=True):
@@ -196,7 +202,7 @@ def try_reconnect(self, get_next=True):
             try:
                 # Attempt to reload and get the next batch.
                 self.reload_cursor()
-                return self.cursor.next() if get_next else True
+                return next(self.cursor) if get_next else True
 
             # Replica set hasn't come online yet.
             except AutoReconnect:
@@ -224,9 +230,8 @@ def try_reconnect(self, get_next=True):
         raise MongoReconnectFailure()
 
     def count(self, with_limit_and_skip=False):
-        while True:
-            try:
-                return self.cursor.count(
-                    with_limit_and_skip=with_limit_and_skip)
-            except AutoReconnect:
-                self.try_reconnect(get_next=False)
+        return self._with_retry(
+            get_next=False,
+            f=self.cursor.count,
+            with_limit_and_skip=with_limit_and_skip,
+        )

mongo_proxy/durable_cursor.py

@@ -17,15 +17,16 @@
 import time
 import pymongo
 
+
 def get_methods(*objs):
     return set(
         attr
         for obj in objs
         for attr in dir(obj)
-        if not attr.startswith('_')
-           and hasattr(getattr(obj, attr), '__call__')
+        if not attr.startswith('_') and hasattr(getattr(obj, attr), '__call__')
     )
 
+
 try:
     # will fail to import from older versions of pymongo
     from pymongo import MongoClient, MongoReplicaSetClient
@@ -57,7 +58,7 @@ def get_connection(obj):
         return None
 
 
-class Executable:
+class Executable(object):
     """ Wrap a MongoDB-method and handle AutoReconnect-exceptions
     using the safe_mongocall decorator.
     """
@@ -118,7 +119,8 @@ def __str__(self):
     def __repr__(self):
         return self.method.__repr__()
 
-class MongoProxy:
+
+class MongoProxy(object):
     """ Proxy for MongoDB connection.
     Methods that are executable, i.e find, insert etc, get wrapped in an
     Executable-instance that handles AutoReconnect-exceptions transparently.

mongo_proxy/mongodb_proxy.py

@@ -16,24 +16,8 @@
 
 from setuptools import setup, find_packages
 
-LONG_DESCRIPTION = None
-README_MARKDOWN = None
-
-with open('README.md') as markdown_source:
-    README_MARKDOWN = markdown_source.read()
-
-try:
-    import pandoc
-    pandoc.core.PANDOC_PATH = 'pandoc'
-    # Converts the README.md file to ReST, since PyPI uses ReST for formatting,
-    # This allows to have one canonical README file, being the README.md
-    doc = pandoc.Document()
-    doc.markdown = README_MARKDOWN
-    LONG_DESCRIPTION = doc.rst
-except ImportError:
-    # If pandoc isn't installed, e.g. when downloading from pip,
-    # just use the regular README.
-    LONG_DESCRIPTION = README_MARKDOWN
+with open('README.rst') as readme_file:
+    LONG_DESCRIPTION = readme_file.read()
 
 setup(
     name='MongoDBProxy',
@@ -45,7 +29,6 @@
     classifiers=[
         'License :: OSI Approved :: Apache Software License',
     ],
-    setup_requires=['pyandoc'],
     install_requires=['pymongo'],
     url="https://github.com/arngarden/MongoDBProxy"
 )