From 35221dcdb6f7e029bb2bc74304bb3263b023b60c Mon Sep 17 00:00:00 2001
From: Garner Fox McCloud <garner@datadoghq.com>
Date: Tue, 28 Jul 2015 19:37:15 +0000
Subject: [PATCH] [mysql] Added logging and testing for custom queries

---
 checks.d/mysql.py                      | 33 ++++++++++++++++++--------
 ci/mysql.rb                            |  5 ++++
 tests/checks/integration/test_mysql.py | 20 +++++++++++++++-
 3 files changed, 47 insertions(+), 11 deletions(-)

diff --git a/checks.d/mysql.py b/checks.d/mysql.py
index 043a8cea18..f3e3209747 100644
--- a/checks.d/mysql.py
+++ b/checks.d/mysql.py
@@ -1,17 +1,17 @@
 # stdlib
-import os
-import re
 import subprocess
+import os
 import sys
+import re
 import traceback
 
-# 3rd party
-import pymysql
-
 # project
 from checks import AgentCheck
 from utils.platform import Platform
 
+# 3rd party
+import pymysql
+
 GAUGE = "gauge"
 RATE = "rate"
 
@@ -52,6 +52,8 @@
 
 class MySql(AgentCheck):
     SERVICE_CHECK_NAME = 'mysql.can_connect'
+    MAX_CUSTOM_QUERIES = 20
+    DEFAULT_TIMEOUT = 5
 
     def __init__(self, name, init_config, agentConfig, instances=None):
         AgentCheck.__init__(self, name, init_config, agentConfig, instances)
@@ -62,8 +64,10 @@ def get_library_versions(self):
         return {"pymysql": pymysql.__version__}
 
     def check(self, instance):
-        host, port, user, password, mysql_sock, defaults_file, tags, options, queries = self._get_config(instance)
+        host, port, user, password, mysql_sock, defaults_file, tags, options, queries = \
+            self._get_config(instance)
 
+        default_timeout = self.init_config.get('default_timeout', self.DEFAULT_TIMEOUT)
 
         if (not host or not user) and not defaults_file:
             raise Exception("Mysql host and user are needed.")
@@ -192,13 +196,20 @@ def _collect_metrics(self, host, db, tags, options, queries):
                 "SHOW SLAVE STATUS", db, tags=tags
             )
 
-    def _collect_metadata(self, db, host):
-        self._get_version(db, host)
-
+        # Collect custom query metrics
+        # Max of 20 queries allowed
         if isinstance(queries, list):
-            for check in queries:
+            for index, check in enumerate(queries[:self.MAX_CUSTOM_QUERIES]):
                 self._collect_dict(check['type'], {check['field']: check['metric']}, check['query'], db, tags=tags)
 
+            if len(queries) > self.MAX_CUSTOM_QUERIES:
+                self.warning("Maximum number (%s) of custom queries reached.  Skipping the rest.",
+                % self.MAX_CUSTOM_QUERIES)
+
+
+    def _collect_metadata(self, db, host):
+        self._get_version(db, host)
+
     def _rate_or_gauge_statuses(self, statuses, dbResults, tags):
         for status, metric in statuses.iteritems():
             metric_name, metric_type = metric
@@ -292,7 +303,9 @@ def _collect_dict(self, metric_type, field_metric_map, query, db, tags):
                     # cursor.description is a tuple of (column_name, ..., ...)
                     try:
                         col_idx = [d[0].lower() for d in cursor.description].index(field.lower())
+                        self.log.debug("Collecting metric: %s" % metric)
                         if result[col_idx] is not None:
+                            self.log.debug("Collecting done, value %s" % result[col_idx])
                             if metric_type == GAUGE:
                                 self.gauge(metric, float(result[col_idx]), tags=tags)
                             elif metric_type == RATE:
diff --git a/ci/mysql.rb b/ci/mysql.rb
index 2a8e8947f3..b2a87d62f0 100644
--- a/ci/mysql.rb
+++ b/ci/mysql.rb
@@ -10,6 +10,11 @@
 
     task before_script: ['ci:common:before_script'] do
       sh %(mysql -e "create user 'dog'@'localhost' identified by 'dog'" -uroot)
+      sh %(mysql -e "CREATE DATABASE testdb;" -uroot)
+      sh %(mysql -e "CREATE TABLE testdb.users (name VARCHAR(20), age INT);" -uroot)
+      sh %(mysql -e "GRANT SELECT ON testdb.users TO 'dog'@'localhost';" -uroot)
+      sh %(mysql -e "INSERT INTO testdb.users (name,age) VALUES('Alice',25);" -uroot)
+      sh %(mysql -e "INSERT INTO testdb.users (name,age) VALUES('Bob',20);" -uroot)
     end
 
     task script: ['ci:common:script'] do
diff --git a/tests/checks/integration/test_mysql.py b/tests/checks/integration/test_mysql.py
index 2d9e18f2f0..729b4a2c59 100644
--- a/tests/checks/integration/test_mysql.py
+++ b/tests/checks/integration/test_mysql.py
@@ -18,7 +18,21 @@ class TestMySql(AgentCheckTest):
         'user': 'dog',
         'pass': 'dog',
         'options': {'replication': True},
-        'tags': METRIC_TAGS
+        'tags': METRIC_TAGS,
+        'queries': [
+            {
+                'query': "SELECT * from testdb.users where name='Alice' limit 1;",
+                'metric': 'alice.age',
+                'type': 'gauge',
+                'field': 'age'
+            },
+            {
+                'query': "SELECT * from testdb.users where name='Bob' limit 1;",
+                'metric': 'bob.age',
+                'type': 'gauge',
+                'field': 'age'
+            }
+        ]
     }]
 
     CONNECTION_FAILURE = [{
@@ -117,6 +131,10 @@ def test_check(self):
         # Assert service metadata
         self.assertServiceMetadata(['version'], count=1)
 
+        #test custom query metrics
+        self.assertMetric('alice.age', value=25)
+        self.assertMetric('bob.age', value=20)
+
         # Raises when COVERAGE=true and coverage < 100%
         self.coverage_report()