From 5c913c2aec8e5462d9c8e7480646534008fba849 Mon Sep 17 00:00:00 2001
From: Liang Guo <liang.guo@airbnb.com>
Date: Fri, 8 Jan 2016 11:43:43 -0800
Subject: [PATCH] Add mysql service connectivity check

It's similar to haproxy option mysql-check, and is used for local
mysql service check. A simple mysql-check uses a user with minimal
USAGE privilege and the user may or may not have a password.
---
 lib/nerve/service_watcher.rb       |  1 +
 lib/nerve/service_watcher/mysql.rb | 40 ++++++++++++++++++++++++++++++
 nerve.gemspec                      |  1 +
 3 files changed, 42 insertions(+)
 create mode 100644 lib/nerve/service_watcher/mysql.rb

diff --git a/lib/nerve/service_watcher.rb b/lib/nerve/service_watcher.rb
index 40283cdc..f0fd6ef8 100644
--- a/lib/nerve/service_watcher.rb
+++ b/lib/nerve/service_watcher.rb
@@ -1,6 +1,7 @@
 require 'nerve/service_watcher/tcp'
 require 'nerve/service_watcher/http'
 require 'nerve/service_watcher/rabbitmq'
+require 'nerve/service_watcher/mysql'
 
 module Nerve
   class ServiceWatcher
diff --git a/lib/nerve/service_watcher/mysql.rb b/lib/nerve/service_watcher/mysql.rb
new file mode 100644
index 00000000..525c2143
--- /dev/null
+++ b/lib/nerve/service_watcher/mysql.rb
@@ -0,0 +1,40 @@
+require 'nerve/service_watcher/base'
+
+module Nerve
+  module ServiceCheck
+    class MysqlServiceCheck < BaseServiceCheck
+      require 'mysql2'
+
+      def initialize(opts={})
+        super
+
+        raise ArgumentError, "missing required argument 'user' in mysql check" unless opts['user']
+
+        @user = opts['user']
+        @pass = opts['password'] || ''
+        @host = opts['host'] || '127.0.0.1'
+        @port = opts['port'] || '3306'
+      end
+
+      def check
+        # the idea of health check is similar to haproxy option mysql-check
+        log.debug "nerve: running mysql health check #{@name}"
+
+        begin
+          log.debug "nerve: mysql connect #{@host}:#{@port} as #{@user}"
+          conn = Mysql2::Client.new(:host => @host, :username => @user, :password => @pass, :port => @port)
+        rescue Mysql2::Error => e
+          log.debug "nerve: mysql check error #{e.errno}: #{e.error}"
+          return false
+        ensure
+          conn.close if conn
+        end
+
+        return true
+      end
+    end
+
+    CHECKS ||= {}
+    CHECKS['mysql'] = MysqlServiceCheck
+  end
+end
diff --git a/nerve.gemspec b/nerve.gemspec
index 6b9b0451..e20613ed 100644
--- a/nerve.gemspec
+++ b/nerve.gemspec
@@ -25,6 +25,7 @@ Gem::Specification.new do |gem|
   gem.add_runtime_dependency "zk", "~> 1.9.2"
   gem.add_runtime_dependency "bunny", "= 1.1.0"
   gem.add_runtime_dependency "etcd", "~> 0.2.3"
+  gem.add_runtime_dependency "mysql2", "~> 0.4.2"
 
   gem.add_development_dependency "rake"
   gem.add_development_dependency "rspec", "~> 3.1.0"