From e5bda2a54539ff3908b4ab36e5901e6d8997e6b2 Mon Sep 17 00:00:00 2001
From: "Z.J. van de Weg" <git@zjvandeweg.nl>
Date: Tue, 1 Nov 2016 20:06:36 +0100
Subject: [PATCH] Add ability to query row counts

This commit allows for queries which only request for row counts to be
gathered and written for Prometheus to gather.

At this time 3 simple queries are added which monitor how many projects
are orphaned and/or pending_delete, though the structure of the code
allows for easily adding more queries which do
SELECT COUNT(<whatever>) FROM <whatever>;
---
 lib/gitlab_monitor/cli.rb                | 43 ++++++++++++++
 lib/gitlab_monitor/database.rb           |  1 +
 lib/gitlab_monitor/database/row_count.rb | 76 ++++++++++++++++++++++++
 3 files changed, 120 insertions(+)
 create mode 100644 lib/gitlab_monitor/database/row_count.rb

diff --git a/lib/gitlab_monitor/cli.rb b/lib/gitlab_monitor/cli.rb
index 6f574dd..9fcbe70 100644
--- a/lib/gitlab_monitor/cli.rb
+++ b/lib/gitlab_monitor/cli.rb
@@ -244,6 +244,48 @@ module GitLab
         end
       end
 
+      # Database row counts query runner.
+      #
+      # This will take the database connection and print the result to STDOUT
+      class DatabaseRowCounts
+        COMMAND_NAME = "row-counts".freeze
+
+        def initialize(args)
+          @options = options(args)
+          @options.parse!
+
+          @target = args.shift || STDOUT
+          @target = File.open(@target, "a") if @target.is_a?(String)
+        end
+
+        def options(args)
+          args.options do |opts|
+            opts.banner = "Usage: #{EXECUTABLE_NAME} #{COMMAND_NAME} [options]"
+            opts.on("--db-conn=\"dbname=test port=5432\"", "Database connection string") do |val|
+              @db_connection_string = val
+            end
+          end
+        end
+
+        def help
+          @options.help
+        end
+
+        def run
+          validate!
+
+          ::GitLab::Monitor::Database::RowCountProber.new(connection_string: @db_connection_string)
+                                                     .probe_db
+                                                     .write_to(@target)
+        end
+
+        private
+
+        def validate!
+          fail InvalidCLICommand.new(help) unless @db_connection_string
+        end
+      end
+
       # Run a web server that exposes the metrics specified in a config file
       class Server
         COMMAND_NAME = "web".freeze
@@ -378,6 +420,7 @@ module GitLab
           DatabaseBlockedQueries,
           DatabaseSlowQueries,
           DatabaseVacuumQueries,
+          DatabaseRowCounts,
           Process,
           SidekiqRunner,
           Server
diff --git a/lib/gitlab_monitor/database.rb b/lib/gitlab_monitor/database.rb
index 17080b2..934cd83 100644
--- a/lib/gitlab_monitor/database.rb
+++ b/lib/gitlab_monitor/database.rb
@@ -7,6 +7,7 @@ module GitLab
       autoload :SlowQueriesProber,    "gitlab_monitor/database/slow_queries"
       autoload :DeadTuplesProber,     "gitlab_monitor/database/dead_tuples"
       autoload :VacuumQueriesProber,  "gitlab_monitor/database/vacuum_queries"
+      autoload :RowCountProber,       "gitlab_monitor/database/row_count"
     end
   end
 end
diff --git a/lib/gitlab_monitor/database/row_count.rb b/lib/gitlab_monitor/database/row_count.rb
new file mode 100644
index 0000000..6896d46
--- /dev/null
+++ b/lib/gitlab_monitor/database/row_count.rb
@@ -0,0 +1,76 @@
+module GitLab
+  module Monitor
+    module Database
+      # A helper class that executes the query its given and returns an int of
+      # the row count
+      # This class works under the assumption you do COUNT(*) queries
+      class RowCountCollector < Base
+        QUERIES = [
+          :orphaned_projects,
+          :soft_deleted_projects,
+          :orphaned_soft_deleted_projects
+        ].freeze
+
+        def run
+          results = Hash.new(0)
+
+          QUERIES.each do |query|
+            results[query] = execute(query)
+          end
+
+          results
+        end
+
+        private
+
+        def execute(query)
+          connection.exec(send(query))[0]["count"]
+        end
+
+        def orphaned_projects
+          <<-SQL
+            SELECT COUNT(*) FROM projects
+            WHERE namespace_id IS NULL;
+          SQL
+        end
+
+        def soft_deleted_projects
+          <<-SQL
+            SELECT COUNT(*) FROM projects
+            WHERE pending_delete=true;
+          SQL
+        end
+
+        def orphaned_soft_deleted_projects
+          <<-SQL
+            SELECT COUNT(*) FROM projects
+            WHERE pending_delete=true AND namespace_id IS NULL;
+          SQL
+        end
+      end
+
+      # The prober which is called when gathering metrics
+      class RowCountProber
+        def initialize(opts, metrics: PrometheusMetrics.new)
+          @metrics = metrics
+          @collector = RowCountCollector.new(connection_string: opts[:connection_string])
+        end
+
+        def probe_db
+          return self unless @collector.connected?
+
+          results = @collector.run
+          results.each do |key, value|
+            @metrics.add("#{key}_count", value.to_i)
+          end
+
+          self
+        end
+
+        def write_to(target)
+          target.write(@metrics.to_s)
+        end
+      end
+    end
+  end
+end
-- 
GitLab