3706: Merge branch 'master' into 3706-keep-warning
[arvados.git] / sdk / python / arvados / commands / ws.py
1 #!/usr/bin/env python
2
3 import sys
4 import logging
5 import argparse
6 import arvados
7 import json
8 from arvados.events import subscribe
9 import signal
10
11 def main(arguments=None):
12     logger = logging.getLogger('arvados.arv-ws')
13
14     parser = argparse.ArgumentParser()
15     parser.add_argument('-u', '--uuid', type=str, default="", help="Filter events on object_uuid")
16     parser.add_argument('-f', '--filters', type=str, default="", help="Arvados query filter to apply to log events (JSON encoded)")
17
18     group = parser.add_mutually_exclusive_group()
19     group.add_argument('--poll-interval', default=15, type=int, help="If websockets is not available, specify the polling interval, default is every 15 seconds")
20     group.add_argument('--no-poll', action='store_false', dest='poll_interval', help="Do not poll if websockets are not available, just fail")
21
22     group = parser.add_mutually_exclusive_group()
23     group.add_argument('-p', '--pipeline', type=str, default="", help="Supply pipeline uuid, print log output from pipeline and its jobs")
24     group.add_argument('-j', '--job', type=str, default="", help="Supply job uuid, print log output from jobs")
25
26     args = parser.parse_args(arguments)
27
28     global filters
29     global known_component_jobs
30     global ws
31
32     filters = []
33     known_component_jobs = set()
34     ws = None
35
36     def update_subscribed_components(components):
37         global known_component_jobs
38         global filters
39         pipeline_jobs = set()
40         for c in components:
41             if "job" in components[c]:
42                 pipeline_jobs.add(components[c]["job"]["uuid"])
43         if known_component_jobs != pipeline_jobs:
44             ws.unsubscribe(filters)
45             filters = [['object_uuid', 'in', [args.pipeline] + list(pipeline_jobs)]]
46             ws.subscribe([['object_uuid', 'in', [args.pipeline] + list(pipeline_jobs)]])
47             known_component_jobs = pipeline_jobs
48
49     api = arvados.api('v1', cache=False)
50
51     if args.uuid:
52         filters += [ ['object_uuid', '=', args.uuid] ]
53
54     if args.filters:
55         filters += json.loads(args.filters)
56
57     if args.job:
58         filters += [ ['object_uuid', '=', args.job] ]
59
60     if args.pipeline:
61         filters += [ ['object_uuid', '=', args.pipeline] ]
62
63     def on_message(ev):
64         global filters
65         global ws
66
67         logger.debug(ev)
68         if 'event_type' in ev and (args.pipeline or args.job):
69             if ev['event_type'] in ('stderr', 'stdout'):
70                 sys.stdout.write(ev["properties"]["text"])
71             elif ev["event_type"] in ("create", "update"):
72                 if ev["object_kind"] == "arvados#pipelineInstance":
73                     update_subscribed_components(ev["properties"]["new_attributes"]["components"])
74         elif 'status' in ev and ev['status'] == 200:
75             pass
76         else:
77             print json.dumps(ev)
78
79     try:
80         ws = subscribe(arvados.api('v1', cache=False), filters, on_message, poll_fallback=args.poll_interval)
81         if ws:
82             if args.pipeline:
83                 c = api.pipeline_instances().get(uuid=args.pipeline).execute()
84                 update_subscribed_components(c["components"])
85
86             while True:
87                 signal.pause()
88     except KeyboardInterrupt:
89         pass
90     except Exception as e:
91         logger.error(e)
92     finally:
93         if ws:
94             ws.close()