Update Twisted integration

rollbar/__init__.py

@@ -96,6 +96,7 @@ def wrap(*args, **kwargs):
         return func(*args, **kwargs)
     return wrap
 
+
 try:
     from tornado.gen import coroutine as tornado_coroutine
     from tornado.httpclient import AsyncHTTPClient as TornadoAsyncHTTPClient
@@ -105,9 +106,10 @@ def wrap(*args, **kwargs):
 
 try:
     import treq
+    from twisted.internet.defer import inlineCallbacks
     from twisted.python import log as twisted_log
 
-    def log_handler(event):
+    def twisted_log_observer(event):
         """
         Default uncaught error handler
         """
@@ -116,20 +118,15 @@ def log_handler(event):
                 return
 
             err = event['failure']
-
             # Don't report Rollbar internal errors to ourselves
             if issubclass(err.type, ApiException):
                 log.error('Rollbar internal error: %s', err.value)
             else:
                 report_exc_info((err.type, err.value, err.getTracebackObject()))
         except:
             log.exception('Error while reporting to Rollbar')
-
-    # Add Rollbar as a log handler which will report uncaught errors
-    twisted_log.addObserver(log_handler)
-
-
 except ImportError:
+    inlineCallbacks = passthrough_decorator
     treq = None
 
 
@@ -256,10 +253,13 @@ def _get_pylons_request():
 _CURRENT_LAMBDA_CONTEXT = None
 
 # Set in init()
+_threads = None
 _transforms = []
 _serialize_transform = None
 
 _initialized = False
+_agent_log_installed = False
+_twisted_log_installed = False
 
 from rollbar.lib.transforms.scrub_redact import REDACT_REF
 
