Merge pull request #1556 from minrk/safe_query

shallow-copy DictDB query results

`db_query()` with default keys pops buffers, and this would result in actually removing the buffers from the record itself.  Later calls to `get_result()` would then raise a KeyError, which would go unhandled.

This PR adds shallow-copy to DictDB results, so that changes to the returned dict are not reflected in the db itself.

The bug revealed a general issue: ZMQStreams are closed when an exception is raised in their `on_recv()` callback (following the IOStream pattern, and on the general principal that a socket with failing handlers should not be allowed to continue), which meant that invoking this bug would prevent any future queries to the Hub.  So in addition to fixing the particular bug, a decorator is added that catches and logs exceptions, and is applied to all `on_recv()` handlers in IPython.parallel, such that other similar bugs will no longer result in the closure of the socket.

Previously failing tests included at the db and client levels.

IPython/parallel/controller/dictdb.py

@@ -46,7 +46,7 b' We support a subset of mongodb operators:'
 #  the file COPYING, distributed as part of this software.
 #-----------------------------------------------------------------------------
 
-
+from copy import copy
 from datetime import datetime
 
 from IPython.config.configurable import LoggingConfigurable
@@ -120,7 +120,7 b' class DictDB(BaseDB):'
 
         for rec in self._records.itervalues():
             if self._match_one(rec, tests):
-                matches.append(rec)
+                matches.append(copy(rec))
         return matches
 
     def _extract_subdict(self, rec, keys):
@@ -139,9 +139,9 b' class DictDB(BaseDB):'
 
     def get_record(self, msg_id):
         """Get a specific Task Record, by msg_id."""
-        if not self._records.has_key(msg_id):
+        if not msg_id in self._records:
             raise KeyError("No such msg_id %r"%(msg_id))
-        return self._records[msg_id]
+        return copy(self._records[msg_id])
 
     def update_record(self, msg_id, rec):
         """Update the data in an existing record."""

IPython/parallel/controller/heartmonitor.py

@@ -25,7 +25,7 b' from zmq.eventloop import ioloop, zmqstream'
 from IPython.config.configurable import LoggingConfigurable
 from IPython.utils.traitlets import Set, Instance, CFloat, Integer
 
-from IPython.parallel.util import asbytes
+from IPython.parallel.util import asbytes, log_errors
 
 class Heart(object):
     """A basic heart object for responding to a HeartMonitor.
@@ -148,6 +148,7 b' class HeartMonitor(LoggingConfigurable):'
         self.hearts.remove(heart)
 
 
+    @log_errors
     def handle_pong(self, msg):
         "a heart just beat"
         current = asbytes(str(self.lifetime))

IPython/parallel/controller/hub.py

@@ -454,6 +454,7 b' class Hub(SessionFactory):'
     #-----------------------------------------------------------------------------
 
 
+    @util.log_errors
     def dispatch_monitor_traffic(self, msg):
         """all ME and Task queue messages come through here, as well as
         IOPub traffic."""
@@ -473,6 +474,7 b' class Hub(SessionFactory):'
             self.log.error("Invalid monitor topic: %r", switch)
 
 
+    @util.log_errors
     def dispatch_query(self, msg):
         """Route registration requests and queries from clients."""
         try:

IPython/parallel/controller/scheduler.py

@@ -43,7 +43,7 b' from IPython.config.application import Application'
 from IPython.config.loader import Config
 from IPython.utils.traitlets import Instance, Dict, List, Set, Integer, Enum, CBytes
 
-from IPython.parallel import error
+from IPython.parallel import error, util
 from IPython.parallel.factory import SessionFactory
 from IPython.parallel.util import connect_logger, local_logger, asbytes
 
@@ -240,6 +240,8 b' class TaskScheduler(SessionFactory):'
     # [Un]Registration Handling
     #-----------------------------------------------------------------------
 
+    
+    @util.log_errors
     def dispatch_notification(self, msg):
         """dispatch register/unregister events."""
         try:
@@ -343,6 +345,9 b' class TaskScheduler(SessionFactory):'
     #-----------------------------------------------------------------------
     # Job Submission
     #-----------------------------------------------------------------------
+    
+    
+    @util.log_errors
     def dispatch_submission(self, raw_msg):
         """Dispatch job submission to appropriate handlers."""
         # ensure targets up to date:
@@ -560,6 +565,9 b' class TaskScheduler(SessionFactory):'
     #-----------------------------------------------------------------------
     # Result Handling
     #-----------------------------------------------------------------------
+    
+    
+    @util.log_errors
     def dispatch_result(self, raw_msg):
         """dispatch method for result replies"""
         try:

IPython/parallel/tests/test_client.py

@@ -238,6 +238,18 b' class TestClient(ClusterTestCase):'
         for rec in found:
             self.assertTrue('msg_id' in rec.keys())
     
+    def test_db_query_get_result(self):
+        """pop in db_query shouldn't pop from result itself"""
+        self.client[:].apply_sync(lambda : 1)
+        found = self.client.db_query({'msg_id': {'$ne' : ''}})
+        rc2 = clientmod.Client(profile='iptest')
+        # If this bug is not fixed, this call will hang:
+        ar = rc2.get_result(self.client.history[-1])
+        ar.wait(2)
+        self.assertTrue(ar.ready())
+        ar.get()
+        rc2.close()
+    
     def test_db_query_in(self):
         """test db query with '$in','$nin' operators"""
         hist = self.client.hub_history()

