#!/usr/bin/env python3
# Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.


from optparse import OptionParser, OptionGroup

import pygraphviz as pgv
import psycopg2
import sys


# Query found here:
# https://stackoverflow.com/a/46594226/105137

def writedeps(conn):
    sql = """WITH RECURSIVE view_deps AS (
SELECT DISTINCT dependent_ns.nspname as dependent_schema
, dependent_view.relname as dependent_view
, source_ns.nspname as source_schema
, source_table.relname as source_table
FROM pg_depend
JOIN pg_rewrite ON pg_depend.objid = pg_rewrite.oid
JOIN pg_class as dependent_view ON pg_rewrite.ev_class = dependent_view.oid
JOIN pg_class as source_table ON pg_depend.refobjid = source_table.oid
JOIN pg_namespace dependent_ns ON dependent_ns.oid = dependent_view.relnamespace
JOIN pg_namespace source_ns ON source_ns.oid = source_table.relnamespace
WHERE NOT (dependent_ns.nspname = source_ns.nspname AND dependent_view.relname = source_table.relname)
UNION
SELECT DISTINCT dependent_ns.nspname as dependent_schema
, dependent_view.relname as dependent_view
, source_ns.nspname as source_schema
, source_table.relname as source_table
FROM pg_depend
JOIN pg_rewrite ON pg_depend.objid = pg_rewrite.oid
JOIN pg_class as dependent_view ON pg_rewrite.ev_class = dependent_view.oid
JOIN pg_class as source_table ON pg_depend.refobjid = source_table.oid
JOIN pg_namespace dependent_ns ON dependent_ns.oid = dependent_view.relnamespace
JOIN pg_namespace source_ns ON source_ns.oid = source_table.relnamespace
INNER JOIN view_deps vd
    ON vd.dependent_schema = source_ns.nspname
    AND vd.dependent_view = source_table.relname
    AND NOT (dependent_ns.nspname = vd.dependent_schema AND dependent_view.relname = vd.dependent_view)
)

SELECT *
FROM view_deps
WHERE dependent_schema = 'public'
ORDER BY source_schema, source_table;"""

    G=pgv.AGraph(directed=True)


    with conn.cursor() as cursor:
        cursor.execute(sql)

        for row in cursor.fetchall():

            dependent_schema, dependent_view, source_schema, source_table = row
            print('"%s" -> "%s";' % (dependent_view, source_table))
            G.add_edge(dependent_view, source_table, color='blue')

    G.layout(prog='dot')
    G.draw('view-dependencies.png')


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("--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' password='%s'"
            % (options.dbname, options.dbuser, options.dbhost, options.dbpass))

        writedeps(conn)

    except psycopg2.OperationalError as e:
        print("Failed to connect to database,",)
        print("perhaps you need to supply auth details:\n %s" % str(e))
        print("Use --help for more info.")
        sys.exit(1)


if __name__ == "__main__":
    main()