@@ -285,7 +285,18 @@ def init(access_token, environment='production', **kw):
                  'staging', 'yourname'
     **kw: provided keyword arguments will override keys in SETTINGS.
     """
-    global SETTINGS, agent_log, _initialized, _transforms, _serialize_transform, _threads
+    global SETTINGS, agent_log, _initialized, _transforms, _serialize_transform,\
+            _threads, _agent_log_installed, _twisted_log_installed
+
+    handler = SETTINGS.get('handler')
+    if 'handler' in kw:
+        handler = kw.get('handler')
+    if handler == 'agent' and not _agent_log_installed:
+        agent_log = _create_agent_log()
+        _agent_log_installed = True
+    elif handler == 'twisted' and treq and not _twisted_log_installed:
+        twisted_log.addObserver(twisted_log_observer)
+        _twisted_log_installed = True
 
     if _initialized:
         # NOTE: Temp solution to not being able to re-init.
@@ -304,8 +315,6 @@ def init(access_token, environment='production', **kw):
     if SETTINGS.get('allow_logging_basic_config'):
         logging.basicConfig()
 
-    if SETTINGS.get('handler') == 'agent':
-        agent_log = _create_agent_log()
 
     # We will perform these transforms in order:
     # 1. Serialize the payload to be all python built-in objects
@@ -450,7 +459,7 @@ def send_payload(payload, access_token):
         _send_payload_appengine(payload_str, access_token)
     elif handler == 'twisted':
         if treq is None:
-            log.error('Unable to find Treq')
+            log.error('treq and twisted are required for the twisted handler')
             return
         _send_payload_twisted(payload_str, access_token)
     else:
@@ -1219,7 +1228,7 @@ def _build_server_data():
     # argv does not always exist in embedded python environments
     argv = getattr(sys, 'argv', None)
     if argv:
-         server_data['argv'] = argv
+        server_data['argv'] = argv
 
     for key in ['branch', 'root']:
         if SETTINGS.get(key):
@@ -1353,25 +1362,43 @@ def _send_payload_twisted(payload_str, access_token):
         log.exception('Exception while posting item %r', e)
 
 
+@inlineCallbacks
 def _post_api_twisted(path, payload_str, access_token=None):
-    def post_data_cb(data, resp):
-        resp._content = data
-        _parse_response(path, SETTINGS['access_token'], payload_str, resp)
-
-    def post_cb(resp):
-        r = requests.Response()
-        r.status_code = resp.code
-        r.headers.update(resp.headers.getAllRawHeaders())
-        return treq.content(resp).addCallback(post_data_cb, r)
+    import datetime
 
     headers = {'Content-Type': ['application/json']}
     if access_token is not None:
         headers['X-Rollbar-Access-Token'] = [access_token]
 
     url = urljoin(SETTINGS['endpoint'], path)
-    d = treq.post(url, payload_str, headers=headers,
-                  timeout=SETTINGS.get('timeout', DEFAULT_TIMEOUT))
-    d.addCallback(post_cb)
+
+    try:
+        resp = yield treq.post(url, payload_str, headers=headers,
+                               timeout=SETTINGS.get('timeout', DEFAULT_TIMEOUT))
+        text = yield treq.content(resp)
+    except Exception as e:
+        # An exception that escapes here will be caught by Deferred
+        # and sent into Twisted's logging system.  This will again
+        # invoke the observer that called this function, starting an
+        # infinite recursion. Catch all exceptions to prevent this,
+        # including exceptions in the logging system itself.
+
+        twisted_log.removeObserver(twisted_log_observer)
+
+        try:
+            now = datetime.datetime.utcnow().isoformat()
+            twisted_log.err(e, '{} - Failed to post to rollbar'.format(now))
+        except:
+            pass
+
+        twisted_log.addObserver(twisted_log_observer)
+    else:
+        r = requests.Response()
+        r._content = text
+        r.status_code = resp.code
+        r.headers.update(resp.headers.getAllRawHeaders())
+
+        _parse_response(path, SETTINGS['access_token'], payload_str, r)
 
 
 def _send_failsafe(message, uuid, host):

rollbar/test/twisted_tests/test_twisted.py

@@ -17,7 +17,9 @@
     from twisted.test import proto_helpers
     from twisted.trial import unittest
     from twisted.internet import protocol
-    from twisted.python import log
+    from twisted.logger import Logger
+
+    log = Logger()
 
     TWISTED_INSTALLED = True
 except ImportError:
@@ -29,7 +31,7 @@ def dataReceived(self, data):
             try:
                 number = int(data)
             except ValueError:
-                log.err()
+                log.failure('error')
                 self.transport.write('error')
             else:
                 self.transport.write(str(number**2))
@@ -39,7 +41,7 @@ class SquareFactory(protocol.Factory):
 
     class TwistedTest(unittest.TestCase):
         def setUp(self):
-            rollbar.init(TOKEN, 'twisted-test')
+            rollbar.init(TOKEN, 'twisted-test', handler='twisted')
             factory = SquareFactory()
             self.proto = factory.buildProtocol(('127.0.0.1', 0))
             self.tr = proto_helpers.StringTransport()
@@ -55,7 +57,7 @@ def test_base_case(self, send_payload):
         @mock.patch('rollbar.send_payload')
         def test_caught_exception(self, send_payload):
             self.proto.dataReceived('rollbar')
-            self.assertEqual(self.tr.value(), "error")
+            self.assertEqual(self.tr.value(), 'error')
             errors = self.flushLoggedErrors(ValueError)
             self.assertEqual(len(errors), 1)
 
@@ -73,7 +75,7 @@ def test_caught_exception(self, send_payload):
         # @mock.patch('rollbar.send_payload')
         # def test_uncaught_exception(self, send_payload):
         #     self.proto.dataReceived([8, 9])
-        #     self.assertEqual(self.tr.value(), "error")
+        #     self.assertEqual(self.tr.value(), 'error')
         #     errors = self.flushLoggedErrors(TypeError)
         #     self.assertEqual(len(errors), 1)
         #