The script shows or sends notification about slow or idle-in-transaction queries

pg_query_inform.py

#!/usr/bin/python3
# pg_query_inform.py - shows or sends notification about slow
# or idle-in-transaction queries.
# (functionality will be added if necessary)
# use pg_query_inform.py -h to show syntax
#
# Author: Andrew Klychkov aaklychkov@mail.ru
# Version: 0.2
# Date: 2019-10-02
# License: Copyleft free software
# Requirements:
# * Python 3+
# * psycopg2 python library
# * postfix or something similar that allows to send mail
# * run it as postgres user
#
# This utility:
# 1) gets query stat from cluster
# 2) sends notification to passed mailboxes if conditions triggered
#
# Examples:
# ./pg_query_inform.py --duration 1000 -m aaklychkov@mail.ru
# If there are queries with duration more than 1000 seconds,
# send a notification (like below for example):
#
# Subject: mytestsrv.lan: query status
#
# Idle-in-transaction queries:
# ----------------
# dur: 4284.930712
# db: postgres
# pid: 11278
# user: postgres
# addr: localhost
# state: idle in transaction
# body: SELECT now();
#
# If -m/--mailboxes arg is not passed, prints messages to console


import argparse
import datetime
import smtplib
import socket
import sys

from email.mime.multipart import MIMEMultipart
from email.mime.text import MIMEText


try:
    import psycopg2
    import psycopg2.extras
except ImportError as e:
    print(e, 'Hint: use pip3 install psycopg2-binary')
    sys.exit(1)


__VERSION__ = '0.2'

HOSTNAME = socket.gethostname()


def parse_cli_args():
    parser = argparse.ArgumentParser(description="Check cluster condition "
                                     "and inform by email")
    parser.add_argument("-q", "--query",
                        action="store_true",
                        help="query state check")
    parser.add_argument("-d", "--duration",
                        dest="duration",
                        metavar="DURATION",
                        help=("max current query duration (in sec), "
                              "alarm if longer"))
    parser.add_argument("-m", "--mailboxes",
                        dest="mailboxes",
                        metavar="MAILBOXES",
                        help="mailboxes to send notifications")
    parser.add_argument("--version",
                        action="version",
                        version=__VERSION__,
                        help="show version and exit")
    return parser.parse_args()


args = parse_cli_args()

# PgSQL conn params:
DBUSER = 'postgres'
DBNAME = 'postgres'

# Mail params:
if args.mailboxes:
    SBJ = '%s: query status' % HOSTNAME
    SMTP_SRV = 'localhost'
    SENDER = 'postgres'


def send_mail_local(sbj, ms, snd, rcv=[]):
    if not isinstance(rcv, (list,)):
        rcv = list(rcv)

    msg = MIMEMultipart()
    msg['Subject'] = sbj
    msg['From'] = snd
    msg['To'] = rcv[0]
    msg.attach(MIMEText(ms, 'plain'))
    smtpconnect = None
    try:
        smtpconnect = smtplib.SMTP('localhost')
        smtpconnect.sendmail(snd, rcv, msg.as_string())
    except SMTPException as e:
        print(e, "Error: unable to send mail")
    finally:
        smtpconnect.quit()


def get_query_info(cursor):
    query = ("SELECT datname, pid, usename, client_addr, "
             "EXTRACT(EPOCH FROM (now() - query_start)) as duration, "
             "state, query FROM pg_stat_activity "
             "WHERE backend_type = 'client backend' "
             "AND pid != pg_backend_pid()")
    cursor.execute(query)
    return cursor.fetchall()

def get_locks_info(cursor):
    query = ("SELECT database, pid, locktype, c.relname, mode "
             "FROM pg_locks l JOIN pg_class c "
             "ON c.oid = l.relation WHERE c.relname NOT LIKE 'pg_%'")
    cursor.execute(query)
    return cursor.fetchall()


def main():
    try:
        db_conn = psycopg2.connect('dbname=%s user=%s' % (DBNAME, DBUSER))
        cursor = db_conn.cursor(cursor_factory=psycopg2.extras.DictCursor)
    except Exception as e:
        print(e, "Error: unable to connect to database")

    if args.duration:
        idle_trans_list = []
        slow_query_list = []

        pg_stat_activity = get_query_info(cursor)
        for query_inf in pg_stat_activity:
            duration = int(query_inf[4])

            if duration > int(args.duration):
                db = query_inf[0]
                pid = query_inf[1]
                user = query_inf[2]
                addr = query_inf[3]
                if addr is None:
                    addr = 'localhost'
                state = query_inf[5]
                body = query_inf[6]

                msg = ('dur: %s sec\ndb: %s \npid: %s \nuser: %s \naddr: %s\n'
                       'state: %s \nbody: %s\n' % (duration, db, pid, user,
                                                   addr, state, body))

                if state == 'active':
                    slow_query_list.append(msg)
                elif state == 'idle in transaction':
                    idle_trans_list.append(msg)

    msg = ''
    if slow_query_list or idle_trans_list:
        delim = 16 * '-'
        if idle_trans_list:
            msg += 'Idle-in-transaction queries:\n%s\n%s' % (
                delim, '\n'.join(idle_trans_list))
        if slow_query_list:
            msg += '%s\nSlow queries:\n%s\n%s' % (
                msg, delim, '\n'.join(slow_query_list))

    pg_locks = get_locks_info(cursor)
    if pg_locks:
        msg += '\nLocks:\n%s' % '\n'.join(pg_locks)

    if msg:
        if args.mailboxes:
            send_mail_local(SBJ, msg, SENDER, [args.mailboxes])
        else:
            print(msg)

    db_conn.close()


if __name__ == '__main__':
    main()