Merge branch '2800-python-global-state' into 2800-pgs

[arvados.git] / sdk / python / arvados / events.py

diff --git a/sdk/python/arvados/events.py b/sdk/python/arvados/events.py
index e61b20c70851f8079137c540680842592a3dcb9b..0dbfe6359dd18b93cb95d6154f22ec7dcbe24177 100644 (file)
--- a/sdk/python/arvados/events.py
+++ b/sdk/python/arvados/events.py
@@ -6,6 +6,9 @@ import time
 import ssl
 import re
 import config
+import logging
+
+_logger = logging.getLogger('arvados.events')
 
 class EventClient(WebSocketClient):
     def __init__(self, url, filters, on_event):
@@ -26,8 +29,22 @@ class EventClient(WebSocketClient):
     def received_message(self, m):
         self.on_event(json.loads(str(m)))
 
+    def close_connection(self):
+        try:
+            self.sock.shutdown(socket.SHUT_RDWR)
+            self.sock.close()
+        except:
+            pass
+
 def subscribe(api, filters, on_event):
-    url = "{}?api_token={}".format(api._rootDesc['websocketUrl'], config.get('ARVADOS_API_TOKEN'))
-    ws = EventClient(url, filters, on_event)
-    ws.connect()
-    return ws
+    ws = None
+    try:
+        url = "{}?api_token={}".format(api._rootDesc['websocketUrl'], config.get('ARVADOS_API_TOKEN'))
+        ws = EventClient(url, filters, on_event)
+        ws.connect()
+        return ws
+    except Exception:
+        _logger.exception('')
+        if (ws):
+          ws.close_connection()
+        raise