mirror of
https://github.com/OpenBB-finance/OpenBB.git
synced 2026-10-28 04:52:07 +00:00
808 lines
31 KiB
Python
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()
|