IPython/parallel/tests/test_db.py

@@ -182,6 +182,36 b' class TestDictBackend(TestCase):'
         query = {'msg_id' : {'$ne' : None}}
         recs = self.db.find_records(query)
         self.assertTrue(len(recs) >= 10)
+    
+    def test_pop_safe_get(self):
+        """editing query results shouldn't affect record [get]"""
+        msg_id = self.db.get_history()[-1]
+        rec = self.db.get_record(msg_id)
+        rec.pop('buffers')
+        rec['garbage'] = 'hello'
+        rec2 = self.db.get_record(msg_id)
+        self.assertTrue('buffers' in rec2)
+        self.assertFalse('garbage' in rec2)
+    
+    def test_pop_safe_find(self):
+        """editing query results shouldn't affect record [find]"""
+        msg_id = self.db.get_history()[-1]
+        rec = self.db.find_records({'msg_id' : msg_id})[0]
+        rec.pop('buffers')
+        rec['garbage'] = 'hello'
+        rec2 = self.db.find_records({'msg_id' : msg_id})[0]
+        self.assertTrue('buffers' in rec2)
+        self.assertFalse('garbage' in rec2)
+
+    def test_pop_safe_find_keys(self):
+        """editing query results shouldn't affect record [find+keys]"""
+        msg_id = self.db.get_history()[-1]
+        rec = self.db.find_records({'msg_id' : msg_id}, keys=['buffers'])[0]
+        rec.pop('buffers')
+        rec['garbage'] = 'hello'
+        rec2 = self.db.find_records({'msg_id' : msg_id})[0]
+        self.assertTrue('buffers' in rec2)
+        self.assertFalse('garbage' in rec2)
 
 
 class TestSQLiteBackend(TestDictBackend):

IPython/parallel/util.py

@@ -39,6 +39,8 b' except:'
 import zmq
 from zmq.log import handlers
 
+from IPython.external.decorator import decorator
+
 # IPython imports
 from IPython.config.application import Application
 from IPython.utils import py3compat
@@ -106,6 +108,19 b' class ReverseDict(dict):'
 # Functions
 #-----------------------------------------------------------------------------
 
+@decorator
+def log_errors(f, self, *args, **kwargs):
+    """decorator to log unhandled exceptions raised in a method.
+    
+    For use wrapping on_recv callbacks, so that exceptions
+    do not cause the stream to be closed.
+    """
+    try:
+        return f(self, *args, **kwargs)
+    except Exception:
+        self.log.error("Uncaught exception in %r" % f, exc_info=True)
+    
+
 def asbytes(s):
     """ensure that an object is ascii bytes"""
     if isinstance(s, unicode):