sensu-plugins/sensu-plugins-postgres

#! /usr/bin/env ruby
# frozen_string_literal: false

#
# metrics-postgres-query
#
# DESCRIPTION:
#   This plugin collects metrics from the results of a postgres query. Can optionally
#   count the number of tuples (rows) returned by the query.
#
# OUTPUT:
#   metric data
#
# PLATFORMS:
#   Linux
#
# DEPENDENCIES:
#   gem: pg
#   gem: sensu-plugin
#
# USAGE:
#   metrics-postgres-query.rb -u db_user -p db_pass -h db_server -d db -q 'select foo from bar'
#
# NOTES:
#
# LICENSE:
#   Copyright 2015, Eric Heydrick <eheydrick@gmail.com>
#   Released under the same terms as Sensu (the MIT license); see LICENSE
#   for details.
#

require 'sensu-plugins-postgres/pgpass'
require 'sensu-plugin/metric/cli'
require 'pg'

class MetricsPostgresQuery < Sensu::Plugin::Metric::CLI::Graphite
  option :pgpass,
         description: 'Pgpass file',
         short: '-f FILE',
         long: '--pgpass',
         default: ENV['PGPASSFILE'] || "#{ENV['HOME']}/.pgpass"

  option :user,
         description: 'Postgres User',
         short: '-u USER',
         long: '--user USER'

  option :password,
         description: 'Postgres Password',
         short: '-p PASS',
         long: '--password PASS'

  option :hostname,
         description: 'Hostname to login to',
         short: '-h HOST',
         long: '--hostname HOST'

  option :port,
         description: 'Database port',
         short: '-P PORT',
         long: '--port PORT'

  option :database,
         description: 'Database name',
         short: '-d DB',
         long: '--db DB'

  option :query,
         description: 'Database query to execute',
         short: '-q QUERY',
         long: '--query QUERY',
         required: true

  option :count_tuples,
         description: 'Count the number of tuples (rows) returned by the query',
         short: '-t',
         long: '--tuples',
         boolean: true,
         default: false

  option :scheme,
         description: 'Metric naming scheme, text to prepend to metric',
         short: '-s SCHEME',
         long: '--scheme SCHEME',
         default: 'postgres'

  option :multirow,
         description: 'Determines if we return first row or all rows',
         short: '-m',
         long: '--multirow',
         boolean: true,
         default: false

  option :timeout,
         description: 'Connection timeout (seconds)',
         short: '-T TIMEOUT',
         long: '--timeout TIMEOUT',
         default: nil

  include Pgpass

  def run
    begin
      pgpass
      con = PG.connect(host: config[:hostname],
                       dbname: config[:database],
                       user: config[:user],
                       password: config[:password],
                       port: config[:port],
                       connect_timeout: config[:timeout])
      res = con.exec(config[:query].to_s)
    rescue PG::Error => e
      unknown "Unable to query PostgreSQL: #{e.message}"
    end

    value = if config[:count_tuples]
              res.ntuples
            else
              res.first&.values&.first
            end

    if config[:multirow] && !config[:count_tuples]
      res.values.each do |row|
        output "#{config[:scheme]}.#{row[0]}", row[1]
      end
    else
      output config[:scheme], value
    end
    ok
  end
end