From 031e04d02ed2897b0997dabc59de3edf07a3618c Mon Sep 17 00:00:00 2001 From: erickson Date: Mon, 20 Dec 2010 01:20:48 +0000 Subject: [PATCH] Python srfsh enhancements * Srfsh plugins can now insert new commands and add words to the tab completion word bank. * Addded support reading script files * Added support for service open/close (connect/disconnect) for stateful connections * Moved to class-based srfsh module for easier state maintenance * More doc strings git-svn-id: svn://svn.open-ils.org/OpenSRF/trunk@2133 9efc2488-bf62-4759-914b-345cdb29e865 --- src/python/srfsh.py | 706 +++++++++++++++++++++++++------------------- 1 file changed, 404 insertions(+), 302 deletions(-) diff --git a/src/python/srfsh.py b/src/python/srfsh.py index d1c3c04..519c12b 100755 --- a/src/python/srfsh.py +++ b/src/python/srfsh.py @@ -39,393 +39,495 @@ srfsh.py - provides a basic shell for issuing OpenSRF requests SRFSH_LOCALE = - request responses to be returned in locale if available """ - -import os, sys, time, readline, atexit, re, pydoc +import os, sys, time, readline, atexit, re, pydoc, traceback import osrf.json, osrf.system, osrf.ses, osrf.conf, osrf.log, osrf.net -router_command_map = { - 'services' : 'opensrf.router.info.class.list', - 'service-stats' : 'opensrf.router.info.stats.class.node.all', - 'service-nodes' : 'opensrf.router.info.stats.class.all' -} - -# List of words to use for readline tab completion -tab_complete_words = [ - 'request', - 'set', - 'router', - 'help', - 'exit', - 'quit', - 'introspect', - 'opensrf.settings', - 'opensrf.math' -] - -# add the router commands to the tab-complete list -for rcommand in router_command_map.keys(): - tab_complete_words.append(rcommand) - -# ------------------------------------------------------------------- -# main listen loop -# ------------------------------------------------------------------- -def do_loop(): - - command_map = { - 'request' : handle_request, - 'router' : handle_router, - 'math_bench' : handle_math_bench, - 'introspect' : handle_introspect, - 'help' : handle_help, - 'set' : handle_set, - 'get' : handle_get, - } - - while True: +class Srfsh(object): - try: - report("", True) - line = raw_input("srfsh# ") - if not len(line): - continue + def __init__(self, script_file=None): - if str.lower(line) == 'exit' or str.lower(line) == 'quit': - break + # used for paging + self.output_buffer = '' - parts = str.split(line) - command = parts.pop(0) + # true if invoked with a script file + self.reading_script = False - if command not in command_map: - report("unknown command: '%s'\n" % command) - continue + # multi-request sessions + self.active_session = None - command_map[command](parts) + # default opensrf request timeout + self.timeout = 120 - except EOFError: # ^-d - sys.exit(0) + # map of command name to handler + self.command_map = {} - except KeyboardInterrupt: # ^-c - report("\n") + if script_file: + self.open_script(script_file) + self.reading_script = True - except Exception, e: - report("%s\n" % e) + # map of router sub-commands to router API calls + self.router_command_map = { + 'services' : 'opensrf.router.info.class.list', + 'service-stats' : 'opensrf.router.info.stats.class.node.all', + 'service-nodes' : 'opensrf.router.info.stats.class.all' + } - cleanup() + # seed the tab completion word bank + self.tab_complete_words = self.router_command_map.keys() + [ + 'exit', + 'quit', + 'opensrf.settings', + 'opensrf.math', + 'opensrf.dbmath', + 'opensrf.py-example' + ] -def handle_introspect(parts): + # add the default commands + for command in ['request', 'router', 'help', 'set', + 'get', 'math_bench', 'introspect', 'connect', 'disconnect' ]: - if len(parts) == 0: - report("usage: introspect [api_prefix]\n") - return + self.add_command(command = command, handler = getattr(Srfsh, 'handle_' + command)) - service = parts.pop(0) - args = [service, 'opensrf.system.method'] + # for compat w/ srfsh.c + self.add_command(command = 'open', handler = Srfsh.handle_connect) + self.add_command(command = 'close', handler = Srfsh.handle_disconnect) - if len(parts) > 0: - api_pfx = parts[0] - if api_pfx[0] != '"': # json-encode if necessary - api_pfx = '"%s"' % api_pfx - args.append(api_pfx) - else: - args[1] += '.all' + def open_script(self, script_file): + ''' Opens the script file and redirects the contents to STDIN for reading. ''' - return handle_request(args) + try: + script = open(script_file, 'r') + os.dup2(script.fileno(), sys.stdin.fileno()) + script.close() + except Exception, e: + self.report_error("Error opening script file '%s': %s" % (script_file, str(e))) + raise e -def handle_router(parts): + def main_loop(self): + ''' Main listen loop. ''' - if len(parts) == 0: - report("usage: router \n") - return + self.set_vars() + self.do_connect() + self.load_plugins() + self.setup_readline() - query = parts[0] + while True: - if query not in router_command_map: - report("router query options: %s\n" % ','.join(router_command_map.keys())) - return + try: + self.report("", True) + line = raw_input("srfsh# ") - return handle_request(['router', router_command_map[query]]) + if not len(line): + continue -# ------------------------------------------------------------------- -# Set env variables to control behavior -# ------------------------------------------------------------------- -def handle_set(parts): - cmd = "".join(parts) - pattern = re.compile('(.*)=(.*)').match(cmd) - key = pattern.group(1) - val = pattern.group(2) - set_var(key, val) - report("%s = %s\n" % (key, val)) + if re.search('^\s*#', line): # ignore lines starting with # + continue -def handle_get(parts): - try: - report("%s=%s\n" % (parts[0], get_var(parts[0]))) - except: - report("\n") + if str.lower(line) == 'exit' or str.lower(line) == 'quit': + break + parts = str.split(line) + command = parts.pop(0) -# ------------------------------------------------------------------- -# Prints help info -# ------------------------------------------------------------------- -def handle_help(foo): - report(__doc__) + if command not in self.command_map: + self.report("unknown command: '%s'\n" % command) + continue -# ------------------------------------------------------------------- -# performs an opensrf request -# ------------------------------------------------------------------- -def handle_request(parts): + self.command_map[command](self, parts) - if len(parts) < 2: - report("usage: request [, , ...]\n") - return + except EOFError: # ctrl-d + break - service = parts.pop(0) - method = parts.pop(0) - locale = __get_locale() - jstr = '[%s]' % "".join(parts) - params = None + except KeyboardInterrupt: # ctrl-c + self.report("\n") - try: - params = osrf.json.to_object(jstr) - except: - report("Error parsing JSON: %s\n" % jstr) - return + except Exception, e: + self.report("%s\n" % traceback.format_exc()) - ses = osrf.ses.ClientSession(service, locale=locale) + self.cleanup() - start = time.time() + def handle_connect(self, parts): + ''' Opens a connected session to an opensrf service ''' - req = ses.request2(method, tuple(params)) + if len(parts) == 0: + self.report("usage: connect ") + return + service = parts.pop(0) - while True: - resp = None + if self.active_session: + if self.active_session['service'] == service: + return # use the existing active session + else: + # currently, we only support one active session at a time + self.handle_disconnect([self.active_session['service']]) - try: - resp = req.recv(timeout=120) - except osrf.net.XMPPNoRecipient: - report("Unable to communicate with %s\n" % service) - total = 0 - break + self.active_session = { + 'ses' : osrf.ses.ClientSession(service, locale = self.__get_locale()), + 'service' : service + } - if not resp: break + self.active_session['ses'].connect() - total = time.time() - start - content = resp.content() + def handle_disconnect(self, parts): + ''' Disconnects the currently active session. ''' - if content is not None: - if get_var('SRFSH_OUTPUT_NET_OBJ_KEYS') == 'true': - report("Received Data: %s\n" % osrf.json.debug_net_object(content)) - else: - if get_var('SRFSH_OUTPUT_FORMAT_JSON') == 'true': - report("Received Data: %s\n" % osrf.json.pprint(osrf.json.to_json(content))) - else: - report("Received Data: %s\n" % osrf.json.to_json(content)) + if len(parts) == 0: + self.report("usage: disconnect ") + return - req.cleanup() - ses.cleanup() + service = parts.pop(0) - report('-'*60 + "\n") - report("Total request time: %f\n" % total) - report('-'*60 + "\n") + if self.active_session: + if self.active_session['service'] == service: + self.active_session['ses'].disconnect() + self.active_session['ses'].cleanup() + self.active_session = None + else: + self.report_error("There is no open connection for service '%s'" % service) + def handle_introspect(self, parts): + ''' Introspect an opensrf service. ''' -def handle_math_bench(parts): + if len(parts) == 0: + self.report("usage: introspect [api_prefix]\n") + return - count = int(parts.pop(0)) - ses = osrf.ses.ClientSession('opensrf.math') - times = [] + service = parts.pop(0) + args = [service, 'opensrf.system.method'] - for cnt in range(100): - if cnt % 10: - sys.stdout.write('.') + if len(parts) > 0: + api_pfx = parts[0] + if api_pfx[0] != '"': # json-encode if necessary + api_pfx = '"%s"' % api_pfx + args.append(api_pfx) else: - sys.stdout.write( str( cnt / 10 ) ) - print "" + args[1] += '.all' - for cnt in range(count): - - starttime = time.time() - req = ses.request('add', 1, 2) - resp = req.recv(timeout=2) - endtime = time.time() - - if resp.content() == 3: - sys.stdout.write("+") - sys.stdout.flush() - times.append( endtime - starttime ) + return handle_request(args) + + + def handle_router(self, parts): + ''' Send requests to the router. ''' + + if len(parts) == 0: + self.report("usage: router \n") + return + + query = parts[0] + + if query not in self.router_command_map: + self.report("router query options: %s\n" % ','.join(self.router_command_map.keys())) + return + + return handle_request(['router', self.router_command_map[query]]) + + def handle_set(self, parts): + ''' Set env variables to control srfsh behavior. ''' + + cmd = "".join(parts) + pattern = re.compile('(.*)=(.*)').match(cmd) + key = pattern.group(1) + val = pattern.group(2) + self.set_var(key, val) + self.report("%s = %s\n" % (key, val)) + + def handle_get(self, parts): + ''' Returns environment variable value ''' + try: + self.report("%s=%s\n" % (parts[0], self.get_var(parts[0]))) + except: + self.report("\n") + + + def handle_help(self, foo): + ''' Prints help info ''' + self.report(__doc__) + + def handle_request(self, parts): + ''' Performs an OpenSRF request and reports the results. ''' + + if len(parts) < 2: + self.report("usage: request [, , ...]\n") + return + + self.report("\n") + + service = parts.pop(0) + method = parts.pop(0) + locale = self.__get_locale() + jstr = '[%s]' % "".join(parts) + params = None + + try: + params = osrf.json.to_object(jstr) + except: + self.report("Error parsing JSON: %s\n" % jstr) + return + + using_active = False + if self.active_session and self.active_session['service'] == service: + # if we have an open connection to the same service, use it + ses = self.active_session['ses'] + using_active = True else: - print "What happened? %s" % str(resp.content()) - - req.cleanup() - if not ( (cnt + 1) % 100): - print ' [%d]' % (cnt + 1) - - ses.cleanup() - total = 0 - for cnt in times: - total += cnt - print "\naverage time %f" % (total / len(times)) + ses = osrf.ses.ClientSession(service, locale=locale) + start = time.time() + req = ses.request2(method, tuple(params)) + last_content = None + while True: + resp = None -# ------------------------------------------------------------------- -# Defines the tab-completion handling and sets up the readline history -# ------------------------------------------------------------------- -def setup_readline(): + try: + resp = req.recv(timeout=self.timeout) + except osrf.net.XMPPNoRecipient: + self.report("Unable to communicate with %s\n" % service) + total = 0 + break - class SrfshCompleter(object): + if not resp: break - def __init__(self, words): - self.words = words - self.prefix = None - - def complete(self, prefix, index): + total = time.time() - start + content = resp.content() - if prefix != self.prefix: + if content is not None: + last_content = content + if self.get_var('SRFSH_OUTPUT_NET_OBJ_KEYS') == 'true': + self.report("Received Data: %s\n" % osrf.json.debug_net_object(content)) + else: + if self.get_var('SRFSH_OUTPUT_FORMAT_JSON') == 'true': + self.report("Received Data: %s\n" % osrf.json.pprint(osrf.json.to_json(content))) + else: + self.report("Received Data: %s\n" % osrf.json.to_json(content)) - self.prefix = prefix + req.cleanup() + if not using_active: + ses.cleanup() - # find all words that start with this prefix - self.matching_words = [ - w for w in self.words if w.startswith(prefix) - ] + self.report("\n" + '-'*60 + "\n") + self.report("Total request time: %f\n" % total) + self.report('-'*60 + "\n") - if len(self.matching_words) == 0: - return None + return last_content + + + def handle_math_bench(self, parts): + ''' Sends a series of request to the opensrf.math service and collects timing stats. ''' + + count = int(parts.pop(0)) + ses = osrf.ses.ClientSession('opensrf.math') + times = [] + + for cnt in range(100): + if cnt % 10: + sys.stdout.write('.') + else: + sys.stdout.write( str( cnt / 10 ) ) + print "" + + for cnt in range(count): + + starttime = time.time() + req = ses.request('add', 1, 2) + resp = req.recv(timeout=2) + endtime = time.time() + + if resp.content() == 3: + sys.stdout.write("+") + sys.stdout.flush() + times.append( endtime - starttime ) + else: + print "What happened? %s" % str(resp.content()) + + req.cleanup() + if not ( (cnt + 1) % 100): + print ' [%d]' % (cnt + 1) + + ses.cleanup() + total = 0 + for cnt in times: + total += cnt + print "\naverage time %f" % (total / len(times)) - if len(self.matching_words) == 1: - return self.matching_words[0] - sys.stdout.write('\n%s\nsrfsh# %s' % - (' '.join(self.matching_words), readline.get_line_buffer())) - return None - completer = SrfshCompleter(tuple(tab_complete_words)) - readline.parse_and_bind("tab: complete") - readline.set_completer(completer.complete) + def setup_readline(self): + ''' Initialize readline history and tab completion. ''' - histfile = os.path.join(get_var('HOME'), ".srfsh_history") - try: - readline.read_history_file(histfile) - except IOError: - pass - atexit.register(readline.write_history_file, histfile) + class SrfshCompleter(object): - readline.set_completer_delims(readline.get_completer_delims().replace('-','')) + def __init__(self, words): + self.words = words + self.prefix = None + + def complete(self, prefix, index): + if prefix != self.prefix: -def do_connect(): - file = os.path.join(get_var('HOME'), ".srfsh.xml") - osrf.system.System.connect(config_file=file, config_context='srfsh') + self.prefix = prefix -def load_plugins(): - # Load the user defined external plugins - # XXX Make this a real module interface, with tab-complete words, commands, etc. - try: - plugins = osrf.conf.get('plugins') + # find all words that start with this prefix + self.matching_words = [ + w for w in self.words if w.startswith(prefix) + ] - except: - # XXX standard srfsh.xml does not yet define element - #print("No plugins defined in /srfsh/plugins/plugin\n") - return + if len(self.matching_words) == 0: + return None - plugins = osrf.conf.get('plugins.plugin') - if not isinstance(plugins, list): - plugins = [plugins] + if len(self.matching_words) == 1: + return self.matching_words[0] + + # re-print the prompt w/ all of the possible word completions + sys.stdout.write('\n%s\nsrfsh# %s' % + (' '.join(self.matching_words), readline.get_line_buffer())) + + return None - for module in plugins: - name = module['module'] - init = module['init'] - print "Loading module %s..." % name + completer = SrfshCompleter(tuple(self.tab_complete_words)) + readline.parse_and_bind("tab: complete") + readline.set_completer(completer.complete) + histfile = os.path.join(self.get_var('HOME'), ".srfsh_history") try: - string = 'import %s\n%s.%s()' % (name, name, init) - exec(string) - print 'OK' + readline.read_history_file(histfile) + except IOError: + pass + atexit.register(readline.write_history_file, histfile) - except Exception, e: - sys.stderr.write("\nError importing plugin %s, with init symbol %s: \n%s\n" % (name, init, e)) + readline.set_completer_delims(readline.get_completer_delims().replace('-','')) -def cleanup(): - osrf.system.System.net_disconnect() - -_output_buffer = '' # collect output for pager -def report(text, flush=False): - global _output_buffer - if get_var('SRFSH_OUTPUT_PAGED') == 'true': - _output_buffer += text + def do_connect(self): + ''' Connects this instance to the OpenSRF network. ''' - if flush and _output_buffer != '': - pipe = os.popen('less -EX', 'w') - pipe.write(_output_buffer) - pipe.close() - _output_buffer = '' + file = os.path.join(self.get_var('HOME'), ".srfsh.xml") + osrf.system.System.connect(config_file=file, config_context='srfsh') - else: - sys.stdout.write(text) - if flush: - sys.stdout.flush() + def add_command(self, **kwargs): + ''' Adds a new command to the supported srfsh commands. -def set_vars(): + Command is also added to the tab-completion word bank. - if not get_var('SRFSH_OUTPUT_NET_OBJ_KEYS'): - set_var('SRFSH_OUTPUT_NET_OBJ_KEYS', 'false') + kwargs : + command : the command name + handler : reference to a two-argument function. + Arguments are Srfsh instance and command arguments. + ''' - if not get_var('SRFSH_OUTPUT_FORMAT_JSON'): - set_var('SRFSH_OUTPUT_FORMAT_JSON', 'true') + command = kwargs['command'] + self.command_map[command] = kwargs['handler'] + self.tab_complete_words.append(command) - if not get_var('SRFSH_OUTPUT_PAGED'): - set_var('SRFSH_OUTPUT_PAGED', 'true') - # XXX Do we need to differ between LANG and LC_MESSAGES? - if not get_var('SRFSH_LOCALE'): - set_var('SRFSH_LOCALE', get_var('LC_ALL')) + def load_plugins(self): + ''' Load plugin modules from the srfsh configuration file ''' -def set_var(key, val): - os.environ[key] = val + try: + plugins = osrf.conf.get('plugins.plugin') + except: + return + + if not isinstance(plugins, list): + plugins = [plugins] + + for plugin in plugins: + module = plugin['module'] + init = plugin.get('init', 'load') + self.report("Loading module %s..." % module, True, True) + + try: + mod = __import__(module, fromlist=' ') + getattr(mod, init)(self, plugin) + self.report("OK.\n", True, True) + + except Exception, e: + self.report_error("Error importing plugin '%s' : %s\n" % (module, traceback.format_exc())) + + def cleanup(self): + ''' Disconnects from opensrf. ''' + osrf.system.System.net_disconnect() + + def report_error(self, msg): + ''' Log to stderr. ''' + sys.stderr.write("%s\n" % msg) + sys.stderr.flush() + + def report(self, text, flush=False, no_page=False): + ''' Logs to the pager or stdout, depending on env vars and context ''' + + if self.reading_script or no_page or self.get_var('SRFSH_OUTPUT_PAGED') != 'true': + sys.stdout.write(text) + if flush: + sys.stdout.flush() + else: + self.output_buffer += text + + if flush and self.output_buffer != '': + pipe = os.popen('less -EX', 'w') + pipe.write(self.output_buffer) + pipe.close() + self.output_buffer = '' + + def set_vars(self): + ''' Set defaults for environment variables. ''' + + if not self.get_var('SRFSH_OUTPUT_NET_OBJ_KEYS'): + self.set_var('SRFSH_OUTPUT_NET_OBJ_KEYS', 'false') + + if not self.get_var('SRFSH_OUTPUT_FORMAT_JSON'): + self.set_var('SRFSH_OUTPUT_FORMAT_JSON', 'true') + + if not self.get_var('SRFSH_OUTPUT_PAGED'): + self.set_var('SRFSH_OUTPUT_PAGED', 'true') + + # XXX Do we need to differ between LANG and LC_MESSAGES? + if not self.get_var('SRFSH_LOCALE'): + self.set_var('SRFSH_LOCALE', self.get_var('LC_ALL')) + + def set_var(self, key, val): + ''' Sets an environment variable's value. ''' + os.environ[key] = val + + def get_var(self, key): + ''' Returns an environment variable's value. ''' + return os.environ.get(key, '') + + def __get_locale(self): + """ + Return the defined locale for this srfsh session. + + A locale in OpenSRF is currently defined as a [a-z]{2}-[A-Z]{2} pattern. + This function munges the LC_ALL setting to conform to that pattern; for + example, trimming en_CA.UTF-8 to en-CA. + + >>> import srfsh + >>> shell = srfsh.Srfsh() + >>> shell.set_var('SRFSH_LOCALE', 'zz-ZZ') + >>> print shell.__get_locale() + zz-ZZ + >>> shell.set_var('SRFSH_LOCALE', 'en_CA.UTF-8') + >>> print shell.__get_locale() + en-CA + """ + + env_locale = self.get_var('SRFSH_LOCALE') + if env_locale: + pattern = re.compile(r'^\s*([a-z]+)[^a-zA-Z]([A-Z]+)').search(env_locale) + lang = pattern.group(1) + region = pattern.group(2) + locale = "%s-%s" % (lang, region) + else: + locale = 'en-US' -def get_var(key): - return os.environ.get(key, '') - -def __get_locale(): - """ - Return the defined locale for this srfsh session. - - A locale in OpenSRF is currently defined as a [a-z]{2}-[A-Z]{2} pattern. - This function munges the LC_ALL setting to conform to that pattern; for - example, trimming en_CA.UTF-8 to en-CA. - - >>> import srfsh - >>> srfsh.set_var('SRFSH_LOCALE', 'zz-ZZ') - >>> print __get_locale() - zz-ZZ - >>> srfsh.set_var('SRFSH_LOCALE', 'en_CA.UTF-8') - >>> print __get_locale() - en-CA - """ - - env_locale = get_var('SRFSH_LOCALE') - if env_locale: - pattern = re.compile(r'^\s*([a-z]+)[^a-zA-Z]([A-Z]+)').search(env_locale) - lang = pattern.group(1) - region = pattern.group(2) - locale = "%s-%s" % (lang, region) - else: - locale = 'en-US' - - return locale + return locale if __name__ == '__main__': - - # Kick it off - set_vars() - setup_readline() - do_connect() - load_plugins() - do_loop() + script = sys.argv[1] if len(sys.argv) > 1 else None + Srfsh(script).main_loop() -- 2.43.2