EVOLUTION-NINJA

Edit File: postgresql.py

import agent_util
import csv
import sys

try:
    import pg8000
except:
    # Fallback for older Python versions that aren't supported
    pg8000 = None
import logging
from agent_util import float

def execute_query(config, query):
    unix_sock = config.get("unix_sock", None)
    hostname = config.get("hostname", "localhost")
    port = int(config.get("port", "5432"))
    user = config.get("username", "postgres")
    password = config.get("password", "postgres")
    database = config.get("database", "postgres")

# make the connection based on either a unix socket or tcp connection
    if unix_sock:
        try:
            pgconn = pg8000.connect(
                user=user,
                host=None,
                unix_sock=unix_sock,
                port=port,
                password=password,
                database=database,
            )
        except Exception:
            logging.exception("Error connecting using provided socket %s" % unix_sock)
            return None
    else:
        try:
            pgconn = pg8000.connect(
                user=user,
                host=hostname,
                port=port,
                password=password,
                database=database,
            )
        except Exception:
            logging.exception("Error connecting using provided PostgreSQL credentials")
            return None

cursor = pgconn.cursor()
    q = cursor.execute(query)
    output = cursor.fetchall()
    cursor.close()
    pgconn.close()
    return output

class PostgreSQLPlugin(agent_util.Plugin):
    """
    If this plugin does not work you'll need to confirm at least two steps
    1. The user has at least SELECT privileges on the database of your choice
    2. Your postgres config has password authentication turned ON. This can be done by putting these settings
        at the top of your pg_hba.conf file:
            local   all     all     trust
            host    all 127.0.0.1/32    trust
        This will allow password authentication to PostgreSQL but still lock it down to only localhost to use password auth
    """

textkey = "postgresql"
    label = "PostgreSQL"

@classmethod
    def get_metadata(self, config):
        status = agent_util.SUPPORTED
        msg = None

if pg8000 is None:
            self.log.info("Postgres plugin not supported with current Python version")
            return {}

# check if postgresql is even installed
        installed = agent_util.which("psql")
        if not installed:
            self.log.info("psql binary not found")
            status = agent_util.UNSUPPORTED
            msg = "psql binary not found"
            return {}

if installed and (
            not config.get("database", "postgres")
            or not config.get("database", "postgres")
            or not config.get("database", "postgres")
        ):
            msg = "Missing config information, please include username, password and database"
            status = agent_util.MISCONFIGURED

if status is agent_util.SUPPORTED:
            try:
                output = execute_query(
                    config,
                    "select datname from pg_database where datistemplate = false",
                )
                if output is None:
                    status = agent_util.MISCONFIGURED
                    msg = "Double check the PostgreSQL credentials provided in the agent config file"
            except:
                self.log.exception("error running postgresql query")
                status = agent_util.MISCONFIGURED
                msg = "Double check the PostgreSQL credentials provided in the agent config file"

db_list = []
        if status is agent_util.SUPPORTED:
            res = execute_query(config, "SELECT datname FROM pg_database")
            for db in res:
                db_list.append(db[0])

