Merge pull request #389 from getsentry/rewrite-context

Updates context() API

6 files changed, 145 insertions, 138 deletions

diff --git a/raven/base.py b/raven/base.py
index ffaf13e..b0620c0 100644
--- a/raven/base.py
+++ b/raven/base.py
@@ -10,7 +10,6 @@ from __future__ import absolute_import
 
 import base64
 import zlib
-import datetime
 import logging
 import os
 import sys
@@ -18,16 +17,17 @@ import time
 import uuid
 import warnings
 
+from datetime import datetime
+
 import raven
 from raven.conf import defaults
 from raven.context import Context
-from raven.utils import json, get_versions, get_auth_header
+from raven.utils import six, json, get_versions, get_auth_header, merge_dicts
 from raven.utils.encoding import to_unicode
 from raven.utils.serializer import transform
 from raven.utils.stacks import get_stack_info, iter_stack_frames, get_culprit
 from raven.utils.urlparse import urlparse
 from raven.utils.compat import HTTPError
-from raven.utils import six
 from raven.transport.registry import TransportRegistry, default_transports
 
 __all__ = ('Client',)
@@ -42,8 +42,8 @@ class ModuleProxyCache(dict):
     def __missing__(self, key):
         module, class_name = key.rsplit('.', 1)
 
-        handler = getattr(__import__(module, {},
-                {}, [class_name]), class_name)
+        handler = getattr(__import__(
+            module, {}, {}, [class_name]), class_name)
 
         self[key] = handler
 
@@ -195,6 +195,8 @@ class Client(object):
         if Raven is None:
             Raven = self
 
+        self._context = Context()
+
     @classmethod
     def register_scheme(cls, scheme, transport_class):
         cls._registry.register_scheme(scheme, transport_class)
@@ -266,12 +268,12 @@ class Client(object):
         # create ID client-side so that it can be passed to application
         event_id = uuid.uuid4().hex
 
-        if data is None:
-            data = {}
-        if extra is None:
-            extra = {}
-        if not date:
-            date = datetime.datetime.utcnow()
+        data = merge_dicts(self.context.data, data)
+
+        data.setdefault('tags', {})
+        data.setdefault('extra', {})
+        data.setdefault('level', logging.ERROR)
+
         if stack is None:
             stack = self.auto_log_stacks
 
@@ -336,22 +338,13 @@ class Client(object):
         if not data.get('modules'):
             data['modules'] = self.get_module_versions()
 
-        data['tags'] = tags or {}
-        data.setdefault('extra', {})
-        data.setdefault('level', logging.ERROR)
-
-        # Add default extra context
-        if self.extra:
-            for k, v in six.iteritems(self.extra):
-                data['extra'].setdefault(k, v)
+        data['tags'] = merge_dicts(self.tags, data['tags'], tags)
+        data['extra'] = merge_dicts(self.extra, data['extra'], extra)
 
-        # Add default tag context
-        if self.tags:
-            for k, v in six.iteritems(self.tags):
-                data['tags'].setdefault(k, v)
-
-        for k, v in six.iteritems(extra):
-            data['extra'][k] = v
+        # Legacy support for site attribute
+        site = data.pop('site', None) or self.site
+        if site:
+            data['tags'].setdefault('site', site)
 
         if culprit:
             data['culprit'] = culprit
@@ -363,27 +356,20 @@ class Client(object):
         if 'message' not in data:
             data['message'] = handler.to_string(data)
 
-        data.setdefault('project', self.project)
-
-        # Legacy support for site attribute
-        site = data.pop('site', None) or self.site
-        if site:
-            data['tags'].setdefault('site', site)
-
+        # tags should only be key=>u'value'
         for key, value in six.iteritems(data['tags']):
             data['tags'][key] = to_unicode(value)
 
-        # Make sure custom data is coerced
+        # extra data can be any arbitrary value
         for k, v in six.iteritems(data['extra']):
             data['extra'][k] = self.transform(v)
 
         # It's important date is added **after** we serialize
-        data.update({
-            'timestamp': date,
-            'time_spent': time_spent,
-            'event_id': event_id,
-            'platform': PLATFORM_NAME,
-        })
+        data.setdefault('project', self.project)
+        data.setdefault('timestamp', date or datetime.utcnow())
+        data.setdefault('time_spent', time_spent)
+        data.setdefault('event_id', event_id)
+        data.setdefault('platform', PLATFORM_NAME)
 
         return data
 
@@ -392,18 +378,39 @@ class Client(object):
             data, list_max_length=self.list_max_length,
             string_max_length=self.string_max_length)
 
