Merge pull request #2828 from alphagov/push-database-metrics-to-prome…

…theus

Add a rake task to push database metrics to prometheus

Gemfile

@@ -19,6 +19,7 @@ gem "json-schema", require: false
 gem "oj"
 gem "pg"
 gem "plek"
+gem "prometheus-client"
 gem "sentry-sidekiq"
 gem "sidekiq-unique-jobs"
 gem "whenever", require: false

Gemfile.lock

@@ -569,6 +569,8 @@ GEM
     parslet (2.0.0)
     pg (1.5.7)
     plek (5.2.0)
+    prometheus-client (4.2.3)
+      base64
     prometheus_exporter (2.1.1)
       webrick
     psych (5.1.2)
@@ -823,6 +825,7 @@ DEPENDENCIES
   pact_broker-client
   pg
   plek
+  prometheus-client
   rails (= 7.1.3.4)
   rspec-rails
   rubocop-govuk

lib/tasks/metrics.rake

@@ -0,0 +1,46 @@
+require "prometheus/client"
+require "prometheus/client/push"
+require "prometheus/client/registry"
+require "prometheus/client/formats/text"
+
+namespace :metrics do
+  desc "Reports metrics about data in the database to prometheus via the pushgateway"
+  task report_to_prometheus: :environment do
+    default_labels = { database: "publishing-api" }
+    edition_count_dimensions = %i[state content_store document_type publishing_app locale]
+
+    prometheus_registry = Prometheus::Client::Registry.new
+    editions_in_database_gauge = prometheus_registry.gauge(
+      :editions_in_database_total,
+      docstring: "Count of editions in various databases labeled by state, document_type etc.",
+      labels: default_labels.keys + edition_count_dimensions,
+    )
+
+    edition_counts = Edition
+      .with_document
+      .where.not(state: "superseded")
+      .where.not(content_store: nil)
+      .group(*edition_count_dimensions)
+      .count
+
+    edition_counts.sort.each do |dimensions, count|
+      labels = default_labels.merge(edition_count_dimensions.zip(dimensions).to_h)
+      editions_in_database_gauge.set(count, labels:)
+    end
+
+    puts "Found #{edition_counts.count} combinations of labels"
+    puts Prometheus::Client::Formats::Text.marshal(prometheus_registry)
+
+    pushgateway_url = ENV["PROMETHEUS_PUSHGATEWAY_URL"]
+    if pushgateway_url.present?
+      puts "Pushing metrics to prometheus via #{pushgateway_url}"
+      Prometheus::Client::Push.new(
+        job: "publishing-api-metrics",
+        gateway: pushgateway_url,
+      ).add(prometheus_registry)
+    end
+  rescue StandardError => e
+    warn e.inspect
+    raise e
+  end
+end

spec/lib/tasks/metrics_spec.rb

@@ -0,0 +1,32 @@
+RSpec.describe "Metrics rake tasks" do
+  describe "metrics:report_to_prometheus" do
+    let(:task) { Rake::Task["metrics:report_to_prometheus"] }
+    before { task.reenable }
+
+    it "outputs empty metrics when there are no editions" do
+      expect { task.invoke }.to output(<<~PROMETHEUS).to_stdout
+        Found 0 combinations of labels
+        # TYPE editions_in_database_total gauge
+        # HELP editions_in_database_total Count of editions in various databases labeled by state, document_type etc.
+      PROMETHEUS
+    end
+
+    it "outputs edition count metrics" do
+      5.times do
+        create(:draft_edition, document_type: "press_release", publishing_app: "whitehall")
+      end
+
+      10.times do
+        create(:live_edition, document_type: "services_and_information", publishing_app: "publisher")
+      end
+
+      expect { task.invoke }.to output(<<~PROMETHEUS).to_stdout
+        Found 2 combinations of labels
+        # TYPE editions_in_database_total gauge
+        # HELP editions_in_database_total Count of editions in various databases labeled by state, document_type etc.
+        editions_in_database_total{database="publishing-api",state="draft",content_store="draft",document_type="press_release",publishing_app="whitehall",locale="en"} 5.0
+        editions_in_database_total{database="publishing-api",state="published",content_store="live",document_type="services_and_information",publishing_app="publisher",locale="en"} 10.0
+      PROMETHEUS
+    end
+  end
+end