diff --git a/openbb_platform/core/pyproject.toml b/openbb_platform/core/pyproject.toml index 2d4244e5ce7..b6218a330d8 100644 --- a/openbb_platform/core/pyproject.toml +++ b/openbb_platform/core/pyproject.toml @@ -8,7 +8,7 @@ readme = "README.md" packages = [{ include = "openbb_core" }] [tool.poetry.dependencies] -python = "^3.9.21" +python = ">=3.9.21,<3.13" uvicorn = "^0.34.0" websockets = "^15.0" pandas = ">=1.5.3" diff --git a/openbb_platform/extensions/commodity/pyproject.toml b/openbb_platform/extensions/commodity/pyproject.toml index e710d506ac6..c33fefc7ef7 100644 --- a/openbb_platform/extensions/commodity/pyproject.toml +++ b/openbb_platform/extensions/commodity/pyproject.toml @@ -8,7 +8,7 @@ readme = "README.md" packages = [{ include = "openbb_commodity" }] [tool.poetry.dependencies] -python = "^3.9" +python = ">=3.9.21,<3.13" openbb-core = "^1.4.3" [build-system] diff --git a/openbb_platform/extensions/crypto/pyproject.toml b/openbb_platform/extensions/crypto/pyproject.toml index fb09bde2398..36dd711ee80 100644 --- a/openbb_platform/extensions/crypto/pyproject.toml +++ b/openbb_platform/extensions/crypto/pyproject.toml @@ -8,7 +8,7 @@ readme = "README.md" packages = [{ include = "openbb_crypto" }] [tool.poetry.dependencies] -python = "^3.9" +python = ">=3.9.21,<3.13" openbb-core = "^1.4.3" [build-system] diff --git a/openbb_platform/extensions/currency/pyproject.toml b/openbb_platform/extensions/currency/pyproject.toml index 9cdadb63da4..c72cc1c3a62 100644 --- a/openbb_platform/extensions/currency/pyproject.toml +++ b/openbb_platform/extensions/currency/pyproject.toml @@ -8,7 +8,7 @@ readme = "README.md" packages = [{ include = "openbb_currency" }] [tool.poetry.dependencies] -python = "^3.9" +python = ">=3.9.21,<3.13" openbb-core = "^1.4.3" [build-system] diff --git a/openbb_platform/extensions/derivatives/pyproject.toml b/openbb_platform/extensions/derivatives/pyproject.toml index 99895bec7b8..a3fdd952192 100644 --- a/openbb_platform/extensions/derivatives/pyproject.toml +++ b/openbb_platform/extensions/derivatives/pyproject.toml @@ -8,7 +8,7 @@ readme = "README.md" packages = [{ include = "openbb_derivatives" }] [tool.poetry.dependencies] -python = "^3.9" +python = ">=3.9.21,<3.13" openbb-core = "^1.4.3" [build-system] diff --git a/openbb_platform/extensions/economy/pyproject.toml b/openbb_platform/extensions/economy/pyproject.toml index c0fc395a015..2b681e2da32 100644 --- a/openbb_platform/extensions/economy/pyproject.toml +++ b/openbb_platform/extensions/economy/pyproject.toml @@ -8,7 +8,7 @@ readme = "README.md" packages = [{ include = "openbb_economy" }] [tool.poetry.dependencies] -python = "^3.9" +python = ">=3.9.21,<3.13" openbb-core = "^1.4.3" [build-system] diff --git a/openbb_platform/extensions/equity/pyproject.toml b/openbb_platform/extensions/equity/pyproject.toml index 0d7921750e5..dd8f841735e 100644 --- a/openbb_platform/extensions/equity/pyproject.toml +++ b/openbb_platform/extensions/equity/pyproject.toml @@ -8,7 +8,7 @@ readme = "README.md" packages = [{ include = "openbb_equity" }] [tool.poetry.dependencies] -python = "^3.9" +python = ">=3.9.21,<3.13" openbb-core = "^1.4.3" [build-system] diff --git a/openbb_platform/extensions/fixedincome/pyproject.toml b/openbb_platform/extensions/fixedincome/pyproject.toml index 69d7c2a4dcf..f766082bc40 100644 --- a/openbb_platform/extensions/fixedincome/pyproject.toml +++ b/openbb_platform/extensions/fixedincome/pyproject.toml @@ -8,7 +8,7 @@ readme = "README.md" packages = [{ include = "openbb_fixedincome" }] [tool.poetry.dependencies] -python = "^3.9" +python = ">=3.9.21,<3.13" openbb-core = "^1.4.3" [build-system] diff --git a/openbb_platform/extensions/index/pyproject.toml b/openbb_platform/extensions/index/pyproject.toml index 9c985a00783..73760b5655e 100644 --- a/openbb_platform/extensions/index/pyproject.toml +++ b/openbb_platform/extensions/index/pyproject.toml @@ -8,7 +8,7 @@ readme = "README.md" packages = [{ include = "openbb_index" }] [tool.poetry.dependencies] -python = "^3.9" +python = ">=3.9.21,<3.13" openbb-core = "^1.4.3" [build-system] diff --git a/openbb_platform/extensions/news/pyproject.toml b/openbb_platform/extensions/news/pyproject.toml index 93cc17cbbf8..0cc939be24e 100644 --- a/openbb_platform/extensions/news/pyproject.toml +++ b/openbb_platform/extensions/news/pyproject.toml @@ -8,7 +8,7 @@ readme = "README.md" packages = [{ include = "openbb_news" }] [tool.poetry.dependencies] -python = "^3.9" +python = ">=3.9.21,<3.13" openbb-core = "^1.4.3" [build-system] diff --git a/openbb_platform/extensions/websockets/openbb_websockets/_client.py b/openbb_platform/extensions/websockets/openbb_websockets/_client.py new file mode 100644 index 00000000000..82b04254c1c --- /dev/null +++ b/openbb_platform/extensions/websockets/openbb_websockets/_client.py @@ -0,0 +1,775 @@ +"""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 queue import Queue + from pathlib import Path + from openbb_core.app.model.abstract.error import OpenBBError + from openbb_core.provider.utils.websockets.database import Database + from openbb_core.provider.utils.websockets.helpers import ( + encrypt_value, + get_logger, + ) + + 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.database = Database( + results_file=self.results_file, + table_name=self.table_name, + limit=self._limit, + keep_results=self.save_results, + logger=self.logger, + data_model=self.data_model, + ) + except Exception as e: # pylint: disable=broad-except + msg = ( + "Unexpected error setting up the SQLite database and table ->" + f" {e.__class__.__name__ if hasattr(e, '__class__') else e.__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 _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_core.provider.utils.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_core.provider.utils.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_core.provider.utils.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) -> list: + """Retrieve the deserialized results from the results file.""" + # pylint: disable=import-outside-toplevel + from openbb_core.app.model.abstract.error import OpenBBError + + try: + return self.database.fetch_all() + except Exception as e: # pylint: disable=broad-except + raise OpenBBError(f"Error retrieving results: {e}") from e + + @results.deleter + def results(self): + """Clear results stored from the WebSocket stream.""" + # pylint: disable=import-outside-toplevel + try: + self.database.clear_results() + 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_core.provider.utils.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() diff --git a/openbb_platform/extensions/websockets/openbb_websockets/app copy 2.py b/openbb_platform/extensions/websockets/openbb_websockets/app copy 2.py new file mode 100644 index 00000000000..5462515e000 --- /dev/null +++ b/openbb_platform/extensions/websockets/openbb_websockets/app copy 2.py @@ -0,0 +1,457 @@ +import asyncio +import contextlib +import json +import logging +import queue +import threading +from datetime import datetime +from typing import Dict, Optional, Set + +import websockets +from fastapi import FastAPI, WebSocket, WebSocketDisconnect +from fastapi.middleware.cors import CORSMiddleware +from websockets.exceptions import ConnectionClosedError, ConnectionClosedOK + +app = FastAPI() + +app.add_middleware( + CORSMiddleware, + allow_origins=["*"], + allow_credentials=True, + allow_methods=["*"], + allow_headers=["*"], +) + +logger = logging.getLogger("uvicorn-error") + + +@app.get("/") +def read_root(): + return {"Info": "Live Grid Widget Example For OpenBB Workspace Custom Backend"} + + +class ConnectionManager: + def __init__(self): + self.active_connections: Set[WebSocket] = set() + self.client_subscriptions: Dict[WebSocket, Set[str]] = {} + self.symbol_subscribers: Dict[str, Set[WebSocket]] = {} + self.external_ws: Optional[websockets.WebSocketClientProtocol] = None + self.external_task: Optional[asyncio.Task] = None + self.message_queue = queue.Queue() + self.processor_thread: Optional[threading.Thread] = None + self.is_running = False + self.lock = asyncio.Lock() + self.thread_running = False + self.loop = None + + async def connect(self, websocket: WebSocket): + await websocket.accept() + self.active_connections.add(websocket) + self.client_subscriptions[websocket] = set() + + # Start the master connection if not already running + if not self.is_running: + await self.start_master_connection() + + async def disconnect(self, websocket: WebSocket): + if websocket in self.active_connections: + self.active_connections.remove(websocket) + + # Remove client subscriptions + if websocket in self.client_subscriptions: + subscribed_symbols = self.client_subscriptions.pop(websocket) + + # Update symbol subscribers + for symbol in subscribed_symbols: + if symbol in self.symbol_subscribers: + self.symbol_subscribers[symbol].discard(websocket) + + # If no more subscribers for this symbol, unsubscribe from Deribit + if not self.symbol_subscribers[symbol] and self.external_ws: + await self.unsubscribe_symbols([symbol]) + + if websocket.application_state == 1: + await websocket.close(reason="Client disconnected") + + # If no more clients, stop the master connection + if not self.active_connections and self.is_running: + await self.stop_master_connection() + + def process_messages_thread(self): + """Thread function to process messages from the queue""" + asyncio.set_event_loop(self.loop) + self.thread_running = True + + while self.thread_running: + try: + # Get message from queue with timeout to allow checking thread_running + try: + message = self.message_queue.get(timeout=0.5) + except queue.Empty: + continue + + # Check if this is a reconnection signal + if isinstance(message, dict) and message.get("_reconnect"): + asyncio.run_coroutine_threadsafe(self.reconnect(), self.loop) + self.message_queue.task_done() + continue + + # Process normal message + try: + message_data = json.loads(message) + + # Process the message + data = message_data.get("result") or message_data.get( + "params", {} + ).get("data", {}) + if data and isinstance(data, dict): + symbol = data.get("index_name") + if symbol: + # Format message + msg = { + "date": datetime.fromtimestamp( + data.get("timestamp") / 1000 + ).strftime("%Y-%m-%d %H:%M:%S"), + "symbol": symbol, + "price": data.get("price"), + } + + # Broadcast to subscribers - must be done in the event loop + subscribers = self.symbol_subscribers.get( + symbol, set() + ).copy() + if subscribers: + asyncio.run_coroutine_threadsafe( + self.broadcast_to_subscribers(symbol, msg), + self.loop, + ) + except json.JSONDecodeError: + logger.error(f"Invalid JSON received: {message}") + except Exception as e: + logger.error(f"Error processing message: {e}") + finally: + # Mark the message as processed + self.message_queue.task_done() + except Exception as e: + logger.error(f"Error in message processor thread: {e}") + + async def broadcast_to_subscribers(self, symbol, msg): + """Broadcast a message to all subscribers of a symbol""" + subscribers = self.symbol_subscribers.get(symbol, set()) + for client in subscribers: + try: + await client.send_json(msg) + except Exception as e: + logger.error(f"Error sending to client: {e}") + + async def start_master_connection(self): + """Start a single master connection to Deribit""" + # Clean up any existing connection first + await self.stop_master_connection() + + try: + # Establish the WebSocket connection + self.external_ws = await websockets.connect( + "wss://test.deribit.com/ws/api/v2" + ) + self.is_running = True + + # Start task to receive messages from Deribit and put them in the queue + self.external_task = asyncio.create_task(self.receive_from_external()) + + # Store the event loop for the thread to use + self.loop = asyncio.get_event_loop() + + # Start thread to process messages from the queue + self.thread_running = True + self.processor_thread = threading.Thread( + target=self.process_messages_thread, daemon=True + ) + self.processor_thread.start() + + logger.info("Master connection to Deribit established") + + # Resubscribe to all active symbols after reconnection + await self.resubscribe_active_symbols() + except Exception as e: + logger.error(f"Failed to establish master connection: {e}", exc_info=True) + self.is_running = False + + async def resubscribe_active_symbols(self): + """Resubscribe to all symbols that have active subscribers after reconnection""" + active_symbols = [ + symbol + for symbol, subscribers in self.symbol_subscribers.items() + if subscribers # Only include symbols with active subscribers + ] + + if active_symbols and self.external_ws: + await self.subscribe_symbols(active_symbols) + logger.info( + f"Resubscribed to {len(active_symbols)} symbols after reconnection" + ) + + async def stop_master_connection(self): + """Stop the master connection to Deribit""" + self.is_running = False + + # Stop the message processing thread + if self.processor_thread and self.processor_thread.is_alive(): + self.thread_running = False + self.processor_thread.join(timeout=2.0) + self.processor_thread = None + + # Cancel the task that receives messages from Deribit + if self.external_task: + self.external_task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await self.external_task + self.external_task = None + + # Close the WebSocket connection + if self.external_ws: + await self.external_ws.close() + self.external_ws = None + + # Clear the message queue + while True: + try: + self.message_queue.get_nowait() + self.message_queue.task_done() + except queue.Empty: + break + + logger.info("Master connection to Deribit closed") + + async def receive_from_external(self): + """Receive messages from the external WebSocket and put them in the queue""" + while self.is_running and self.external_ws: + try: + response = await self.external_ws.recv() + # Put the raw message in the queue for processing + self.message_queue.put(response) + except (ConnectionClosedOK, ConnectionClosedError, WebSocketDisconnect): + logger.warning("Deribit connection closed, attempting to reconnect...") + # Add a special reconnect message to the queue + self.message_queue.put({"_reconnect": True}) + break + except Exception as e: + logger.error(f"Error in master connection handler: {e}", exc_info=True) + await asyncio.sleep(1) # Prevent tight loop on errors + + async def reconnect(self): + """Handle reconnection after a connection failure""" + if not self.is_running: + return + + self.is_running = False + + # Wait a bit before reconnecting + await asyncio.sleep(5) + + # Only reconnect if we still have clients + if self.active_connections: + await self.start_master_connection() + + async def update_client_symbols(self, websocket: WebSocket, symbols: list): + """Update the symbols a client is subscribed to""" + async with self.lock: + if websocket not in self.active_connections: + return + + old_symbols = self.client_subscriptions.get(websocket, set()) + new_symbols = set(symbols) + + # Symbols to subscribe to + to_subscribe = new_symbols - old_symbols + + # Symbols to unsubscribe from + to_unsubscribe = old_symbols - new_symbols + + # Update client subscriptions + self.client_subscriptions[websocket] = new_symbols + + # Update symbol subscribers and subscribe/unsubscribe as needed + to_add_to_deribit = [] + to_remove_from_deribit = [] + + for symbol in to_subscribe: + # Check if symbol exists but has no subscribers + if ( + symbol in self.symbol_subscribers + and not self.symbol_subscribers[symbol] + ): + to_add_to_deribit.append(symbol) + # Check if symbol doesn't exist at all + elif symbol not in self.symbol_subscribers: + self.symbol_subscribers[symbol] = set() + to_add_to_deribit.append(symbol) + self.symbol_subscribers[symbol].add(websocket) + + for symbol in to_unsubscribe: + if symbol in self.symbol_subscribers: + self.symbol_subscribers[symbol].discard(websocket) + if not self.symbol_subscribers[symbol]: + to_remove_from_deribit.append(symbol) + # Don't delete the key here, we'll do it after unsubscribing + + # Make the actual subscription changes on Deribit + if to_add_to_deribit and self.external_ws: + await self.subscribe_symbols(to_add_to_deribit) + + if to_remove_from_deribit and self.external_ws: + await self.unsubscribe_symbols(to_remove_from_deribit) + # Now we can safely remove empty subscriber sets + for symbol in to_remove_from_deribit: + if ( + symbol in self.symbol_subscribers + and not self.symbol_subscribers[symbol] + ): + del self.symbol_subscribers[symbol] + + async def subscribe_symbols(self, symbols: list): + """Subscribe to symbols on Deribit""" + if not symbols or not self.external_ws: + return + + channels = [f"deribit_price_index.{sym}" for sym in symbols] + + msg = { + "jsonrpc": "2.0", + "id": 3600, + "method": "public/subscribe", + "params": { + "channels": channels, + }, + } + + try: + await self.external_ws.send(json.dumps(msg)) + logger.info(f"Subscribed to: {symbols}") + except Exception as e: + logger.error(f"Failed to subscribe: {e}", exc_info=True) + + async def unsubscribe_symbols(self, symbols: list): + """Unsubscribe from symbols on Deribit""" + if not symbols or not self.external_ws: + return + + channels = [f"deribit_price_index.{sym}" for sym in symbols] + + msg = { + "jsonrpc": "2.0", + "id": 3600, + "method": "public/unsubscribe", + "params": { + "channels": channels, + }, + } + + try: + await self.external_ws.send(json.dumps(msg)) + logger.info(f"Unsubscribed from: {symbols}") + except Exception as e: + logger.error(f"Failed to unsubscribe: {e}", exc_info=True) + + +# Initialize the connection manager +manager = ConnectionManager() + + +async def shutdown_event(): + await manager.stop_master_connection() + + +app.add_event_handler("shutdown", shutdown_event) + + +@app.websocket("/ws") +async def websocket_endpoint(websocket: WebSocket): + await manager.connect(websocket) + + try: + while True: + try: + data = await websocket.receive_text() + symbols = json.loads(data).get("params", {}).get("symbol", []) + await manager.update_client_symbols(websocket, symbols) + except (WebSocketDisconnect, ConnectionClosedOK): + break + except Exception as e: + logger.error(f"Error handling client message: {e}", exc_info=True) + finally: + await manager.disconnect(websocket) + + +@app.get("/get_ws_data") +async def get_messages(symbol: str = None) -> list: + await asyncio.sleep(1) + symbols = symbol.split(",") if isinstance(symbol, str) else symbol + msg = [] + if symbols: + for sym in symbols: + msg.append({"date": None, "symbol": sym, "price": 0}) + + return msg + + +@app.get("/widgets.json") +async def get_widgets(): + """Get widgets.""" + return { + "live_grid_example": { + "name": "Live Grid", + "description": "Live Grid", + "type": "live_grid", + "endpoint": "get_ws_data", + "wsEndpoint": "ws", + "data": { + "wsRowIdColumn": "symbol", + "table": { + "showAll": True, + "columnsDefs": [ + {"field": "date", "headerName": "Date", "type": "date"}, + {"field": "symbol", "headerName": "Symbol", "type": "text"}, + {"field": "price", "headerName": "Price", "type": "number"}, + ], + }, + }, + "params": [ + { + "paramName": "symbol", + "description": "The symbol to get details for", + "value": None, + "label": "Symbol", + "type": "text", + "multiSelect": True, + "options": [ + {"label": "BTC/USD", "value": "btc_usd"}, + {"label": "ETH/USD", "value": "eth_usd"}, + {"label": "SOL/USD", "value": "sol_usd"}, + {"label": "XRP/USD", "value": "xrp_usd"}, + {"label": "BNB/USD", "value": "bnb_usd"}, + {"label": "MATIC/USD", "value": "matic_usd"}, + {"label": "STETH/USD", "value": "steth_usd"}, + {"label": "USYC/USD", "value": "usyc_usd"}, + {"label": "USDC/USD", "value": "usdc_usd"}, + {"label": "USDT/USD", "value": "usdt_usd"}, + {"label": "USDE/USD", "value": "usde_usd"}, + {"label": "PAXG/USD", "value": "paxg_usd"}, + ], + } + ], + "gridData": {"w": 15, "h": 9}, + } + } + + +@app.get("/templates.json") +async def get_templates(): + """Get templates.""" + return [] + + +if __name__ == "__main__": + import uvicorn + + uvicorn.run(app, port=6940) diff --git a/openbb_platform/extensions/websockets/openbb_websockets/app copy.py b/openbb_platform/extensions/websockets/openbb_websockets/app copy.py new file mode 100644 index 00000000000..b09a4e7a0f5 --- /dev/null +++ b/openbb_platform/extensions/websockets/openbb_websockets/app copy.py @@ -0,0 +1,260 @@ +import asyncio +import contextlib +import importlib.util +import json +import logging +import sys +from datetime import datetime +from typing import Annotated + +import websockets +from fastapi import Depends, FastAPI, WebSocket, WebSocketDisconnect +from sqlalchemy import select +from sqlalchemy.sql import text +from websockets.exceptions import ConnectionClosedOK + +# Dynamically import the database module +database_path = "/Users/darrenlee/github/OpenBB/openbb_platform/extensions/websockets/openbb_websockets/database.py" +spec = importlib.util.spec_from_file_location("database", database_path) +database = importlib.util.module_from_spec(spec) +sys.modules["database"] = database +spec.loader.exec_module(database) + +app = FastAPI() +logger = logging.getLogger("uvicorn-error") + +message_queue = asyncio.Queue() + +SUBSCRIBED_SYMBOLS = [] + + +class ConnectedClients: + + def __init__(self): + self.clients = set() + + def add_client(self, client): + self.clients.add(client) + + def remove_client(self, client): + self.clients.remove(client) + + +ConnectedWebsocketClients = ConnectedClients() + + +def get_connected_clients(): + return ConnectedWebsocketClients + + +ActiveClients = Annotated[ConnectedClients, Depends(get_connected_clients)] + + +DbSessionLocal = database.WebSocketDatabaseSession( + host="sqlite+aiosqlite:///sqlite.db", + batch_size=100, + batch_interval=1, + logger=logger, +) + + +@app.on_event("startup") +async def startup_event(): + # Create the database tables + await DbSessionLocal.create_all() + + # Start the task to write messages to the database + asyncio.create_task(await asyncio.to_thread(DbSessionLocal.start_writer_task)) + + +@app.on_event("shutdown") +async def shutdown_event(): + # Dispose of the database connection + await DbSessionLocal.stop_writer_task() + await DbSessionLocal.close() + + +async def symbol_handler(symbols: str, external_ws): + unsubscribe_channels = [] + channels = [] + for sym in list(set(symbols)): + if sym not in SUBSCRIBED_SYMBOLS: + SUBSCRIBED_SYMBOLS.append(sym) + channels.append(f"deribit_price_index.{sym}") + + UNSUBSCRIBE_SYMBOLS = list(set(SUBSCRIBED_SYMBOLS) - set(symbols)) + print(f"Subscribed Symbols: {SUBSCRIBED_SYMBOLS}") + print(f"Unsubscribed Symbols: {UNSUBSCRIBE_SYMBOLS}") + + for sym in UNSUBSCRIBE_SYMBOLS: + SUBSCRIBED_SYMBOLS.remove(sym) + unsubscribe_channels.append(f"deribit_price_index.{sym}") + + if unsubscribe_channels: + msg = { + "jsonrpc": "2.0", + "id": 3600, + "method": "public/unsubscribe", + "params": { + "channels": unsubscribe_channels, + }, + } + print(msg) + await external_ws.send(json.dumps(msg)) + + print(f"Updating subscription to: {channels}") + + msg = { + "jsonrpc": "2.0", + "id": 3600, + "method": "public/subscribe", + "params": { + "channels": channels, + }, + } + + await external_ws.send(json.dumps(msg)) + + +async def message_handler(external_ws, client_ws): + while True: + try: + message = await message_queue.get() + if not message: + asyncio.sleep(0.1) + continue + message = message.get("result") or message.get("params", {}).get("data", {}) + if message and isinstance(message, dict): + msg = { + "date": datetime.fromtimestamp( + message.get("timestamp") / 1000 + ).strftime("%Y-%m-%d %H:%M:%S"), + "symbol": message.get("index_name"), + "price": message.get("price"), + } + await DbSessionLocal.queue.put(msg) + + try: + await client_ws.send_json(msg) + except RuntimeError: + break + + except (WebSocketDisconnect, ConnectionClosedOK): + return # Exit the entire function on disconnect + else: + continue + + except Exception as e: + logger.error(f"Error in message handler: {e}", exc_info=True) + # Continue to the next iteration of the outer loop + + +@app.websocket("/ws") +async def websocket_endpoint(websocket: WebSocket): + await websocket.accept() + + msg = { + "jsonrpc": "2.0", + "id": 3600, + "method": "public/subscribe", + "params": {"channels": []}, + } + try: + async with websockets.connect( + "wss://test.deribit.com/ws/api/v2" + ) as external_ws: + + handler_task = asyncio.create_task( + await asyncio.to_thread(message_handler, external_ws, websocket) + ) + + async def receive_from_external_ws(): + while True: + try: + response = await external_ws.recv() + await message_queue.put(json.loads(response)) + except (WebSocketDisconnect, ConnectionClosedOK): + break + + async def receive_from_client_ws(): + while True: + try: + data = await websocket.receive_text() + print(data) + _symbols = json.loads(data).get("params", {}).get("symbol", []) + await symbol_handler(_symbols, external_ws) + + except (WebSocketDisconnect, ConnectionClosedOK): + break + + try: + await external_ws.send(json.dumps(msg)) + + external_task = asyncio.create_task(receive_from_external_ws()) + client_task = asyncio.create_task(receive_from_client_ws()) + + done, pending = await asyncio.wait( + [external_task, client_task], + return_when=asyncio.FIRST_COMPLETED, + ) + for task in pending: + task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await task + except (WebSocketDisconnect, KeyboardInterrupt, ConnectionClosedOK): + logger.info("WebSocket disconnected") + pass + except Exception as e: + logger.error(f"Error: {e}") + finally: + handler_task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await handler_task + for task in pending: + task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await task + if websocket.application_state == 1: + await websocket.close(reason="WebSocket Disconnected") + + +@app.get("/get_ws_data") +async def get_messages(symbol: str = None) -> list: + global SUBSCRIBED_SYMBOLS + + symbols = symbol.split(",") if isinstance(symbol, str) else symbol + msg = [] + + async with DbSessionLocal.get_read_session() as session: + try: + if symbols: + for sym in symbols: + # Use JSON filtering with SQLite's json operator + if sym not in SUBSCRIBED_SYMBOLS: + SUBSCRIBED_SYMBOLS.append(sym) + result = await session.execute( + select(database.WebSocketMessages) + .where(text("message->>'symbol' = :sym")) + .limit(1) + .offset(-1) + .params(sym=sym) + ) + messages = result.scalars().all() + msg.extend( + [m.message for m in messages] + if messages + else [{"date": None, "symbol": sym, "price": None}] + ) + else: + msg.extend({}) + + return msg + except Exception as e: + logger.error(f"Error querying database: {e}", exc_info=True) + return [] + + +if __name__ == "__main__": + import uvicorn + + uvicorn.run(app, host="0.0.0.0", port=8000) diff --git a/openbb_platform/extensions/websockets/openbb_websockets/app.py b/openbb_platform/extensions/websockets/openbb_websockets/app.py new file mode 100644 index 00000000000..a304a4a12c4 --- /dev/null +++ b/openbb_platform/extensions/websockets/openbb_websockets/app.py @@ -0,0 +1,875 @@ +import asyncio +import contextlib +import json +import logging +import queue +import subprocess +import sys +import threading +import uuid +from datetime import datetime +from pathlib import Path +from typing import Any, Callable, Dict, List, Optional, Set + +import uvicorn +import websockets +from fastapi import BackgroundTasks, FastAPI, WebSocket, WebSocketDisconnect +from fastapi.middleware.cors import CORSMiddleware +from fastapi.responses import StreamingResponse +from websockets.exceptions import ConnectionClosedError, ConnectionClosedOK + +logger = logging.getLogger("uvicorn-error") + + +class ConnectionManager: + def __init__( + self, + ws_uri: str = "wss://test.deribit.com/ws/api/v2", + connection_msg: Optional[Dict] = None, + process_message: Optional[Callable[[Dict], List[Dict]]] = None, + format_subscribe: Optional[Callable[[List[str]], Dict]] = None, + format_unsubscribe: Optional[Callable[[List[str]], Dict]] = None, + connection_timeout: float = 15.0, + retry_delay: float = 5.0, + max_retries: int = 3, + ): + """Initialize a WebSocket connection manager + + Args: + app: The FastAPI app instance to configure + ws_uri: URI for the external WebSocket service + connection_msg: Optional message to send on connection + process_message: Function to process incoming messages + format_subscribe: Function to format subscription requests + format_unsubscribe: Function to format unsubscription requests + connection_timeout: Timeout in seconds for connection attempts + retry_delay: Delay in seconds between connection retries + max_retries: Maximum number of connection retry attempts + """ + # Create the FastAPI app + self.app = FastAPI() + + # Add CORS middleware + self.app.add_middleware( + CORSMiddleware, + allow_origins=["*"], + allow_credentials=True, + allow_methods=["*"], + allow_headers=["*"], + ) + + # Client connections and subscriptions + self.active_connections: Set[WebSocket] = set() + self.client_subscriptions: Dict[WebSocket, Set[str]] = {} + self.symbol_subscribers: Dict[str, Set[WebSocket]] = {} + + # External WebSocket connection + self.external_ws: Optional[websockets.WebSocketClientProtocol] = None + self.external_task: Optional[asyncio.Task] = None + + # Threading and queuing + self.message_queue = queue.Queue() + self.processor_thread: Optional[threading.Thread] = None + self.is_running = False + self.lock = asyncio.Lock() + self.thread_running = False + self.loop = None + + # Configuration + self.ws_uri = ws_uri + self.connection_msg = connection_msg + self.connection_timeout = connection_timeout + self.retry_delay = retry_delay + self.max_retries = max_retries + + # Message handling functions + self.process_message = process_message or self._default_process_message + self.format_subscribe = format_subscribe or self._default_format_subscribe + self.format_unsubscribe = format_unsubscribe or self._default_format_unsubscribe + + # Register endpoints with the app + self._setup_app_routes() + + def _setup_app_routes(self): + """Configure the FastAPI app with routes and event handlers""" + + @self.app.get("/") + def read_root(): + return { + "Info": "Live Grid Widget Example For OpenBB Workspace Custom Backend" + } + + @self.app.on_event("shutdown") + async def shutdown_event(): + await self.stop_master_connection() + + @self.app.websocket("/ws") + async def websocket_endpoint(websocket: WebSocket): + await self.connect(websocket) + try: + while True: + try: + data = await websocket.receive_text() + symbols = json.loads(data).get("params", {}).get("symbol", []) + await self.update_client_symbols(websocket, symbols) + except (WebSocketDisconnect, ConnectionClosedOK): + break + except Exception as e: + logger.error( + f"Error handling client message: {e}", exc_info=True + ) + finally: + await self.disconnect(websocket) + + @self.app.get("/control/subscribe") + async def control_subscribe(symbols: str): + """Subscribe to symbols from external control""" + symbols = symbols.split(",") if isinstance(symbols, str) else symbols + if not symbols: + return {"success": False, "message": "No symbols provided"} + + if not self.external_ws: + # Start the master connection if not already running + if not self.is_running: + await self.start_master_connection() + + if not self.external_ws: + return { + "success": False, + "message": "Failed to establish external connection", + } + + # Add the symbols to a special control subscriber + for symbol in symbols: + if symbol not in self.symbol_subscribers: + self.symbol_subscribers[symbol] = set() + + # Subscribe to the symbols on the external service + await self.subscribe_symbols(symbols) + + return {"success": True, "subscribed": symbols} + + @self.app.get("/master_stream") + async def master_stream_endpoint(): + """Stream raw data directly from the master connection""" + + async def event_generator(): + async for message in self.get_master_stream(): + # Format as Server-Sent Events + yield f"data: {json.dumps(message)}\n\n" + + return StreamingResponse( + event_generator(), + media_type="text/event-stream", + headers={ + "Cache-Control": "no-cache", + "Connection": "keep-alive", + "Content-Type": "text/event-stream", + }, + ) + + @self.app.get("/control/status") + async def control_status(): + """Get the status of the WebSocket connection""" + return { + "connected": self.is_running and self.external_ws is not None, + "active_clients": len(self.active_connections), + "subscribed_symbols": list(self.symbol_subscribers.keys()), + } + + @self.app.get("/get_ws_data") + async def get_messages(symbol: str = None) -> list: + await asyncio.sleep(0.1) # Quick response + symbols = symbol.split(",") if isinstance(symbol, str) else symbol + msg = [] + if symbols: + for sym in symbols: + msg.append({"date": None, "symbol": sym, "price": 0}) + return msg + + @self.app.get("/widgets.json") + async def get_widgets(): + """Get widgets.""" + return { + "live_grid_example": { + "name": "Live Grid", + "description": "Live Grid", + "type": "live_grid", + "endpoint": "get_ws_data", + "wsEndpoint": "ws", + "data": { + "wsRowIdColumn": "symbol", + "table": { + "showAll": True, + "columnsDefs": [ + {"field": "date", "headerName": "Date", "type": "date"}, + { + "field": "symbol", + "headerName": "Symbol", + "type": "text", + }, + { + "field": "price", + "headerName": "Price", + "type": "number", + }, + ], + }, + }, + "params": [ + { + "paramName": "symbol", + "description": "The symbol to get details for", + "value": None, + "label": "Symbol", + "type": "text", + "multiSelect": True, + "options": [ + {"label": "BTC/USD", "value": "btc_usd"}, + {"label": "ETH/USD", "value": "eth_usd"}, + {"label": "SOL/USD", "value": "sol_usd"}, + {"label": "XRP/USD", "value": "xrp_usd"}, + {"label": "BNB/USD", "value": "bnb_usd"}, + {"label": "MATIC/USD", "value": "matic_usd"}, + {"label": "STETH/USD", "value": "steth_usd"}, + {"label": "USYC/USD", "value": "usyc_usd"}, + {"label": "USDC/USD", "value": "usdc_usd"}, + {"label": "USDT/USD", "value": "usdt_usd"}, + {"label": "USDE/USD", "value": "usde_usd"}, + {"label": "PAXG/USD", "value": "paxg_usd"}, + ], + } + ], + "gridData": {"w": 15, "h": 9}, + } + } + + @self.app.get("/templates.json") + async def get_templates(): + """Get templates.""" + return [] + + @self.app.get("/control/unsubscribe") + async def control_unsubscribe(symbols: str): + """Unsubscribe from symbols from external control""" + symbols = symbols.split(",") if isinstance(symbols, str) else symbols + if not symbols or not self.external_ws: + return { + "success": False, + "message": "No symbols or no external connection", + } + + to_unsubscribe = [] + for symbol in symbols: + if ( + symbol in self.symbol_subscribers + and not self.symbol_subscribers[symbol] + ): + to_unsubscribe.append(symbol) + + if to_unsubscribe: + await self.unsubscribe_symbols(to_unsubscribe) + + # Clean up the symbol subscribers dictionary + for symbol in to_unsubscribe: + if ( + symbol in self.symbol_subscribers + and not self.symbol_subscribers[symbol] + ): + del self.symbol_subscribers[symbol] + + return {"success": True, "unsubscribed": to_unsubscribe} + + async def get_master_stream(self): + """ + Creates an async generator that yields messages directly from the master connection. + This can be used to create streaming endpoints. + + Yields: + Dict: Raw messages from the external WebSocket after parsing from JSON + """ + if not self.external_ws: + # Start the master connection if not already running + if not self.is_running: + await self.start_master_connection() + + if not self.external_ws: + raise RuntimeError("Failed to establish external connection") + + # Create a message queue specifically for this stream + stream_queue = asyncio.Queue() + stream_id = str(uuid.uuid4()) + + # Register this stream in a dictionary of active streams + if not hasattr(self, "active_streams"): + self.active_streams = {} + self.active_streams[stream_id] = stream_queue + + try: + while self.is_running and self.external_ws: + try: + # Wait for a message to be added to this stream's queue + message = await asyncio.wait_for(stream_queue.get(), timeout=30.0) + + # Parse and yield the message + if isinstance(message, str): + try: + yield json.loads(message) + except json.JSONDecodeError: + yield {"error": "Invalid JSON", "raw": message} + else: + yield message + + # Mark as processed + stream_queue.task_done() + + except asyncio.TimeoutError: + # Send a heartbeat to keep the connection alive + yield {"type": "heartbeat", "timestamp": datetime.now().isoformat()} + except Exception as e: + logger.error(f"Error in master stream: {e}", exc_info=True) + yield {"error": str(e)} + await asyncio.sleep(1) + finally: + # Clean up when the stream ends + if hasattr(self, "active_streams") and stream_id in self.active_streams: + del self.active_streams[stream_id] + + def _default_process_message(self, message_data: Dict) -> List[Dict]: + """Default message processor for Deribit format""" + results = [] + + # Extract data from Deribit message format + data = message_data.get("result") or message_data.get("params", {}).get( + "data", {} + ) + if data and isinstance(data, dict): + symbol = data.get("index_name") + if symbol: + # Format message + msg = { + "date": datetime.fromtimestamp( + data.get("timestamp") / 1000 + ).strftime("%Y-%m-%d %H:%M:%S"), + "symbol": symbol, + "price": data.get("price"), + } + results.append(msg) + + return results + + def _default_format_subscribe(self, symbols: List[str]) -> Dict: + """Default subscription formatter for Deribit""" + channels = [f"deribit_price_index.{sym}" for sym in symbols] + return { + "jsonrpc": "2.0", + "id": 3600, + "method": "public/subscribe", + "params": { + "channels": channels, + }, + } + + def _default_format_unsubscribe(self, symbols: List[str]) -> Dict: + """Default unsubscription formatter for Deribit""" + channels = [f"deribit_price_index.{sym}" for sym in symbols] + return { + "jsonrpc": "2.0", + "id": 3600, + "method": "public/unsubscribe", + "params": { + "channels": channels, + }, + } + + async def connect(self, websocket: WebSocket): + await websocket.accept() + self.active_connections.add(websocket) + self.client_subscriptions[websocket] = set() + + # Start the master connection if not already running + if not self.is_running: + await self.start_master_connection() + + async def disconnect(self, websocket: WebSocket): + if websocket in self.active_connections: + self.active_connections.remove(websocket) + + # Remove client subscriptions + if websocket in self.client_subscriptions: + subscribed_symbols = self.client_subscriptions.pop(websocket) + + # Update symbol subscribers + for symbol in subscribed_symbols: + if symbol in self.symbol_subscribers: + self.symbol_subscribers[symbol].discard(websocket) + + # If no more subscribers for this symbol, unsubscribe from external service + if not self.symbol_subscribers[symbol] and self.external_ws: + await self.unsubscribe_symbols([symbol]) + + if websocket.application_state == 1: + await websocket.close(reason="Client disconnected") + + # If no more clients, stop the master connection + if not self.active_connections and self.is_running: + await self.stop_master_connection() + + def process_messages_thread(self): + """Thread function to process messages from the queue""" + asyncio.set_event_loop(self.loop) + self.thread_running = True + + while self.thread_running: + try: + # Get message from queue with timeout to allow checking thread_running + try: + message = self.message_queue.get(timeout=0.5) + except queue.Empty: + continue + + # Check if this is a reconnection signal + if isinstance(message, dict) and message.get("_reconnect"): + asyncio.run_coroutine_threadsafe(self.reconnect(), self.loop) + self.message_queue.task_done() + continue + + # Process normal message + try: + # Parse JSON message + message_data = ( + json.loads(message) if isinstance(message, str) else message + ) + + # Process the message using the custom processor + processed_messages = self.process_message(message_data) + + # Send each processed message to its subscribers + for msg in processed_messages: + symbol = msg.get("symbol") + if symbol: + subscribers = self.symbol_subscribers.get( + symbol, set() + ).copy() + if subscribers: + asyncio.run_coroutine_threadsafe( + self.broadcast_to_subscribers(subscribers, msg), + self.loop, + ) + except json.JSONDecodeError: + logger.error(f"Invalid JSON received: {message}") + except Exception as e: + logger.error(f"Error processing message: {e}") + finally: + # Mark the message as processed + self.message_queue.task_done() + except Exception as e: + logger.error(f"Error in message processor thread: {e}") + + async def broadcast_to_subscribers(self, subscribers: Set[WebSocket], msg: Dict): + """Broadcast a message to all specified subscribers""" + for client in subscribers: + try: + await client.send_json(msg) + except Exception as e: + logger.error(f"Error sending to client: {e}") + + async def start_master_connection(self): + """Start a single master connection to the external WebSocket service""" + # Clean up any existing connection first + await self.stop_master_connection() + + retries = 0 + while retries <= self.max_retries: + try: + # Establish the WebSocket connection with timeout + logger.info( + f"Connecting to {self.ws_uri} (attempt {retries + 1}/{self.max_retries + 1})" + ) + self.external_ws = await asyncio.wait_for( + websockets.connect(self.ws_uri), timeout=self.connection_timeout + ) + self.is_running = True + + # Send connection message if provided + if self.connection_msg and self.external_ws: + await self.external_ws.send(json.dumps(self.connection_msg)) + + # Start task to receive messages from external service + self.external_task = asyncio.create_task(self.receive_from_external()) + + # Store the event loop for the thread to use + self.loop = asyncio.get_event_loop() + + # Start thread to process messages from the queue + self.thread_running = True + self.processor_thread = threading.Thread( + target=self.process_messages_thread, daemon=True + ) + self.processor_thread.start() + + logger.info(f"Master connection to {self.ws_uri} established") + + # Resubscribe to all active symbols after reconnection + await self.resubscribe_active_symbols() + + # Connection successful, exit retry loop + break + + except asyncio.TimeoutError: + retries += 1 + if retries <= self.max_retries: + logger.warning( + f"Connection attempt timed out, retrying in {self.retry_delay} seconds..." + ) + await asyncio.sleep(self.retry_delay) + else: + logger.error( + f"Failed to connect after {self.max_retries + 1} attempts, giving up" + ) + self.is_running = False + + except Exception as e: + retries += 1 + if retries <= self.max_retries: + logger.warning( + f"Connection error: {e}, retrying in {self.retry_delay} seconds..." + ) + await asyncio.sleep(self.retry_delay) + else: + logger.error( + f"Failed to establish master connection: {e}", exc_info=True + ) + self.is_running = False + break + + async def resubscribe_active_symbols(self): + """Resubscribe to all symbols that have active subscribers after reconnection""" + active_symbols = [ + symbol + for symbol, subscribers in self.symbol_subscribers.items() + if subscribers # Only include symbols with active subscribers + ] + + if active_symbols and self.external_ws: + await self.subscribe_symbols(active_symbols) + logger.info( + f"Resubscribed to {len(active_symbols)} symbols after reconnection" + ) + + async def stop_master_connection(self): + """Stop the master connection to the external WebSocket service""" + self.is_running = False + + # Stop the message processing thread + if self.processor_thread and self.processor_thread.is_alive(): + self.thread_running = False + self.processor_thread.join(timeout=2.0) + self.processor_thread = None + + # Cancel the task that receives messages from external service + if self.external_task: + self.external_task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await self.external_task + self.external_task = None + + # Close the WebSocket connection + if self.external_ws: + with contextlib.suppress(Exception): + await self.external_ws.close() + self.external_ws = None + + # Clear the message queue + while True: + try: + self.message_queue.get_nowait() + self.message_queue.task_done() + except queue.Empty: + break + + logger.info(f"Master connection to {self.ws_uri} closed") + + async def receive_from_external(self): + """Receive messages from the external WebSocket and put them in the queue""" + while self.is_running and self.external_ws: + try: + response = await asyncio.wait_for( + self.external_ws.recv(), timeout=30.0 # Add reasonable timeout + ) + # Put the raw message in the queue for processing + self.message_queue.put(response) + + if hasattr(self, "active_streams") and self.active_streams: + for stream_queue in self.active_streams.values(): + await stream_queue.put(response) + + except asyncio.TimeoutError: + # Just a timeout, not necessarily an error, try again + continue + except (ConnectionClosedOK, ConnectionClosedError, WebSocketDisconnect): + logger.warning( + f"Connection to {self.ws_uri} closed, attempting to reconnect..." + ) + # Add a special reconnect message to the queue + self.message_queue.put({"_reconnect": True}) + break + except Exception as e: + logger.error(f"Error in master connection handler: {e}", exc_info=True) + await asyncio.sleep(1) # Prevent tight loop on errors + + async def reconnect(self): + """Handle reconnection after a connection failure""" + if not self.is_running: + return + + self.is_running = False + + # Wait a bit before reconnecting + await asyncio.sleep(self.retry_delay) + + # Only reconnect if we still have clients + if self.active_connections: + await self.start_master_connection() + + async def update_client_symbols(self, websocket: WebSocket, symbols: list): + """Update the symbols a client is subscribed to""" + async with self.lock: + if websocket not in self.active_connections: + return + + old_symbols = self.client_subscriptions.get(websocket, set()) + new_symbols = set(symbols) + + # Symbols to subscribe to + to_subscribe = new_symbols - old_symbols + + # Symbols to unsubscribe from + to_unsubscribe = old_symbols - new_symbols + + # Update client subscriptions + self.client_subscriptions[websocket] = new_symbols + + # Update symbol subscribers and subscribe/unsubscribe as needed + to_add_to_external = [] + to_remove_from_external = [] + + for symbol in to_subscribe: + # Check if symbol exists but has no subscribers + if ( + symbol in self.symbol_subscribers + and not self.symbol_subscribers[symbol] + ): + to_add_to_external.append(symbol) + # Check if symbol doesn't exist at all + elif symbol not in self.symbol_subscribers: + self.symbol_subscribers[symbol] = set() + to_add_to_external.append(symbol) + self.symbol_subscribers[symbol].add(websocket) + + for symbol in to_unsubscribe: + if symbol in self.symbol_subscribers: + self.symbol_subscribers[symbol].discard(websocket) + if not self.symbol_subscribers[symbol]: + to_remove_from_external.append(symbol) + # Don't delete the key here, we'll do it after unsubscribing + + # Make the actual subscription changes on external service + if to_add_to_external and self.external_ws: + await self.subscribe_symbols(to_add_to_external) + + if to_remove_from_external and self.external_ws: + await self.unsubscribe_symbols(to_remove_from_external) + # Now we can safely remove empty subscriber sets + for symbol in to_remove_from_external: + if ( + symbol in self.symbol_subscribers + and not self.symbol_subscribers[symbol] + ): + del self.symbol_subscribers[symbol] + + async def subscribe_symbols(self, symbols: list): + """Subscribe to symbols on external service""" + if not symbols or not self.external_ws: + return + + # Format the subscription message using the configured formatter + msg = self.format_subscribe(symbols) + + try: + await self.external_ws.send(json.dumps(msg)) + logger.info(f"Subscribed to: {symbols}") + except Exception as e: + logger.error(f"Failed to subscribe: {e}", exc_info=True) + + async def unsubscribe_symbols(self, symbols: list): + """Unsubscribe from symbols on external service""" + if not symbols or not self.external_ws: + return + + # Format the unsubscription message using the configured formatter + msg = self.format_unsubscribe(symbols) + + try: + await self.external_ws.send(json.dumps(msg)) + logger.info(f"Unsubscribed from: {symbols}") + except Exception as e: + logger.error(f"Failed to unsubscribe: {e}", exc_info=True) + + +def create_deribit_manager(): + """Create a ConnectionManager configured for Deribit""" + + # Create the Deribit message processor + def process_deribit_message(message_data): + results = [] + + # Extract data from Deribit message format + data = message_data.get("result") or message_data.get("params", {}).get( + "data", {} + ) + if data and isinstance(data, dict): + symbol = data.get("index_name") + if symbol: + # Format message + msg = { + "date": datetime.fromtimestamp( + data.get("timestamp") / 1000 + ).strftime("%Y-%m-%d %H:%M:%S"), + "symbol": symbol, + "price": data.get("price"), + } + results.append(msg) + + return results + + # Create and return the connection manager + return ConnectionManager( + ws_uri="wss://test.deribit.com/ws/api/v2", + process_message=process_deribit_message, + connection_timeout=60, + retry_delay=3, + max_retries=2, + ) + + +def create_binance_manager(): + """Create a ConnectionManager configured for Binance""" + + # Create the Binance message processor and formatters + def process_binance_message(message_data): + results = [] + + # For kline/candlestick data + if "k" in message_data: + symbol = message_data.get("s", "").lower() + candle = message_data.get("k", {}) + + if symbol and candle: + msg = { + "date": datetime.fromtimestamp( + message_data.get("E", 0) / 1000 + ).strftime("%Y-%m-%d %H:%M:%S"), + "symbol": symbol, + "price": float(candle.get("c", 0)), + } + results.append(msg) + + return results + + def format_binance_subscribe(symbols): + streams = [f"{sym.lower()}@kline_1m" for sym in symbols] + return {"method": "SUBSCRIBE", "params": streams, "id": 1} + + def format_binance_unsubscribe(symbols): + streams = [f"{sym.lower()}@kline_1m" for sym in symbols] + return {"method": "UNSUBSCRIBE", "params": streams, "id": 1} + + # Create and return the connection manager + return ConnectionManager( + ws_uri="wss://stream.binance.com:9443/ws", + process_message=process_binance_message, + format_subscribe=format_binance_subscribe, + format_unsubscribe=format_binance_unsubscribe, + connection_timeout=10.0, + retry_delay=3.0, + max_retries=2, + ) + + +def run_websocket_server(manager_factory_func=create_deribit_manager, port: int = 6940): + """Run a WebSocket server with the specified manager factory function""" + manager = manager_factory_func() + uvicorn.run(manager.app, host="0.0.0.0", port=port) + + +def start_server_subprocess( + manager_factory_name: str = "create_deribit_manager", port: int = 6940 +) -> tuple[subprocess.Popen, str]: + """Start a WebSocket server in a subprocess and return a controller for it + + Args: + manager_factory_name: Name of the manager factory function to use + port: Port number to run the server on + + Returns: + Tuple of (subprocess object, base URL for server) + """ + import time + + # Get the path to this script + script_path = Path(__file__).resolve() + + # Start the server in a subprocess + cmd = [ + sys.executable, + str(script_path), + "--factory", + manager_factory_name, + "--port", + str(port), + ] + + process = subprocess.Popen( + cmd, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + + # Give the server a moment to start + time.sleep(1.0) + + # Return the process and the base URL for communicating with it + base_url = f"http://localhost:{port}" + return process, base_url + + +# When running this file directly +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser(description="Start a WebSocket server") + parser.add_argument( + "--factory", + default="create_deribit_manager", + help="Name of the manager factory function to use", + ) + parser.add_argument( + "--port", type=int, default=6940, help="Port to run the server on" + ) + args = parser.parse_args() + + # Select the factory function based on the argument + factory_map = { + "create_deribit_manager": create_deribit_manager, + "create_binance_manager": create_binance_manager, + } + + selected_factory = factory_map.get(args.factory) + if not selected_factory: + print(f"Unknown factory function: {args.factory}") + sys.exit(1) + + # Run the server + run_websocket_server(selected_factory, args.port) diff --git a/openbb_platform/extensions/websockets/openbb_websockets/database.py b/openbb_platform/extensions/websockets/openbb_websockets/database.py new file mode 100644 index 00000000000..9daa132971d --- /dev/null +++ b/openbb_platform/extensions/websockets/openbb_websockets/database.py @@ -0,0 +1,234 @@ +import asyncio +import contextlib +import logging +from typing import Any, AsyncIterator, Optional + +from sqlalchemy import Index +from sqlalchemy.ext.asyncio import ( + AsyncConnection, + AsyncSession, + async_sessionmaker, + create_async_engine, +) +from sqlalchemy.orm import DeclarativeMeta, Mapped, declarative_base, mapped_column +from sqlalchemy.pool import NullPool +from sqlalchemy.sql import insert, text +from sqlalchemy.types import JSON + +DATABASE_URL = "sqlite+aiosqlite:///" + __file__.replace("database.py", "sqlite.db") + +JSON_Base: DeclarativeMeta = declarative_base( + type_annotation_map={ + dict[str, Any]: JSON, + } +) + +logger = logging.getLogger(__name__) + + +class WebSocketMessages(JSON_Base): + """Table to store messages from the WebSocket server.""" + + __tablename__ = "records" + + id: Mapped[Optional[int]] = mapped_column( + primary_key=True, autoincrement=True, nullable=False + ) + message: Mapped[dict[str, Any]] = mapped_column(type_=JSON, nullable=False) + + __table_args__ = ( + Index("idx_symbol", text("message->>'symbol'")), + Index("idx_date", text("message->>'date'")), + ) + + +class DatabaseSessionManager: + """Manages the database connection and session creation.""" + + def __init__(self, host: str, **engine_kwargs: Any): + """ + Initialize the DatabaseSessionManager. + + Parameters + ---------- + host : str + The database URL. + **engine_kwargs : Any + Additional keyword arguments to pass to the engine. + """ + self._engine = create_async_engine(host, **engine_kwargs) + self._sessionmaker = async_sessionmaker(self._engine, expire_on_commit=False) + + async def close(self): + """Close the database connection.""" + if self._engine is None: + raise Exception("DatabaseSessionManager is not initialized") + await self._engine.dispose() + + self._engine = None + self._sessionmaker = None + + async def create_all(self): + """Create all tables in the database.""" + base = [JSON_Base] + async with self._engine.begin() as conn: + for b in base: + await conn.run_sync(b.metadata.create_all) + + async def set_pragmas(self, connection: AsyncConnection, read_only: bool = False): + """Set the PRAGMA statements for the connection.""" + pragmas = [ + "PRAGMA journal_mode=WAL", + "PRAGMA synchronous=off", + ] + if read_only: + await connection.execute(text("PRAGMA query_only=ON")) + else: + for pragma in pragmas: + await connection.execute(text(pragma)) + + @contextlib.asynccontextmanager + async def connect(self) -> AsyncIterator[AsyncConnection]: + """Connect to the database.""" + if self._engine is None: + raise Exception("DatabaseSessionManager is not initialized") + + async with self._engine.begin() as connection: + await self.set_pragmas(connection) + yield connection + + @contextlib.asynccontextmanager + async def read_session(self) -> AsyncIterator[AsyncSession]: + """Create a read-only session.""" + if self._sessionmaker is None: + raise Exception("DatabaseSessionManager is not initialized") + + async with self._engine.connect() as connection: + await self.set_pragmas(connection, read_only=True) + async with self._sessionmaker(bind=connection) as session: + yield session + + @contextlib.asynccontextmanager + async def write_session(self) -> AsyncIterator[AsyncSession]: + """Create a write session.""" + if self._sessionmaker is None: + raise Exception("DatabaseSessionManager is not initialized") + + async with self._engine.begin() as connection: + await self.set_pragmas(connection) + async with self._sessionmaker(bind=connection) as session: + try: + yield session + await session.flush() # Explicitly flush before the transaction commits + except Exception: + # Let the engine.begin() context manager handle rollback + raise + + +class WebSocketDatabaseSession: + """Manages the database connection and session creation for the WebSocket client.""" + + def __init__( + self, + host: str, + batch_size: int = 100, + batch_interval: int = 1, + echo=False, + poolclass=NullPool, + queue: asyncio.Queue = None, + logger: logging.Logger = None, + **kw: Any, + ): + """ + Initialize the WebSocketDatabaseSession. + + Parameters + ---------- + host : str + The database URL. + batch_size : int, optional + The number of messages to write to the database at once, by default 100. + batch_interval : int, optional + The interval in seconds to write messages to the database, by default 1. + echo : bool, optional + Echo SQL statements to the console, by default False. + poolclass : Pool, optional + The pool class to use for the database connection, by default NullPool. + queue : asyncio.Queue, optional + The queue to store messages, by default None. + logger : logging.Logger, optional + The logger to use for logging, by default None. + **kw : Any + Additional keyword arguments to pass to the DatabaseSessionManager. + """ + self.batch_size = batch_size + self.batch_interval = batch_interval + self.queue = queue or asyncio.Queue() + self.writer_task_running = False + self.prune_task_running = False + + if "sqlite+aiosqlite:///" not in host: + raise ValueError("Only SQLite databases are supported") + + self._session_manager = DatabaseSessionManager( + host, echo=echo, poolclass=poolclass, **kw + ) + self.get_read_session = self._session_manager.read_session + self.get_write_session = self._session_manager.write_session + + async def create_all(self): + """Create all tables in the database.""" + await self._session_manager.create_all() + + async def close(self): + """Close the database connection.""" + await self._session_manager.close() + + async def stop_writer_task(self): + """Stop the writer task.""" + self.writer_task_running = False + + async def start_writer_task(self): + """Start the writer task.""" + self.writer_task_running = True + while True: + messages: list = [] + if self.writer_task_running is False: + break + try: + # Collect messsages in batches + while len(messages) < self.batch_size: + try: + message = await asyncio.wait_for( + self.queue.get(), timeout=self.batch_interval + ) + if message: + messages.append(message) + except asyncio.TimeoutError: + break + except Exception as e: + logger.error(f"Error in write_messages_to_db: {e}") + # Continue to the next iteration of the outer loop + if messages: + try: + async with self._session_manager.write_session() as session: + try: + # Batch insert for more efficiency + await session.execute( + insert(WebSocketMessages), + [{"message": msg} for msg in messages], + ) + # The context manager will handle the commit + except Exception as e: + logger.error(f"Error writing messages to database: {e}") + logger.error(f"Messages: {messages}") + # The context manager will handle rollback + raise + except Exception as e: + logger.error(f"Error in database transaction: {e}") + + self.writer_task_running = False + + async def prune_database(self): + """Prune the database at a set interval.""" + pass diff --git a/openbb_platform/extensions/websockets/openbb_websockets/sqlite.db b/openbb_platform/extensions/websockets/openbb_websockets/sqlite.db new file mode 100644 index 00000000000..c3609679785 Binary files /dev/null and b/openbb_platform/extensions/websockets/openbb_websockets/sqlite.db differ diff --git a/openbb_platform/extensions/websockets/openbb_websockets/widgets.json b/openbb_platform/extensions/websockets/openbb_websockets/widgets.json new file mode 100644 index 00000000000..8137e1cb8f9 --- /dev/null +++ b/openbb_platform/extensions/websockets/openbb_websockets/widgets.json @@ -0,0 +1,96 @@ +{ + "live_grid_example": { + "name": "Live Grid", + "description": "Live Grid", + "type": "live_grid", + "endpoint": "get_ws_data", + "wsEndpoint": "ws", + "data": { + "wsRowIdColumn": "symbol", + "table": { + "showAll": true, + "columnsDefs": [ + { + "field": "date", + "headerName": "Date", + "type": "date" + }, + { + "field": "symbol", + "headerName": "Symbol", + "type": "text" + }, + { + "field": "price", + "headerName": "Price", + "type": "number" + } + ] + } + }, + "params": [ + { + "paramName": "symbol", + "description": "The symbol to get details for", + "value": null, + "label": "Symbol", + "type": "text", + "multiSelect": true, + "options": [ + { + "label": "BTC/USD", + "value": "btc_usd" + }, + { + "label": "ETH/USD", + "value": "eth_usd" + }, + { + "label": "SOL/USD", + "value": "sol_usd" + }, + { + "label": "XRP/USD", + "value": "xrp_usd" + }, + { + "label": "BNB/USD", + "value": "bnb_usd" + }, + { + "label": "MATIC/USD", + "value": "matic_usd" + }, + { + "label": "STETH/USD", + "value": "steth_usd" + }, + { + "label": "USYC/USD", + "value": "usyc_usd" + }, + { + "label": "USDC/USD", + "value": "usdc_usd" + }, + { + "label": "USDT/USD", + "value": "usdt_usd" + }, + { + "label": "USDE/USD", + "value": "usde_usd" + }, + { + "label": "PAXG/USD", + "value": "paxg_usd" + } + ] + } + ], + "gridData": { + "w": 20, + "h": 9 + } + } + } \ No newline at end of file diff --git a/openbb_platform/providers/alpha_vantage/pyproject.toml b/openbb_platform/providers/alpha_vantage/pyproject.toml index e0df0f44fd0..5f456670340 100644 --- a/openbb_platform/providers/alpha_vantage/pyproject.toml +++ b/openbb_platform/providers/alpha_vantage/pyproject.toml @@ -8,7 +8,7 @@ readme = "README.md" packages = [{ include = "openbb_alpha_vantage" }] [tool.poetry.dependencies] -python = "^3.9.21" +python = ">=3.9.21,<3.13" openbb-core = "^1.4.3" [build-system] diff --git a/openbb_platform/providers/benzinga/pyproject.toml b/openbb_platform/providers/benzinga/pyproject.toml index ecfd44a2b9c..14c9e0069f5 100644 --- a/openbb_platform/providers/benzinga/pyproject.toml +++ b/openbb_platform/providers/benzinga/pyproject.toml @@ -8,7 +8,7 @@ readme = "README.md" packages = [{ include = "openbb_benzinga" }] [tool.poetry.dependencies] -python = "^3.9.21" +python = ">=3.9.21,<3.13" openbb-core = "^1.4.3" [build-system] diff --git a/openbb_platform/providers/biztoc/pyproject.toml b/openbb_platform/providers/biztoc/pyproject.toml index 8544cf4ffda..4cbc1df3484 100644 --- a/openbb_platform/providers/biztoc/pyproject.toml +++ b/openbb_platform/providers/biztoc/pyproject.toml @@ -8,7 +8,7 @@ readme = "README.md" packages = [{ include = "openbb_biztoc" }] [tool.poetry.dependencies] -python = "^3.9.21" +python = ">=3.9.21,<3.13" openbb-core = "^1.4.3" [build-system] diff --git a/openbb_platform/providers/bls/pyproject.toml b/openbb_platform/providers/bls/pyproject.toml index 4e27b4b1432..4ec3ae26747 100644 --- a/openbb_platform/providers/bls/pyproject.toml +++ b/openbb_platform/providers/bls/pyproject.toml @@ -8,7 +8,7 @@ readme = "README.md" packages = [{ include = "openbb_bls" }] [tool.poetry.dependencies] -python = "^3.9.21" +python = ">=3.9.21,<3.13" openbb-core = "^1.4.3" [build-system] diff --git a/openbb_platform/providers/cboe/pyproject.toml b/openbb_platform/providers/cboe/pyproject.toml index 499fdb79983..d74724f2751 100644 --- a/openbb_platform/providers/cboe/pyproject.toml +++ b/openbb_platform/providers/cboe/pyproject.toml @@ -8,7 +8,7 @@ readme = "README.md" packages = [{ include = "openbb_cboe" }] [tool.poetry.dependencies] -python = "^3.9.21" +python = ">=3.9.21,<3.13" aiohttp-client-cache = "^0.11.0" aiosqlite = "^0.20.0" openbb-core = "^1.4.3" diff --git a/openbb_platform/providers/cftc/pyproject.toml b/openbb_platform/providers/cftc/pyproject.toml index 859084afe68..86dc83d59e0 100644 --- a/openbb_platform/providers/cftc/pyproject.toml +++ b/openbb_platform/providers/cftc/pyproject.toml @@ -8,7 +8,7 @@ readme = "README.md" packages = [{ include = "openbb_cftc" }] [tool.poetry.dependencies] -python = "^3.9.21" +python = ">=3.9.21,<3.13" openbb-core = "^1.4.3" [build-system] diff --git a/openbb_platform/providers/deribit/pyproject.toml b/openbb_platform/providers/deribit/pyproject.toml index 2a5465a9f1d..3f671c39496 100644 --- a/openbb_platform/providers/deribit/pyproject.toml +++ b/openbb_platform/providers/deribit/pyproject.toml @@ -8,7 +8,7 @@ readme = "README.md" packages = [{ include = "openbb_deribit" }] [tool.poetry.dependencies] -python = "^3.9.21" +python = ">=3.9.21,<3.13" openbb-core = "^1.4.3" async-lru = "^2.0.4" diff --git a/openbb_platform/providers/ecb/pyproject.toml b/openbb_platform/providers/ecb/pyproject.toml index bc913c9d74e..454cc0c3d59 100644 --- a/openbb_platform/providers/ecb/pyproject.toml +++ b/openbb_platform/providers/ecb/pyproject.toml @@ -8,7 +8,7 @@ readme = "README.md" packages = [{ include = "openbb_ecb" }] [tool.poetry.dependencies] -python = "^3.9.21" +python = ">=3.9.21,<3.13" openbb-core = "^1.4.3" xmltodict = "^0.13.0" diff --git a/openbb_platform/providers/econdb/pyproject.toml b/openbb_platform/providers/econdb/pyproject.toml index dce3785cd53..2079a10ddaf 100644 --- a/openbb_platform/providers/econdb/pyproject.toml +++ b/openbb_platform/providers/econdb/pyproject.toml @@ -8,7 +8,7 @@ readme = "README.md" packages = [{ include = "openbb_econdb" }] [tool.poetry.dependencies] -python = "^3.9.21" +python = ">=3.9.21,<3.13" openbb-core = "^1.4.3" aiohttp-client-cache = "^0.11.0" aiosqlite = "^0.20.0" diff --git a/openbb_platform/providers/eia/pyproject.toml b/openbb_platform/providers/eia/pyproject.toml index 295be83a9b2..6eb73e62baa 100644 --- a/openbb_platform/providers/eia/pyproject.toml +++ b/openbb_platform/providers/eia/pyproject.toml @@ -8,7 +8,7 @@ readme = "README.md" packages = [{ include = "openbb_us_eia" }] [tool.poetry.dependencies] -python = "^3.9.21" +python = ">=3.9.21,<3.13" openbb-core = "^1.4.3" async-lru = "^2.0.4" openpyxl = "^3.1.5" diff --git a/openbb_platform/providers/federal_reserve/pyproject.toml b/openbb_platform/providers/federal_reserve/pyproject.toml index e358ee68a79..362be23284e 100644 --- a/openbb_platform/providers/federal_reserve/pyproject.toml +++ b/openbb_platform/providers/federal_reserve/pyproject.toml @@ -8,7 +8,7 @@ readme = "README.md" packages = [{ include = "openbb_federal_reserve" }] [tool.poetry.dependencies] -python = "^3.9.21" +python = ">=3.9.21,<3.13" openbb-core = "^1.4.3" [build-system] diff --git a/openbb_platform/providers/finra/pyproject.toml b/openbb_platform/providers/finra/pyproject.toml index d316be9dd16..ff90fe2926e 100644 --- a/openbb_platform/providers/finra/pyproject.toml +++ b/openbb_platform/providers/finra/pyproject.toml @@ -8,7 +8,7 @@ readme = "README.md" packages = [{ include = "openbb_finra" }] [tool.poetry.dependencies] -python = "^3.9.21" +python = ">=3.9.21,<3.13" openbb-core = "^1.4.3" [build-system] diff --git a/openbb_platform/providers/finviz/pyproject.toml b/openbb_platform/providers/finviz/pyproject.toml index 45d4098ea55..e4c7a763389 100644 --- a/openbb_platform/providers/finviz/pyproject.toml +++ b/openbb_platform/providers/finviz/pyproject.toml @@ -8,7 +8,7 @@ readme = "README.md" packages = [{ include = "openbb_finviz" }] [tool.poetry.dependencies] -python = "^3.9.21" +python = ">=3.9.21,<3.13" openbb-core = "^1.4.3" finvizfinance = "^1.1.0" diff --git a/openbb_platform/providers/fmp/pyproject.toml b/openbb_platform/providers/fmp/pyproject.toml index 2861576abed..c748a3972e8 100644 --- a/openbb_platform/providers/fmp/pyproject.toml +++ b/openbb_platform/providers/fmp/pyproject.toml @@ -8,7 +8,7 @@ readme = "README.md" packages = [{ include = "openbb_fmp" }] [tool.poetry.dependencies] -python = "^3.9.21" +python = ">=3.9.21,<3.13" openbb-core = "^1.4.3" [build-system] diff --git a/openbb_platform/providers/fred/pyproject.toml b/openbb_platform/providers/fred/pyproject.toml index 61b7819adb9..49e79cf243c 100644 --- a/openbb_platform/providers/fred/pyproject.toml +++ b/openbb_platform/providers/fred/pyproject.toml @@ -8,7 +8,7 @@ readme = "README.md" packages = [{ include = "openbb_fred" }] [tool.poetry.dependencies] -python = "^3.9.21" +python = ">=3.9.21,<3.13" openbb-core = "^1.4.3" [build-system] diff --git a/openbb_platform/providers/government_us/pyproject.toml b/openbb_platform/providers/government_us/pyproject.toml index efcd1cc9091..815f0a61187 100644 --- a/openbb_platform/providers/government_us/pyproject.toml +++ b/openbb_platform/providers/government_us/pyproject.toml @@ -8,7 +8,7 @@ readme = "README.md" packages = [{ include = "openbb_government_us" }] [tool.poetry.dependencies] -python = "^3.9.21" +python = ">=3.9.21,<3.13" openbb-core = "^1.4.3" random-user-agent = "^1.0.1" diff --git a/openbb_platform/providers/imf/pyproject.toml b/openbb_platform/providers/imf/pyproject.toml index d06326065c5..e6f98dd42fa 100644 --- a/openbb_platform/providers/imf/pyproject.toml +++ b/openbb_platform/providers/imf/pyproject.toml @@ -8,7 +8,7 @@ readme = "README.md" packages = [{ include = "openbb_imf" }] [tool.poetry.dependencies] -python = "^3.9.21" +python = ">=3.9.21,<3.13" openbb-core = "^1.4.3" [build-system] diff --git a/openbb_platform/providers/intrinio/pyproject.toml b/openbb_platform/providers/intrinio/pyproject.toml index b4906d6f65d..360fb64ad49 100644 --- a/openbb_platform/providers/intrinio/pyproject.toml +++ b/openbb_platform/providers/intrinio/pyproject.toml @@ -8,7 +8,7 @@ readme = "README.md" packages = [{ include = "openbb_intrinio" }] [tool.poetry.dependencies] -python = "^3.9.21" +python = ">=3.9.21,<3.13" requests-cache = "^1.1.0" openbb-core = "^1.4.3" diff --git a/openbb_platform/providers/nasdaq/pyproject.toml b/openbb_platform/providers/nasdaq/pyproject.toml index c206f5e39e4..28e00b62915 100644 --- a/openbb_platform/providers/nasdaq/pyproject.toml +++ b/openbb_platform/providers/nasdaq/pyproject.toml @@ -8,7 +8,7 @@ readme = "README.md" packages = [{ include = "openbb_nasdaq" }] [tool.poetry.dependencies] -python = "^3.9.21" +python = ">=3.9.21,<3.13" openbb-core = "^1.4.3" random-user-agent = "^1.0.1" nasdaq-data-link = "^1.0.4" diff --git a/openbb_platform/providers/oecd/pyproject.toml b/openbb_platform/providers/oecd/pyproject.toml index 73e924febd8..9e5fdc49f9b 100644 --- a/openbb_platform/providers/oecd/pyproject.toml +++ b/openbb_platform/providers/oecd/pyproject.toml @@ -8,7 +8,7 @@ readme = "README.md" packages = [{ include = "openbb_oecd" }] [tool.poetry.dependencies] -python = "^3.9.21" +python = ">=3.9.21,<3.13" openbb-core = "^1.4.3" urllib3 = ">1.26.16" defusedxml = "^0.8.0rc2" diff --git a/openbb_platform/providers/polygon/pyproject.toml b/openbb_platform/providers/polygon/pyproject.toml index 3968e9705f0..c8f59c4f977 100644 --- a/openbb_platform/providers/polygon/pyproject.toml +++ b/openbb_platform/providers/polygon/pyproject.toml @@ -8,7 +8,7 @@ readme = "README.md" packages = [{ include = "openbb_polygon" }] [tool.poetry.dependencies] -python = "^3.9.21" +python = ">=3.9.21,<3.13" openbb-core = "^1.4.3" [build-system] diff --git a/openbb_platform/providers/seeking_alpha/pyproject.toml b/openbb_platform/providers/seeking_alpha/pyproject.toml index 24ebee62d3a..298f1e91c8f 100644 --- a/openbb_platform/providers/seeking_alpha/pyproject.toml +++ b/openbb_platform/providers/seeking_alpha/pyproject.toml @@ -8,7 +8,7 @@ readme = "README.md" packages = [{ include = "openbb_seeking_alpha" }] [tool.poetry.dependencies] -python = "^3.9.21" +python = ">=3.9.21,<3.13" openbb-core = "^1.4.3" [build-system] diff --git a/openbb_platform/providers/stockgrid/pyproject.toml b/openbb_platform/providers/stockgrid/pyproject.toml index 5d7f10fd3d9..d0b04f08b90 100644 --- a/openbb_platform/providers/stockgrid/pyproject.toml +++ b/openbb_platform/providers/stockgrid/pyproject.toml @@ -8,7 +8,7 @@ readme = "README.md" packages = [{ include = "openbb_stockgrid" }] [tool.poetry.dependencies] -python = "^3.9.21" +python = ">=3.9.21,<3.13" openbb-core = "^1.4.3" pytest-freezegun = "^0.4.2" diff --git a/openbb_platform/providers/tiingo/pyproject.toml b/openbb_platform/providers/tiingo/pyproject.toml index 7840fe3d4db..89899f76189 100644 --- a/openbb_platform/providers/tiingo/pyproject.toml +++ b/openbb_platform/providers/tiingo/pyproject.toml @@ -8,7 +8,7 @@ readme = "README.md" packages = [{ include = "openbb_tiingo" }] [tool.poetry.dependencies] -python = "^3.9.21" +python = ">=3.9.21,<3.13" openbb-core = "^1.4.3" [build-system] diff --git a/openbb_platform/providers/tmx/pyproject.toml b/openbb_platform/providers/tmx/pyproject.toml index a40dda67b8c..8fc02e26894 100644 --- a/openbb_platform/providers/tmx/pyproject.toml +++ b/openbb_platform/providers/tmx/pyproject.toml @@ -8,7 +8,7 @@ readme = "README.md" packages = [{ include = "openbb_tmx" }] [tool.poetry.dependencies] -python = "^3.9.21" +python = ">=3.9.21,<3.13" aiohttp-client-cache = "^0.11.0" aiosqlite = "^0.20.0" random-user-agent = "^1.0.1" diff --git a/openbb_platform/providers/tradier/pyproject.toml b/openbb_platform/providers/tradier/pyproject.toml index 3c9904fa245..96caf3a8099 100644 --- a/openbb_platform/providers/tradier/pyproject.toml +++ b/openbb_platform/providers/tradier/pyproject.toml @@ -8,7 +8,7 @@ readme = "README.md" packages = [{ include = "openbb_tradier" }] [tool.poetry.dependencies] -python = "^3.9.21" +python = ">=3.9.21,<3.13" openbb-core = "^1.4.3" [build-system] diff --git a/openbb_platform/providers/tradingeconomics/pyproject.toml b/openbb_platform/providers/tradingeconomics/pyproject.toml index aae1feb78bb..41144de3c1e 100644 --- a/openbb_platform/providers/tradingeconomics/pyproject.toml +++ b/openbb_platform/providers/tradingeconomics/pyproject.toml @@ -8,7 +8,7 @@ readme = "README.md" packages = [{ include = "openbb_tradingeconomics" }] [tool.poetry.dependencies] -python = "^3.9.21" +python = ">=3.9.21,<3.13" openbb-core = "^1.4.3" [build-system] diff --git a/openbb_platform/providers/wsj/pyproject.toml b/openbb_platform/providers/wsj/pyproject.toml index 57187dfada9..1082b807d6f 100644 --- a/openbb_platform/providers/wsj/pyproject.toml +++ b/openbb_platform/providers/wsj/pyproject.toml @@ -8,7 +8,7 @@ readme = "README.md" packages = [{ include = "openbb_wsj" }] [tool.poetry.dependencies] -python = "^3.9.21" +python = ">=3.9.21,<3.13" openbb-core = "^1.4.3" [build-system] diff --git a/openbb_platform/providers/yfinance/pyproject.toml b/openbb_platform/providers/yfinance/pyproject.toml index 1b3e1958e85..0871ed6e80b 100644 --- a/openbb_platform/providers/yfinance/pyproject.toml +++ b/openbb_platform/providers/yfinance/pyproject.toml @@ -8,7 +8,7 @@ readme = "README.md" packages = [{ include = "openbb_yfinance" }] [tool.poetry.dependencies] -python = "^3.9.21" +python = ">=3.9.21,<3.13" yfinance = "^0.2.55" openbb-core = "^1.4.3"