from Cheetah.Template import Template
import re, sys

from getopt import getopt

#try: import config
#except:
class c: pass
config = c()
config.dbparams = {'autocommit': True}
config.dbstring = ''
sys.modules['config'] = config

#try:    from pgdb  import Query, make_connection
#except: from pgdb2 import Query, make_connection
from pgdb2 import Query

def help():
    print """\
    -h,--help    This help text.
    -d,--dsn     Data Source Name, e.g. "host=h user=u dbname=frob pass=gr0nk"
    -t,--tmpl    Supply a template file to fill
    -o,--out     Write output to this file
    """
    sys.exit(0)

def error(s):
    print >>sys.stderr, s
    sys.exit(0)

opt = {'d:': 'dsn=', 't:': 'tmpl=', 'o:': 'out=','h':'help'}
try:    opts, junk = getopt(sys.argv[1:], ''.join(opt.keys()), opt.values())
except Exception, x:
    print x
    help()

tmpl, out = 'moin.tmpl', 'moin.wok'

for k,v in opts:
    if   k in ('--help', '-h'): help()  # does not return
    elif k in ('--dsn',  '-d'): config.dbstring = v
    elif k in ('--tmpl', '-t'): tmpl = v
    elif k in ('--out',  '-o'): out = v

if opts: print opts

acl_chars = 'rwadRxtXUCT'

acl_perms = { 'a': 'Insert',
              'r': 'Select',
              'w': 'Update',
              'd': 'Delete',
              'R': 'Rule',
              'x': 'References',
              't': 'Trigger',
              'X': 'Execute',
              'U': 'Usage',
              'C': 'Create',
              'T': 'Temporary',
              '*': 'Grant' }

acl_pat = re.compile('[{,"](?P<user>(?:[^="]|(?:\\\\"))*)='
                     '(?P<Insert>a?)' '(?P<GrantInsert>[*]?)'
                     '(?P<Select>r?)' '(?P<GrantSelect>[*]?)'
                     '(?P<Update>w?)' '(?P<GrantUpdate>[*]?)'
                     '(?P<Delete>d?)' '(?P<GrantDelete>[*]?)'
                     '(?P<Rule>R?)' '(?P<GrantRule>[*]?)'
                     '(?P<References>x?)' '(?P<GrantReferences>[*]?)'
                     '(?P<Trigger>t?)' '(?P<GrantTrigger>[*]?)'
                     '(?P<Execute>X?)' '(?P<GrantExecute>[*]?)'
                     '(?P<Usage>U?)' '(?P<GrantUsage>[*]?)'
                     '(?P<Create>C?)' '(?P<GrantCreate>[*]?)'
                     '(?P<Temporary>T?)' '(?P<GrantTemporary>[*]?)'
                     '/(?P<grantor>(?:[^,}"]|(?:\\\\"))+)')

def munge_acl(s):
    if not s: return []
    acl = []
    for z in acl_pat.finditer(s):
        d = z.groupdict()
        d['user'] = d['user'] or 'public'
        acl.append(d)
    return acl

typ = [ dict(t) for t in Query("SELECT oid, * FROM pg_catalog.pg_type")() ]

typd = dict([ (t['oid'], t) for t in typ ])

usr = [ dict(u) for u in Query("""
SELECT   usename AS uname, usesysid AS uid,
         usecreatedb AS createdb, usesuper AS su, usecatupd AS catupd
FROM     pg_user
ORDER BY uid;
""")() ]
usrd = dict([ (u['uid'], u) for u in usr ])

def fun():
    fun = Query("""
SELECT    p.oid, nspname AS schema, proname AS function,
          proowner AS own, proacl AS acl, prosecdef AS setuid,
          lanname AS language, proisstrict AS strict, provolatile AS volatile,
          proisagg AS aggregate, proretset AS multiret, ret.typname AS ret,
          pronargs AS argn, proargtypes AS aty, prosrc AS src
FROM      pg_proc p
LEFT JOIN pg_language l    ON l.oid     = prolang
LEFT JOIN pg_type ret      ON ret.oid   = prorettype
LEFT JOIN pg_namespace n   ON n.oid     = pronamespace
WHERE     nspname !~ 'pg_.*|information_schema'
ORDER BY  proname, nspname;
""")
    ff = fun()
    for f in ff:
        f['owner'] = usrd[f['own']]
        f['argt'] = [ typ[int(x)] for x in f['aty'].split() ]
        f['argn'] = ', '.join([ t['typname'] for t in f['argt'] ])
        f['proto'] = ''
        f['proto'] += '%(language)s ' % f
        f['proto'] += 'setuid ' * bool(f['setuid'])
        f['proto'] += 'strict ' * bool(f['strict'])
        f['proto'] += 'volatile ' * bool(f['volatile'])
        f['proto'] += 'aggregate ' * bool(f['aggregate'])
        f['proto'] += '%(ret)s' % f
        f['proto'] += '{}' * bool(f['multiret']) +' '
        f['proto'] += "%(schema)s.%(function)s(%(argn)s)" % f
        f['acl'] = munge_acl(f['acl'])
    return ff

