jerry/OpenBB
1
0
Fork 0
mirror of https://github.com/OpenBB-finance/OpenBB.git synced 2026-10-28 04:52:07 +00:00
Code Issues Packages Projects Releases Wiki Activity

bracket all the extensions

Browse Source
This commit is contained in:
Danglewood
2025-03-31 16:50:51 -07:00
parent ed63e172bd
commit edaccee104
46 changed files with 2736 additions and 39 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
openbb_platform/core/pyproject.toml vendored
View File

@@ -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"

2
openbb_platform/extensions/commodity/pyproject.toml vendored
View File

@@ -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]

2
openbb_platform/extensions/crypto/pyproject.toml vendored
View File

@@ -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]

2
openbb_platform/extensions/currency/pyproject.toml vendored
View File

@@ -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]

2
openbb_platform/extensions/derivatives/pyproject.toml vendored
View File

@@ -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]

2
openbb_platform/extensions/economy/pyproject.toml vendored
View File

@@ -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]

2
openbb_platform/extensions/equity/pyproject.toml vendored
View File

@@ -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]

2
openbb_platform/extensions/fixedincome/pyproject.toml vendored
View File

@@ -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]

2
openbb_platform/extensions/index/pyproject.toml vendored
View File

@@ -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]

2
openbb_platform/extensions/news/pyproject.toml vendored
View File

@@ -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]

775
openbb_platform/extensions/websockets/openbb_websockets/_client.py Normal file
View File

@@ -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()

457
openbb_platform/extensions/websockets/openbb_websockets/app copy 2.py Normal file
View File

@@ -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)

260
openbb_platform/extensions/websockets/openbb_websockets/app copy.py Normal file
View File

@@ -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)

875
openbb_platform/extensions/websockets/openbb_websockets/app.py Normal file
View File

@@ -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)

234
openbb_platform/extensions/websockets/openbb_websockets/database.py Normal file
View File

@@ -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

BIN
openbb_platform/extensions/websockets/openbb_websockets/sqlite.db vendored Normal file
View File

Binary file not shown.

96
openbb_platform/extensions/websockets/openbb_websockets/widgets.json vendored Normal file
View File

@@ -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
}
}
}

2
openbb_platform/providers/alpha_vantage/pyproject.toml vendored
View File

@@ -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]

2
openbb_platform/providers/benzinga/pyproject.toml vendored
View File

@@ -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]

2
openbb_platform/providers/biztoc/pyproject.toml vendored
View File

@@ -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]

2
openbb_platform/providers/bls/pyproject.toml vendored
View File

@@ -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]

2
openbb_platform/providers/cboe/pyproject.toml vendored
View File

@@ -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"

2
openbb_platform/providers/cftc/pyproject.toml vendored
View File

@@ -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]

2
openbb_platform/providers/deribit/pyproject.toml vendored
View File

@@ -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"

2
openbb_platform/providers/ecb/pyproject.toml vendored
View File

@@ -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"

2
openbb_platform/providers/econdb/pyproject.toml vendored
View File

@@ -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"

2
openbb_platform/providers/eia/pyproject.toml vendored
View File

@@ -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"

2
openbb_platform/providers/federal_reserve/pyproject.toml vendored
View File

@@ -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]

2
openbb_platform/providers/finra/pyproject.toml vendored
View File

@@ -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]

2
openbb_platform/providers/finviz/pyproject.toml vendored
View File

@@ -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"

2
openbb_platform/providers/fmp/pyproject.toml vendored
View File

@@ -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]

2
openbb_platform/providers/fred/pyproject.toml vendored
View File

@@ -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]

2
openbb_platform/providers/government_us/pyproject.toml vendored
View File

@@ -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"

2
openbb_platform/providers/imf/pyproject.toml vendored
View File

@@ -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]

2
openbb_platform/providers/intrinio/pyproject.toml vendored
View File

@@ -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"

2
openbb_platform/providers/nasdaq/pyproject.toml vendored
View File

@@ -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"

2
openbb_platform/providers/oecd/pyproject.toml vendored
View File

@@ -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"

2
openbb_platform/providers/polygon/pyproject.toml vendored
View File

@@ -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]

2
openbb_platform/providers/seeking_alpha/pyproject.toml vendored
View File

@@ -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]

2
openbb_platform/providers/stockgrid/pyproject.toml vendored
View File

@@ -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"

2
openbb_platform/providers/tiingo/pyproject.toml vendored
View File

@@ -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]

2
openbb_platform/providers/tmx/pyproject.toml vendored
View File

@@ -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"

2
openbb_platform/providers/tradier/pyproject.toml vendored
View File

@@ -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]

2
openbb_platform/providers/tradingeconomics/pyproject.toml vendored
View File

@@ -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]

2
openbb_platform/providers/wsj/pyproject.toml vendored
View File

@@ -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]

2
openbb_platform/providers/yfinance/pyproject.toml vendored
View File

@@ -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"