Fix bug whereby stack traces were not preserved when a Node's process…

… method raised an exception

featureflow/__init__.py

@@ -1,4 +1,4 @@
-__version__ = '1.16.11'
+__version__ = '1.16.12'
 
 from model import BaseModel
 

featureflow/model.py

@@ -30,7 +30,7 @@ def _ensure_persistence_settings(cls):
         if not issubclass(cls, PersistenceSettings):
             raise NoPersistenceSettingsError(
                     'The class {cls} is not a PersistenceSettings subclass'
-                    .format(cls=cls.__name__))
+                        .format(cls=cls.__name__))
 
     def __iter__(cls):
         cls._ensure_persistence_settings(cls)
@@ -73,6 +73,17 @@ def _build_extractor(cls, _id):
             feature._build_extractor(_id, g, cls)
         return g
 
+    @classmethod
+    def _rollback(cls, _id):
+        for f in cls.features.itervalues():
+            if not f.store:
+                continue
+            key = cls.key_builder.build(_id, f.key, f.version)
+            try:
+                del cls.database[key]
+            except:
+                pass
+
     @classmethod
     def process(cls, **kwargs):
         BaseModel._ensure_persistence_settings(cls)
@@ -82,13 +93,6 @@ def process(cls, **kwargs):
         try:
             graph.process(**kwargs)
             return _id
-        except Exception as e:
-            for f in cls.features.itervalues():
-                if not f.store:
-                    continue
-                key = cls.key_builder.build(_id, f.key, f.version)
-                try:
-                    del cls.database[key]
-                except:
-                    pass
-            raise e
+        except Exception:
+            cls._rollback(_id)
+            raise

featureflow/test_integration.py

@@ -19,6 +19,7 @@
 from persistence import PersistenceSettings
 from tempfile import mkdtemp
 from shutil import rmtree
+import traceback
 
 data_source = {
     'mary': 'mary had a little lamb little lamb little lamb',
@@ -150,6 +151,18 @@ def _process(self, data):
         yield data
 
 
+class Broken(Node):
+
+    MESSAGE = uuid4().hex
+
+    def __init__(self, needs=None):
+        super(Broken, self).__init__(needs=needs)
+
+    def _process(self, data):
+        raise Exception(Broken.MESSAGE)
+        yield data
+
+
 class TheLastWord(Node):
     def __init__(self, needs=None):
         super(TheLastWord, self).__init__(needs=needs)
@@ -329,6 +342,24 @@ class MultipleRoots(BaseModel):
 
 
 class BaseTest(object):
+
+    def test_get_sane_stack_trace_when_node_raises(self):
+        class D(BaseModel, self.Settings):
+            stream = Feature(TextStream, store=True)
+            words = Feature(Tokenizer, needs=stream, store=False)
+            broken = Feature(Broken, needs=words, store=False)
+            count = JSONFeature(WordCount, needs=broken, store=True)
+
+        try:
+            D.process(stream='mary')
+        except Exception:
+            _, _, tb = sys.exc_info()
+            items = traceback.extract_tb(tb)
+            self.assertEqual('raise Exception(Broken.MESSAGE)', items[-1][-1])
+            return
+
+        self.fail('Exception should have been raised')
+
     def test_can_iter_over_document_class(self):
         class D(BaseModel, self.Settings):
             stream = Feature(TextStream, store=True)