fun = fun()
fund = dict([ (f['oid'], f) for f in fun ])

def con():
    con = Query("""
SELECT    nf.nspname AS from_schema, cf.relname AS from_table,
          nt.nspname AS   to_schema, ct.relname AS   to_table,
          conrelid, confrelid,
          conkey     AS from_col,
          confkey    AS to_col,
          conname AS constraint,
          contype AS "type",
          confupdtype AS update,
          confdeltype AS delete,
          confmatchtype AS match,
          pg_get_constraintdef(c.oid) AS src,
          description, c.oid
FROM      pg_constraint c
LEFT JOIN pg_class cf      ON cf.oid = conrelid
LEFT JOIN pg_class ct      ON ct.oid = confrelid
LEFT JOIN pg_namespace nf  ON nf.oid = cf.relnamespace
LEFT JOIN pg_namespace nt  ON nt.oid = ct.relnamespace
LEFT JOIN pg_description x ON x.objoid = c.oid
WHERE     conrelid != 0
ORDER BY  conname, from_table;
""")
    cc = [ dict(c) for c in con() ]
    for c in cc:
        if  c['from_col'] and c['from_col'][0] == '{' and c['from_col'][-1]:
            c['from_col'] = [ int(k) for k in c['from_col'][1:-1].split(',') ]
        if  c[  'to_col'] and c[  'to_col'][0] == '{' and c[  'to_col'][-1]:
            c[  'to_col'] = [ int(k) for k in c[  'to_col'][1:-1].split(',') ]
    return cc

con = con()

def col():
    col = Query("""
SELECT    nspname AS schema, relname AS table, attname AS column,
          attnum AS colnum, atttypid AS typ_id,
          typname, atttypmod AS "mod", attnotnull AS notnull,
          adsrc AS default, typlen AS size,
          attstorage AS stor,
          CASE WHEN attstorage = 'p' THEN 'Plain'
               WHEN attstorage = 'e' THEN 'External'
               WHEN attstorage = 'm' THEN 'Main'
               WHEN attstorage = 'x' THEN 'Extended'
                                     ELSE 'Unknown' END AS storage,
          pg_catalog.col_description(attrelid, attnum) as description
FROM      pg_attribute a
LEFT JOIN pg_type t        ON t.oid     = atttypid
LEFT JOIN pg_class c       ON c.oid     = attrelid
LEFT JOIN pg_namespace n   ON n.oid     = relnamespace
LEFT JOIN pg_attrdef d     ON d.adrelid = c.oid AND d.adnum = a.attnum
WHERE     nspname !~ 'pg_.*|information_schema'
AND       relkind IN ('r','s','v')
AND       attnum > 0
AND       attisdropped IS FALSE
ORDER BY  attnum
;
""")
    cc = [ dict(c) for c in col() ]
    for c in cc:
        c['type'] = typd[c['typ_id']]
        c['constraints'] = [ k for k in con
                             if  c['schema'] == k['from_schema']
                             and c['table' ] == k['from_table' ]
                             and c['colnum'] in k['from_col'   ] ]
        c['constr_in'  ] = [ k for k in con
                             if  c['schema'] == k['to_schema']
                             and c['table' ] == k['to_table' ]
                             and c['colnum'] in k['to_col'   ] ]
        c['default'] = c['default'] or ''
    return cc

col = col()

def ind():
    ind = Query("""
SELECT    nspname AS schema, c.relname AS table, i.relname AS index,
          t.typname, --i.relam AS access, --z.indpred AS expr,
          z.indisunique AS unique, z.indisprimary AS primary,
          z.indkey AS cols,
          a.amname AS method, i.relpages AS pages, i.reltuples AS rows,
          pg_get_indexdef(i.oid) AS def,
          substring(pg_get_indexdef(i.oid) FROM
                    strpos(pg_get_indexdef(i.oid), '(')+1 FOR
                    strpos(pg_get_indexdef(i.oid), ')') -
                    strpos(pg_get_indexdef(i.oid), '(')-1
                    ) AS deflist,
          CASE WHEN strpos(pg_get_indexdef(i.oid), 'WHERE') = 0 THEN ''
               ELSE substring(pg_get_indexdef(i.oid) FROM
                              strpos(pg_get_indexdef(i.oid), 'WHERE')+6)
          END AS pred,
          description
FROM      pg_index       z
JOIN      pg_class       c ON c.oid = z.indrelid
JOIN      pg_class       i ON i.oid = z.indexrelid
LEFT JOIN pg_am          a ON a.oid = i.relam
LEFT JOIN pg_type        t ON t.oid = i.relam
LEFT JOIN pg_namespace   n ON n.oid = c.relnamespace
LEFT JOIN pg_description x ON x.objoid = pg_get_indexdef(i.oid)
WHERE     c.relkind = 'r'
AND       i.relkind = 'i'
ORDER BY  i.relname, c.relname
""")
    ii = [ dict(i) for i in ind() ]
    for i in ii:
        i['def'] = i['deflist'].split(', ')
    return ii
