Commit 329d7e76 authored by Ahmad Sherif's avatar Ahmad Sherif

Allow RowCountProber to execute selected queries

parent 29cce2af
......@@ -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:
......
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
......
......@@ -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
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment