Allow RowCountProber to execute selected queries

config/gitlab-monitor.yml.example

@@ -46,6 +46,12 @@ probes:
     rows_count:
       class_name: Database::RowCountProber
       <<: *db_common
+      opts:
+        <<: *db_common_opts
+        selected_queries:
+        - soft_deleted_projects
+        - orphaned_projects
+        - uploads
 
   process:
     methods:

lib/gitlab_monitor/database/row_count.rb

+require "set"
+
 module GitLab
   module Monitor
     module Database
@@ -25,11 +27,18 @@ module GitLab
           uploads: { select: :uploads }
         }.freeze
 
+        def initialize(args)
+          super(args)
+
+          @selected_queries = Set.new(args[:selected_queries].map(&:to_sym)) unless args[:selected_queries].nil?
+        end
+
         def run
           results = Hash.new(0)
 
           QUERIES.each do |key, query_hash|
             next if query_hash[:check] && !successful_check?(query_hash[:check])
+            next if !@selected_queries.nil? && !@selected_queries.include?(key)
 
             results[key] = count_from_query_hash(query_hash)
           end
@@ -74,7 +83,10 @@ module GitLab
       class RowCountProber
         def initialize(opts, metrics: PrometheusMetrics.new)
           @metrics = metrics
-          @collector = RowCountCollector.new(connection_string: opts[:connection_string])
+          @collector = RowCountCollector.new(
+            connection_string: opts[:connection_string],
+            selected_queries: opts[:selected_queries]
+          )
         end
 
         def probe_db

spec/database/row_count_spec.rb

@@ -2,18 +2,30 @@ require "spec_helper"
 require "gitlab_monitor/database/row_count"
 
 describe GitLab::Monitor::Database::RowCountCollector do
-  let(:query)     { { project_1: { select: :projects, where: "id=1" } } }
+  let(:query) {
+    { project_1: { select: :projects, where: "id=1" },
+      project_2: { select: :projects, where: "id=2" } }
+  }
   let(:collector) { described_class.new(connection_string: "host=localhost") }
 
   describe "#run" do
     before do
       stub_const("GitLab::Monitor::Database::RowCountCollector::QUERIES", query)
-    end
 
-    it "executes the query" do
       allow(collector).to receive(:count_from_query_hash).with(query[:project_1]).and_return(3)
+      allow(collector).to receive(:count_from_query_hash).with(query[:project_2]).and_return(6)
+    end
+
+    it "executes all the queries" do
+      expect(collector.run).to eq(project_1: 3, project_2: 6)
+    end
+
+    context "when selected_queries is passed" do
+      let(:collector) { described_class.new(connection_string: "host=localhost", selected_queries: ["project_2"]) }
 
-      expect(collector.run).to eq(project_1: 3)
+      it "executes the selected queries" do
+        expect(collector.run).to eq(project_2: 6)
+      end
     end
   end