From 355cce4aa87afbc688cc806d75ec0484be76e04a Mon Sep 17 00:00:00 2001
From: Xavier Vello <xavier.vello@datadoghq.com>
Date: Wed, 2 Aug 2017 16:30:55 +0200
Subject: [PATCH] add type_overrides member to PrometheusCheck class

Allows to override metric type if the provided type is invalid or absent
---
 checks/prometheus_check.py | 24 ++++++++++++++++++++++++
 1 file changed, 24 insertions(+)

diff --git a/checks/prometheus_check.py b/checks/prometheus_check.py
index 20b8d70149..e4335a5e41 100644
--- a/checks/prometheus_check.py
+++ b/checks/prometheus_check.py
@@ -65,6 +65,13 @@ def __init__(self, name, init_config, agentConfig, instances=None):
         # will just not be added as tags when submitting the metric.
         self.exclude_labels = []
 
+        # `type_overrides` is a dictionnary where the keys are prometheus metric names
+        # and the values are a metric type (name as string) to use instead of the one
+        # listed in the payload. It can be used to force a type on untyped metrics.
+        # Note: it is empty in the mother class but will need to be
+        # overloaded/hardcoded in the final check not to be counted as custom metric.
+        self.type_overrides = {}
+
     def check(self, instance):
         """
         check should take care of getting the url and other params
@@ -102,6 +109,14 @@ def parse_metric_family(self, buf, content_type):
 
                 message = metrics_pb2.MetricFamily()
                 message.ParseFromString(msg_buf)
+
+                # Lookup type overrides:
+                if self.type_overrides and message.name in self.type_overrides:
+                    new_type = self.type_overrides[message.name]
+                    if new_type in self.METRIC_TYPES:
+                        message.type = self.METRIC_TYPES.index(new_type)
+                    else:
+                        self.log.debug("type override %s for %s is not a valid type name" % (new_type, message.name))
                 yield message
         elif 'text/plain' in content_type:
             messages = {}  # map with the name of the element (before the labels) and the list of occurrences with labels and values
@@ -109,6 +124,15 @@ def parse_metric_family(self, buf, content_type):
             obj_help = {}  # help for the metrics
             for line in buf.splitlines():
                 self._extract_metrics_from_string(line, messages, obj_map, obj_help)
+
+            # Add type overrides:
+            for m_name, m_type in self.type_overrides.iteritems():
+                if m_type in self.METRIC_TYPES:
+                    obj_map[m_name] = m_type
+                else:
+                    self.log.debug("type override %s for %s is not a valid type name" % (m_type,m_name))
+
+
             for _m in obj_map:
                 if _m in messages or (obj_map[_m] == 'histogram' and '{}_bucket'.format(_m) in messages):
                     yield self._extract_metric_from_map(_m, messages, obj_map, obj_help)