#!/usr/bin/env python # # Copyright 2007 Google Inc. # # Licensed 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. # """Command line client for Google SQL Service. This module makes use of the sqlcmd Python library (with some minor modifications to support alternate commands and syntax), and defines the following classes: DatabaseConfig: A small container object to pass around database configuration settings, like the Google SQL instance ID and database name. The sqlcmd APIs don't allow for extra, database specific configuration variables, so we combine the instance ID with the database name, and pass this object as the value for database. GoogleSqlDriver: Grizzled DB Driver for Google SQL Service. sqlcmd makes use of the Grizzled Python library to handle database connectivity, and we provide an appropriate driver for connecting to a Google SQL instance. GoogleSqlCmd: A subclass of the sqlcmd.SQLCmd class, which is the main class of the application. It overrides a few pieces of the parent class to change some UI elements, command syntax, and support for additional commands. """ import logging import optparse import os import sys from google.storage.speckle.python.api import rdbms_googleapi from grizzled import db from grizzled.db import mysql import prettytable import sqlcmd from sqlcmd import config logging.basicConfig(level=logging.WARNING) sqlcmd.log = logging.getLogger('google_sql') sqlcmd.DEFAULT_CONFIG_DIR = os.path.join(os.environ.get('HOME', os.getcwd()), '.googlesql') sqlcmd.RC_FILE = os.path.join(sqlcmd.DEFAULT_CONFIG_DIR, 'config') sqlcmd.HISTORY_FILE_FORMAT = os.path.join(sqlcmd.DEFAULT_CONFIG_DIR, '%s.hist') sqlcmd.INTRO = 'Google SQL Client\n\nType "help" or "?" for help.\n' DEFAULT_ENCODING = 'utf-8' USAGE = '%prog [options] instance [database]' class DatabaseConfig(object): """A small configuration object used to pass around database settings. Attributes: instance: The Google SQL instance ID. name: The database name to use. oauth_credentials_path: A location to use for oauth credentials storage instead of the default """ def __init__(self, instance, name=None, oauth_credentials_path=None): self.instance = instance self.name = name self.oauth_credentials_path = oauth_credentials_path def __str__(self): result = self.instance if self.name: result = '%s|%s' % (result, self.name) return result class GoogleSqlDriver(mysql.MySQLDriver): """Grizzled DB Driver for Google SQL Service.""" NAME = 'googlesql' def get_import(self): return rdbms_googleapi def get_display_name(self): return 'Google SQL' def do_connect(self, host, port, user, password, database): """Connect to the actual underlying database, using the driver. Args: host: The host where the database lives. port: The TCP port to use when connecting. (UNUSED) user: The user to use when connecting. password: The password to use when connecting. database: A DatabaseConfig instance containing the instance id and database name. Returns: The DB API Connection object. """ dbi = self.get_import() return dbi.connect( host, database.instance, database=database.name, user=user, password=password, oauth_credentials_path=database.oauth_credentials_path) class GoogleSqlCmd(sqlcmd.SQLCmd): """The SQLCmd command interpreter for Google SQL.""" sqlcmd.SQLCmd.MAIN_PROMPT = 'sql> ' sqlcmd.SQLCmd.CONTINUATION_PROMPT = ' -> ' sqlcmd.SQLCmd.NO_SEMI_NEEDED.update( ['about', 'connect', 'desc', 'describe', 'echo', 'exit', 'h', 'hist', 'history', 'load', 'quit', 'run', 'set', 'show', 'var', 'vars']) for method in [x for x in dir(sqlcmd.SQLCmd) if x.startswith('do_dot_')]: setattr(sqlcmd.SQLCmd, method.replace('dot_', ''), getattr( sqlcmd.SQLCmd, method)) delattr(sqlcmd.SQLCmd, method) def __init__(self, *args, **kwargs): sqlcmd.SQLCmd.__init__(self, *args, **kwargs) self.prompt = sqlcmd.SQLCmd.MAIN_PROMPT self.output_encoding = DEFAULT_ENCODING def set_output_encoding(self, encoding): self.output_encoding = encoding def do_quit(self, args): """Quit the google_sql. Same as exit.""" return self.do_exit(args) def _BuildTable(self, cursor): """Builds an output PrettyTable from the results in the given cursor.""" if not cursor.description: return None column_names = [column[0] for column in cursor.description] table = prettytable.PrettyTable(column_names) rows = cursor.fetchall() if not rows: return table for i, col in enumerate(rows[0]): table.set_field_align( column_names[i], isinstance(col, basestring) and 'l' or 'r') for row in rows: table.add_row(row) return table def _SQLCmd__handle_select(self, args, cursor, command='select'): """Overrides SQLCmd.__handle_select to display output with prettytable.""" self._SQLCmd__exec_SQL(cursor, command, args) table = self._BuildTable(cursor) if table: output = table.get_string() if isinstance(output, unicode): print output.encode(self.output_encoding) else: print output def _CreateConfigDir(): """Creates the sqlcmd config directory if necessary.""" directory = sqlcmd.DEFAULT_CONFIG_DIR if not os.access(directory, os.R_OK | os.W_OK | os.X_OK): old_umask = os.umask(077) os.makedirs(sqlcmd.DEFAULT_CONFIG_DIR) os.umask(old_umask) def main(argv): parser = optparse.OptionParser(usage=USAGE) parser.add_option('-e', '--output_encoding', dest='output_encoding', default=DEFAULT_ENCODING, help='Output encoding. Defaults to %s.' % DEFAULT_ENCODING) parser.add_option('--oauth_credentials_path', dest='oauth_credentials_path', default=None, help=optparse.SUPPRESS_HELP) (options, args) = parser.parse_args(argv[1:]) if len(args) < 1 or len(args) > 2: parser.print_help(sys.stderr) sys.exit(1) instance = args[0] instance_alias = instance.replace(':', '#') database = None if len(args) == 2: database = args[1] _CreateConfigDir() database = DatabaseConfig(instance, database, options.oauth_credentials_path) db.add_driver(GoogleSqlDriver.NAME, GoogleSqlDriver) sql_cmd_config = config.SQLCmdConfig(None) sql_cmd_config.add('__googlesql__', instance_alias, None, None, database, GoogleSqlDriver.NAME, None, None) sql_cmd = GoogleSqlCmd(sql_cmd_config) sql_cmd.set_output_encoding(options.output_encoding) sql_cmd.set_database(instance_alias) sql_cmd.cmdloop() if __name__ == '__main__': main(sys.argv)