diff --git a/rollbar/__init__.py b/rollbar/__init__.py index 6285150b..20e91160 100644 --- a/rollbar/__init__.py +++ b/rollbar/__init__.py @@ -528,7 +528,12 @@ def send_payload(payload, access_token): if payload is False: return - handler = SETTINGS.get('handler') + if sys.version_info >= (3, 6): + from rollbar.lib._async import get_current_handler + handler = get_current_handler() + else: + handler = SETTINGS.get('handler') + if handler == 'twisted': payload['data']['framework'] = 'twisted' diff --git a/rollbar/lib/_async.py b/rollbar/lib/_async.py index 19df4881..1fd80173 100644 --- a/rollbar/lib/_async.py +++ b/rollbar/lib/_async.py @@ -22,6 +22,28 @@ ) +if sys.version_info[:2] == (3, 6): + # Backport PEP 567 + try: + import aiocontextvars + except ImportError: + log.warning( + 'Python3.6 does not provide the `contextvars` module.' + ' Some advanced features may not work as expected.' + ' Please upgrade Python or install `aiocontextvars`.' + ) + +try: + from contextvars import ContextVar +except ImportError: + ContextVar = None + +if ContextVar: + _ctx_handler = ContextVar('handler', default=None) +else: + _ctx_handler = None + + class RollbarAsyncError(Exception): ... @@ -142,21 +164,48 @@ async def try_report( ) -@contextlib.contextmanager -def async_handler(): - original_handler = rollbar.SETTINGS.get('handler') +class async_handler: + def __init__(self): + self.global_handler = None + self.token = None - if original_handler not in ALLOWED_HANDLERS: - log.warning( - 'Running coroutines requires async compatible handler. Switching to default async handler.' - ) - rollbar.SETTINGS['handler'] = 'async' + def with_ctx_handler(self): + if self.global_handler in ALLOWED_HANDLERS: + self.token = _ctx_handler.set(self.global_handler) + else: + log.warning( + 'Running coroutines requires async compatible handler. Switching to default async handler.' + ) + self.token = _ctx_handler.set('async') - try: - yield rollbar.SETTINGS['handler'] - finally: - if original_handler is not None: - rollbar.SETTINGS['handler'] = original_handler + return _ctx_handler.get() + + def with_global_handler(self): + return self.global_handler + + def __enter__(self): + self.global_handler = rollbar.SETTINGS.get('handler') + + if _ctx_handler: + return self.with_ctx_handler() + else: + return self.with_global_handler() + + def __exit__(self, exc_type, exc_value, traceback): + if _ctx_handler and self.token: + _ctx_handler.reset(self.token) + + +def get_current_handler(): + if _ctx_handler is None: + return rollbar.SETTINGS.get('handler') + + handler = _ctx_handler.get() + + if handler is None: + return rollbar.SETTINGS.get('handler') + + return handler def call_later(coro): diff --git a/rollbar/test/async_tests/test_async.py b/rollbar/test/async_tests/test_async.py index 556403d9..11296bde 100644 --- a/rollbar/test/async_tests/test_async.py +++ b/rollbar/test/async_tests/test_async.py @@ -146,8 +146,11 @@ def test_report_message_should_use_async_handler_regardless_of_settings( 'Running coroutines requires async compatible handler. Switching to default async handler.' ) + @mock.patch('logging.Logger.warning') @mock.patch('rollbar._send_payload_async') - def test_report_exc_info_should_allow_async_handler(self, mock__send_payload_async): + def test_report_exc_info_should_allow_async_handler( + self, mock__send_payload_async, mock_log + ): import rollbar from rollbar.lib._async import report_exc_info, run @@ -158,9 +161,13 @@ def test_report_exc_info_should_allow_async_handler(self, mock__send_payload_asy self.assertEqual(rollbar.SETTINGS['handler'], 'async') mock__send_payload_async.assert_called_once() + mock_log.assert_not_called() + @mock.patch('logging.Logger.warning') @mock.patch('rollbar._send_payload_async') - def test_report_message_should_allow_async_handler(self, mock__send_payload_async): + def test_report_message_should_allow_async_handler( + self, mock__send_payload_async, mock_log + ): import rollbar from rollbar.lib._async import report_message, run @@ -171,9 +178,13 @@ def test_report_message_should_allow_async_handler(self, mock__send_payload_asyn self.assertEqual(rollbar.SETTINGS['handler'], 'async') mock__send_payload_async.assert_called_once() + mock_log.assert_not_called() + @mock.patch('logging.Logger.warning') @mock.patch('rollbar._send_payload_httpx') - def test_report_exc_info_should_allow_httpx_handler(self, mock__send_payload_httpx): + def test_report_exc_info_should_allow_httpx_handler( + self, mock__send_payload_httpx, mock_log + ): import rollbar from rollbar.lib._async import report_exc_info, run @@ -184,9 +195,13 @@ def test_report_exc_info_should_allow_httpx_handler(self, mock__send_payload_htt self.assertEqual(rollbar.SETTINGS['handler'], 'httpx') mock__send_payload_httpx.assert_called_once() + mock_log.assert_not_called() + @mock.patch('logging.Logger.warning') @mock.patch('rollbar._send_payload_httpx') - def test_report_message_should_allow_httpx_handler(self, mock__send_payload_httpx): + def test_report_message_should_allow_httpx_handler( + self, mock__send_payload_httpx, mock_log + ): import rollbar from rollbar.lib._async import report_message, run @@ -197,62 +212,297 @@ def test_report_message_should_allow_httpx_handler(self, mock__send_payload_http self.assertEqual(rollbar.SETTINGS['handler'], 'httpx') mock__send_payload_httpx.assert_called_once() + mock_log.assert_not_called() @mock.patch('logging.Logger.warning') - def test_ctx_manager_should_temporary_set_async_handler(self, mock_log): + def test_ctx_manager_should_use_async_handler(self, mock_log): import rollbar from rollbar.lib._async import async_handler - rollbar.SETTINGS['handler'] = 'threading' - self.assertEqual(rollbar.SETTINGS['handler'], 'threading') - + rollbar.SETTINGS['handler'] = 'thread' with async_handler() as handler: self.assertEqual(handler, 'async') - self.assertEqual(rollbar.SETTINGS['handler'], handler) mock_log.assert_called_once_with( 'Running coroutines requires async compatible handler.' ' Switching to default async handler.' ) + rollbar.SETTINGS['handler'] = 'thread' + + @mock.patch('logging.Logger.warning') + def test_ctx_manager_should_use_global_handler_if_contextvar_is_not_supported( + self, mock_log + ): + import rollbar + import rollbar.lib._async + from rollbar.lib._async import async_handler - self.assertEqual(rollbar.SETTINGS['handler'], 'threading') + try: + # simulate missing `contextvars` module + _ctx_handler = rollbar.lib._async._ctx_handler + rollbar.lib._async._ctx_handler = None + + rollbar.SETTINGS['handler'] = 'thread' + self.assertEqual(rollbar.SETTINGS['handler'], 'thread') + + with async_handler() as handler: + self.assertEqual(handler, 'thread') + mock_log.assert_not_called() + + self.assertEqual(rollbar.SETTINGS['handler'], 'thread') + finally: + # restore original _ctx_handler + rollbar.lib._async._ctx_handler = _ctx_handler - def test_ctx_manager_should_not_substitute_async_handler(self): + @mock.patch('logging.Logger.warning') + def test_ctx_manager_should_not_substitute_global_handler(self, mock_log): import rollbar from rollbar.lib._async import async_handler - rollbar.SETTINGS['handler'] = 'async' - self.assertEqual(rollbar.SETTINGS['handler'], 'async') + rollbar.SETTINGS['handler'] = 'thread' + self.assertEqual(rollbar.SETTINGS['handler'], 'thread') with async_handler() as handler: self.assertEqual(handler, 'async') - self.assertEqual(rollbar.SETTINGS['handler'], handler) + self.assertEqual(rollbar.SETTINGS['handler'], 'thread') - self.assertEqual(rollbar.SETTINGS['handler'], 'async') + self.assertEqual(rollbar.SETTINGS['handler'], 'thread') + mock_log.assert_called_once_with( + 'Running coroutines requires async compatible handler.' + ' Switching to default async handler.' + ) - def test_ctx_manager_should_not_substitute_httpx_handler(self): + @mock.patch('rollbar._send_payload_httpx') + @mock.patch('rollbar._send_payload_async') + def test_report_exc_info_message_should_allow_multiple_async_handlers( + self, mock__send_payload_async, mock__send_payload_httpx + ): + import asyncio import rollbar - from rollbar.lib._async import async_handler + from rollbar.lib._async import report_exc_info, run - rollbar.SETTINGS['handler'] = 'httpx' - self.assertEqual(rollbar.SETTINGS['handler'], 'httpx') + async def report(handler): + rollbar.SETTINGS['handler'] = handler + try: + raise Exception('foo') + except: + await report_exc_info() - with async_handler() as handler: - self.assertEqual(handler, 'httpx') - self.assertEqual(rollbar.SETTINGS['handler'], handler) + async def send_reports(): + await asyncio.gather(report('async'), report('httpx')) - self.assertEqual(rollbar.SETTINGS['handler'], 'httpx') + run(send_reports()) + + mock__send_payload_async.assert_called_once() + mock__send_payload_httpx.assert_called_once() + + @mock.patch('rollbar._send_payload_httpx') + @mock.patch('rollbar._send_payload_async') + def test_report_message_should_allow_multiple_async_handlers( + self, mock__send_payload_async, mock__send_payload_httpx + ): + import asyncio + import rollbar + from rollbar.lib._async import report_message, run + + async def report(handler): + rollbar.SETTINGS['handler'] = handler + await report_message('foo') + + async def send_reports(): + await asyncio.gather(report('async'), report('httpx')) + + run(send_reports()) + + mock__send_payload_async.assert_called_once() + mock__send_payload_httpx.assert_called_once() + + @mock.patch('rollbar._send_payload') + @mock.patch('rollbar._send_payload_thread') + @mock.patch('rollbar._send_payload_httpx') + @mock.patch('rollbar._send_payload_async') + def test_report_exc_info_should_allow_multiple_handlers( + self, + mock__send_payload_async, + mock__send_payload_httpx, + mock__send_payload_thread, + mock__send_payload, + ): + import asyncio + import rollbar + from rollbar.lib._async import report_exc_info, run + + async def async_report(handler): + rollbar.SETTINGS['handler'] = handler + try: + raise Exception('foo') + except: + await report_exc_info() + + async def sync_report(handler): + rollbar.SETTINGS['handler'] = handler + try: + raise Exception('foo') + except: + rollbar.report_exc_info() + + async def send_reports(): + await asyncio.gather( + sync_report('thread'), + async_report('httpx'), + sync_report('blocking'), + async_report('async'), + ) + + run(send_reports()) + + mock__send_payload_async.assert_called_once() + mock__send_payload_httpx.assert_called_once() + mock__send_payload_thread.assert_called_once() + mock__send_payload.assert_called_once() + + @mock.patch('rollbar._send_payload') + @mock.patch('rollbar._send_payload_thread') + @mock.patch('rollbar._send_payload_httpx') + @mock.patch('rollbar._send_payload_async') + def test_report_message_should_allow_multiple_handlers( + self, + mock__send_payload_async, + mock__send_payload_httpx, + mock__send_payload_thread, + mock__send_payload, + ): + import asyncio + import rollbar + from rollbar.lib._async import report_message, run + + async def async_report(handler): + rollbar.SETTINGS['handler'] = handler + await report_message('foo') + + async def sync_report(handler): + rollbar.SETTINGS['handler'] = handler + rollbar.report_message('foo') + + async def send_reports(): + await asyncio.gather( + sync_report('thread'), + async_report('httpx'), + sync_report('blocking'), + async_report('async'), + ) + + run(send_reports()) + + mock__send_payload_async.assert_called_once() + mock__send_payload_httpx.assert_called_once() + mock__send_payload_thread.assert_called_once() + mock__send_payload.assert_called_once() + + @mock.patch('logging.Logger.warning') + @mock.patch('rollbar._send_payload_thread') + @mock.patch('rollbar._send_payload_async') + def test_report_exc_info_should_allow_multiple_handlers_with_threads( + self, + mock__send_payload_async, + mock__send_payload_thread, + mock_log, + ): + import time + import threading + import rollbar + from rollbar.lib._async import report_exc_info, run + + async def async_report(): + try: + raise Exception('foo') + except: + await report_exc_info() + + def sync_report(): + # give a chance to execute async_report() first + time.sleep(0.1) + + try: + raise Exception('foo') + except: + rollbar.report_exc_info() + + rollbar.SETTINGS['handler'] = 'thread' + t1 = threading.Thread(target=run, args=(async_report(),)) + t2 = threading.Thread(target=sync_report) + + t1.start() + t2.start() + t1.join() + t2.join() + + mock__send_payload_async.assert_called_once() + mock__send_payload_thread.assert_called_once() + mock_log.assert_called_once_with( + 'Running coroutines requires async compatible handler. Switching to default async handler.' + ) + + @mock.patch('logging.Logger.warning') + @mock.patch('rollbar._send_payload_thread') + @mock.patch('rollbar._send_payload_async') + def test_report_message_should_allow_multiple_handlers_with_threads( + self, + mock__send_payload_async, + mock__send_payload_thread, + mock_log, + ): + import time + import threading + import rollbar + from rollbar.lib._async import report_message, run + + async def async_report(): + await report_message('foo') + + def sync_report(): + # give a chance to execute async_report() first + time.sleep(0.1) + rollbar.report_message('foo') + + rollbar.SETTINGS['handler'] = 'thread' + t1 = threading.Thread(target=run, args=(async_report(),)) + t2 = threading.Thread(target=sync_report) + + t1.start() + t2.start() + t1.join() + t2.join() + + mock__send_payload_async.assert_called_once() + mock__send_payload_thread.assert_called_once() + mock_log.assert_called_once_with( + 'Running coroutines requires async compatible handler. Switching to default async handler.' + ) @mock.patch('rollbar.lib._async.report_exc_info', new_callable=AsyncMock) def test_should_try_report_with_async_handler(self, async_report_exc_info): import rollbar from rollbar.lib._async import run, try_report + rollbar.SETTINGS['handler'] = 'async' self.assertEqual(rollbar.SETTINGS['handler'], 'async') run(try_report()) async_report_exc_info.assert_called_once() + @mock.patch('rollbar.lib._async.report_exc_info', new_callable=AsyncMock) + def test_should_try_report_if_default_handler(self, async_report_exc_info): + import rollbar + from rollbar.lib._async import run, try_report + + rollbar.SETTINGS['handler'] = 'default' + self.assertEqual(rollbar.SETTINGS['handler'], 'default') + + run(try_report()) + + async_report_exc_info.assert_called_once() + @mock.patch('rollbar.lib._async.report_exc_info', new_callable=AsyncMock) def test_should_not_try_report_with_async_handler_if_non_async_handler( self, async_report_exc_info diff --git a/setup.py b/setup.py index 0f11c3df..e642e3db 100644 --- a/setup.py +++ b/setup.py @@ -29,6 +29,8 @@ tests_require.append('enum34') if sys.version_info >= (3, 6): tests_require.append('httpx') +if sys.version_info[:2] == (3, 6): + tests_require.append('aiocontextvars') setup( name='rollbar',