data = {
            "rows_select_idx": {
                "label": "Rows returned",
                "options": db_list,
                "status": status,
                "error_message": msg,
                "unit": "rows/min",
            },
            "rows_select_scan": {
                "label": "Rows scanned",
                "options": db_list,
                "status": status,
                "error_message": msg,
                "unit": "rows/min",
            },
            "rows_insert": {
                "label": "INSERTS",
                "options": db_list,
                "status": status,
                "error_message": msg,
                "unit": "rows/min",
            },
            "rows_update": {
                "label": "UPDATES",
                "options": db_list,
                "status": status,
                "error_message": msg,
                "unit": "rows/min",
            },
            "rows_delete": {
                "label": "DELETES",
                "options": db_list,
                "status": status,
                "error_message": msg,
                "unit": "rows/min",
            },
            "rows_total": {
                "label": "Total queries",
                "options": db_list,
                "status": status,
                "error_message": msg,
                "unit": "rows/min",
            },
            "numbackends": {
                "label": "Number of active backends",
                "options": db_list,
                "status": status,
                "error_message": msg,
                "unit": "backends",
            },
            "xact_commit": {
                "label": "Transactions committed",
                "options": db_list,
                "status": status,
                "error_message": msg,
                "unit": "txn/min",
            },
            "xact_rollback": {
                "label": "Transactions rolled back",
                "options": db_list,
                "status": status,
                "error_message": msg,
                "unit": "txn/min",
            },
            "xact_total": {
                "label": "Total transactions",
                "options": db_list,
                "status": status,
                "error_message": msg,
                "unit": "txn/min",
            },
            "blks_read": {
                "label": "Blocks read from disk",
                "options": db_list,
                "status": status,
                "error_message": msg,
                "unit": "blocks/min",
            },
            "blks_hit": {
                "label": "Blocks read from buffer cache",
                "options": db_list,
                "status": status,
                "error_message": msg,
                "unit": "blocks/min",
            },
            "blks_cache_pc": {
                "label": "Buffer cache hit rate (%)",
                "options": db_list,
                "status": status,
                "error_message": msg,
                "unit": "percent",
            },
            "total_checkpoints": {
                "label": "Total database checkpoints",
                "options": None,
                "status": status,
                "error_message": msg,
                "unit": "checkpoints",
            },
            "minutes_between_checkpoints": {
                "label": "Minutes between database checkpoints",
                "options": None,
                "status": status,
                "error_message": msg,
                "unit": "minutes",
            },
            "total_connections": {
                "label": "Total connections",
                "options": db_list,
                "status": status,
                "error_message": msg,
                "unit": "connections",
            },
            "active_connections": {
                "label": "Active connections",
                "options": db_list,
                "status": status,
                "error_message": msg,
                "unit": "connections",
            },
            "idle_connections": {
                "label": "Idle connections",
                "options": db_list,
                "status": status,
                "error_message": msg,
                "unit": "connections",
            },
            "idle_in_txn_connections": {
                "label": "Idle in transaction connections",
                "options": db_list,
                "status": status,
                "error_message": msg,
                "unit": "connections",
            },
            "connections_waiting_for_lock": {
                "label": "Connections waiting for lock",
                "options": db_list,
                "status": status,
                "error_message": msg,
                "unit": "connections",
            },
            "max_txn_age": {
                "label": "Max transaction age",
                "options": db_list,
                "status": status,
                "error_message": msg,
                "unit": "seconds",
            },
            "temp_files_created": {
                "label": "Temp files generated",
                "options": db_list,
                "status": status,
                "error_message": msg,
                "unit": "files",
            },
            "temp_bytes_created": {
                "label": "Temp bytes generated",
                "options": db_list,
                "status": status,
                "error_message": msg,
                "unit": "bytes",
            },
            "checkpoints_timed": {
                "label": "Timed Checkpoints",
                "options": None,
                "status": status,
                "error_message": msg,
                "unit": "checkpoints",
            },
            "checkpoints_req": {
                "label": "Requested Checkpoints",
                "options": None,
                "status": status,
                "error_message": msg,
                "unit": "checkpoints",
            },
        }
        return data

def check(self, textkey, data, config):
        if data:
            db_name = data
        else:
            db_name = config.get("database", "postgres")

try:
            #######################################
            # these are point-in-time metric values
            #######################################
            if textkey == "minutes_between_checkpoints":
                res = execute_query(
                    config,
                    "SELECT seconds_since_start / total_checkpoints / 60 AS minutes_between_checkpoints FROM (SELECT EXTRACT(EPOCH FROM (now() - pg_postmaster_start_time())) AS seconds_since_start, (checkpoints_timed+checkpoints_req) AS total_checkpoints FROM pg_stat_bgwriter) AS sub;",
                )
                row = res[0]
                self.log.debug("Minutes between database checkpoints: %s" % str(row[0]))
                return int(row[0])

elif textkey == "checkpoints_timed":
                res = execute_query(
                    config, "SELECT checkpoints_timed from pg_stat_bgwriter;"
                )
                row = res[0]
                self.log.debug("Timed checkpoints: %s" % str(row[0]))
                return int(row[0])

