Refactor of caching & data storage part 9 #118

providers/spot/sse_spot_provider.py

@@ -0,0 +1,79 @@
+import logging
+from datetime import datetime
+from threading import Thread
+from time import sleep
+
+import pytz
+from requests_sse import EventSource
+
+from core.constants import HTTP_HEADERS
+from providers.spot.spot_provider import SpotProvider
+
+
+class SSESpotProvider(SpotProvider):
+    """Spot provider using Server-Sent Events."""
+
+    def __init__(self, name, provider_config, url):
+        super().__init__(name, provider_config)
+        self._url = url
+        self._event_source = None
+        self._thread = None
+        self._stopped = False
+        self._last_event_id = None
+
+    def start(self):
+        logging.info("Set up SSE connection to " + self.name + " spot API.")
+        self._stopped = False
+        self._thread = Thread(target=self._run)
+        self._thread.daemon = True
+        self._thread.start()
+
+    def stop(self):
+        self._stopped = True
+        if self._event_source:
+            self._event_source.close()
+        if self._thread:
+            self._thread.join()
+
+    def _on_open(self):
+        self.status = "Waiting for Data"
+
+    def _on_error(self):
+        self.status = "Connecting"
+
+    def _run(self):
+        while not self._stopped:
+            try:
+                logging.debug("Connecting to " + self.name + " spot API...")
+                self.status = "Connecting"
+                with EventSource(self._url, headers=HTTP_HEADERS, latest_event_id=self._last_event_id, timeout=30,
+                                 on_open=self._on_open, on_error=self._on_error) as event_source:
+                    self._event_source = event_source
+                    for event in self._event_source:
+                        if event.type == 'message':
+                            try:
+                                self._last_event_id = event.last_event_id
+                                new_spot = self._sse_message_to_spot(event.data)
+                                if new_spot:
+                                    self._submit(new_spot)
+
+                                self.status = "OK"
+                                self.last_update_time = datetime.now(pytz.UTC)
+                                logging.debug("Received data from " + self.name + " spot API.")
+
+                            except Exception:
+                                logging.exception(
+                                    "Exception processing message from SSE Spot Provider (" + self.name + ")")
+
+            except Exception:
+                self.status = "Error"
+                logging.exception("Exception in SSE Spot Provider (" + self.name + ")")
+            else:
+                self.status = "Disconnected"
+            sleep(5)  # Wait before trying to reconnect
+
+    def _sse_message_to_spot(self, message_data):
+        """Convert an SSE message received from the API into a spot. The whole message data is provided here so the subclass
+        implementations can handle the message as JSON, XML, text, whatever the API actually provides."""
+
+        raise NotImplementedError("Subclasses must implement this method")