From da61da0d503b64267f4da97be2a34d89c53b7337 Mon Sep 17 00:00:00 2001
From: Mark McClain <mark.mcclain@shootq.com>
Date: Thu, 27 Oct 2011 13:38:59 -0400
Subject: [PATCH] adding after_rollback functionality

---
 pecan/decorators.py | 36 +++++++++++++++++++++++++++++-----
 pecan/hooks.py      | 16 +++++++++------
 tests/test_hooks.py | 48 +++++++++++++++++++++++++++++++++++++++++++--
 3 files changed, 87 insertions(+), 13 deletions(-)

diff --git a/pecan/decorators.py b/pecan/decorators.py
index 95e3cef..7670489 100644
--- a/pecan/decorators.py
+++ b/pecan/decorators.py
@@ -2,7 +2,7 @@ from inspect import getargspec, getmembers, isclass, ismethod
 from util import _cfg
 
 __all__ = [
-    'expose', 'transactional', 'accept_noncanonical', 'after_commit'
+    'expose', 'transactional', 'accept_noncanonical', 'after_commit', 'after_rollback'
 ]
 
 
@@ -105,6 +105,24 @@ def transactional(ignore_redirects=True):
     return deco
 
 
+def after_action(action_type, action):
+    '''
+    If utilizing the :mod:`pecan.hooks` ``TransactionHook``, allows you
+    to flag a controller method to perform a callable action after the
+    action_type is successfully issued.
+
+    :param action: The callable to call after the commit is successfully issued.
+    '''
+
+    if action_type not in ('commit', 'rollback'):
+        raise Exception, 'action_type (%s) is not valid' % action_type
+
+
+    def deco(func):
+        _cfg(func).setdefault('after_%s' % action_type, []).append(action)
+        return func
+    return deco
+
 def after_commit(action):
     '''
     If utilizing the :mod:`pecan.hooks` ``TransactionHook``, allows you
@@ -113,10 +131,18 @@ def after_commit(action):
 
     :param action: The callable to call after the commit is successfully issued.
     '''
-    def deco(func):
-        _cfg(func).setdefault('after_commit', []).append(action)
-        return func
-    return deco
+    return  after_action('commit', action)
+
+
+def after_rollback(action):
+    '''
+    If utilizing the :mod:`pecan.hooks` ``TransactionHook``, allows you
+    to flag a controller method to perform a callable action after the
+    rollback is successfully issued.
+
+    :param action: The callable to call after the rollback is successfully issued.
+    '''
+    return  after_action('rollback', action)
 
 
 def accept_noncanonical(func):
diff --git a/pecan/hooks.py b/pecan/hooks.py
index 535c2bb..90c5b63 100644
--- a/pecan/hooks.py
+++ b/pecan/hooks.py
@@ -156,19 +156,23 @@ class TransactionHook(PecanHook):
     
     def after(self, state):
         if state.request.transactional:
+            action_name = None
             if state.request.error:
+                action_name = 'after_rollback'
                 self.rollback()
             else:
+                action_name = 'after_commit'
                 self.commit()
 
-                #
-                # If a controller was routed to, find any
-                # after_commit actions it may have registered, and perform
-                # them.
-                #
+            #
+            # If a controller was routed to, find any
+            # after_* actions it may have registered, and perform
+            # them.
+            #
+            if action_name:
                 controller = getattr(state, 'controller', None)
                 if controller is not None:
-                    actions = _cfg(controller).get('after_commit', [])
+                    actions = _cfg(controller).get(action_name, [])
                     for action in actions:
                         action()
 
diff --git a/tests/test_hooks.py b/tests/test_hooks.py
index f40aa44..3ba4022 100644
--- a/tests/test_hooks.py
+++ b/tests/test_hooks.py
@@ -1,9 +1,9 @@
 from cStringIO           import StringIO
-from pecan               import make_app, expose, request, redirect
+from pecan               import make_app, expose, request, redirect, abort
 from pecan.core          import state
 from pecan.hooks         import PecanHook, TransactionHook, HookController, RequestViewerHook
 from pecan.configuration import Config
-from pecan.decorators    import transactional, after_commit
+from pecan.decorators    import transactional, after_commit, after_rollback
 from copy                import copy
 from formencode          import Schema, validators
 from webtest             import TestApp
@@ -559,6 +559,17 @@ class TestTransactionHook(object):
             def decorated(self):
                 run_hook.append('inside')
                 return 'Decorated Method!'
+
+            @expose()
+            @after_rollback(action('action-three'))
+            def rollback(self):
+                abort(500)
+
+            @expose()
+            @transactional()
+            @after_rollback(action('action-four'))
+            def rollback_decorated(self):
+                abort(500)
         
         def gen(event):
             return lambda: run_hook.append(event)
@@ -612,6 +623,39 @@ class TestTransactionHook(object):
 
         run_hook = []
 
+        response = app.get('/rollback', expect_errors=True)
+        assert response.status_int == 500
+
+        assert len(run_hook) == 2
+        assert run_hook[0] == 'start_ro'
+        assert run_hook[1] == 'clear'
+
+        run_hook = []
+
+        response = app.post('/rollback', expect_errors=True)
+        assert response.status_int == 500
+
+        assert len(run_hook) == 4 
+        assert run_hook[0] == 'start'
+        assert run_hook[1] == 'rollback'
+        assert run_hook[2] == 'action-three'
+        assert run_hook[3] == 'clear'
+        
+        run_hook = []
+    
+        response = app.get('/rollback_decorated', expect_errors=True)
+        assert response.status_int == 500
+        
+        assert len(run_hook) == 6 
+        assert run_hook[0] == 'start_ro'
+        assert run_hook[1] == 'clear'
+        assert run_hook[2] == 'start'
+        assert run_hook[3] == 'rollback'
+        assert run_hook[4] == 'action-four'
+        assert run_hook[5] == 'clear'
+
+        run_hook = []
+
         response = app.get('/fourohfour', status=404)
         assert response.status_int == 404