Generating a dependency graph for Postgres DB

postgres-deps.py

# From https://sigterm.sh/2010/07/09/generating-a-dependency-graph-for-a-postgresql-database/
# Converted to Python3
# Add config fort dbport
# Increase PNG Size
# pylinting

from optparse import OptionParser, OptionGroup

import sys
import psycopg2


def writedeps(cursor, tbl):
    sql = """SELECT
        tc.constraint_name, tc.table_name, kcu.column_name,
        ccu.table_name AS foreign_table_name,
        ccu.column_name AS foreign_column_name
    FROM
        information_schema.table_constraints AS tc
    JOIN information_schema.key_column_usage AS kcu ON
        tc.constraint_name = kcu.constraint_name
    JOIN information_schema.constraint_column_usage AS ccu ON
        ccu.constraint_name = tc.constraint_name
    WHERE constraint_type = 'FOREIGN KEY' AND tc.table_name = '%s'"""
    cursor.execute(sql % tbl)
    for row in cursor.fetchall():
        constraint, table, column, foreign_table, foreign_column = row
        print('"%s" -> "%s" [label="%s"];' % (tbl, foreign_table, constraint))


def get_tables(cursor):
    cursor.execute("SELECT tablename FROM pg_tables WHERE schemaname='public'")
    for row in cursor.fetchall():
        yield row[0]


def main():
    parser = OptionParser()

    group = OptionGroup(parser, "Database Options")
    group.add_option("--dbname", action="store", dest="dbname",
            help="The database name.")
    group.add_option("--dbhost", action="store", dest="dbhost",
            default="localhost", help="The database host.")
    group.add_option("--dbport", action="store", dest="dbport",
            default="10000", help="The database port.")
    group.add_option("--dbuser", action="store", dest="dbuser",
            help="The database username.")
    group.add_option("--dbpass", action="store", dest="dbpass",
            help="The database password.")
    parser.add_option_group(group)

    (options, args) = parser.parse_args()

    if not options.dbname:
        print("Please supply a database name, see --help for more info.")
        sys.exit(1)

    try:
        conn = psycopg2.connect("dbname='%s' user='%s' host='%s' port='%s' password='%s'"
            % (options.dbname, options.dbuser, options.dbhost, options.dbport, options.dbpass))
    except psycopg2.OperationalError as e:
        print("Failed to connect to database,", end=' ')
        print("perhaps you need to supply auth details:\n %s" % str(e))
        print("Use --help for more info.")
        sys.exit(1)

    cursor = conn.cursor()

    print("Digraph F {\n")
    print('ranksep=1.0; size="50"; rankdir=LR;')
    for i in get_tables(cursor):
        writedeps(cursor, i)
    print("}")

    sys.exit(0)


if __name__ == "__main__":
    main()