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

Allow RowCountProber to execute selected queries

parent 29cce2af
...@@ -46,6 +46,12 @@ probes: ...@@ -46,6 +46,12 @@ probes:
rows_count: rows_count:
class_name: Database::RowCountProber class_name: Database::RowCountProber
<<: *db_common <<: *db_common
opts:
<<: *db_common_opts
selected_queries:
- soft_deleted_projects
- orphaned_projects
- uploads
process: process:
methods: methods:
......
require "set"
module GitLab module GitLab
module Monitor module Monitor
module Database module Database
...@@ -25,11 +27,18 @@ module GitLab ...@@ -25,11 +27,18 @@ module GitLab
uploads: { select: :uploads } uploads: { select: :uploads }
}.freeze }.freeze
def initialize(args)
super(args)
@selected_queries = Set.new(args[:selected_queries].map(&:to_sym)) unless args[:selected_queries].nil?
end
def run def run
results = Hash.new(0) results = Hash.new(0)
QUERIES.each do |key, query_hash| QUERIES.each do |key, query_hash|
next if query_hash[:check] && !successful_check?(query_hash[:check]) 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) results[key] = count_from_query_hash(query_hash)
end end
...@@ -74,7 +83,10 @@ module GitLab ...@@ -74,7 +83,10 @@ module GitLab
class RowCountProber class RowCountProber
def initialize(opts, metrics: PrometheusMetrics.new) def initialize(opts, metrics: PrometheusMetrics.new)
@metrics = metrics @metrics = metrics
@collector = RowCountCollector.new(connection_string: opts[:connection_string]) @collector = RowCountCollector.new(
connection_string: opts[:connection_string],
selected_queries: opts[:selected_queries]
)
end end
def probe_db def probe_db
......
...@@ -2,18 +2,30 @@ require "spec_helper" ...@@ -2,18 +2,30 @@ require "spec_helper"
require "gitlab_monitor/database/row_count" require "gitlab_monitor/database/row_count"
describe GitLab::Monitor::Database::RowCountCollector do 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") } let(:collector) { described_class.new(connection_string: "host=localhost") }
describe "#run" do describe "#run" do
before do before do
stub_const("GitLab::Monitor::Database::RowCountCollector::QUERIES", query) 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_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
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