ind = ind()

def trg():
    trg = Query("""
SELECT   n.nspname AS schema, rel.relname AS table, tgname AS trigger,
         trg.*, trg.oid
FROM     pg_trigger   trg
JOIN     pg_class     rel ON rel.oid = tgrelid
JOIN     pg_namespace n   ON n.oid   = relnamespace
WHERE    tgisconstraint IS FALSE
AND      tgname !~ 'pg_.*'
ORDER BY tgname, rel.relname
""")
    tt = [ dict(t) for t in trg() ]
    for t in tt:
        t['function'] = fund[ t['tgfoid'] ]
    return tt

trg = trg()

def tbl():
    tbl = Query("""
SELECT    nspname AS schema, relname AS table, relowner AS own, relacl AS acl,
          c.relpages AS pages, c.reltuples AS rows, relkind AS kind,
          CASE WHEN relkind = 'v'
               THEN pg_get_viewdef(c.oid)
               ELSE NULL END AS view,
          relhasoids AS has_oids, description, c.oid
FROM      pg_class c
LEFT JOIN pg_namespace   n ON n.oid     = relnamespace
LEFT JOIN pg_description x ON x.objoid  = c.oid AND x.objsubid = 0
WHERE     nspname !~ 'pg_.*|information_schema'
AND       relkind IN ('r','s','v')
ORDER BY  2
""")
    tt = [ dict(t) for t in tbl() ]

    vref = '(?:FROM |JOIN )([@A-Za-z][@A-Za-z0-9]+).([@A-Za-z][@A-Za-z0-9]+)'

    for t in tt: t['refs'], t['refd'] = [], []

    for t in tt:
        t['owner']       = usrd[t['own']]
        t['columns']     = [ c for c in col
                             if  t['schema'] == c['schema']
                             and t['table' ] == c['table' ] ]
        t['indices']     = [ i for i in ind
                             if  t['schema'] == i['schema']
                             and t['table' ] == i['table' ] ]
        t['constraints'] = [ c for c in con
                             if  t['schema'] == c['from_schema']
                             and t['table' ] == c['from_table' ] ]
        t['constr_in']   = [ c for c in con
                             if  t['schema'] == c['to_schema']
                             and t['table' ] == c['to_table' ] ]
        t['triggers']    = [ u for u in trg
                             if  t['schema'] == u['schema']
                             and t['table' ] == u['table' ] ]
        t['acl'] = munge_acl(t['acl'])
        if t['view']:
           rr = set([ (z, u) for z,u in re.findall(vref, t['view']) ])
           for schema, table in rr:
               for q in tt:
                   if q['schema'] == schema and q['table'] == table:
                       t['refs'].append(q)
                       q['refd'].append(t)
           t['view'] = re.sub(' *(FULL (?:OUTER )?|LEFT |RIGHT |'
                              'INNER |CROSS |)?JOIN',
                              '\n\g<1>JOIN',
                              (t['view']
                               .replace(',',',\n')
                               .replace('SELECT ','SELECT\n ')
                               .replace('CASE','\nCASE')
                               .replace('THEN','\n     THEN')
                               .replace('ELSE','\n     ELSE')))
    return tt

tbl = tbl()
tbld = dict([ (t['oid'], t) for t in tbl ])

def sch():
    sch = Query("""
SELECT    nspname AS schema, nspowner AS own, nspacl AS acl, description,
          nspacl AS rawacl
FROM      pg_namespace n
LEFT JOIN pg_description x ON x.objoid = n.oid
WHERE     nspname !~ 'pg_.*|information_schema'
ORDER BY  nspname;
""")
    ss = [ dict(s) for s in sch() ]
    for s in ss:
        s['owner']     = usrd[s['own']]
        s['tables']    = [ t for t in tbl if t['schema'] == s['schema'] ]
        s['functions'] = [ f for f in fun if f['schema'] == s['schema'] ]
        s['acl'] = munge_acl(s['acl'])
    return ss

sch = sch()

def db():
    db = Query("""
SELECT    datname AS db, datdba AS own, datacl AS acl, description, *
FROM      pg_database d
LEFT JOIN pg_description x ON x.objoid = d.oid
WHERE     datname !~ 'postgres|template.*';
""")
    dd = [ dict(d) for d in db() ]
    for d in dd:
        d['owner'] = usrd[d['own']]
        d['acl'] = munge_acl(d['acl'])
    return dd

db = db()[0]

#w = wiki(searchList = [locals()] )
#del wiki
#w = str(w)

file(out,'w').write(str(Template(file=tmpl, namespaces=[locals()])))