Merge pull request #987 from arnavk/arnav/sample_rate

Add provision to sample per message

4 files changed, 34 insertions, 3 deletions

diff --git a/raven/base.py b/raven/base.py
index 3627f33..a41508b 100644
--- a/raven/base.py
+++ b/raven/base.py
@@ -566,7 +566,8 @@ class Client(object):
         })
 
     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 @@ class Client(object):
             **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']
diff --git a/raven/handlers/logging.py b/raven/handlers/logging.py
index 3764c07..39f60d8 100644
--- a/raven/handlers/logging.py
+++ b/raven/handlers/logging.py
@@ -174,7 +174,9 @@ class SentryHandler(logging.Handler, object):
         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)
diff --git a/tests/base/tests.py b/tests/base/tests.py
index 736cf1c..5644db0 100644
--- a/tests/base/tests.py
+++ b/tests/base/tests.py
@@ -551,6 +551,20 @@ class ClientTest(TestCase):
             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:secret@example.com/1',
                         transport=HTTPTransport)
diff --git a/tests/handlers/logging/tests.py b/tests/handlers/logging/tests.py
index 163b8ab..5b6dc5b 100644
--- a/tests/handlers/logging/tests.py
+++ b/tests/handlers/logging/tests.py
@@ -262,6 +262,17 @@ class LoggingIntegrationTest(TestCase):
         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):