-    def context(self, **kwargs):
+    @property
+    def context(self):
         """
-        Create default context around a block of code for exception management.
-
-        >>> with client.context(tags={'key': 'value'}) as raven:
-        >>>     # use the context manager's client reference
-        >>>     raven.captureMessage('hello!')
-        >>>
-        >>>     # uncaught exceptions also contain the context
-        >>>     1 / 0
+        Updates this clients thread-local context for future events.
+
+        >>> def view_handler(view_func, *args, **kwargs):
+        >>>     client.context.merge(tags={'key': 'value'})
+        >>>     try:
+        >>>         return view_func(*args, **kwargs)
+        >>>     finally:
+        >>>         client.context.clear()
         """
-        return Context(self, **kwargs)
+        return self._context
+
+    def user_context(self, data):
+        return self.context.merge({
+            'sentry.interfaces.User': data,
+        })
+
+    def http_context(self, data, **kwargs):
+        return self.context.merge({
+            'sentry.interfaces.Http': data,
+        })
+
+    def extra_context(self, data, **kwargs):
+        return self.context.merge({
+            'extra': data,
+        })
+
+    def tags_context(self, data, **kwargs):
+        return self.context.merge({
+            'tags': data,
+        })
 
     def capture(self, event_type, data=None, date=None, time_spent=None,
                 extra=None, stack=None, tags=None, **kwargs):
diff --git a/raven/context.py b/raven/context.py
index 5d1e068..3fde3ac 100644
--- a/raven/context.py
+++ b/raven/context.py
@@ -7,41 +7,51 @@ raven.context
 """
 from __future__ import absolute_import
 
+from collections import Mapping, Iterable
+from threading import local
+
 from raven.utils import six
 
 
-class Context(object):
+class Context(local, Mapping, Iterable):
     """
-    Create default context around a block of code for exception management.
-
-    >>> with Context(client, tags={'key': 'value'}) as raven:
-    >>>     # use the context manager's client reference
-    >>>     raven.captureMessage('hello!')
-    >>>
-    >>>     # uncaught exceptions also contain the context
-    >>>     1 / 0
+    Stores context until cleared.
+
+    >>> def view_handler(view_func, *args, **kwargs):
+    >>>     context = Context()
+    >>>     context.merge(tags={'key': 'value'})
+    >>>     try:
+    >>>         return view_func(*args, **kwargs)
+    >>>     finally:
+    >>>         context.clear()
     """
-    def __init__(self, client, **defaults):
-        self.client = client
-        self.defaults = defaults
-        self.result = None
+    def __init__(self):
+        self.data = {}
+
+    def __getitem__(self, key):
+        return self.data[key]
 
-    def __enter__(self):
-        return self
+    def __iter__(self):
+        return iter(self.data)
 
-    def __exit__(self, *exc_info):
-        if all(exc_info):
-            self.result = self.captureException(exc_info)
+    def __repr__(self):
+        return '<%s: %s>' % (type(self).__name__, self.data)
 
-    def __call(self, function, *args, **kwargs):
-        for key, value in six.iteritems(self.defaults):
-            if key not in kwargs:
-                kwargs[key] = value
+    def merge(self, data):
+        d = self.data
+        for key, value in six.iteritems(data):
+            if key in ('tags', 'extra'):
+                d.setdefault(key, {})
+                for t_key, t_value in six.iteritems(value):
+                    d[key][t_key] = t_value
+            else:
+                d[key] = value
 
-        return function(*args, **kwargs)
+    def set(self, data):
+        self.data = data
 
-    def captureException(self, *args, **kwargs):
-        return self.__call(self.client.captureException, *args, **kwargs)
+    def get(self):
+        return self.data
 
-    def captureMessage(self, *args, **kwargs):
-        return self.__call(self.client.captureMessage, *args, **kwargs)
+    def clear(self):
+        self.data = {}
diff --git a/raven/middleware.py b/raven/middleware.py
index 3f1e420..dd96a45 100644
--- a/raven/middleware.py
+++ b/raven/middleware.py
@@ -48,6 +48,9 @@ class Sentry(object):
                 except Exception:
                     self.handle_exception(environ)
 
+    def process_response(self, request, response):
+        self.client.context.clear()
+
     def handle_exception(self, environ):
         event_id = self.client.captureException(
             data={
diff --git a/raven/utils/__init__.py b/raven/utils/__init__.py
index d624ebc..abbf5b8 100644
--- a/raven/utils/__init__.py
+++ b/raven/utils/__init__.py
@@ -18,6 +18,17 @@ import sys
 logger = logging.getLogger('raven.errors')
 
 
+def merge_dicts(*dicts):
+    out = {}
+    for d in dicts:
+        if not d:
+            continue
+
+        for k, v in six.iteritems(d):
+            out[k] = v
+    return out
+
+
 def varmap(func, var, context=None, name=None):
     """
     Executes ``func(key_name, value)`` on all values
