Add option to record websocket received time

This commit is contained in:
Alan Rominger 2021-05-03 11:34:33 -04:00
parent faded278e3
commit b919befc90
No known key found for this signature in database
GPG Key ID: C2D7EAAA12B63559

View File

@ -3,6 +3,7 @@ import logging
import atexit
import json
import ssl
import datetime
from queue import Queue, Empty
from urllib.parse import urlparse
@ -50,7 +51,7 @@ class WSClient(object):
# Subscription group types
def __init__(self, token=None, hostname='', port=443, secure=True, session_id=None, csrftoken=None):
def __init__(self, token=None, hostname='', port=443, secure=True, session_id=None, csrftoken=None, add_received_time=False):
# delay this import, because this is an optional dependency
import websocket
@ -90,6 +91,7 @@ class WSClient(object):
self._message_cache = []
self._should_subscribe_to_pending_job = False
self._pending_unsubscribe = threading.Event()
self._add_received_time = add_received_time
def connect(self):
wst = threading.Thread(target=self._ws_run_forever, args=(self.ws, {"cert_reqs": ssl.CERT_NONE}))
@ -195,6 +197,8 @@ class WSClient(object):
def _on_message(self, message):
message = json.loads(message)
log.debug('received message: {}'.format(message))
if self._add_received_time:
message['received_time'] = datetime.datetime.utcnow()
if all([message.get('group_name') == 'jobs', message.get('status') == 'pending', message.get('unified_job_id'), self._should_subscribe_to_pending_job]):
if bool(message.get('project_id')) == (self._should_subscribe_to_pending_job['events'] == 'project_update_events'):