src/python/srfsh.py

   1 #!/usr/bin/python
   2 # vim:et:ts=4
   3 """
   4 srfsh.py - provides a basic shell for issuing OpenSRF requests
   5
   6   help
   7     - show this menu
   8
   9   math_bench <count>
  10     - runs <count> opensrf.math requests and prints the average time
  11
  12   request <service> <method> [<param1>, <param2>, ...]
  13     - performs an opensrf request
  14
  15   router <query>
  16     - Queries the router.  Query options: services service-stats service-nodes
  17
  18   set VAR=<value>
  19     - sets an environment variable
  20
  21   get VAR
  22     - Returns the value for the environment variable
  23
  24   Environment variables:
  25     SRFSH_OUTPUT_NET_OBJ_KEYS  = true - If a network object is array-encoded and key registry exists for the object type, annotate the object with field names
  26                                = false - Print JSON
  27
  28     SRFSH_OUTPUT_FORMAT_JSON = true - Use JSON pretty printer
  29                              = false - Print raw JSON
  30
  31     SRFSH_OUTPUT_PAGED = true - Paged output.  Uses "less -EX"
  32                        = false - Output is not paged
  33
  34     SRFSH_LOCALE = <locale> - request responses to be returned in locale <locale> if available
  35
  36 """
  37
  38 import os, sys, time, readline, atexit, re, pydoc
  39 import osrf.json, osrf.system, osrf.ses, osrf.conf, osrf.log, osrf.net
  40
  41 router_command_map = {
  42     'services' : 'opensrf.router.info.class.list',
  43     'service-stats' : 'opensrf.router.info.stats.class.node.all',
  44     'service-nodes' : 'opensrf.router.info.stats.class.all'
  45 }
  46
  47 # List of words to use for readline tab completion
  48 tab_complete_words = [
  49     'request',
  50     'set',
  51     'router',
  52     'help',
  53     'exit',
  54     'quit',
  55     'opensrf.settings',
  56     'opensrf.math'
  57 ]
  58
  59 # add the router commands to the tab-complete list
  60 for rcommand in router_command_map.keys():
  61     tab_complete_words.append(rcommand)
  62
  63 # -------------------------------------------------------------------
  64 # main listen loop
  65 # -------------------------------------------------------------------
  66 def do_loop():
  67
  68     command_map = {
  69         'request' : handle_request,
  70         'router' : handle_router,
  71         'math_bench' : handle_math_bench,
  72         'help' : handle_help,
  73         'set' : handle_set,
  74         'get' : handle_get,
  75     }
  76
  77     while True:
  78
  79         try:
  80             report("", True)
  81             line = raw_input("srfsh# ")
  82
  83             if not len(line):
  84                 continue
  85
  86             if str.lower(line) == 'exit' or str.lower(line) == 'quit':
  87                 break
  88
  89             parts = str.split(line)
  90             command = parts.pop(0)
  91
  92             if command not in command_map:
  93                 report("unknown command: '%s'\n" % command)
  94                 continue
  95
  96             command_map[command](parts)
  97
  98         except EOFError: # ^-d
  99             sys.exit(0)
 100
 101         except KeyboardInterrupt: # ^-c
 102             report("\n")
 103
 104         except Exception, e:
 105             report("%s\n" % e)
 106
 107     cleanup()
 108
 109 def handle_router(parts):
 110
 111     if len(parts) == 0:
 112         report("usage: router <query>\n")
 113         return
 114
 115     query = parts[0]
 116
 117     if query not in router_command_map:
 118         report("router query options: %s\n" % ','.join(router_command_map.keys()))
 119         return
 120
 121     return handle_request(['router', router_command_map[query]])
 122
 123 # -------------------------------------------------------------------
 124 # Set env variables to control behavior
 125 # -------------------------------------------------------------------
 126 def handle_set(parts):
 127     cmd = "".join(parts)
 128     pattern = re.compile('(.*)=(.*)').match(cmd)
 129     key = pattern.group(1)
 130     val = pattern.group(2)
 131     set_var(key, val)
 132     report("%s = %s\n" % (key, val))
 133
 134 def handle_get(parts):
 135     try:
 136         report("%s=%s\n" % (parts[0], get_var(parts[0])))
 137     except:
 138         report("\n")
 139
 140
 141 # -------------------------------------------------------------------
 142 # Prints help info
 143 # -------------------------------------------------------------------
 144 def handle_help(foo):
 145     report(__doc__)
 146
 147 # -------------------------------------------------------------------
 148 # performs an opensrf request
 149 # -------------------------------------------------------------------
 150 def handle_request(parts):
 151
 152     if len(parts) < 2:
 153         report("usage: request <service> <api_name> [<param1>, <param2>, ...]\n")
 154         return
 155
 156     service = parts.pop(0)
 157     method = parts.pop(0)
 158     locale = __get_locale()
 159     jstr = '[%s]' % "".join(parts)
 160     params = None
 161
 162     try:
 163         params = osrf.json.to_object(jstr)
 164     except:
 165         report("Error parsing JSON: %s\n" % jstr)
 166         return
 167
 168     ses = osrf.ses.ClientSession(service, locale=locale)
 169
 170     start = time.time()
 171
 172     req = ses.request2(method, tuple(params))
 173
 174
 175     while True:
 176         resp = None
 177
 178         try:
 179             resp = req.recv(timeout=120)
 180         except osrf.net.XMPPNoRecipient:
 181             report("Unable to communicate with %s\n" % service)
 182             total = 0
 183             break
 184
 185         if not resp: break
 186
 187         total = time.time() - start
 188         content = resp.content()
 189
 190         if content is not None:
 191             if get_var('SRFSH_OUTPUT_NET_OBJ_KEYS') == 'true':
 192                 report("Received Data: %s\n" % osrf.json.debug_net_object(content))
 193             else:
 194                 if get_var('SRFSH_OUTPUT_FORMAT_JSON') == 'true':
 195                     report("Received Data: %s\n" % osrf.json.pprint(osrf.json.to_json(content)))
 196                 else:
 197                     report("Received Data: %s\n" % osrf.json.to_json(content))
 198
 199     req.cleanup()
 200     ses.cleanup()
 201
 202     report('-'*60 + "\n")
 203     report("Total request time: %f\n" % total)
 204     report('-'*60 + "\n")
 205
 206
 207 def handle_math_bench(parts):
 208
 209     count = int(parts.pop(0))
 210     ses = osrf.ses.ClientSession('opensrf.math')
 211     times = []
 212
 213     for cnt in range(100):
 214         if cnt % 10:
 215             sys.stdout.write('.')
 216         else:
 217             sys.stdout.write( str( cnt / 10 ) )
 218     print ""
 219
 220     for cnt in range(count):
 221
 222         starttime = time.time()
 223         req = ses.request('add', 1, 2)
 224         resp = req.recv(timeout=2)
 225         endtime = time.time()
 226
 227         if resp.content() == 3:
 228             sys.stdout.write("+")
 229             sys.stdout.flush()
 230             times.append( endtime - starttime )
 231         else:
 232             print "What happened? %s" % str(resp.content())
 233
 234         req.cleanup()
 235         if not ( (cnt + 1) % 100):
 236             print ' [%d]' % (cnt + 1)
 237
 238     ses.cleanup()
 239     total = 0
 240     for cnt in times:
 241         total += cnt
 242     print "\naverage time %f" % (total / len(times))
 243
 244
 245
 246
 247 # -------------------------------------------------------------------
 248 # Defines the tab-completion handling and sets up the readline history
 249 # -------------------------------------------------------------------
 250 def setup_readline():
 251
 252     class SrfshCompleter(object):
 253
 254         def __init__(self, words):
 255             self.words = words
 256             self.prefix = None
 257
 258         def complete(self, prefix, index):
 259
 260             if prefix != self.prefix:
 261
 262                 self.prefix = prefix
 263
 264                 # find all words that start with this prefix
 265                 self.matching_words = [
 266                     w for w in self.words if w.startswith(prefix)
 267                 ]
 268
 269                 if len(self.matching_words) == 0:
 270                     return None
 271
 272                 if len(self.matching_words) == 1:
 273                     return self.matching_words[0]
 274
 275                 sys.stdout.write('\n%s\nsrfsh# %s' %
 276                     (' '.join(self.matching_words), readline.get_line_buffer()))
 277
 278                 return None
 279
 280     completer = SrfshCompleter(tuple(tab_complete_words))
 281     readline.parse_and_bind("tab: complete")
 282     readline.set_completer(completer.complete)
 283
 284     histfile = os.path.join(get_var('HOME'), ".srfsh_history")
 285     try:
 286         readline.read_history_file(histfile)
 287     except IOError:
 288         pass
 289     atexit.register(readline.write_history_file, histfile)
 290
 291     readline.set_completer_delims(readline.get_completer_delims().replace('-',''))
 292
 293
 294 def do_connect():
 295     file = os.path.join(get_var('HOME'), ".srfsh.xml")
 296     osrf.system.System.connect(config_file=file, config_context='srfsh')
 297
 298 def load_plugins():
 299     # Load the user defined external plugins
 300     # XXX Make this a real module interface, with tab-complete words, commands, etc.
 301     try:
 302         plugins = osrf.conf.get('plugins')
 303
 304     except:
 305         # XXX standard srfsh.xml does not yet define <plugins> element
 306         #print("No plugins defined in /srfsh/plugins/plugin\n")
 307         return
 308
 309     plugins = osrf.conf.get('plugins.plugin')
 310     if not isinstance(plugins, list):
 311         plugins = [plugins]
 312
 313     for module in plugins:
 314         name = module['module']
 315         init = module['init']
 316         print "Loading module %s..." % name
 317
 318         try:
 319             string = 'import %s\n%s.%s()' % (name, name, init)
 320             exec(string)
 321             print 'OK'
 322
 323         except Exception, e:
 324             sys.stderr.write("\nError importing plugin %s, with init symbol %s: \n%s\n" % (name, init, e))
 325
 326 def cleanup():
 327     osrf.system.System.net_disconnect()
 328
 329 _output_buffer = '' # collect output for pager
 330 def report(text, flush=False):
 331     global _output_buffer
 332
 333     if get_var('SRFSH_OUTPUT_PAGED') == 'true':
 334         _output_buffer += text
 335
 336         if flush and _output_buffer != '':
 337             pipe = os.popen('less -EX', 'w')
 338             pipe.write(_output_buffer)
 339             pipe.close()
 340             _output_buffer = ''
 341
 342     else:
 343         sys.stdout.write(text)
 344         if flush:
 345             sys.stdout.flush()
 346
 347 def set_vars():
 348
 349     if not get_var('SRFSH_OUTPUT_NET_OBJ_KEYS'):
 350         set_var('SRFSH_OUTPUT_NET_OBJ_KEYS', 'false')
 351
 352     if not get_var('SRFSH_OUTPUT_FORMAT_JSON'):
 353         set_var('SRFSH_OUTPUT_FORMAT_JSON', 'true')
 354
 355     if not get_var('SRFSH_OUTPUT_PAGED'):
 356         set_var('SRFSH_OUTPUT_PAGED', 'true')
 357
 358     # XXX Do we need to differ between LANG and LC_MESSAGES?
 359     if not get_var('SRFSH_LOCALE'):
 360         set_var('SRFSH_LOCALE', get_var('LC_ALL'))
 361
 362 def set_var(key, val):
 363     os.environ[key] = val
 364
 365 def get_var(key):
 366     return os.environ.get(key, '')
 367
 368 def __get_locale():
 369     """
 370     Return the defined locale for this srfsh session.
 371
 372     A locale in OpenSRF is currently defined as a [a-z]{2}-[A-Z]{2} pattern.
 373     This function munges the LC_ALL setting to conform to that pattern; for
 374     example, trimming en_CA.UTF-8 to en-CA.
 375
 376     >>> import srfsh
 377     >>> srfsh.set_var('SRFSH_LOCALE', 'zz-ZZ')
 378     >>> print __get_locale()
 379     zz-ZZ
 380     >>> srfsh.set_var('SRFSH_LOCALE', 'en_CA.UTF-8')
 381     >>> print __get_locale()
 382     en-CA
 383     """
 384
 385     env_locale = get_var('SRFSH_LOCALE')
 386     if env_locale:
 387         pattern = re.compile(r'^\s*([a-z]+)[^a-zA-Z]([A-Z]+)').search(env_locale)
 388         lang = pattern.group(1)
 389         region = pattern.group(2)
 390         locale = "%s-%s" % (lang, region)
 391     else:
 392         locale = 'en-US'
 393
 394     return locale
 395
 396 if __name__ == '__main__':
 397
 398     # Kick it off
 399     set_vars()
 400     setup_readline()
 401     do_connect()
 402     load_plugins()
 403     do_loop()
 404