Merge pull request #987 from arnavk/arnav/sample_rate

Add provision to sample per message

Author: dcramer

raven/base.py

@@ -566,7 +566,8 @@ def tags_context(self, data, **kwargs):
         })
 
     def capture(self, event_type, data=None, date=None, time_spent=None,
-                extra=None, stack=None, tags=None, **kwargs):
+                extra=None, stack=None, tags=None, sample_rate=None,
+                **kwargs):
         """
         Captures and processes an event and pipes it off to SentryClient.send.
 
@@ -636,7 +637,10 @@ def capture(self, event_type, data=None, date=None, time_spent=None,
             **kwargs)
 
         # should this event be sampled?
-        if self._random.random() < self.sample_rate:
+        if sample_rate is None:
+            sample_rate = self.sample_rate
+
+        if self._random.random() < sample_rate:
             self.send(**data)
 
         self._local_state.last_event_id = data['event_id']

raven/handlers/logging.py

@@ -174,7 +174,9 @@ def _emit(self, record, **kwargs):
         tags.update(getattr(record, 'tags', {}))
 
         kwargs.update(handler_kwargs)
+        sample_rate = extra.pop('sample_rate', None)
 
         return self.client.capture(
             event_type, stack=stack, data=data,
-            extra=extra, date=date, **kwargs)
+            extra=extra, date=date, sample_rate=sample_rate,
+            **kwargs)

tests/base/tests.py

@@ -551,6 +551,20 @@ def test_client_extra_context(self):
             expected = {'logger': "u'test'", 'foo': "u'bar'"}
         self.assertEquals(event['extra'], expected)
 
+    def test_sample_rate(self):
+        self.client.sample_rate = 0.0
+        self.client.captureMessage(message='test')
+        self.assertEquals(len(self.client.events), 0)
+
+    def test_sample_rate_per_message(self):
+        self.client.sample_rate = 1
+        self.client.captureMessage(message='test', sample_rate=0.0)
+        self.assertEquals(len(self.client.events), 0)
+
+        self.client.sample_rate = 0
+        self.client.captureMessage(message='test', sample_rate=1.0)
+        self.assertEquals(len(self.client.events), 1)
+
     def test_transport_registration(self):
         client = Client('http://public:[email protected]/1',
                         transport=HTTPTransport)

tests/handlers/logging/tests.py

@@ -262,6 +262,17 @@ def test_server_name_on_event(self):
         event = self.client.events.pop(0)
         assert event['server_name'] == 'foo'
 
+    def test_sample_rate(self):
+        record = self.make_record('Message', extra={'sample_rate': 0.0})
+        self.handler.emit(record)
+
+        self.assertEqual(len(self.client.events), 0)
+
+        record = self.make_record('Message', extra={'sample_rate': 1.0})
+        self.handler.emit(record)
+
+        self.assertEqual(len(self.client.events), 1)
+
 
 class LoggingHandlerTest(TestCase):
     def test_client_arg(self):