3609: Add test case for PollClient. arvados.event.subscribe() starts a new
[arvados.git] / sdk / python / arvados / events.py
1 from ws4py.client.threadedclient import WebSocketClient
2 import threading
3 import json
4 import os
5 import time
6 import ssl
7 import re
8 import config
9 import logging
10
11 _logger = logging.getLogger('arvados.events')
12
13 class EventClient(WebSocketClient):
14     def __init__(self, url, filters, on_event):
15         ssl_options = None
16         if re.match(r'(?i)^(true|1|yes)$',
17                     config.get('ARVADOS_API_HOST_INSECURE', 'no')):
18             ssl_options={'cert_reqs': ssl.CERT_NONE}
19         else:
20             ssl_options={'cert_reqs': ssl.CERT_REQUIRED}
21         super(EventClient, self).__init__(url, ssl_options=ssl_options)
22         self.filters = filters
23         self.on_event = on_event
24
25     def opened(self):
26         self.subscribe(self.filters)
27
28     def received_message(self, m):
29         self.on_event(json.loads(str(m)))
30
31     def close_connection(self):
32         try:
33             self.sock.shutdown(socket.SHUT_RDWR)
34             self.sock.close()
35         except:
36             pass
37
38     def subscribe(self, filters, last_log_id=None):
39         m = {"method": "subscribe", "filters": filters}
40         if last_log_id is not None:
41             m["last_log_id"] = last_log_id
42         self.send(json.dumps(m))
43
44     def unsubscribe(self, filters):
45         self.send(json.dumps({"method": "unsubscribe", "filters": filters}))
46
47 class PollClient(threading.Thread):
48     def __init__(self, api, filters, on_event, poll_time):
49         super(PollClient, self).__init__()
50         self.api = api
51         self.filters = filters
52         self.on_event = on_event
53         items = self.api.logs().list(limit=1, order=["id desc"], filters=filters).execute()['items']
54         if len(items) > 0:
55             self.id = items[0]["id"]
56         else:
57             self.id = 0
58         self.poll_time = poll_time
59         self.loop = True
60         self.on_event({'status': 200})
61
62     def run(self):
63         while self.loop:
64             time.sleep(self.poll_time)
65             items = self.api.logs().list(order=["id asc"], filters=self.filters+[["id", ">", str(self.id)]]).execute()['items']
66             for i in items:
67                 self.id = i['id']
68                 self.on_event(i)
69
70     def close(self):
71         self.loop = False
72         self.join()
73
74     def subscribe(self, filters):
75         self.on_event({'status': 200})
76         self.filters += filters
77
78     def unsubscribe(self, filters):
79         del self.filters[self.filters.index(filters)]
80
81 def subscribe(api, filters, on_event, poll_fallback=15):
82     ws = None
83     if 'websocketUrl' in api._rootDesc:
84         try:
85             url = "{}?api_token={}".format(api._rootDesc['websocketUrl'], api.api_token)
86             ws = EventClient(url, filters, on_event)
87             ws.connect()
88             return ws
89         except Exception as e:
90             _logger.warn("Got exception %s trying to connect to web sockets" % e)
91             if ws:
92                 ws.close()
93     if poll_fallback:
94         _logger.warn("Web sockets not available, falling back to log table polling")
95         p = PollClient(api, filters, on_event, poll_fallback)
96         p.start()
97         return p
98     else:
99         _logger.error("Web sockets not available")
100         return None