FEATURE: Add ActiveRecord 6.1 instrumentation support (#144)

FEATURE: Add ActiveRecord 6.1 instrumentation support (#144)

diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 61dc12f..dcc2d26 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -28,9 +28,13 @@ jobs:
             ${{ runner.os }}-${{ matrix.ruby }}-gems-
       - name: Setup gems
         run: |
+          gem install bundle
+          # for Ruby <= 2.6 , details https://github.com/rubygems/rubygems/issues/3284
+          gem update --system 3.0.8 && gem update --system
           bundle config path vendor/bundle
           bundle install --jobs 4
+          bundle exec appraisal install
       - name: Rubocop
         run: bundle exec rubocop
       - name: Run tests
-        run: bundle exec rake
+        run: bundle exec appraisal rake
diff --git a/.gitignore b/.gitignore
index 9b7ea1b..91f0bf0 100644
--- a/.gitignore
+++ b/.gitignore
@@ -7,5 +7,7 @@
 /spec/reports/
 /tmp/
 Gemfile.lock
+/gemfiles/*.gemfile.lock
+
 
 .rubocop-https---raw-githubusercontent-com-discourse-discourse-master--rubocop-yml
diff --git a/.rubocop.yml b/.rubocop.yml
index d46296c..3f1b32d 100644
--- a/.rubocop.yml
+++ b/.rubocop.yml
@@ -1,2 +1,7 @@
 inherit_gem:
   rubocop-discourse: default.yml
+
+AllCops:
+  Exclude:
+    - 'gemfiles/**/*'
+    - 'vendor/**/*'
\ No newline at end of file
diff --git a/Appraisals b/Appraisals
new file mode 100644
index 0000000..e322ffd
--- /dev/null
+++ b/Appraisals
@@ -0,0 +1,10 @@
+# frozen_string_literal: true
+
+appraise "ar-60" do
+  # we are using this version as default in gemspec
+  # gem "activerecord", "~> 6.0.0"
+end
+
+appraise "ar-61" do
+  gem "activerecord", "~> 6.1.0.rc2"
+end
diff --git a/gemfiles/.bundle/config b/gemfiles/.bundle/config
new file mode 100644
index 0000000..c127f80
--- /dev/null
+++ b/gemfiles/.bundle/config
@@ -0,0 +1,2 @@
+---
+BUNDLE_RETRY: "1"
diff --git a/gemfiles/ar_60.gemfile b/gemfiles/ar_60.gemfile
new file mode 100644
index 0000000..095e660
--- /dev/null
+++ b/gemfiles/ar_60.gemfile
@@ -0,0 +1,5 @@
+# This file was generated by Appraisal
+
+source "https://rubygems.org"
+
+gemspec path: "../"
diff --git a/gemfiles/ar_61.gemfile b/gemfiles/ar_61.gemfile
new file mode 100644
index 0000000..3b665c0
--- /dev/null
+++ b/gemfiles/ar_61.gemfile
@@ -0,0 +1,7 @@
+# This file was generated by Appraisal
+
+source "https://rubygems.org"
+
+gem "activerecord", "~> 6.1.0.rc2"
+
+gemspec path: "../"
diff --git a/lib/prometheus_exporter/instrumentation/active_record.rb b/lib/prometheus_exporter/instrumentation/active_record.rb
index 03b21df..8d4201a 100644
--- a/lib/prometheus_exporter/instrumentation/active_record.rb
+++ b/lib/prometheus_exporter/instrumentation/active_record.rb
@@ -67,21 +67,30 @@ module PrometheusExporter::Instrumentation
       ObjectSpace.each_object(::ActiveRecord::ConnectionAdapters::ConnectionPool) do |pool|
         next if pool.connections.nil?
 
-        labels_from_config = pool.spec.config
-          .select { |k, v| @config_labels.include? k }
-          .map { |k, v| [k.to_s.dup.prepend("dbconfig_"), v] }
-
-        labels = @metric_labels.merge(pool_name: pool.spec.name).merge(Hash[labels_from_config])
-
         metric = {
           pid: pid,
           type: "active_record",
           hostname: ::PrometheusExporter.hostname,
-          metric_labels: labels
+          metric_labels: labels(pool)
         }
         metric.merge!(pool.stat)
         metrics << metric
       end
     end
+
+    private
+
+    def labels(pool)
+      if pool.respond_to?(:spec) # ActiveRecord <= 6.0
+        @metric_labels.merge(pool_name: pool.spec.name).merge(pool.spec.config
+          .select { |k, v| @config_labels.include? k }
+          .map { |k, v| [k.to_s.dup.prepend("dbconfig_"), v] }.to_h)
+      elsif pool.respond_to?(:db_config) # ActiveRecord >= 6.1.rc1
+        @metric_labels.merge(pool_name: pool.db_config.name).merge(
+          @config_labels.each_with_object({}) { |l, acc| acc["dbconfig_#{l}"] = pool.db_config.public_send(l) })
+      else
+        raise "Unsupported connection pool"
+      end
+    end
   end
 end
diff --git a/prometheus_exporter.gemspec b/prometheus_exporter.gemspec
index 20bc60f..9bb0d26 100644
--- a/prometheus_exporter.gemspec
+++ b/prometheus_exporter.gemspec
@@ -35,6 +35,8 @@ Gem::Specification.new do |spec|
   spec.add_development_dependency "rack-test", "~> 0.8.3"
   spec.add_development_dependency "minitest-stub-const", "~> 0.6"
   spec.add_development_dependency "rubocop-discourse", ">2"
+  spec.add_development_dependency "appraisal", "~> 2.3"
+  spec.add_development_dependency "activerecord", "~> 6.0.0"
   if !RUBY_ENGINE == 'jruby'
     spec.add_development_dependency "raindrops", "~> 0.19"
   end
diff --git a/test/instrumentation/active_record_test.rb b/test/instrumentation/active_record_test.rb
new file mode 100644
index 0000000..46f6246
--- /dev/null
+++ b/test/instrumentation/active_record_test.rb
@@ -0,0 +1,62 @@
+# frozen_string_literal: true
+
+require 'test_helper'
+require 'prometheus_exporter/instrumentation'
+require 'active_record'
+
+class PrometheusInstrumentationActiveRecordTest < Minitest::Test
+
+  def setup
+    super
+
+    # With this trick this variable with be accessible with ::ObjectSpace
+    @pool = if active_record_version >= Gem::Version.create('6.1.0.rc1')
+      active_record61_pool
+    elsif active_record_version >= Gem::Version.create('6.0.0')
+      active_record60_pool
+    else
+      raise 'unsupported active_record version'
+    end
+  end
+
+  def metric_labels
+    { foo: :bar }
+  end
+
+  def config_labels
+    %i[database username]
+  end
+
+  def collector
+    @collector ||= PrometheusExporter::Instrumentation::ActiveRecord.new(metric_labels, config_labels)
+  end
+
+  %i[size connections busy dead idle waiting checkout_timeout type metric_labels].each do |key|
+    define_method("test_collecting_metrics_contain_#{key}_key") do
+      assert_includes collector.collect.first, key
+    end
+  end
+
+  def test_metrics_labels
+    assert_includes collector.collect.first[:metric_labels], :foo
+  end
+
+  def test_type
+    assert_equal collector.collect.first[:type], 'active_record'
+  end
+
+  private
+
+  def active_record_version
+    Gem.loaded_specs["activerecord"].version
+  end
+
+  def active_record60_pool
+    ::ActiveRecord::ConnectionAdapters::ConnectionPool.new(OpenStruct.new(config: {}))
+  end
+
+  def active_record61_pool
+    ::ActiveRecord::ConnectionAdapters::ConnectionPool.new(
+      OpenStruct.new(db_config: OpenStruct.new(checkout_timeout: 0, idle_timeout: 0, pool: 5)))
+  end
+end

GitHub sha: bb3a0f7c

This commit appears in #144 which was merged by SamSaffron.