diff --git a/tests/base/tests.py b/tests/base/tests.py
index 4d99538..1d8bd1d 100644
--- a/tests/base/tests.py
+++ b/tests/base/tests.py
@@ -270,35 +270,20 @@ class ClientTest(TestCase):
         self.assertFalse('sentry.interfaces.Stacktrace' in event)
         self.assertTrue('timestamp' in event)
 
-    def test_exception_context_manager(self):
-        cm = self.client.context(tags={'foo': 'bar'})
+    def test_context(self):
+        self.client.context.merge({
+            'tags': {'foo': 'bar'},
+        })
         try:
-            with cm:
-                raise ValueError('foo')
+            raise ValueError('foo')
         except:
-            pass
+            self.client.captureException()
         else:
             self.fail('Exception should have been raised')
 
-        self.assertNotEquals(cm.result, None)
-
-        self.assertEquals(len(self.client.events), 1)
+        assert len(self.client.events) == 1
         event = self.client.events.pop(0)
-        self.assertEquals(event['message'], 'ValueError: foo')
-        self.assertTrue('sentry.interfaces.Exception' in event)
-        exc = event['sentry.interfaces.Exception']
-        self.assertEquals(exc['type'], 'ValueError')
-        self.assertEquals(exc['value'], 'foo')
-        self.assertEquals(exc['module'], ValueError.__module__)  # this differs in some Python versions
-        self.assertTrue('sentry.interfaces.Stacktrace' in event)
-        frames = event['sentry.interfaces.Stacktrace']
-        self.assertEquals(len(frames['frames']), 1)
-        frame = frames['frames'][0]
-        self.assertEquals(frame['abs_path'], __file__.replace('.pyc', '.py'))
-        self.assertEquals(frame['filename'], 'tests/base/tests.py')
-        self.assertEquals(frame['module'], __name__)
-        self.assertEquals(frame['function'], 'test_exception_context_manager')
-        self.assertTrue('timestamp' in event)
+        assert event['tags'] == {'foo': 'bar'}
 
     def test_stack_explicit_frames(self):
         def bar():
diff --git a/tests/context/tests.py b/tests/context/tests.py
index 04a1427..cc65564 100644
--- a/tests/context/tests.py
+++ b/tests/context/tests.py
@@ -1,46 +1,37 @@
-import mock
-import sys
-from exam import fixture
 from raven.utils.testutils import TestCase
 
 from raven.context import Context
 
 
 class ContextTest(TestCase):
-    @fixture
-    def client(self):
-        return mock.Mock()
-
-    def context(self, **kwargs):
-        return Context(self.client, **kwargs)
-
-    def test_capture_exception(self):
-        with self.context(tags={'foo': 'bar'}) as client:
-            result = client.captureException('exception')
-            self.assertEquals(result, self.client.captureException.return_value)
-            self.client.captureException.assert_called_once_with('exception', tags={
+    def test_simple(self):
+        context = Context()
+        context.merge({'foo': 'bar'})
+        context.merge({'biz': 'baz'})
+        context.merge({'biz': 'boz'})
+        assert context.get() == {
+            'foo': 'bar',
+            'biz': 'boz',
+        }
+
+    def test_tags(self):
+        context = Context()
+        context.merge({'tags': {'foo': 'bar'}})
+        context.merge({'tags': {'biz': 'baz'}})
+        assert context.get() == {
+            'tags': {
                 'foo': 'bar',
-            })
-
-    def test_capture_message(self):
-        with self.context(tags={'foo': 'bar'}) as client:
-            result = client.captureMessage('exception')
-            self.assertEquals(result, self.client.captureMessage.return_value)
-            self.client.captureMessage.assert_called_once_with('exception', tags={
+                'biz': 'baz',
+            }
+        }
+
+    def test_extra(self):
+        context = Context()
+        context.merge({'extra': {'foo': 'bar'}})
+        context.merge({'extra': {'biz': 'baz'}})
+        assert context.get() == {
+            'extra': {
                 'foo': 'bar',
-            })
-
-    def test_implicit_exception_handling(self):
-        try:
-            with self.context(tags={'foo': 'bar'}):
-                try:
-                    1 / 0
-                except Exception:
-                    exc_info = sys.exc_info()
-                    raise
-        except Exception:
-            pass
-
-        self.client.captureException.assert_called_once_with(exc_info, tags={
-            'foo': 'bar',
-        })
+                'biz': 'baz',
+            }
+        }