From 4cba70a604e789a455befde24eee2e1435039f69 Mon Sep 17 00:00:00 2001
From: Scott Geary <scoot@scoot.me.uk>
Date: Tue, 9 Aug 2016 04:03:16 +1200
Subject: [PATCH] [mysql] Allow `connection_timeout` to be set (#2729)

* Allow connection_timeout to be set for pymysql instances

Defaults to the original 'None' behaviour: falling back to OS socket defaults on sock.settimeout()

* Pass in connect_timeout when using default file too
---
 checks.d/mysql.py         | 24 ++++++++++++++++--------
 conf.d/mysql.yaml.example |  1 +
 2 files changed, 17 insertions(+), 8 deletions(-)

diff --git a/checks.d/mysql.py b/checks.d/mysql.py
index 20aabe7ab6..0a0badf08b 100644
--- a/checks.d/mysql.py
+++ b/checks.d/mysql.py
@@ -294,7 +294,7 @@ def check(self, instance):
             procfs_path = self.agentConfig.get('procfs_path', '/proc').rstrip('/')
             psutil.PROCFS_PATH = procfs_path
 
-        host, port, user, password, mysql_sock, defaults_file, tags, options, queries, ssl = \
+        host, port, user, password, mysql_sock, defaults_file, tags, options, queries, ssl, connect_timeout = \
             self._get_config(instance)
 
         self._set_qcache_stats()
@@ -303,7 +303,7 @@ def check(self, instance):
             raise Exception("Mysql host and user are needed.")
 
         with self._connect(host, port, mysql_sock, user,
-                           password, defaults_file, ssl) as db:
+                           password, defaults_file, ssl, connect_timeout) as db:
             try:
                 # Metadata collection
                 self._collect_metadata(db, host)
@@ -330,9 +330,10 @@ def _get_config(self, instance):
         options = instance.get('options', {})
         queries = instance.get('queries', [])
         ssl = instance.get('ssl', {})
+        connect_timeout = instance.get('connect_timeout', None)
 
         return (self.host, self.port, user, password, self.mysql_sock,
-                self.defaults_file, tags, options, queries, ssl)
+                self.defaults_file, tags, options, queries, ssl, connect_timeout)
 
     def _set_qcache_stats(self):
         host_key = self._get_host_key()
@@ -363,7 +364,7 @@ def _get_host_key(self):
         return hostkey
 
     @contextmanager
-    def _connect(self, host, port, mysql_sock, user, password, defaults_file, ssl):
+    def _connect(self, host, port, mysql_sock, user, password, defaults_file, ssl, connect_timeout):
         self.service_check_tags = [
             'server:%s' % (mysql_sock if mysql_sock != '' else host),
             'port:%s' % ('unix_socket' if port == 0 else port)
@@ -374,7 +375,11 @@ def _connect(self, host, port, mysql_sock, user, password, defaults_file, ssl):
             ssl = dict(ssl) if ssl else None
 
             if defaults_file != '':
-                db = pymysql.connect(read_default_file=defaults_file, ssl=ssl)
+                db = pymysql.connect(
+                    read_default_file=defaults_file,
+                    ssl=ssl,
+                    connect_timeout=connect_timeout
+                )
             elif mysql_sock != '':
                 self.service_check_tags = [
                     'server:{0}'.format(mysql_sock),
@@ -383,7 +388,8 @@ def _connect(self, host, port, mysql_sock, user, password, defaults_file, ssl):
                 db = pymysql.connect(
                     unix_socket=mysql_sock,
                     user=user,
-                    passwd=password
+                    passwd=password,
+                    connect_timeout=connect_timeout
                 )
             elif port:
                 db = pymysql.connect(
@@ -391,14 +397,16 @@ def _connect(self, host, port, mysql_sock, user, password, defaults_file, ssl):
                     port=port,
                     user=user,
                     passwd=password,
-                    ssl=ssl
+                    ssl=ssl,
+                    connect_timeout=connect_timeout
                 )
             else:
                 db = pymysql.connect(
                     host=host,
                     user=user,
                     passwd=password,
-                    ssl=ssl
+                    ssl=ssl,
+                    connect_timeout=connect_timeout
                 )
             self.log.debug("Connected to MySQL")
             self.service_check(self.SERVICE_CHECK_NAME, AgentCheck.OK,
diff --git a/conf.d/mysql.yaml.example b/conf.d/mysql.yaml.example
index f4ede9b518..90d04c3820 100644
--- a/conf.d/mysql.yaml.example
+++ b/conf.d/mysql.yaml.example
@@ -7,6 +7,7 @@ instances:
     # port: 3306             # Optional
     # sock: /path/to/sock    # Connect via Unix Socket
     # defaults_file: my.cnf  # Alternate configuration mechanism
+    # connect_timeout: None  # Optional integer seconds
     # tags:                  # Optional
     #   - optional_tag1
     #   - optional_tag2