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