Support optional health check queries in MultiActiveRecordCheck

lib/easymon/checks/multi_active_record_check.rb

@@ -1,6 +1,7 @@
 module Easymon
   class MultiActiveRecordCheck
     attr_accessor :block
+    attr_accessor :query
     attr_accessor :results
 
     # Here we pass a block so we get a fresh instance of ActiveRecord::Base or
@@ -26,25 +27,50 @@ class MultiActiveRecordCheck
     #    }
     # }
     # Easymon::Repository.add("multi-database", check)
+    #
+    # Optionally, a health check query can be given instead of relying on
+    # connection.active?. The query is run against each connection and must
+    # return a single truthy/falsey value (e.g. 1/0 from MySQL):
+    # check = Easymon::MultiActiveRecordCheck.new(query: "SELECT @@read_only") {
+    #    { "Primary": ActiveRecord::Base.connection }
+    # }
+    #
+    # A per-connection query can also be given by pairing the connection with
+    # its query. It takes precedence over the default query: option above:
+    # check = Easymon::MultiActiveRecordCheck.new {
+    #    {
+    #      "Primary": [ActiveRecord::Base.connection, "SELECT @@read_only"],
+    #      "PrimaryReplica": [Easymon::PrimaryReplica.connection, "SELECT TIMESTAMPDIFF(MICROSECOND, MAX(ts), NOW(6)) / 1000000 > 1 FROM percona.heartbeat"],
+    #      "OtherReplica": Easymon::OtherReplica.connection # plain connection.active?
+    #    }
+    # }
 
-    def initialize(&block)
+    def initialize(query: nil, &block)
+      self.query = query
       self.block = block
     end
 
     def check
       connections = Hash(@block.call)
 
-      results = connections.transform_values { |connection| database_up?(connection) }
+      results = connections.transform_values do |value|
+        connection, connection_query = value.is_a?(Array) ? value : [ value, nil ]
+        database_up?(connection, query: connection_query || query)
+      end
 
       status = results.map { |db_name, result| "#{db_name}: #{result ? 'Up' : 'Down'}" }.join(" - ")
 
       [ (results.any? && results.values.all?), status ]
     end
 
     private
-      def database_up?(connection)
+      def database_up?(connection, query: nil)
         connection.connect!
-        connection.active?
+        if query
+          ActiveModel::Type::Boolean.new.cast(connection.select_value(query)) || false
+        else
+          connection.active?
+        end
       rescue
         false
       end

test/unit/checks/multi_active_record_check_test.rb

@@ -53,6 +53,100 @@ class MultiActiveRecordCheckTest < ActiveSupport::TestCase
     assert_equal(false, results[0])
   end
 
+  test "#run uses the health check query when given" do
+    primary = ActiveRecord::Base.connection
+    primary_replica = Easymon::PrimaryReplica.connection
+
+    check = Easymon::MultiActiveRecordCheck.new(query: "SELECT 1") do
+      { "Primary": primary, "PrimaryReplica": primary_replica }
+    end
+    results = check.check
+
+    assert_equal("Primary: Up - PrimaryReplica: Up", results[1])
+    assert_equal(true, results[0])
+  end
+
+  test "#run sets failed conditions when the health check query returns false" do
+    primary = ActiveRecord::Base.connection
+    primary_replica = Easymon::PrimaryReplica.connection
+
+    primary_replica.stubs(:select_value).with("SELECT 1").returns(0)
+
+    check = Easymon::MultiActiveRecordCheck.new(query: "SELECT 1") do
+      { "Primary": primary, "PrimaryReplica": primary_replica }
+    end
+    results = check.check
+
+    assert_equal("Primary: Up - PrimaryReplica: Down", results[1])
+    assert_equal(false, results[0])
+  end
+
+  test "#run sets failed conditions when the health check query returns null" do
+    primary = ActiveRecord::Base.connection
+    primary_replica = Easymon::PrimaryReplica.connection
+
+    check = Easymon::MultiActiveRecordCheck.new(query: "SELECT NULL") do
+      { "Primary": primary, "PrimaryReplica": primary_replica }
+    end
+    results = check.check
+
+    assert_equal("Primary: Down - PrimaryReplica: Down", results[1])
+    assert_equal(false, results[0])
+  end
+
+  test "#run sets failed conditions when the health check query raises" do
+    primary = ActiveRecord::Base.connection
+    primary_replica = Easymon::PrimaryReplica.connection
+
+    primary_replica.stubs(:select_value).raises("boom")
+
+    check = Easymon::MultiActiveRecordCheck.new(query: "SELECT 1") do
+      { "Primary": primary, "PrimaryReplica": primary_replica }
+    end
+    results = check.check
+
+    assert_equal("Primary: Up - PrimaryReplica: Down", results[1])
+    assert_equal(false, results[0])
+  end
+
+  test "#run uses per-connection queries when given" do
+    primary = ActiveRecord::Base.connection
+    primary_replica = Easymon::PrimaryReplica.connection
+    other_replica = Easymon::OtherReplica.connection
+
+    primary_replica.stubs(:select_value).with("SELECT 0").returns(0)
+
+    check = Easymon::MultiActiveRecordCheck.new do
+      {
+        "Primary": [ primary, "SELECT 1" ],
+        "PrimaryReplica": [ primary_replica, "SELECT 0" ],
+        "OtherReplica": other_replica
+      }
+    end
+    results = check.check
+
+    assert_equal("Primary: Up - PrimaryReplica: Down - OtherReplica: Up", results[1])
+    assert_equal(false, results[0])
+  end
+
+  test "#run prefers the per-connection query over the default query" do
+    primary = ActiveRecord::Base.connection
+    primary_replica = Easymon::PrimaryReplica.connection
+
+    primary_replica.stubs(:select_value).with("SELECT 0").returns(0)
+
+    check = Easymon::MultiActiveRecordCheck.new(query: "SELECT 1") do
+      {
+        "Primary": primary,
+        "PrimaryReplica": [ primary_replica, "SELECT 0" ]
+      }
+    end
+    results = check.check
+
+    assert_equal("Primary: Up - PrimaryReplica: Down", results[1])
+    assert_equal(false, results[0])
+  end
+
   test "given nil as a config" do
     check = Easymon::MultiActiveRecordCheck.new { }
     results = check.check