FEATURE: Custom labels support in Activerecord Connection Pool and Sidekiq Queue (#119)

FEATURE: Custom labels support in Activerecord Connection Pool and Sidekiq Queue (#119)

collectors.

Closes #118.

diff --git a/lib/prometheus_exporter/server/active_record_collector.rb b/lib/prometheus_exporter/server/active_record_collector.rb
index 9cdb9ac..4f9a228 100644
--- a/lib/prometheus_exporter/server/active_record_collector.rb
+++ b/lib/prometheus_exporter/server/active_record_collector.rb
@@ -27,6 +27,7 @@ module PrometheusExporter::Server
 
       @active_record_metrics.map do |m|
         metric_key = (m["metric_labels"] || {}).merge("pid" => m["pid"])
+        metric_key.merge!(m["custom_labels"]) if m["custom_labels"]
 
         ACTIVE_RECORD_GAUGES.map do |k, help|
           k = k.to_s
diff --git a/lib/prometheus_exporter/server/sidekiq_queue_collector.rb b/lib/prometheus_exporter/server/sidekiq_queue_collector.rb
index c2be01f..81134f2 100644
--- a/lib/prometheus_exporter/server/sidekiq_queue_collector.rb
+++ b/lib/prometheus_exporter/server/sidekiq_queue_collector.rb
@@ -37,6 +37,7 @@ module PrometheusExporter::Server
       now = ::Process.clock_gettime(::Process::CLOCK_MONOTONIC)
       object['queues'].each do |queue|
         queue["created_at"] = now
+        queue["labels"].merge!(object['custom_labels']) if object['custom_labels']
         sidekiq_metrics.delete_if { |metric| metric['created_at'] + MAX_SIDEKIQ_METRIC_AGE < now }
         sidekiq_metrics << queue
       end
diff --git a/test/server/active_record_collector_test.rb b/test/server/active_record_collector_test.rb
index e3284b2..91e13c6 100644
--- a/test/server/active_record_collector_test.rb
+++ b/test/server/active_record_collector_test.rb
@@ -45,4 +45,28 @@ class PrometheusActiveRecordCollectorTest < Minitest::Test
     assert_equal 6, metrics.size
     assert(metrics.first.metric_text.include?('active_record_connection_pool_connections{service="service1",pid="1000"} 50'))
   end
+
+  def test_collecting_metrics_with_client_default_labels
+
+    collector.collect(
+      "type" => "active_record",
+      "pid" => "1000",
+      "connections" => 50,
+      "busy" => 20,
+      "dead" => 10,
+      "idle" => 20,
+      "waiting" => 0,
+      "size" => 120,
+      "metric_labels" => {
+        "service" => "service1"
+      },
+      "custom_labels" => {
+        "environment" => "test"
+      }
+    )
+
+    metrics = collector.metrics
+    assert_equal 6, metrics.size
+    assert(metrics.first.metric_text.include?('active_record_connection_pool_connections{service="service1",pid="1000",environment="test"} 50'))
+  end
 end
diff --git a/test/server/sidekiq_queue_collector_test.rb b/test/server/sidekiq_queue_collector_test.rb
index fc783a6..703980a 100644
--- a/test/server/sidekiq_queue_collector_test.rb
+++ b/test/server/sidekiq_queue_collector_test.rb
@@ -31,6 +31,27 @@ class PrometheusSidekiqQueueCollectorTest < Minitest::Test
     assert_equal expected, metrics.map(&:metric_text)
   end
 
+  def test_collecting_metrics_with_client_default_labels
+    collector.collect(
+      'queues' => [
+        'backlog_total' => 16,
+        'latency_seconds' => 7,
+        'labels' => { 'queue' => 'default' }
+      ],
+      'custom_labels' => {
+        'environment' => 'test'
+      }
+    )
+
+    metrics = collector.metrics
+
+    expected = [
+      'sidekiq_queue_backlog_total{queue="default",environment="test"} 16',
+      'sidekiq_queue_latency_seconds{queue="default",environment="test"} 7',
+    ]
+    assert_equal expected, metrics.map(&:metric_text)
+  end
+
   def test_only_fresh_metrics_are_collected
     Process.stub(:clock_gettime, 1.0) do
       collector.collect(

GitHub sha: 4cd634e7