cassandra/drivers/py/cqlsh

#!/usr/bin/env python

# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from optparse import OptionParser
from StringIO import StringIO

import cmd
import sys
import readline
import os
import re

try:
    import cql
except ImportError:
    sys.path.append(os.path.abspath(os.path.dirname(__file__)))
    import cql
from cql.cursor import _COUNT_DESCRIPTION

HISTORY = os.path.join(os.path.expanduser('~'), '.cqlsh')
CQLTYPES = ("bytes", "ascii", "utf8", "timeuuid", "uuid", "long", "int")

RED = "\033[1;31m%s\033[0m"
GREEN = "\033[1;32m%s\033[0m"
BLUE = "\033[1;34m%s\033[0m"
YELLOW = "\033[1;33m%s\033[0m"
CYAN = "\033[1;36m%s\033[0m"
MAGENTA = "\033[1;35m%s\033[0m"

def startswith(words, text):
    return [i for i in words if i.startswith(text)]

class Shell(cmd.Cmd):
    default_prompt  = "cqlsh> "
    continue_prompt = "   ... "

    def __init__(self, hostname, port, color=False, username=None,
            password=None):
        cmd.Cmd.__init__(self)
        self.conn = cql.connect(hostname, port, user=username, password=password)
        self.cursor = self.conn.cursor()

        if os.path.exists(HISTORY):
            readline.read_history_file(HISTORY)

        if sys.stdin.isatty():
            self.prompt = Shell.default_prompt
        else:
            self.prompt = ""
        
        self.statement = StringIO()
        self.color = color
        self.in_comment = False
        self.in_batch = False
    
    def reset_statement(self):
        self.set_prompt(Shell.default_prompt)
        self.statement.truncate(0)
        
    def get_statement(self, line):
        if self.in_comment:
            if "*/" in line:
                fragment = line[line.index("*/")+2:]
                if fragment.strip():
                    line = fragment
                    self.in_comment = False
                else:
                    self.in_comment = False
                    self.set_prompt(Shell.default_prompt)
                    return None
            else:
                return None
        
        if "/*" in line and (not self.in_comment):
            self.in_comment = True
            self.set_prompt(Shell.continue_prompt)
            if line.lstrip().index("/*") != 0:
                self.statement.write(line[:line.lstrip().index("/*")])
            return None
        
        if ["BEGIN", "BATCH"] == line.upper().split()[:2]:
            self.set_prompt(Shell.continue_prompt)
            self.in_batch = True
            self.statement.write("%s\n" % line)
            return None
        elif (["APPLY", "BATCH"] == line.upper().split()) and self.in_batch:
            self.in_batch = False
            self.statement.write("%s\n" % line)
            try:
                return self.statement.getvalue()
            finally:
                self.reset_statement()
        
        self.statement.write("%s\n" % line)
        
        # The end of a statement    
        if not line.endswith(";") or self.in_batch:
            self.set_prompt(Shell.continue_prompt)
            return None
        
        try:    
            return self.statement.getvalue()
        finally:
            self.reset_statement()

    def default(self, arg):
        def scrub_oneline_comments(s):
            res = re.sub(r'\/\*.*\*\/', '', s)
            res = re.sub(r'--.*$', '', res)
            return res
        
        input = scrub_oneline_comments(arg)
        if not input.strip(): return
        statement = self.get_statement(input)
        if not statement: return

        self.cursor.execute(statement)

        if self.cursor.description is _COUNT_DESCRIPTION:
            if self.cursor.result: print self.cursor.result[0]
        else:
            for x in range(self.cursor.rowcount):
                row = self.cursor.fetchone()
                self.printout(repr(row[0]), BLUE, False)
                for (i, value) in enumerate(row[1:]):
                    name = self.cursor.description[i+1][0]
                    self.printout(" | ", newline=False)
                    self.printout(repr(name), MAGENTA, False)
                    self.printout(",", newline=False)
                    self.printout(repr(value), YELLOW, False)
                self.printout("")

    def emptyline(self):
        pass

    def complete_select(self, text, line, begidx, endidx):
        keywords = ('FIRST', 'REVERSED', 'FROM', 'WHERE', 'KEY')
        return startswith(keywords, text.upper())
    complete_SELECT = complete_select

    def complete_update(self, text, line, begidx, endidx):
        keywords = ('WHERE', 'KEY', 'SET')
        return startswith(keywords, text.upper())
    complete_UPDATE = complete_update

    def complete_create(self, text, line, begidx, endidx):
        words = line.split()
        if len(words) < 3:
            return startswith(['COLUMNFAMILY', 'KEYSPACE'], text.upper())

        common = ['WITH', 'AND']

        if words[1].upper() == 'COLUMNFAMILY':
            types = startswith(CQLTYPES, text)
            keywords = startswith(('KEY', 'PRIMARY'), text.upper())
            props = startswith(("comparator",
                                "comment",
                                "row_cache_size",
                                "key_cache_size",
                                "read_repair_chance",
                                "gc_grace_seconds",
                                "default_validation",
                                "min_compaction_threshold",
                                "max_compaction_threshold",
                                "row_cache_save_period_in_seconds",
                                "key_cache_save_period_in_seconds",
                                "memtable_flush_after_mins",
                                "memtable_throughput_in_mb",
                                "memtable_operations_in_millions",
                                "replicate_on_write"), text)
            return startswith(common, text.upper()) + types + keywords + props

        if words[1].upper() == 'KEYSPACE':
            props = ("replication_factor", "strategy_options", "strategy_class")
            return startswith(common, text.upper()) + startswith(props, text)
    complete_CREATE = complete_create

    def complete_drop(self, text, line, begidx, endidx):
        words = line.split()
        if len(words) < 3:
            return startswith(['COLUMNFAMILY', 'KEYSPACE'], text.upper())
    complete_DROP = complete_drop

    def completenames(self, text, *ignored):
        cmds = startswith(('USE', 'SELECT', 'UPDATE', 'DELETE', 'CREATE', 'DROP'),
                          text.upper())
        return cmd.Cmd.completenames(self, text, ignored) + cmds

    def set_prompt(self, prompt):
        if sys.stdin.isatty():
            self.prompt = prompt

    def do_EOF(self, arg):
        if sys.stdin.isatty(): print
        self.do_exit(None)

    def do_exit(self, arg):
        sys.exit()
    do_quit = do_exit

    def printout(self, text, color=None, newline=True, out=sys.stdout):
        if not color or not self.color:
            out.write(text)
        else:
            out.write(color % text)

        if newline:
            out.write("\n");

    def printerr(self, text, color=None, newline=True):
        self.printout(text, color, newline, sys.stderr)

if __name__ == '__main__':
    parser = OptionParser(usage = "Usage: %prog [host [port]]")
    parser.add_option("-C",
                      "--color",
                      action="store_true",
                      help="Enable color output.")
    parser.add_option("-u", "--username", help="Authenticate as user.")
    parser.add_option("-p", "--password", help="Authenticate using password.")
    (options, arguments) = parser.parse_args()

    hostname = len(arguments) > 0 and arguments[0] or "localhost"

    if len(arguments) > 1:
        try:
            port = int(arguments[1])
        except ValueError:
            print >>sys.stderr, "%s is not a valid port number" % arguments[1]
            parser.print_help(file=sys.stderr)
            sys.exit(1)
    else:
        port = 9160


    shell = Shell(hostname,
                  port,
                  color=options.color,
                  username=options.username,
                  password=options.password)
    while(True):
        try:
            shell.cmdloop()
        except SystemExit:
            readline.write_history_file(HISTORY)
            break
        except cql.Error, cqlerr:
            shell.printerr(str(cqlerr), color=RED)
        except KeyboardInterrupt:
            shell.reset_statement()
            print
        except Exception, err:
            shell.printerr("Exception: %s" % err, color=RED)