lib/pgq_prometheus/processor.rb in pgq_prometheus-0.2.0 vs lib/pgq_prometheus/processor.rb in pgq_prometheus-0.2.1

@@ -3,11 +3,15 @@
 require 'prometheus_exporter/client'
 
 module PgqPrometheus
   class Processor
     class << self
-      attr_accessor :sql_caller, :logger, :on_error
+      attr_accessor :sql_caller,
+                    :logger,
+                    :on_error,
+                    :before_collect,
+                    :after_collect
 
       def start(client: nil, frequency: 30, labels: nil)
         raise ArgumentError, "#{name}.sql_caller must be defined" if sql_caller.nil?
 
         stop
@@ -20,17 +24,19 @@
           wrap_thread_loop(name) do
             sql_caller.release_connection
             logger&.info { "Start #{name}" }
             while true
               begin
+                before_collect&.call
                 metrics = process_collector.collect
                 metrics.each do |metric|
                   client.send_json metric
                 end
+                after_collect&.call
               rescue => e
                 STDERR.puts "#{self.class} Failed To Collect Stats #{e.class} #{e.message}"
-                log(:error) { "#{e.class} #{e.message} #{e.backtrace.join("\n")}" }
+                logger&.error { "#{e.class} #{e.message} #{e.backtrace.join("\n")}" }
                 on_error&.call(e)
               end
               sleep frequency
             end
           end
@@ -95,15 +101,9 @@
 
     private
 
     [:sql_caller, :logger].each do |meth|
       define_method(meth) { |*args, &block| self.class.public_send(meth, *args, &block) }
-    end
-
-    def log(severity)
-      return yield if logger.nil?
-
-      logger.public_send(severity) { yield }
     end
 
     def queue_metric_opts
       Config._metrics.select { |_, opts| opts[:from] == :queue }
     end