June 21, 2014

How to Access Postgresql from the Command Line

Tonight, I wrote myself a postgres interface in python. For this script, I've taken to incorporating readline for vi or emacs keybindings, completion, and command-line history. Still working on it, but I've learned a lot so far:

#!python
# -*- coding: utf-8 -*-
import argparse
import curses
import posix
import logging
import sys
import readline

import psycopg2


class SimpleCompleter(object):

    def __init__(self, options):
        self.options = sorted(options)
        return

    def complete(self, text, state):
        response = None
        if state == 0:
            # This is the first time for this text, so build a match list.
            if text:
                self.matches = [s
                                for s in self.options
                                if s and s.startswith(text)]
                logging.debug('%s matches: %s', repr(text), self.matches)
            else:
                self.matches = self.options[:]
                logging.debug('(empty input) matches: %s', self.matches)

        # Return the state'th item from the match list,
        # if we have that many.
        try:
            response = self.matches[state]
        except IndexError:
            response = None
        logging.debug('complete(%s, %s) => %s',
                      repr(text), state, repr(response))
        return response


if __name__ == '__main__':
    argparser = argparse.ArgumentParser(description='PostgreSQL console client, by Hasan Diwan')
    argparser.add_argument('-H', '--host', default='', help='Database Hostname', type=unicode)
    argparser.add_argument('-u', '--user', default=posix.getlogin(), help='Database username', type=unicode)
    argparser.add_argument('-p', '--password', default='', help='Database password', type=unicode)
    argparser.add_argument('-d', '--database', default=posix.getlogin(), help='Database to connect to', type=unicode)
    argparser.add_argument('-v', '--verbose', help='verbose logging', action='store_false')
    argparser.add_argument('-e', '--editor', default='emacs', help='Use vi or emacs keybindings', type=unicode)

    parsed = argparser.parse_args()

    if parsed.verbose is False:
        logging.basicConfig(level=logging.DEBUG)
    else:
        logging.basicConfig(level=logging.FATAL)

    authstring = ''
    if parsed.password:
        authstring = 'password={}'.format(parsed.password)
    if parsed.host:
        authstring = 'host={} {}'.format(parsed.host, authstring)

    # Readline configuration -- cribbed from http://pymotw.com/2/readline/
    readline.parse_and_bind('set editing-mode {}'.format(parsed.editor))
    readline.parse_and_bind('tab: complete')
    try:
        while True:
            conn = psycopg2.connect('dbname={} user={} {}'.format(parsed.database, parsed.user, authstring))
            cursor = conn.cursor()
            cursor.execute("select relname from pg_class where relkind='r' and relname !~ '^(pg_|sql_)';")
            # tablenames
            tables = [t[0] for t in cursor.fetchall()]
            readline.set_completer(SimpleCompleter(tables).complete)
            sql = raw_input('SQL (EOF to exit)> ')
            cursor = conn.cursor()
            try:
                cursor.execute(sql)
            except Exception, e:
                print str(e)
                conn.rollback()
                conn.close()
                continue
            # Column names
            columns = []
            map((lambda(column): columns.append(column[0])), cursor.description)
            print [x for x in columns]
            print [r for r in cursor.fetchall()[0]]
            conn.commit()
            conn.close()
    except EOFError, e:
        exit()

No comments:

Post a Comment