elif textkey == "checkpoints_req":
                res = execute_query(
                    config, "SELECT checkpoints_req from pg_stat_bgwriter;"
                )
                row = res[0]
                self.log.debug("Timed checkpoints: %s" % str(row[0]))
                return int(row[0])

elif textkey == "active_connections":
                res = execute_query(
                    config,
                    "SELECT count(*) FROM pg_stat_activity WHERE state='active' AND datname='%s';"
                    % db_name,
                )
                row = res[0]
                self.log.debug("Active connections: %s" % str(row[0]))
                return int(row[0])

elif textkey == "idle_connections":
                res = execute_query(
                    config,
                    "SELECT count(*) FROM pg_stat_activity WHERE state='idle' AND datname='%s';"
                    % db_name,
                )
                row = res[0]
                self.log.debug("Idle connections: %s" % str(row[0]))
                return int(row[0])

elif textkey == "idle_in_txn_connections":
                res = execute_query(
                    config,
                    "SELECT count(*) FROM pg_stat_activity WHERE state='idle in transaction' AND datname='%s';"
                    % db_name,
                )
                row = res[0]
                self.log.debug("Idle in transaction connections: %s" % str(row[0]))
                return int(row[0])

elif textkey == "total_connections":
                res = execute_query(
                    config,
                    "SELECT count(*) FROM pg_stat_activity WHERE datname='%s';"
                    % db_name,
                )
                row = res[0]
                self.log.debug("Total connections: %s" % str(row[0]))
                return int(row[0])

elif textkey == "connections_waiting_for_lock":
                res = execute_query(
                    config,
                    "SELECT count(*) FROM pg_locks pgl, pg_stat_activity pgsa WHERE pgl.pid = pgsa.pid and pgl.granted = false and pgsa.datname='%s';"
                    % db_name,
                )
                row = res[0]
                self.log.debug("Connections waiting for lock: %s" % str(row[0]))
                return int(row[0])

elif textkey == "max_txn_age":
                res = execute_query(
                    config,
                    "SELECT max(now() - xact_start) FROM pg_stat_activity WHERE state IN ('idle in transaction', 'active') AND datname='%s';"
                    % db_name,
                )
                row = res[0]
                self.log.debug("Max transaction age: %s" % str(row[0]))
                return int(row[0].seconds)

elif textkey == "numbackends":
                res = execute_query(
                    config,
                    "SELECT numbackends FROM pg_stat_database where datname='%s';"
                    % db_name,
                )
                row = res[0]
                self.log.debug("Number of active backends: %s" % str(row[0]))
                return int(row[0])

################################################################
            # these are guages which require the different to be calculated
            ################################################################
            cached_result = self.get_cache_results(textkey, db_name)

if textkey == "rows_select_idx":
                res = execute_query(
                    config,
                    "SELECT tup_fetched FROM pg_stat_database where datname='%s';"
                    % db_name,
                )
                row = res[0]
                self.log.debug("SELECTs (from scans) (rows/second): %s" % str(row[0]))
                queried_value = int(row[0])

elif textkey == "rows_select_scan":
                res = execute_query(
                    config,
                    "SELECT tup_returned FROM pg_stat_database where datname='%s';"
                    % db_name,
                )
                row = res[0]
                self.log.debug("SELECTs (from scans) (rows/second): %s" % str(row[0]))
                queried_value = int(row[0])

elif textkey == "rows_insert":
                res = execute_query(
                    config,
                    "SELECT tup_inserted FROM pg_stat_database where datname='%s';"
                    % db_name,
                )
                row = res[0]
                self.log.debug("INSERTs (rows/second): %s" % str(row[0]))
                queried_value = int(row[0])

elif textkey == "rows_update":
                res = execute_query(
                    config,
                    "SELECT tup_updated FROM pg_stat_database where datname='%s';"
                    % db_name,
                )
                row = res[0]
                self.log.debug("UPDATEs (rows/second): %s" % str(row[0]))
                queried_value = int(row[0])

elif textkey == "rows_delete":
                res = execute_query(
                    config,
                    "SELECT tup_deleted FROM pg_stat_database where datname='%s';"
                    % db_name,
                )
                row = res[0]
                self.log.debug("DELETEs (rows/second): %s" % str(row[0]))
                queried_value = int(row[0])

