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 reports the average time
  11
  12   request <service> <method> [<param1>, <param2>, ...]
  13     - performs an opensrf request
  14
  15   set VAR=<value>
  16     - sets an environment variable
  17
  18   Environment variables:
  19     SRFSH_OUTPUT = pretty - print pretty JSON and key/value pairs for network objects
  20                  = raw - print formatted JSON
  21
  22     SRFSH_LOCALE = <locale> - request responses to be returned in locale <locale> if available
  23 """
  24
  25 import os, sys, time, readline, atexit, re
  26 import osrf.json, osrf.system, osrf.ses, osrf.conf, osrf.log, osrf.net
  27
  28 # -------------------------------------------------------------------
  29 # main listen loop
  30 # -------------------------------------------------------------------
  31 def do_loop():
  32     while True:
  33
  34         try:
  35             line = raw_input("srfsh% ")
  36             if not len(line):
  37                 continue
  38             if str.lower(line) == 'exit' or str.lower(line) == 'quit':
  39                 break
  40             parts = str.split(line)
  41
  42             command = parts[0]
  43
  44             if command == 'request':
  45                 parts.pop(0)
  46                 handle_request(parts)
  47                 continue
  48
  49             if command == 'math_bench':
  50                 parts.pop(0)
  51                 handle_math_bench(parts)
  52                 continue
  53
  54             if command == 'help':
  55                 handle_help()
  56                 continue
  57
  58             if command == 'set':
  59                 parts.pop(0)
  60                 handle_set(parts)
  61
  62             if command == 'get':
  63                 parts.pop(0)
  64                 handle_get(parts)
  65
  66
  67
  68         except KeyboardInterrupt:
  69             print ""
  70
  71         except EOFError:
  72             print "exiting..."
  73             sys.exit(0)
  74
  75
  76 # -------------------------------------------------------------------
  77 # Set env variables to control behavior
  78 # -------------------------------------------------------------------
  79 def handle_set(parts):
  80     cmd = "".join(parts)
  81     pattern = re.compile('(.*)=(.*)').match(cmd)
  82     key = pattern.group(1)
  83     val = pattern.group(2)
  84     set_var(key, val)
  85     print "%s = %s" % (key, val)
  86
  87 def handle_get(parts):
  88     try:
  89         print get_var(parts[0])
  90     except:
  91         print ""
  92
  93
  94 # -------------------------------------------------------------------
  95 # Prints help info
  96 # -------------------------------------------------------------------
  97 def handle_help():
  98     print __doc__
  99
 100 # -------------------------------------------------------------------
 101 # performs an opensrf request
 102 # -------------------------------------------------------------------
 103 def handle_request(parts):
 104     service = parts.pop(0)
 105     method = parts.pop(0)
 106     locale = __get_locale()
 107     jstr = '[%s]' % "".join(parts)
 108     params = None
 109
 110     try:
 111         params = osrf.json.to_object(jstr)
 112     except:
 113         print "Error parsing JSON: %s" % jstr
 114         return
 115
 116     ses = osrf.ses.ClientSession(service, locale=locale)
 117
 118     start = time.time()
 119
 120     req = ses.request2(method, tuple(params))
 121
 122
 123     while True:
 124         resp = None
 125         try:
 126             resp = req.recv(timeout=120)
 127         except osrf.net.XMPPNoRecipient:
 128             print "Unable to communicate with %s" % service
 129             total = 0
 130             break
 131
 132         osrf.log.log_internal("Looping through receive request")
 133         if not resp:
 134             break
 135         total = time.time() - start
 136
 137         otp = get_var('SRFSH_OUTPUT')
 138         if otp == 'pretty':
 139             print "\n" + osrf.json.debug_net_object(resp.content())
 140         else:
 141             print osrf.json.pprint(osrf.json.to_json(resp.content()))
 142
 143     req.cleanup()
 144     ses.cleanup()
 145
 146     print '-'*60
 147     print "Total request time: %f" % total
 148     print '-'*60
 149
 150
 151 def handle_math_bench(parts):
 152
 153     count = int(parts.pop(0))
 154     ses = osrf.ses.ClientSession('opensrf.math')
 155     times = []
 156
 157     for cnt in range(100):
 158         if cnt % 10:
 159             sys.stdout.write('.')
 160         else:
 161             sys.stdout.write( str( cnt / 10 ) )
 162     print ""
 163
 164
 165     for cnt in range(count):
 166
 167         starttime = time.time()
 168         req = ses.request('add', 1, 2)
 169         resp = req.recv(timeout=2)
 170         endtime = time.time()
 171
 172         if resp.content() == 3:
 173             sys.stdout.write("+")
 174             sys.stdout.flush()
 175             times.append( endtime - starttime )
 176         else:
 177             print "What happened? %s" % str(resp.content())
 178
 179         req.cleanup()
 180         if not ( (cnt + 1) % 100):
 181             print ' [%d]' % (cnt + 1)
 182
 183     ses.cleanup()
 184     total = 0
 185     for cnt in times:
 186         total += cnt
 187     print "\naverage time %f" % (total / len(times))
 188
 189
 190
 191
 192 # -------------------------------------------------------------------
 193 # Defines the tab-completion handling and sets up the readline history
 194 # -------------------------------------------------------------------
 195 def setup_readline():
 196     class SrfshCompleter(object):
 197         def __init__(self, words):
 198             self.words = words
 199             self.prefix = None
 200
 201         def complete(self, prefix, index):
 202             if prefix != self.prefix:
 203                 # find all words that start with this prefix
 204                 self.matching_words = [
 205                     w for w in self.words if w.startswith(prefix)
 206                 ]
 207                 self.prefix = prefix
 208                 try:
 209                     return self.matching_words[index]
 210                 except IndexError:
 211                     return None
 212
 213     words = 'request', 'help', 'exit', 'quit', 'opensrf.settings', 'opensrf.math', 'set'
 214     completer = SrfshCompleter(words)
 215     readline.parse_and_bind("tab: complete")
 216     readline.set_completer(completer.complete)
 217
 218     histfile = os.path.join(get_var('HOME'), ".srfsh_history")
 219     try:
 220         readline.read_history_file(histfile)
 221     except IOError:
 222         pass
 223     atexit.register(readline.write_history_file, histfile)
 224
 225 def do_connect():
 226     file = os.path.join(get_var('HOME'), ".srfsh.xml")
 227     print_stdout("Connecting to opensrf...")
 228     osrf.system.System.connect(config_file=file, config_context='srfsh')
 229     print_stdout('OK\n')
 230
 231 def load_plugins():
 232     # Load the user defined external plugins
 233     # XXX Make this a real module interface, with tab-complete words, commands, etc.
 234     try:
 235         plugins = osrf.conf.get('plugins')
 236
 237     except:
 238         # XXX standard srfsh.xml does not yet define <plugins> element
 239         #print_stdout("No plugins defined in /srfsh/plugins/plugin\n")
 240         return
 241
 242     plugins = osrf.conf.get('plugins.plugin')
 243     if not isinstance(plugins, list):
 244         plugins = [plugins]
 245
 246     for module in plugins:
 247         name = module['module']
 248         init = module['init']
 249         print_stdout("Loading module %s..." % name)
 250
 251         try:
 252             string = 'import %s\n%s.%s()' % (name, name, init)
 253             exec(string)
 254             print_stdout('OK\n')
 255
 256         except Exception, e:
 257             sys.stderr.write("\nError importing plugin %s, with init symbol %s: \n%s\n" % (name, init, e))
 258
 259 def set_vars():
 260     if not get_var('SRFSH_OUTPUT'):
 261         set_var('SRFSH_OUTPUT', 'raw')
 262
 263     # XXX Do we need to differ between LANG and LC_MESSAGES?
 264     if not get_var('SRFSH_LOCALE'):
 265         set_var('SRFSH_LOCALE', get_var('LC_ALL'))
 266
 267 def set_var(key, val):
 268     os.environ[key] = val
 269
 270 def get_var(key):
 271     return os.environ.get(key, '')
 272
 273 def __get_locale():
 274     """
 275     Return the defined locale for this srfsh session.
 276
 277     A locale in OpenSRF is currently defined as a [a-z]{2}-[A-Z]{2} pattern.
 278     This function munges the LC_ALL setting to conform to that pattern; for
 279     example, trimming en_CA.UTF-8 to en-CA.
 280
 281     >>> import srfsh
 282     >>> srfsh.set_var('SRFSH_LOCALE', 'zz-ZZ')
 283     >>> print __get_locale()
 284     zz-ZZ
 285     >>> srfsh.set_var('SRFSH_LOCALE', 'en_CA.UTF-8')
 286     >>> print __get_locale()
 287     en-CA
 288     """
 289
 290     env_locale = get_var('SRFSH_LOCALE')
 291     if env_locale:
 292         pattern = re.compile(r'^\s*([a-z]+)[^a-zA-Z]([A-Z]+)').search(env_locale)
 293         lang = pattern.group(1)
 294         region = pattern.group(2)
 295         locale = "%s-%s" % (lang, region)
 296     else:
 297         locale = 'en-US'
 298
 299     return locale
 300
 301 def print_stdout(string):
 302     sys.stdout.write(string)
 303     sys.stdout.flush()
 304
 305
 306 if __name__ == '__main__':
 307
 308     # Kick it off
 309     set_vars()
 310     setup_readline()
 311     do_connect()
 312     load_plugins()
 313     do_loop()
 314