Merge pull request #379 from alpacahq/advanced_streamconn_usage

Advanced streamconn usage

Author: Shlomi Kushchi

README.md

@@ -288,6 +288,14 @@ import logging
 logging.basicConfig(format='%(asctime)s %(message)s', level=logging.INFO)
 ```
 
+## Websocket best practices
+Under the examples folder you could find several examples to do the following:
+* Different subscriptions(channels) usage with alpaca/polygon streams
+* pause / resume connection
+* change subscriptions/channels of existing connection
+* ws disconnections handler (make sure we reconnect when the internal mechanism fails)
+
+
 # Polygon API Service
 
 Alpaca's API key ID can be used to access Polygon API, the documentation for

alpaca_trade_api/polygon/streamconn.py

@@ -190,12 +190,15 @@ def run(self, initial_channels=[]):
 
     async def close(self):
         '''Close any open connections'''
-        if self._consume_task:
-            self._consume_task.cancel()
+        await self.cancel_task()
         if self._ws is not None:
             await self._ws.close()
         self._ws = None
 
+    async def cancel_task(self):
+        if self._consume_task:
+            self._consume_task.cancel()
+
     def _cast(self, subject, data):
         if subject == 'T':
             return Trade({trade_mapping[k]: v for k,

alpaca_trade_api/stream2.py

@@ -4,6 +4,7 @@
 import re
 import traceback
 from asyncio import CancelledError
+import queue
 
 import websockets
 from .common import get_base_url, get_data_url, get_credentials, URL
@@ -140,12 +141,15 @@ async def unsubscribe(self, channels):
             }))
 
     async def close(self):
-        if self._consume_task:
-            self._consume_task.cancel()
+        await self.cancel_task()
         if self._ws:
             await self._ws.close()
             self._ws = None
 
+    async def cancel_task(self):
+        if self._consume_task:
+            self._consume_task.cancel()
+
     def _cast(self, channel, msg):
         if channel == 'account_updates':
             return Account(msg)
@@ -227,6 +231,7 @@ def __init__(
         self._data_stream = _data_stream
         self._debug = debug
         self._raw_data = raw_data
+        self._stop_stream_queue = queue.Queue()
 
         self.trading_ws = _StreamConn(self._key_id,
                                       self._secret_key,
@@ -337,6 +342,9 @@ def run(self, initial_channels: List[str] = []):
                 logging.error(f"error while consuming ws messages: {m}")
                 if self._debug:
                     traceback.print_exc()
+                if not self._stop_stream_queue.empty():
+                    self._stop_stream_queue.get()
+                    should_renew = False
                 loop.run_until_complete(self.close(should_renew))
                 if loop.is_running():
                     loop.close()
@@ -370,6 +378,18 @@ async def close(self, renew):
                                            self._oauth,
                                            raw_data=self._raw_data)
 
+    async def stop_ws(self):
+        """
+        Signal the ws connections to stop listenning to api stream.
+        """
+        self._stop_stream_queue.put_nowait({"should_stop": True})
+        if self.trading_ws is not None:
+            logging.info("Stopping the trading websocket connection")
+            await self.trading_ws.cancel_task()
+        if self.data_ws is not None:
+            logging.info("Stopping the data websocket connection")
+            await self.data_ws.cancel_task()
+
     def on(self, channel_pat, symbols=None):
         def decorator(func):
             self.register(channel_pat, func, symbols)

examples/README.md

@@ -17,3 +17,13 @@ This shows a basic approach to opening a streaming connection for Polygon market
 ## Martingale
 
 This trading algorithm explores a strategy based on a gambling technique. Trading every few seconds, it maintains a position in the $SPY symbol of a size determined by the number of up or down candles it's experienced in a row. For a more complete explanation, please see [this post](https://forum.alpaca.markets/t/martingale-day-trading-with-the-alpaca-trading-api/).
+
+## Websocket Best practices
+Under this folder you could find several examples to do the following:
+* Different subscriptions(channels) usage with alpaca/polygon streams
+* pause / resume connection
+* change subscriptions of existing connection
+* ws disconnections handler (make sure we reconnect)
+
+
+Use it to integrate with your own code.

examples/websockets/avoid_server_connection_errors.py

@@ -0,0 +1,57 @@
+"""
+In this example code we wrap the ws connection to make sure we reconnect
+in case of ws disconnection.
+"""
+
+import logging
+import threading
+import asyncio
+import time
+from alpaca_trade_api import StreamConn
+from alpaca_trade_api.common import URL
+
+logging.basicConfig(format='%(asctime)s %(message)s', level=logging.INFO)
+
+ALPACA_API_KEY = "<YOUR-API-KEY>"
+ALPACA_SECRET_KEY = "<YOUR-SECRET-KEY>"
+USE_POLYGON = False
+
+
+def run_connection(conn, channels):
+    try:
+        conn.run(channels)
+    except Exception as e:
+        print(f'Exception from websocket connection: {e}')
+    finally:
+        send_msg(f"Trying to re-establish connection")
+        time.sleep(3)
+        run_connection(conn, channels)
+
+if __name__ == '__main__':
+    channels = ['alpacadatav1/Q.GOOG']
+
+    conn = StreamConn(
+        ALPACA_API_KEY,
+        ALPACA_SECRET_KEY,
+        base_url=URL('https://paper-api.alpaca.markets'),
+        data_url=URL('https://data.alpaca.markets'),
+        # data_url=URL('http://127.0.0.1:8765'),
+        data_stream='polygon' if USE_POLYGON else 'alpacadatav1'
+    )
+
+
+    @conn.on(r'^AM\..+$')
+    async def on_minute_bars(conn, channel, bar):
+        print('bars', bar)
+
+
+    @conn.on(r'Q\..+')
+    async def on_quotes(conn, channel, quote):
+        print('quote', quote)
+
+
+    @conn.on(r'T\..+')
+    async def on_trades(conn, channel, trade):
+        print('trade', trade)
+
+    run_connection(conn, channels)

examples/websockets/dynamic_subscription_example.py

@@ -0,0 +1,69 @@
+"""
+In this example code we will show a pattern that allows a user to change
+the websocket subscriptions as they please.
+"""
+import logging
+import threading
+import asyncio
+import time
+from alpaca_trade_api import StreamConn
+from alpaca_trade_api.common import URL
+
+ALPACA_API_KEY = "<YOUR-API-KEY>"
+ALPACA_SECRET_KEY = "<YOUR-SECRET-KEY>"
+USE_POLYGON = False
+
+conn: StreamConn = None
+
+def consumer_thread():
+
+    try:
+        # make sure we have an event loop, if not create a new one
+        loop = asyncio.get_event_loop()
+        loop.set_debug(True)
+    except RuntimeError:
+        asyncio.set_event_loop(asyncio.new_event_loop())
+
+    global conn
+    conn = StreamConn(
+        ALPACA_API_KEY,
+        ALPACA_SECRET_KEY,
+        base_url=URL('https://paper-api.alpaca.markets'),
+        data_url=URL('https://data.alpaca.markets'),
+        # data_url=URL('http://127.0.0.1:8765'),
+        data_stream='polygon' if USE_POLYGON else 'alpacadatav1'
+    )
+
+    @conn.on(r'^AM\..+$')
+    async def on_minute_bars(conn, channel, bar):
+        print('bars', bar)
+
+
+    @conn.on(r'Q\..+')
+    async def on_quotes(conn, channel, quote):
+        print('quote', quote)
+
+
+    @conn.on(r'T\..+')
+    async def on_trades(conn, channel, trade):
+        print('trade', trade)
+
+    conn.run(['alpacadatav1/Q.GOOG'])
+
+if __name__ == '__main__':
+    logging.basicConfig(format='%(asctime)s %(message)s', level=logging.INFO)
+    threading.Thread(target=consumer_thread).start()
+
+    loop = asyncio.get_event_loop()
+
+    time.sleep(5)  # give the initial connection time to be established
+    subscriptions = [['alpacadatav1/AM.TSLA'], ['alpacadatav1/Q.GOOG'],
+                     ['alpacadatav1/T.AAPL']]
+
+    while 1:
+        for channels in subscriptions:
+            loop.run_until_complete(conn.subscribe(channels))
+            if "AM." in channels[0]:
+                time.sleep(60)  # aggs are once every minute. give it time
+            else:
+                time.sleep(20)

examples/websockets/streamconn_on_and_off.py

@@ -0,0 +1,63 @@
+"""
+In this example code we will show how to shut the streamconn websocket
+connection down and then up again. it's the ability to stop/start the
+connection
+"""
+import logging
+import threading
+import asyncio
+import time
+from alpaca_trade_api import StreamConn
+from alpaca_trade_api.common import URL
+
+ALPACA_API_KEY = "<YOUR-API-KEY>"
+ALPACA_SECRET_KEY = "<YOUR-SECRET-KEY>"
+USE_POLYGON = False
+
+conn: StreamConn = None
+
+def consumer_thread():
+
+    try:
+        # make sure we have an event loop, if not create a new one
+        loop = asyncio.get_event_loop()
+        loop.set_debug(True)
+    except RuntimeError:
+        asyncio.set_event_loop(asyncio.new_event_loop())
+
+    global conn
+    conn = StreamConn(
+        ALPACA_API_KEY,
+        ALPACA_SECRET_KEY,
+        base_url=URL('https://paper-api.alpaca.markets'),
+        data_url=URL('https://data.alpaca.markets'),
+        # data_url=URL('http://127.0.0.1:8765'),
+        data_stream='polygon' if USE_POLYGON else 'alpacadatav1'
+    )
+
+    @conn.on(r'^AM\..+$')
+    async def on_minute_bars(conn, channel, bar):
+        print('bars', bar)
+
+
+    @conn.on(r'Q\..+')
+    async def on_quotes(conn, channel, quote):
+        print('quote', quote)
+
+
+    @conn.on(r'T\..+')
+    async def on_trades(conn, channel, trade):
+        print('trade', trade)
+
+    conn.run(['alpacadatav1/Q.GOOG'])
+
+if __name__ == '__main__':
+    logging.basicConfig(format='%(asctime)s %(message)s', level=logging.INFO)
+
+    loop = asyncio.get_event_loop()
+
+    while 1:
+        threading.Thread(target=consumer_thread).start()
+        time.sleep(5)
+        loop.run_until_complete(conn.stop_ws())
+        time.sleep(20)