elif textkey == "rows_total":
                res = execute_query(
                    config,
                    "SELECT sum(tup_fetched + tup_returned + tup_inserted + tup_updated + tup_deleted) from pg_stat_database where datname='%s';"
                    % db_name,
                )
                row = res[0]
                self.log.debug("Total queries (rows/second): %s" % str(row[0]))
                queried_value = int(row[0])

elif textkey == "xact_commit":
                res = execute_query(
                    config,
                    "SELECT xact_commit FROM pg_stat_database where datname='%s';"
                    % db_name,
                )
                row = res[0]
                self.log.debug("Transactions committed (txn/second): %s" % str(row[0]))
                queried_value = int(row[0])

elif textkey == "xact_rollback":
                res = execute_query(
                    config,
                    "SELECT xact_rollback FROM pg_stat_database where datname='%s';"
                    % db_name,
                )
                row = res[0]
                self.log.debug(
                    "Transactions rolled back (txn/second): %s" % str(row[0])
                )
                queried_value = int(row[0])

elif textkey == "xact_total":
                res = execute_query(
                    config,
                    "SELECT sum(xact_commit + xact_rollback) FROM pg_stat_database where datname='%s';"
                    % db_name,
                )
                row = res[0]
                self.log.debug("Total transactions (txn/second): %s" % str(row[0]))
                queried_value = int(row[0])

elif textkey == "blks_read":
                res = execute_query(
                    config,
                    "SELECT blks_read FROM pg_stat_database where datname='%s';"
                    % db_name,
                )
                row = res[0]
                self.log.debug(
                    "Blocks read from disk (blocks/second): %s" % str(row[0])
                )
                queried_value = int(row[0])

elif textkey == "blks_hit":
                res = execute_query(
                    config,
                    "SELECT blks_hit FROM pg_stat_database where datname='%s';"
                    % db_name,
                )
                row = res[0]
                self.log.debug(
                    "Blocks read from buffer cache (blocks/second): %s" % str(row[0])
                )
                queried_value = int(row[0])

elif textkey == "blks_cache_pc":
                res = execute_query(
                    config,
                    "SELECT blks_read, blks_hit FROM pg_stat_database where datname='%s';"
                    % db_name,
                )
                row = list(map(float, res[0]))
                if row[0] or row[1]:
                    result = int(row[1] / (row[0] + row[1]) * 100)
                else:
                    result = 0
                self.log.debug("Buffer cache hit rate (%%): %s" % str(result))
                queried_value = int(result)

elif textkey == "total_checkpoints":
                res = execute_query(
                    config,
                    "SELECT total_checkpoints FROM (SELECT EXTRACT(EPOCH FROM (now() - pg_postmaster_start_time())) AS seconds_since_start, (checkpoints_timed+checkpoints_req) AS total_checkpoints FROM pg_stat_bgwriter) AS sub;",
                )
                row = res[0]
                self.log.debug("Total database checkpoints: %s" % str(row[0]))
                queried_value = int(row[0])

elif textkey == "temp_files_created":
                res = execute_query(
                    config,
                    "select temp_files from pg_stat_database where datname='%s';"
                    % db_name,
                )
                row = res[0]
                self.log.debug("Temp files generated: %s" % str(row[0]))
                queried_value = int(row[0])

elif textkey == "temp_bytes_created":
                res = execute_query(
                    config,
                    "select temp_bytes from pg_stat_database where datname='%s';"
                    % db_name,
                )
                row = res[0]
                self.log.debug("Temp bytes generated: %s" % str(row[0]))
                queried_value = int(row[0])
        except:
            self.log.error(
                "Unable to gather metric - double check the PostgreSQL credentials provided in the agent config file"
            )
            return None

# we have a cached result, so do the math and return the differential
        if cached_result:
            cached_value = cached_result[0][1]
            self.cache_result(textkey, db_name, queried_value, replace=True)
            return int(queried_value - cached_value)

# no cached value in the agent's cache. This could be the first run of the metric, so return 0 and let the next run report the correct value
        else:
            self.cache_result(textkey, db_name, queried_value, replace=True)
            return 0