pycassaShell

#!/usr/bin/env python

"""
interactive Cassandra Python shell

"""

try:
    from IPython.Shell import IPShellEmbed
except ImportError:
    print "[I]: IPython not found, falling back to default interpreter."

def runshell():
    import os
    os.environ['PYTHONINSPECT'] = '1'

try:
    runshell = IPShellEmbed([])
except NameError:
    pass

import pycassa, optparse
from pycassa.system_manager import *
from sys import stdout, stderr, exit

_pool = None

def _update_cf(cf, delete=False):
    if _pool is not None:
        if not delete:
            existed = cf.upper() in globals()
            globals()[cf.upper()] = pycassa.ColumnFamily(_pool, cf)
            if not existed:
                print "\nLoaded %s as %s" % (cf, cf.upper())
            else:
                print "\nReloaded %s" % cf.upper()
        else:
            globals().pop(cf.upper())
            print "\nDropped %s" % cf.upper()


parser = optparse.OptionParser(usage='Usage: %prog [OPTIONS]')
parser.add_option('-k', '--keyspace', help='Cassandra keyspace name.')
parser.add_option('-H', '--host', help='Hostname.')
parser.add_option('-p', '--port', type="int", help='Thrift port number.')
parser.add_option('-u', '--user', help='Username (for simple auth).')
parser.add_option('-P', '--passwd', help='Password (for simple auth).')
parser.add_option('-S', '--streaming', help='Using streaming transport.',
                  action="store_false", dest='framed')
parser.add_option('-F', '--framed', 
                  help='Use framed transport.  Default transport.',
                  action="store_true", dest='framed')
parser.add_option('-f', '--file', help='Run a script after startup')

(options, args) = parser.parse_args()


hostname = options.host and options.host or 'localhost'
port = options.port and options.port or 9160
framed = True if options.framed is None else options.framed
credentials = None

if options.user or options.passwd:
    if options.user and (not options.passwd):
        print >>stderr, "You must supply a password for username", options.user
        exit(1)
    if options.passwd and (not options.user):
        print >>stderr, "You need a user to go with that password!"
        exit(1)
    credentials = {'username': options.user, 'password': options.passwd}

SYSTEM_MANAGER = pycassa.SystemManager('%s:%d' % (hostname, port), credentials, framed, _update_cf)

print "----------------------------------"
print "Cassandra Interactive Python Shell"
print "----------------------------------"
print "Keyspace: %s" % options.keyspace
print "Host: %s:%d" % (hostname, port)

if options.keyspace:
    _pool = pycassa.QueuePool(keyspace=options.keyspace,
                              server_list=['%s:%d' % (hostname, port)],
                              credentials=credentials,
                              framed_transport=framed,
                              timeout=5)

    print "\nAvailable ColumnFamily instances:"
    for cfname in SYSTEM_MANAGER.get_keyspace_description(options.keyspace).keys():
        cfinstance = pycassa.ColumnFamily(_pool, cfname)
        exec('%s = cfinstance' % cfname.upper())
        spaces = " " * (25 - len(cfname))
        print " *", cfname.upper(), spaces, "(", cfname, ")"
else:
    print "\nColumnFamily instances are only available if a keyspace is specified with -k/--keyspace"

print "\nSchema definition tools and cluster information are available through SYSTEM_MANAGER."

if (options.file):
    print "\nExecuting script ...",
    stdout.flush()
    f = open(options.file, 'r')
    line = f.readline()
    while line != '':
        exec(line);
        line = f.readline()
    print " done."

runshell()