diff --git a/examples/influxdb.py b/examples/influxdb.py index 741f48d..88b0d0f 100644 --- a/examples/influxdb.py +++ b/examples/influxdb.py @@ -32,6 +32,16 @@ class SimpleHandler(metrics.InfluxDBMixin, web.RequestHandler): """ + @gen.coroutine + def prepare(self): + maybe_future = super(SimpleHandler, self).prepare() + if gen.is_future(maybe_future): + yield maybe_future + + if 'Correlation-ID' in self.request.headers: + self.set_metric_tag('correlation_id', + self.request.headers['Correlation-ID']) + @gen.coroutine def get(self): with self.execution_timer('sleepytime'): diff --git a/tests.py b/tests.py index f266789..3b91eed 100644 --- a/tests.py +++ b/tests.py @@ -1,6 +1,7 @@ import logging import socket import time +import uuid from tornado import gen, testing, web import mock @@ -190,3 +191,18 @@ class InfluxDbTests(testing.AsyncHTTPTestCase): response = self.fetch('/') self.assertEqual(response.code, 204) self.assertIs(cfg['db_connection'], conn) + + def test_that_metric_tag_is_tracked(self): + cid = str(uuid.uuid4()) + response = self.fetch('/', headers={'Correlation-ID': cid}) + self.assertEqual(response.code, 204) + + for key, fields, timestamp in self.influx_messages: + if key.startswith('my-service,'): + tag_dict = dict(a.split('=') for a in key.split(',')[1:]) + self.assertEqual(tag_dict['correlation_id'], + '"{}"'.format(cid)) + break + else: + self.fail('Expected to find "request" metric in {!r}'.format( + list(self.application.influx_db['requests'])))