sqltool/sql.py

#!/usr/bin/python3

import os
import sys
import sqlalchemy
import readline
import traceback


def softpad(s, l, padchar=" "):
    padding = padchar * l
    if len(s) >= l:
        return s + padchar
    return (s + padding)[:l]


def pad(s, l, padchar=" "):
    padding = padchar * l
    return (s + padding)[:l]


rowbuffer = None


def execute_cmdline_query(v):
    global rowbuffer
    rowbuffer = None
    try:
        v = v.rstrip(';')
        cur.execute(v)
        print_row()
    except:
        try:
            c2 = db.cursor()
            c2.execute(
                "SELECT 1 FROM INFORMATION_SCHEMA.TABLES WHERE 1 IS NULL")
        except:
            print("Database connection lost!")
            sys.exit(1)
        traceback.print_exc()
        sys.stdout.write('\n')


def print_row():
    global rowbuffer, cur, db
    if not cur or not cur.description:
        return False
    try:
        page = int(opts['page'])
        colhdrs = [x[0] for x in cur.description]
        allcols = []
        lens = []
        for i in range(len(colhdrs)):
            l = min(len(colhdrs[i]), 10)
            lens.append(l)
        if rowbuffer:
            rowbuffer = [repr(x) for x in rowbuffer]
            allcols = [rowbuffer]
            for i in range(len(rowbuffer)):
                l = len(rowbuffer[i])
                if i > len(lens):
                    print(lens)
                    print(rowbuffer)
                if l > lens[i]:
                    lens[i] = l
            rowbuffer = None
        while len(allcols) < page:
            cols = cur.fetchone()
            if not cols:
                break
            cols = [repr(x) for x in cols]
            allcols.append(cols)
            for i in range(len(cols)):
                l = len(cols[i])
                if l > lens[i]:
                    lens[i] = l
        if not allcols:
            return False
        for i in range(len(colhdrs)):
            colhdrs[i] = pad(colhdrs[i], lens[i])

        for i in range(len(allcols)):
            cols = allcols[i]
            for j in range(len(cols)):
                cols[j] = pad(cols[j], lens[j])

        colhdrrow = ' | '.join(colhdrs)
        seprow = '-' * len(colhdrrow)
        print(colhdrrow)
        print(seprow)
        for cols in allcols:
            colsrow = ' | '.join(cols)
            print(colsrow)

        rowbuffer = cur.fetchone()
        if rowbuffer:
            print("  ( more )  ")
        return True
    except:
        traceback.print_exc()
        return False


opts = {
    'page': '20',
    'hide_null_tablespace': '1',
    'hide_tablespaces': 'SYSTEM;SYSAUX;SDE;SPATIAL_DATA'
}

connstr = ""
if os.path.exists('settings.ini'):
    try:
        from cecil.core import inifile
        config = inifile.inifile('settings.ini')
        connstr = config['settings']['db']
    except None:
        pass
cmdline_query = None
if not connstr:
    if len(sys.argv) < 2:
        if not connstr:
            connstr = input('Connect string: ')
    else:
        connstr = sys.argv[-1]
elif len(sys.argv) > 1:
    cmdline_query = sys.argv[-1]

e = sqlalchemy.create_engine(connstr)
pool = e.connect()
db = pool.connection
cur = db.cursor()
c2 = db.cursor()

if cmdline_query != None:
    execute_cmdline_query(cmdline_query)
    db.commit()
    sys.exit(0)

do_quit = False
while True:
    try:
        v = input('# ')
    except EOFError:
        do_quit = True

    if do_quit:
        break
    v = v.strip()
    vl = v.lower()

    if vl == 'quit' or vl == 'exit':
        break
    elif vl == 'commit':
        db.commit()
    elif vl == 'rollback':
        db.rollback()
    elif vl[:4] == 'copy' or vl[:5] == 'copyi':
        vs = v.split(' ')
        try:
            configsrc = inifile.inifile(vs[1])
            connstrsrc = configsrc['settings']['db']
            esrc = sqlalchemy.create_engine(connstrsrc)
            poolsrc = e.connect()
            dbsrc = pool.connection
            cursrc = db.cursor()

            c2 = db.cursor()
            #c2.execute("select * from %s where 1 is null" % (vs[2],))
            cursrc.execute("select * from %s" % (vs[2],))
            colnames = []
            dstart = 1
            if vs[0] == "copyi":
                dstart = 0
            for coldata in cursrc.description[dstart:]:
                colnames.append(coldata[0])
            rowload = 0
            rowerror = 0
            for r in cursrc:
                try:
                    colnamesql = ", ".join(colnames)
                    colmarks = ", ".join("%s" for cn in colnames)
                    srcq = "insert into %s (%s) values (%s)" % (vs[2], colnamesql, colmarks)
                    print(srcq)
                    print(r[0])
                    print(r[1])
                    print(r[2])
                    print(len(r))
                    c2.execute(srcq, tuple(r[dstart:]))
                except:
                    traceback.print_exc()
                    rowerror += 1
            print(f"Rows loaded {rowload}, errors {rowerror}")
        except:
            traceback.print_exc()
    elif vl == 'tables':
        rowbuffer = None
        cur.execute("SELECT TABLE_NAME FROM INFORMATION_SCHEMA.TABLES")
        print_row()
    elif vl[:7] == 'tables ':
        v2 = v[7:]
        rowbuffer = None
        cur.execute("SELECT TABLE_NAME FROM INFORMATION_SCHEMA.TABLES WHERE TABLE_NAME LIKE %s" % (v2,))
        print_row()
    elif vl[:5] == 'cols ':
        rowbuffer = None
        vs = v.split(' ')
        c2 = db.cursor()
        try:
            c2.execute("select * from %s where 1 is null" % (vs[1],))

            for coldata in c2.description:
                typedata = coldata[1]
                try:
                    import pytds.tds_base
                    tdlist = []
                    for x in dir(pytds.tds_base):
                        if getattr(pytds.tds_base, x) == typedata:
                            tdlist.append(x)
                    if tdlist:
                        tdlist2 = [x.replace("NTYPE", "").replace(
                            "TYPE", "") for x in tdlist if "TYPE" in x]
                        if tdlist2:
                            tdlist = tdlist2
                    if tdlist:
                        tdlist.sort()
                        tdlist.sort(key=len)
                        typedata = " ".join(tdlist)

                except:
                    pass

                print("%s%s" % (softpad(coldata[0], 40), typedata))
        except:
            traceback.print_exc()
            sys.stdout.write('\n')
    elif vl[:4] == 'set ':
        vs = v[4:]
        vi = vs.find('=')
        vn = vs[:vi]
        vv = vs[vi+1:]
        opts[vn] = vv
    elif vl == '':
        print_row()
    elif vl == 'help':
        print("""
commands:
  tables            lists tables
  cols <table>      lists columns of table
  commit            commits any pending updates or inserts
  rollback          aborts and pending updates or inserts
  set <key>=<val>   sets internal options, like "set page=40" to page at 40 rows
  help              displays this help
  copy <path> <tbl> reads the config from <path> connects to that database and copies the contents of <tbl>
  quit              exits sql.py
  <empty line>      next page if there are (MORE) pages, just hit "ENTER"
  <anything else>   executes SQL query on the server
""")
    else:
        execute_cmdline_query(v)