Files
OpenBB/openbb_platform/extensions/websockets/openbb_websockets/client.py
2024-12-09 19:31:18 -08:00

808 lines
31 KiB
Python

"""WebSocket Client module for interacting with a provider websocket in a non-blocking pattern."""
# pylint: disable=too-many-statements,protected-access
# flake8: noqa: PLR0915
import logging
from typing import TYPE_CHECKING, Any, Literal, Optional, Union
if TYPE_CHECKING:
from openbb_core.provider.abstract.data import Data
class WebSocketClient: # pylint: disable=too-many-instance-attributes
"""Client for interacting with a websocket server in a non-blocking pattern.
Parameters
----------
name : str
Name to assign the WebSocket connection. Used to identify and manage multiple instances.
module : str
The Python module for the provider websocket_client module. Runs in a separate thread.
Example: 'openbb_fmp.utils.websocket_client'. Pass additional keyword arguments by including kwargs.
symbol : Optional[str]
The symbol(s) requested to subscribe. Enter multiple symbols separated by commas without spaces.
limit : Optional[int]
The limit of records to hold in memory. Once the limit is reached, the oldest records are removed.
Default is 1000. Set to None to keep all records.
results_file : Optional[str]
Absolute path to the file for continuous writing. By default, a temporary file is created.
table_name : Optional[str]
SQL table name to store serialized data messages. By default, 'records'.
save_results : bool
Whether to persist the results after the main Python session ends. Default is False.
data_model : Optional[Data]
Pydantic data model to validate the results before storing them in the database.
Also used to deserialize the results from the database.
auth_token : Optional[str]
The authentication token to use for the WebSocket connection. Default is None.
Only used for API and Python application endpoints.
logger : Optional[logging.Logger]
The pre-configured logger instance to use for this connection. By default, a new logger is created.
kwargs : Optional[dict]
Additional keyword arguments to pass to the target provider module. Keywords and values must not contain spaces.
To pass items to 'websocket.connect()', include them in the 'kwargs' dictionary as,
{'connect_kwargs': {'key': 'value'}}.
Properties
----------
symbol : str
Symbol(s) requested to subscribe.
module : str
Path to the provider connection script.
is_running : bool
Check if the provider connection process is running.
is_broadcasting : bool
Check if the broadcast server process is running.
broadcast_address : str
URI address for the results broadcast server.
results : list[Data]
All stored results from the provider's WebSocket stream.
Results are stored in a SQLite database as a serialized JSON string, this property deserializes the results.
Clear the results by deleting the property. e.g., del client.results
Methods
-------
connect
Connect to the provider WebSocket stream.
disconnect
Disconnect from the provider WebSocket.
subscribe
Subscribe to a new symbol or list of symbols.
unsubscribe
Unsubscribe from a symbol or list of symbols.
start_broadcasting
Start the broadcast server to stream results over a network connection.
stop_broadcasting
Stop the broadcast server and disconnect all listening clients.
send_message
Send a message to the WebSocket process. Messages can be sent to "provider" or "broadcast" targets.
"""
def __init__( # noqa: PLR0913 # pylint: disable=too-many-arguments,too-many-positional-arguments,too-many-locals
self,
name: str,
module: str,
symbol: Optional[str] = None,
limit: Optional[int] = 1000,
results_file: Optional[str] = None,
table_name: Optional[str] = None,
save_results: bool = False,
data_model: Optional["Data"] = None,
auth_token: Optional[str] = None,
logger: Optional[logging.Logger] = None,
**kwargs,
) -> None:
"""Initialize the WebSocketClient class."""
# pylint: disable=import-outside-toplevel
import atexit # noqa
import os
import tempfile
import threading
from aiosqlite import DatabaseError
from queue import Queue
from pathlib import Path
from openbb_core.app.model.abstract.error import OpenBBError
from openbb_websockets.helpers import get_logger
from openbb_websockets.helpers import encrypt_value
self.name = name
self.module = module.replace(".py", "") # type: ignore
self.results_file = results_file if results_file else None
self.table_name = table_name if table_name else "records"
self._limit = limit
self.data_model = data_model
self._symbol = symbol
self._key = os.urandom(32)
self._iv = os.urandom(16)
self._auth_token = (
encrypt_value(self._key, self._iv, auth_token) if auth_token else None
)
# strings in kwargs are encrypted before storing in the class but unencrypted when passed to the provider module.
if kwargs:
for k, v in kwargs.items():
if isinstance(v, str):
kwargs[k] = encrypt_value(self._key, self._iv, v)
else:
kwargs[k] = v
self._kwargs = kwargs if kwargs else {}
self._process: Any = None
self._psutil_process: Any = None
self._thread: Any = None
self._log_thread: Any = None
self._provider_message_queue: Queue = Queue()
self._stop_log_thread_event: threading.Event = threading.Event()
self._stop_broadcasting_event: threading.Event = threading.Event()
self._broadcast_address: Any = None
self._broadcast_process: Any = None
self._psutil_broadcast_process: Any = None
self._broadcast_thread: Any = None
self._broadcast_log_thread: Any = None
self._broadcast_message_queue: Queue = Queue()
self._exception: Any = None
if not results_file:
with tempfile.NamedTemporaryFile(delete=False) as temp_file:
temp_file_path = temp_file.name
self.results_path = Path(temp_file_path).absolute()
self.results_file = temp_file_path
self.results_path = Path(self.results_file).absolute() # type: ignore
self.save_results = save_results
self.logger = logger if logger else get_logger("openbb.websocket.client")
atexit.register(self._atexit)
try:
self._setup_database()
except DatabaseError as e:
msg = f"Unexpected error setting up the SQLite database and table -> {e.__class___.__name__}: {e}"
self.logger.error(msg)
self._exception = OpenBBError(msg)
def _atexit(self) -> None:
"""Clean up the WebSocket client processes at exit."""
# pylint: disable=import-outside-toplevel
import os
self._exception = None
if self.is_running:
self.disconnect()
if self.is_broadcasting:
self.stop_broadcasting()
if self.save_results:
self.logger.info("Websocket results saved to, %s\n", str(self.results_path))
if os.path.exists(self.results_file) and not self.save_results: # type: ignore
os.remove(self.results_file) # type: ignore
def _setup_database(self) -> None:
"""Set up the SQLite database and table."""
# pylint: disable=import-outside-toplevel
from openbb_core.provider.utils.helpers import run_async # noqa
from openbb_websockets.helpers import setup_database
run_async(setup_database, self.results_path, self.table_name)
def _log_provider_output(self, output_queue) -> None:
"""Log output from the provider logger, handling exceptions, errors, and messages that are not data."""
# pylint: disable=import-outside-toplevel
import json # noqa
import queue
import sys
from openbb_core.provider.utils.errors import UnauthorizedError
from openbb_websockets.helpers import clean_message
from pydantic import ValidationError
while not self._stop_log_thread_event.is_set():
try:
output = output_queue.get(timeout=1)
if output:
# Handle raised exceptions from the provider connection thread, killing the process if required.
# UnauthorizedError should be raised by the parent thread, but we kill the process here.
if "UnauthorizedError" in output:
self._psutil_process.kill()
self._process.wait()
self._thread.join()
err = UnauthorizedError(output)
self._exception = err
sys.stdout.write(output + "\n")
sys.stdout.flush()
break
# ValidationError may occur after the provider connection is established.
# We write to stdout in case the exception can't be raised before the main function returns.
# We kill the connection here.
if "ValidationError" in output:
self._psutil_process.kill()
self._process.wait()
self._thread.join()
title, errors = output.split(" -> ")[-1].split(": ")
line_errors = json.loads(errors.strip())
err = ValidationError.from_exception_data(
title=title.strip(), line_errors=line_errors
)
self._exception = err
msg = (
"PROVIDER ERROR: Disconnecting because a ValidatonError was raised"
+ " by the provider while processing data."
+ f"\n\n{str(err)}\n"
)
sys.stdout.write(msg + "\n")
sys.stdout.flush()
break
# We don't kill the process on SymbolError, but raise the exception in the main thread instead.
# This is likely a subscribe event and the connection is already streaming.
if "SymbolError" in output:
err = ValueError(output)
self._exception = err
sys.stdout.write(output + "\n")
sys.stdout.flush()
continue
# Other errors are logged to stdout and the process is killed.
# If the exception is raised by the parent thread, it will be treated as an unexpected error.
if (
"server rejected" in output.lower()
or "PROVIDER ERROR" in output
or "unexpected error" in output.lower()
):
self._psutil_process.kill()
self._process.wait()
self._thread.join()
err = ChildProcessError(output)
self._exception = err
output = output + "\n"
sys.stdout.write(output)
sys.stdout.flush()
break
output = clean_message(output)
if output.startswith("ERROR:"):
output = output.replace("ERROR:", "PROVIDER ERROR:")
elif output.startswith("INFO:"):
output = output.replace("INFO:", "PROVIDER INFO:")
sys.stdout.write(output + "\n")
sys.stdout.flush()
except queue.Empty:
continue
def _log_broadcast_output(self, output_queue) -> None:
"""Log output from the broadcast server queue."""
# pylint: disable=import-outside-toplevel
import queue # noqa
import sys
from openbb_websockets.helpers import clean_message
while not self._stop_broadcasting_event.is_set():
try:
output = output_queue.get(timeout=1)
if output and "Uvicorn running" in output:
address = (
output.split("Uvicorn running on ")[-1]
.strip()
.replace(" (Press CTRL+C to quit)", "")
.replace("http", "ws")
)
output = "INFO: " + f"Stream results from {address}"
self._broadcast_address = address
if output and "Started server process" in output:
output = None
if output and "Waiting for application startup." in output:
output = None
if output and "Application startup complete." in output:
output = None
if output:
if output.startswith("ERROR:"):
output = output.replace("ERROR:", "BROADCAST ERROR:") + "\n"
elif output.startswith("INFO:"):
output = output.replace("INFO:", "BROADCAST INFO:") + "\n"
output = output[0] if isinstance(output, tuple) else output
output = clean_message(output)
sys.stdout.write(output + "\n")
sys.stdout.flush()
except queue.Empty:
continue
def connect(self) -> None: # pylint: disable=too-many-locals
"""Connect to the provider WebSocket."""
# pylint: disable=import-outside-toplevel
import os # noqa
import psutil
import queue
import subprocess
import threading
import time
from openbb_core.app.model.abstract.error import OpenBBError
from openbb_websockets.helpers import decrypt_value
if self.is_running:
self.logger.info("Provider connection already running.")
return
symbol = self.symbol
if not symbol:
self.logger.info("No subscribed symbols.")
return
command = self.module
command.extend([f"symbol={symbol}"])
command.extend([f"results_file={self.results_file}"])
command.extend([f"table_name={self.table_name}"])
if self.limit:
command.extend([f"limit={self.limit}"])
kwargs = self._kwargs.copy()
if kwargs:
for k, v in kwargs.items():
if isinstance(v, str):
unencrypted_value = decrypt_value(
self._key, self._iv, v # pylint: disable=protected-access
)
kwargs[k] = unencrypted_value
else:
kwargs[k] = v
_kwargs = (
[f"{k}={str(v).strip().replace(' ', '_')}" for k, v in kwargs.items()]
if kwargs
else None
)
if _kwargs is not None:
for kwarg in _kwargs:
if kwarg not in command:
command.extend([kwarg])
self._process = subprocess.Popen( # noqa # pylint: disable=consider-using-with
command,
stdout=subprocess.PIPE,
stderr=subprocess.STDOUT,
stdin=subprocess.PIPE,
env=os.environ,
text=True,
bufsize=1,
)
self._psutil_process = psutil.Process(self._process.pid)
log_output_queue: queue.Queue = queue.Queue()
self._thread = threading.Thread(
target=non_blocking_websocket,
args=(
self,
log_output_queue,
self._provider_message_queue,
),
)
self._thread.daemon = True
self._thread.start()
self._log_thread = threading.Thread(
target=self._log_provider_output,
args=(log_output_queue,),
)
self._log_thread.daemon = True
self._log_thread.start()
# Give it some startup time to allow the connection to be established and for exceptions to populate.
time.sleep(2)
if self._exception is not None:
exc = getattr(self, "_exception", None)
self._exception = None
raise OpenBBError(exc)
if not self.is_running:
self.logger.error("The provider server failed to start.")
def send_message(
self, message, target: Literal["provider", "broadcast"] = "provider"
) -> None:
"""Send a message to the WebSocket process."""
if target == "provider":
self._provider_message_queue.put(message)
read_message_queue(self, self._provider_message_queue)
elif target == "broadcast":
self._broadcast_message_queue.put(message)
read_message_queue(self, self._broadcast_message_queue, target="broadcast")
def disconnect(self) -> None:
"""Disconnect from the provider WebSocket."""
self._stop_log_thread_event.set()
if self._process is None or self.is_running is False:
self.logger.info("Not connected to the provider WebSocket.")
return
if (
self._psutil_process is not None
and hasattr(self._psutil_process, "is_running")
and self._psutil_process.is_running()
):
self._psutil_process.kill()
self._process.wait()
self._thread.join()
self._log_thread.join()
self._stop_log_thread_event.clear()
self.logger.info("Disconnected from the provider WebSocket.")
if hasattr(self, "_exception") and self._exception:
raise self._exception
return
def subscribe(self, symbol) -> None:
"""Subscribe to a new symbol or list of symbols."""
# pylint: disable=import-outside-toplevel
import json # noqa
import time
from openbb_core.app.model.abstract.error import OpenBBError
if not self.is_running:
raise OpenBBError("Provider connection is not running.")
ticker = symbol if isinstance(symbol, list) else symbol.split(",")
msg = {"event": "subscribe", "symbol": ticker}
self.send_message(json.dumps(msg))
time.sleep(0.1)
if self._exception:
exc = getattr(self, "_exception", None)
self._exception = None
raise OpenBBError(exc)
old_symbols = self.symbol.split(",") if self.symbol is not None else []
new_symbols = list(set(old_symbols + ticker))
self._symbol = ",".join(new_symbols)
def unsubscribe(self, symbol) -> None:
"""Unsubscribe from a symbol or list of symbols."""
# pylint: disable=import-outside-toplevel
import json # noqa
import time
from openbb_core.app.model.abstract.error import OpenBBError
if not self.symbol:
self.logger.info("No subscribed symbols.")
return
if not self.is_running:
raise OpenBBError("Provider connection is not running.")
ticker = symbol if isinstance(symbol, list) else symbol.split(",")
msg = {"event": "unsubscribe", "symbol": ticker}
self.send_message(json.dumps(msg))
time.sleep(0.1)
old_symbols = self.symbol.split(",")
new_symbols = list(set(old_symbols) - set(ticker))
self._symbol = ",".join(new_symbols)
@property
def is_running(self) -> bool:
"""Check if the provider connection is running."""
if hasattr(self._psutil_process, "is_running"):
return self._psutil_process.is_running()
return False
@property
def is_broadcasting(self) -> bool:
"""Check if the broadcast server is running."""
if hasattr(self._psutil_broadcast_process, "is_running"):
return self._psutil_broadcast_process.is_running()
return False
@property
def results(self) -> Union[list[dict], list["Data"], None]:
"""Retrieve the deserialized results from the results file."""
# pylint: disable=import-outside-toplevel
import json # noqa
import sqlite3
from openbb_core.app.model.abstract.error import OpenBBError
output: list = []
file_path = self.results_path
if file_path.exists():
with sqlite3.connect(file_path) as conn:
try:
cursor = conn.execute(f"SELECT * FROM {self.table_name}") # noqa
for row in cursor:
_, message = row
if self.data_model:
message = json.loads(message)
if isinstance(message, (str, bytes)):
output.append(
self.data_model.model_validate_json(message)
)
elif isinstance(message, dict):
output.append(self.data_model(**message))
else:
output.append(json.loads(json.loads(message)))
except Exception as e:
raise OpenBBError(f"Error retrieving results: {e}") from e
if output:
return output
self.logger.info("No results found in %s", self.results_file)
return None
@results.deleter
def results(self):
"""Clear results stored from the WebSocket stream."""
# pylint: disable=import-outside-toplevel
import sqlite3 # noqa
try:
with sqlite3.connect(self.results_path) as conn:
conn.execute(f"DELETE FROM {self.table_name}") # noqa
conn.commit()
self._setup_database()
self.logger.info(
"Results cleared from table %s in %s",
self.table_name,
self.results_file,
)
except Exception as e: # pylint: disable=broad-except
msg = f"Error clearing results: {e.__class__.__name__}: {e}"
self.logger.error(msg)
@property
def module(self) -> list:
"""Path to the provider connection script."""
return self._module
@module.setter
def module(self, module):
"""Set the path to the provider connection script."""
# pylint: disable=import-outside-toplevel
import sys
self._module = [
sys.executable,
"-m",
module,
]
@property
def symbol(self) -> Union[str, None]:
"""Symbol(s) requested to subscribe."""
return self._symbol
@property
def limit(self) -> Union[int, None]:
"""Get the limit of records to hold in memory."""
return self._limit
@limit.setter
def limit(self, limit):
"""Set the limit of records to hold in memory."""
self._limit = limit
@property
def broadcast_address(self) -> Union[str, None]:
"""Get the WebSocket broadcast address."""
return (
self._broadcast_address
if self._broadcast_address and self.is_broadcasting
else None
)
def start_broadcasting( # pylint: disable=too-many-locals
self,
host: str = "127.0.0.1",
port: int = 6666,
**kwargs,
) -> None:
"""Broadcast results over a network connection."""
# pylint: disable=import-outside-toplevel
import os # noqa
import subprocess
import sys
import threading
import psutil
import queue
from openbb_platform_api.utils.api import check_port
from openbb_websockets.helpers import decrypt_value
if (
self._broadcast_process is not None
and self._broadcast_process.poll() is None
):
msg = f"WebSocket broadcast already running on: {self._broadcast_address}"
self.logger.info(msg)
return
open_port = check_port(host, port)
if open_port != port:
msg = f"Port {port} is already in use. Using {open_port} instead."
self.logger.warning(msg)
command = [
sys.executable,
"-m",
"openbb_websockets.broadcast",
f"host={host}",
f"port={open_port}",
f"results_file={self.results_file}",
f"table_name={self.table_name}",
f"auth_token={decrypt_value(self._key, self._iv, self._auth_token) if self._auth_token else None}",
]
if kwargs:
for k, v in kwargs.items():
command.extend([f"{k}={v}"])
self._broadcast_process = (
subprocess.Popen( # noqa # pylint: disable=consider-using-with
command,
stdout=subprocess.PIPE,
stderr=subprocess.STDOUT,
stdin=subprocess.PIPE,
env=os.environ,
text=True,
bufsize=1,
)
)
self._psutil_broadcast_process = psutil.Process(self._broadcast_process.pid)
output_queue: queue.Queue = queue.Queue()
self._broadcast_thread = threading.Thread(
target=non_blocking_broadcast,
args=(
self,
output_queue,
self._broadcast_message_queue,
),
)
self._broadcast_thread.daemon = True
self._broadcast_thread.start()
self._broadcast_log_thread = threading.Thread(
target=self._log_broadcast_output,
args=(output_queue,),
)
self._broadcast_log_thread.daemon = True
self._broadcast_log_thread.start()
if not self.is_broadcasting:
self.logger.error(
"The broadcast server failed to start on: %s",
self._broadcast_address,
)
def stop_broadcasting(self):
"""Stop the broadcast server."""
broadcast_address = self._broadcast_address
self._stop_broadcasting_event.set()
if self._broadcast_process is None or self.is_broadcasting is False:
self.logger.info("Not currently broadcasting.")
return
if (
self._psutil_broadcast_process is not None
and hasattr(self._psutil_broadcast_process, "is_running")
and self._psutil_broadcast_process.is_running()
):
self._psutil_broadcast_process.kill()
if broadcast_address:
self.logger.info("Stopped broadcasting to: %s", broadcast_address)
self._broadcast_process.wait()
self._broadcast_thread.join()
self._broadcast_log_thread.join()
self._broadcast_process = None
self._psutil_broadcast_process = None
self._broadcast_address = None
self._stop_broadcasting_event.clear()
return
def __repr__(self):
"""Return the WebSocketClient representation."""
return (
f"WebSocketClient(module={self.module}, symbol={self.symbol}, "
f"is_running={self.is_running}, provider_pid: "
f"{self._psutil_process.pid if self._psutil_process else ''}, is_broadcasting={self.is_broadcasting}, "
f"broadcast_address={self.broadcast_address}, "
f"broadcast_pid: {self._psutil_broadcast_process.pid if self._psutil_broadcast_process else ''}, "
f"results_file={self.results_file}, table_name={self.table_name}, "
f"save_results={self.save_results})"
)
def non_blocking_websocket(client, output_queue, provider_message_queue) -> None:
"""Communicate with the threaded process."""
try:
while not client._stop_log_thread_event.is_set():
while not provider_message_queue.empty():
read_message_queue(client, provider_message_queue)
output = client._process.stdout.readline()
if output == "" and client._process.poll() is not None:
break
if output:
output_queue.put(output.strip())
except Exception as e:
msg = (
f"Unexpected error in non_blocking_websocket: {e.__class__.__name__} -> {e}"
)
client.logger.error(msg)
raise e from e
finally:
client._process.stdout.close()
client._process.wait()
def send_message(
client, message, target: Literal["provider", "broadcast"] = "provider"
) -> None:
"""Send a message to the WebSocketConnection process."""
# pylint: disable=import-outside-toplevel
import json
if isinstance(message, (dict, list)):
message = json.dumps(message)
try:
if target == "provider":
if client._process and client._process.stdin:
client._process.stdin.write(message + "\n")
client._process.stdin.flush()
else:
client.logger.error("Provider process is not running.")
elif target == "broadcast":
if client._broadcast_process and client._broadcast_process.stdin:
client._broadcast_process.stdin.write(message + "\n")
client._broadcast_process.stdin.flush()
else:
client.logger.error("Broadcast process is not running.")
except Exception as e: # pylint: disable=broad-except
msg = f"Error sending message to the {target} process: {e.__class__.__name__} -> {e}"
client.logger.error(msg)
def read_message_queue(
client, message_queue, target: Literal["provider", "broadcast"] = "provider"
):
"""Read messages from the queue and send them to the WebSocketConnection process."""
while not message_queue.empty():
message = message_queue.get(timeout=1)
if message:
try:
if target == "provider" and not client._stop_log_thread_event.is_set():
send_message(client, message, target="provider")
elif (
target == "broadcast"
and not client._stop_broadcasting_event.is_set()
):
send_message(client, message, target="broadcast")
except Exception as e: # pylint: disable=broad-except
err = (
f"Error while attempting to transmit from the outgoing message queue: {e.__class__.__name__} "
f"-> {e} -> {message}"
)
client.logger.error(err)
def non_blocking_broadcast(client, output_queue, broadcast_message_queue) -> None:
"""Continuously read the output from the broadcast process and log it to the main thread."""
try:
while not client._stop_broadcasting_event.is_set():
while not broadcast_message_queue.empty():
read_message_queue(client, broadcast_message_queue, target="broadcast")
output = client._broadcast_process.stdout.readline()
if output == "" and client._broadcast_process.poll() is not None:
break
if output:
output_queue.put(output.strip())
except Exception as e: # pylint: disable=broad-except
err = (
f"Unexpected error in non_blocking_broadcast: {e.__class__.__name__} -> {e}"
)
client.logger.error(err)
finally:
client._broadcast_process.stdout.close()
client._broadcast_process.wait()