pglockviz

#!/usr/bin/env python

# Copyright (c) 2009 Peter Eisentraut
# See LICENSE file for conditions.

"""pglockviz

This program displays the current locking situation in a PostgreSQL
database as a graph.

Synopsis: pglockviz [-o OUTFILE] [CONN]

CONN -- A PostgreSQL database connection string, e.g., dbname=test.

OUTFILE -- The file format is determined based on the extension.

If no output file is specified, the dot graph is printed to the
standard output.  You could pipe the output into fdp -Txlib to display
the graph.
"""

import getopt
import psycopg2
import pygraphviz
import sys


def main():
    """Main function"""
    outfile = None

    opts, rest = getopt.gnu_getopt(sys.argv[1:], 'o:')
    for opt, arg in opts:
        if opt == '-o':
            outfile = arg

    conn = psycopg2.connect(rest[0] if rest else '')
    cursor = conn.cursor()
    if conn.server_version >= 90200:
        cursor.execute("""
SELECT n.nspname AS schema,
       c.relname AS table,
       l.pid,
       l.mode,
       l.granted,
       l.locktype,
       a.query
FROM pg_locks l JOIN pg_class c ON (l.relation = c.oid)
                JOIN pg_namespace n ON (c.relnamespace = n.oid)
                JOIN pg_stat_activity a ON (l.pid = a.pid)
WHERE locktype IN ( 'relation', 'tuple' )
  AND l.pid <> pg_backend_pid()
""")
    else:
        cursor.execute("""
SELECT n.nspname AS schema,
       c.relname AS table,
       l.pid,
       l.mode,
       l.granted,
       l.locktype,
       a.current_query
FROM pg_locks l JOIN pg_class c ON (l.relation = c.oid)
                JOIN pg_namespace n ON (c.relnamespace = n.oid)
                JOIN pg_stat_activity a ON (l.pid = a.procpid)
WHERE locktype IN ( 'relation', 'tuple' )
  AND l.pid <> pg_backend_pid()
""")

    graph = pygraphviz.AGraph(directed=True, name='Locks')

    for row in cursor.fetchall():
        owner = str(row[2]) + r"\n" + row[6][:25]
        obj = row[0] + "." + row[1] + r"\n locktype:" + row[5]

        graph.add_node(owner, shape='box')

        if row[4]:
            graph.add_node(obj, shape='ellipse')
            graph.add_edge(obj, owner,
                           label=(r"holds\n" + row[3]))
        else:
            graph.add_node(obj, shape='ellipse', color='red')
            graph.add_edge(owner, obj,
                           label=(r"wants\n" + row[3]),
                           color='red')

    if outfile is not None:
        graph.draw(outfile, prog='fdp')
    else:
        print(graph.to_string())


if __name__ == "__main__":
    main()