import psycopg2
from itertools import chain, permutations
from graphviz import Digraph


DB_SETTINGS = dict(
    host='localhost',
    database='database_name',
    user='user',
    password='*****',
)

def get_relations():
    # db
    conn = psycopg2.connect(**DB_SETTINGS)
    cur = conn.cursor()

    # fetch database objects (views, functions, tables)
    cur.execute('''
            select proname, prosrc
            from pg_proc
            where proname ~* '_vs_';
    ''')
    functions = cur.fetchall()

    cur.execute('''
            select viewname, definition
            from pg_views
            where schemaname in ('api', 'public');
    ''')
    views = cur.fetchall()

    # find relations between database objects.
    relations = set()
    for a, b in permutations(chain(views, functions), 2):
        a_name, a_source = a
        b_name, b_source = b
        if a_name == b_name:  # ignore self references.
            continue
        if (a_name in b_source):
            # "a" is referenced in "b" definition.
            relations.add((a_name, b_name))
    return relations


def make_graph_file(relations):
    dot = Digraph()
    dot.format = 'png'
    dot.graph_attr['rankdir'] = 'LR'
    dot.graph_attr['engine'] = 'dot'
    dot.node_attr['style'] = 'filled'
    dot.node_attr['shape'] = 'box'

    nodes = set(chain(*relations))
    for n in nodes:
        # highlight some nodes based on their names
        dot.node(n, fillcolor=("yellow" if '_vs_' in n else None))

    dot.edges(relations)

    # saves gv and png files
    dot.render()


relations = get_relations()
make_graph_file(relations)