#! /usr/bin/python
#
# File: sql2dot.py
# Author: Dimitri Fontaine <dim@tapoueh.org>
#
# This code is licensed under the GPL.
# Get yourself a version here : http://www.gnu.org/copyleft/gpl.html
#
# This will generate a .dot file for use with Graphviz, the input file
# is a sql table creation script.
#
# Usage example:
#   ./sql2dot.py create.sql | dot -Tps -o create.ps
#
# If you wanna create png files, you may find the following tip useful :
# ttf=`for f in $(locate *ttf); do dirname $f; done | sort -u | tr -s '\n' ':'`
# export DOTFONTPATH=$ttf

import re
import sys
import string

def begin_dot():
    """ begins the dot file """
    str = """
digraph g {
  graph [
    rankdir = "LR"
  ];
  node [
    fontsize = "12"
    fontname = "Helvetica"
    shape = "ellipse"
  ];
  edge [
  ];
"""
    return str

def add_node(name, labels):
    """ add a node to the dot file """

    # build the label string
    lstr = ""
    for l in labels:
        if lstr == "":
            #lstr = "<" + l + "> " + l
            t = l
            lstr = l
        else:
            if l != "":
                lstr = lstr + "|" +  "<" + l +  "> " + l
    
    str = """
  "%s" [
    label = "%s"
    shape = "record"
  ];
""" % (name, lstr)
    return str


def add_relation(relations, n1, l1, n2, l2):
    """ add a relation between nodes """
    # relations needs to be initialized as an empty list
    return relations.append([n1, l1, n2, l2])


def print_relations(relations):
    """ print the relations given in a list """
    id  = 0
    str = ""

    for rel in relations:
        str = str + '  "%s":%s -> "%s":%s [\n' \
              % (rel[0], rel[1], rel[2], rel[3])
        str = str + '    id = %s\n'      % id
        str = str + '  ];\n'

        id = id + 1
        
    return str

def end_dot():
    """ end the dot file """
    return "}";


if __name__ == "__main__":
    # Here we gonna do the job

    if len(sys.argv) == 1:
        print "no file to parse"
        sys.exit(0)
        
    filename = sys.argv[1]
    
    file = open(filename, "r")
    buffer = file.read()
    file.close()
    lines = string.split(buffer, "\n")

    relations = []
    curtable  = ""
    curfields = []

    print begin_dot()

    for line in lines:
        # create table
        mobj = re.match("create table \"?([^\(\s\"]*)", line, re.I)
        if mobj is not None:
            curtable = mobj.group(1)
            # Ignore pgaccess work tables
            mobj = re.match("pga_", curtable)
            if mobj is None:
                curfields = [curtable]
            else:
                curtable = ""
            continue

        # reference
        mobj = re.match(
            "^\s*([a-z0-9_]*).*references\s*([a-z0-9_]*)\s*\(([a-z0-9_]*)\)",
            line, re.I)
        
        if mobj is not None:
            l1 = mobj.group(1)
            n2 = mobj.group(2)
            l2 = mobj.group(3)
            add_relation(relations, curtable, l1, n2, l2)

        # skip primary keys
        mobj = re.search("^\s*primary key", line, re.I)
        if mobj is not None:
            continue
            
        # skip constraints
        mobj = re.search("^\s*constraint", line, re.I)
        if mobj is not None:
            continue
            
        # field
        mobj = re.match("^\s*\"?([a-z0-9_]*)\"?\s*[^\(]*", line, re.I)
        if curtable != "" and mobj is not None:
            field = mobj.group(1)
            curfields.append(field)

        # end create table
        mobj = re.match("\);", line, re.I)
        if mobj is not None:
            if curtable != "":
                print add_node(curtable, curfields)
                curfields = []
                curtable  = ""

    print print_relations(relations)
    print end_dot()