#!/usr/bin/env python
"""
sqlminus -- sqlplus minus. the features? the suck? you be the judge!

usage:
    sqlminus connstr          - interactive shell
    sqlminus connstr  cmds... - execute cmds
    sqlminus connstr  @file   - execute cmd in file.sql

configuration:
    ~/.sqlminus hold connstr aliases.  format:   foo = 'foo/bar@baz'

special commands:
    tables - show user tabes
    desc   - describe a table
    rehash - refresh user table information
    mono   - set monochrome output
    color  - set color output

features:
    readline editing
    semi-intelligent tab completion for tables, columns
    nice table formatting
    single file, easy to install

author:
    Mark Harrison, mh@pixar.com

license and download:
    bsd-ish.  see orapig package for complete text.
    sqlminus is part of the orapig package: http://code.google.com/p/orapig
"""
import sys,os,re,cmd,collections,readline
import cx_Oracle

class OracleCmd(cmd.Cmd):
    #-------------------------------------------------------------------
    def __init__(self,connstr):
        cmd.Cmd.__init__(self)
        self.conn = cx_Oracle.connect(connstr)
        self.curs = self.conn.cursor()
        self.cmds=['desc','tables ','color','mono','rehash',
                   'select','update','delete','insert']
        self.cmd=''
        self.do_rehash()
        self.do_mono()
    #-------------------------------------------------------------------
    def clearinput(self):
        """clear the input state"""
        self.cmd=''         
    #-------------------------------------------------------------------
    def oraprint(self,desc,rows):
        """nicely print a query result set"""
        # get the max width and type of each column,
        # use that to build fmt strings to print the header and rows

        maxlen=[len(i[0]) for i in desc]
        types=[i[1] for i in desc]
        for r in rows:
            for i in range(len(desc)):
                tmpl=len(str(r[i]))
                if maxlen[i]<tmpl:
                    maxlen[i]=tmpl
        fmt0=''
        fmt1=''
        for (tt,mm) in zip(types,maxlen):
            if tt==cx_Oracle.NUMBER:
                minus=''
            else:
                minus='-'
            fmt0+='%%-%ds '%(mm)
            fmt1+='%%%s%ds '%(minus,mm)

        # header
        line=fmt0%tuple([i[0].lower() for i in desc])
        print line
        print re.sub('[^ ]','-',line)

        # rows
        x=0
        for r in rows:
            x=1-x
            line=fmt1%tuple(r)
            print '%s%s'%(self.colors[x],line)
        print self.colors[2]
    #-------------------------------------------------------------------
    def run(self):
        """execute a command and display results"""
        try:
            self.cmd=re.sub('; *$','',self.cmd)
            rc=self.curs.execute(self.cmd)
            if rc==None:
                print 'ok'
            else:
                self.oraprint(self.curs.description,self.curs.fetchall())
        except cx_Oracle.DatabaseError,e:
            print str(e).strip()
        self.cmd=''
    #-------------------------------------------------------------------
    def default0(self,s):
        """process a line of input"""
        if self.cmd=='':   # first line, assigns
            self.cmd=s
        else:
            self.cmd+='\n' # remaining lines, append
            self.cmd+=s
        if re.search('.*; *$',self.cmd):
            self.run()
    #-------------------------------------------------------------------
    def default(self,s):
        """preprocess a line of input"""
        if re.search('^ *@',s):
            fn=s.split()[0] # get fn
            fn=fn[1:]        # chop @
            try:
                fd=open(fn)
                ll=[x.strip() for x in fd]
                fd.close()
            except IOError,e:
                print '%s: %s'%(fn, str(e))
                ll=['']
        else:
            ll=[s]
        for line in ll:
            self.default0(line)
    #-------------------------------------------------------------------
    def emptyline(self):
        """ignore empty lines"""
        pass
    #-------------------------------------------------------------------
    def complete(self, text, state):
        """cmd line completion.  couldn't get completedefault() to
           work, so jus overwrite the entire completer.  by context,
           it tries to figure out when your typing sql, columns,
           and tables.
        """
        if state == 0:
            buf=readline.get_line_buffer()
            buf=buf[:readline.get_begidx()]
            fullcmd=self.cmd+buf
            #print '<%s>'%(fullcmd)
            if re.search(r'^\s*$',fullcmd,re.I|re.S):
                #cmds
                self.tmpcomp=[i for i in self.cmds if i.startswith(text)]
            elif re.search(r'^.*\s+from\s+$',fullcmd,re.I|re.S):
                #tables
                #add order by, where
                self.tmpcomp=[i for i in self.ltabs if i.startswith(text)]
            elif re.search(r'^\s*desc\s+$',fullcmd,re.I|re.S):
                #tables
                #add order by, where
                self.tmpcomp=[i for i in self.ltabs if i.startswith(text)]
            else:
                #columns
                #add from
                self.tmpcomp=[i for i in self.lcols if i.startswith(text)]
        if state < len(self.tmpcomp):
            rv=self.tmpcomp[state]
        else:
            rv=None
        return rv
    #-----------------------------------------------------------------------
    def oneshot(self,cmds):
        """execute a set of commands one time each"""
        for cmd in cmds:
            self.curs.execute(cmd)
            self.oraprint(self.curs.description,self.curs.fetchall())
    #-------------------------------------------------------------------
    def do_tables(self,s):
        "print a list of the tables"
        print self.ltabs
    #-------------------------------------------------------------------
    def do_desc(self,s):
        "describe a table"
        print "(todo: fix order)"
        print s
        cols=[i for i in self.lcols if self.xcols[i].count(s) > 0]
        print '   ',cols
    #-------------------------------------------------------------------
    def do_EOF(self,s):
        """goodbye -- surely there's a better way to catch eof??"""
        print ''
        if len(self.cmd) > 0:
            pass
            #print 'unexpected EOF!'
        sys.exit(0);
    #-------------------------------------------------------------------
    def do_mono(self,s=None):
        """set monochrome output"""
        self.colors=['','','']
    #-------------------------------------------------------------------
    def do_color(self,s=None):
        """set color output"""
        self.colors=['\033[32m','\033[36m','\033[32m']
    #-------------------------------------------------------------------
    def do_rehash(self,s=None):
        """(re)populate the user's tables/columns"""
        self.xtabs=collections.defaultdict(list)
        self.xcols=collections.defaultdict(list)

        self.curs.execute("""select lower(table_name),lower(column_name)
                             from user_tab_cols""")
        for (tt,cc) in self.curs:
            self.xtabs[tt].append(cc)
            self.xcols[cc].append(tt)
        self.ltabs=[i for i in self.xtabs.keys()]; self.ltabs.sort()
        self.lcols=[i for i in self.xcols.keys()]; self.lcols.sort()
        #self.desc[tblname]=...

#-----------------------------------------------------------------------
def lookup(s):
    """resolve an alias"""
    gg={}
    ll={}
    try:
        execfile(os.environ['HOME']+'/.sqlminus',gg,ll)
    except IOError:
        pass
    if ll.has_key(s):
        s=ll[s]
    return s

#-----------------------------------------------------------------------
def main():
    if len(sys.argv) < 2:
        print >>sys.stderr,'usage: sqlminus connstr [cmds...]'
        sys.exit(1)
    connstr=lookup(sys.argv[1])
    connstr2=re.sub('/.*@','@',connstr)

    print 'connecting to %s...'%(connstr2)
    try:
        cc=OracleCmd(connstr)
    except cx_Oracle.DatabaseError,e:
        print e.args[0]
        sys.exit(1)
    if os.isatty(sys.stdout.fileno()):
        cc.prompt=connstr2+'> '
        cc.do_color()
    if len(sys.argv) >= 3:
        for aa in sys.argv[2:]:
            cc.default(aa)
    else:
        while True:
            try:
                cc.cmdloop()
                break
            except KeyboardInterrupt:
                cc.clearinput()
                print '^C',
                print '(input cleared)'

if __name__=='__main__':
    main()