Reorganize database-related classes

lib/gitlab_monitor.rb

@@ -6,7 +6,6 @@ module GitLab
     autoload :PrometheusMetrics, "gitlab_monitor/prometheus"
     autoload :Git,               "gitlab_monitor/git"
     autoload :GitProber,         "gitlab_monitor/git"
-    autoload :DatabaseDeadTuplesCollector, "gitlab_monitor/database_dead_tuples"
-    autoload :DatabaseDeadTuplesProber,    "gitlab_monitor/database_dead_tuples"
+    autoload :Database,          "gitlab_monitor/database"
   end
 end

lib/gitlab_monitor/cli.rb

@@ -104,9 +104,9 @@ module GitLab
         def run
           validate!
 
-          ::GitLab::Monitor::DatabaseDeadTuplesProber.new(connection_string: @db_connection_string)
-                                                     .probe_db
-                                                     .write_to(@target)
+          ::GitLab::Monitor::Database::DeadTuplesProber.new(connection_string: @db_connection_string)
+                                                       .probe_db
+                                                       .write_to(@target)
         end
 
         private

lib/gitlab_monitor/database.rb

+module GitLab
+  module Monitor
+    # Database-related classes
+    module Database
+      autoload :Base,                 "gitlab_monitor/database/base"
+      autoload :DeadTuplesProber,     "gitlab_monitor/database/dead_tuples"
+    end
+  end
+end

lib/gitlab_monitor/database/base.rb

+require "pg"
+
+module GitLab
+  module Monitor
+    module Database
+      # An abstract class for interacting with DB
+      #
+      # It takes a connection string (e.g. "dbname=test port=5432")
+      class Base
+        def initialize(args)
+          @connection_string = args[:connection_string]
+        end
+
+        def run
+          fail NotImplemented
+        end
+
+        def connected?
+          !connection.nil?
+        end
+
+        private
+
+        def connection
+          @connection ||= PG.connect(@connection_string)
+        rescue PG::ConnectionBad # rubocop:disable Lint/HandleExceptions
+          # Do nothing, we could be on the slave machine
+        end
+      end
+    end
+  end
+end

lib/gitlab_monitor/database/dead_tuples.rb

+module GitLab
+  module Monitor
+    module Database
+      # A helper class to collect dead tuples stats from the database
+      #
+      # It takes a connection string (e.g. "dbname=test port=5432")
+      class DeadTuplesCollector < Base
+        def run
+          table_names = connection.exec("SELECT tablename FROM pg_tables WHERE tableowner = 'gitlab'").values.flatten
+          stats       = {}
+
+          table_names.each do |table_name|
+            result = connection.exec("SELECT n_dead_tup FROM pg_stat_user_tables WHERE relname = '#{table_name}'")
+            stats[table_name] = result[0]["n_dead_tup"].to_i
+          end
+
+          stats
+        end
+      end
+
+      # Probes the DB specified by opts[:connection_string] for dead tubles stats, then converts them to metrics
+      class DeadTuplesProber
+        def initialize(opts, metrics = PrometheusMetrics.new)
+          @metrics   = metrics
+          @collector = DeadTuplesCollector.new(connection_string: opts[:connection_string])
+        end
+
+        def probe_db
+          return self unless @collector.connected?
+
+          result = @collector.run
+
+          result.each do |table_name, dead_tuples_count|
+            @metrics.add("pg_dead_tuples", dead_tuples_count, table_name: table_name)
+          end
+
+          self
+        end
+
+        def write_to(target)
+          target.write(@metrics.to_s)
+        end
+      end
+    end
+  end
+end

lib/gitlab_monitor/database_dead_tuples.rb

-require "pg"
-
-module GitLab
-  module Monitor
-    # A helper class to collect dead tuples stats from the database
-    #
-    # It takes a connection string (e.g. "dbname=test port=5432")
-    class DatabaseDeadTuplesCollector
-      attr_reader :result
-
-      def initialize(args)
-        @connection_string = args[:connection_string]
-      end
-
-      def run
-        table_names = connection.exec("SELECT tablename FROM pg_tables WHERE tableowner = 'gitlab'").values.flatten
-        stats       = {}
-
-        table_names.each do |table_name|
-          result = connection.exec("SELECT n_dead_tup FROM pg_stat_user_tables WHERE relname = '#{table_name}'")
-          stats[table_name] = result[0]["n_dead_tup"].to_i
-        end
-
-        stats
-      end
-
-      private
-
-      def connection
-        @connection ||= PG.connect(@connection_string)
-      end
-    end
-
-    # Probes the DB specified by opts[:connection_string] for dead tubles stats, then converts them to metrics
-    class DatabaseDeadTuplesProber
-      def initialize(opts, metrics = PrometheusMetrics.new)
-        @metrics   = metrics
-        @collector = DatabaseDeadTuplesCollector.new(connection_string: opts[:connection_string])
-      end
-
-      def probe_db
-        result = @collector.run
-
-        result.each do |table_name, dead_tuples_count|
-          @metrics.add("pg_dead_tuples", dead_tuples_count, table_name: table_name)
-        end
-
-        self
-      end
-
-      def write_to(target)
-        target.write(@metrics.to_s)
-      end
-    end
-  end
-end