jerry/OpenBB
1
0
Fork 0
mirror of https://github.com/OpenBB-finance/OpenBB.git synced 2026-10-27 12:40:35 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix some imports and make default message queue 10000

Browse Source
This commit is contained in:
Danglewood
2024-12-17 16:36:08 -08:00
parent 14a96cc3ed
commit f7faf07ea8
9 changed files with 77 additions and 119 deletions
Download Patch File Download Diff File Expand all files Collapse all files

19
openbb_platform/core/openbb_core/provider/utils/websockets/database.py
View File

@@ -41,8 +41,6 @@ class Database:
A Pydantic model to validate the JSON data. Default is None.
limit: Optional[int]
The maximum number of records to keep in the database. Default is None.
keep_results: Optional[bool]
Whether to persist the file on disk after the Python session is over. Default is False.
logger : Optional[logging.Logger]
A custom logger to use. If not provided, a new logger will be created.
loop: Optional[asyncio.AbstractEventLoop]
@@ -75,8 +73,7 @@ class Database:
):
"""Initialize the ResultsDB class."""
# pylint: disable=import-outside-toplevel
import atexit # noqa
import tempfile
import tempfile # noqa
from pathlib import Path
from openbb_core.provider.utils.websockets.helpers import get_logger
@@ -94,11 +91,9 @@ class Database:
self.results_path = Path(results_file).absolute()
self.results_file = results_file
self.keep_results = keep_results
self.table_name = table_name if table_name else "records"
self.limit = limit
self.loop = loop
atexit.register(self._atexit)
self.kwargs = kwargs if kwargs else {}
run_async(self._setup_database)
self.data_model = data_model
@@ -136,7 +131,6 @@ class Database:
f"""
CREATE TABLE IF NOT EXISTS {self.table_name} (
id INTEGER PRIMARY KEY AUTOINCREMENT,
symbol TEXT,
message TEXT NOT NULL
)
"""
@@ -150,17 +144,6 @@ class Database:
self._exception = e
raise OpenBBError(msg) from e
def _atexit(self) -> None:
"""Clean up the WebSocket client processes at exit."""
# pylint: disable=import-outside-toplevel
import os
self._exception = None
if self.keep_results:
self.logger.info("Results database saved to, %s\n", str(self.results_file))
if os.path.exists(self.results_file) and not self.keep_results: # type: ignore
os.remove(self.results_file) # type: ignore
async def _write_to_db(self, message) -> None:
"""Write the WebSocket message to the SQLite database."""
# pylint: disable=import-outside-toplevel

4
openbb_platform/core/openbb_core/provider/utils/websockets/message_queue.py
View File

@@ -9,7 +9,7 @@ class MessageQueue:
def __init__(
self,
max_size: int = 5000,
max_size: int = 10000,
max_retries=5,
backoff_factor=0.5,
logger: Optional[logging.Logger] = None,
@@ -52,7 +52,7 @@ class MessageQueue:
async def process_queue(self, handler):
"""Process the message queue."""
while not self.queue.empty():
while True:
message = await self.queue.get()
await self._process_message(message, handler)
self.queue.task_done()

11
openbb_platform/extensions/websockets/openbb_websockets/broadcast.py
View File

@@ -8,10 +8,10 @@ from typing import Optional
import uvicorn
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from openbb_core.provider.utils.websockets.database import Database
from openbb_core.provider.utils.websockets.helpers import get_logger, parse_kwargs
from starlette.websockets import WebSocketState
from openbb_websockets.helpers import get_logger, parse_kwargs, setup_database
connected_clients: set = set()
kwargs = parse_kwargs()
@@ -25,6 +25,7 @@ TABLE_NAME = kwargs.pop("table_name", None) or "records"
SLEEP_TIME = kwargs.pop("sleep_time", None) or 0.25
AUTH_TOKEN = kwargs.pop("auth_token", None)
DATABASE = Database(results_file=RESULTS_FILE, table_name=TABLE_NAME)
app = FastAPI()
@@ -151,14 +152,14 @@ class BroadcastServer: # pylint: disable=too-many-instance-attributes
def _encrypt_value(self, value: str) -> str:
"""Encrypt the value for storage."""
# pylint: disable=import-outside-toplevel
from openbb_websockets.helpers import encrypt_value
from openbb_core.provider.utils.websockets.helpers import encrypt_value
return encrypt_value(self._key, self._iv, value)
def _decrypt_value(self, value: str) -> str:
"""Decrypt the value for use."""
# pylint: disable=import-outside-toplevel
from openbb_websockets.helpers import decrypt_value
from openbb_core.provider.utils.websockets.helpers import decrypt_value
return decrypt_value(self._key, self._iv, value)
@@ -294,5 +295,5 @@ if __name__ == "__main__":
# pylint: disable=import-outside-toplevel
from openbb_core.provider.utils.helpers import run_async
run_async(setup_database, RESULTS_FILE, TABLE_NAME)
run_async(DATABASE._setup_database)
main()

84
openbb_platform/extensions/websockets/openbb_websockets/client.py
View File

@@ -99,12 +99,14 @@ class WebSocketClient: # pylint: disable=too-many-instance-attributes
import os
import tempfile
import threading
from aiosqlite import DatabaseError
from queue import Queue
from pathlib import Path
from openbb_core.app.model.abstract.error import OpenBBError
from openbb_websockets.helpers import get_logger
from openbb_websockets.helpers import encrypt_value
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
@@ -156,8 +158,15 @@ class WebSocketClient: # pylint: disable=too-many-instance-attributes
atexit.register(self._atexit)
try:
self._setup_database()
except DatabaseError as e:
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 = f"Unexpected error setting up the SQLite database and table -> {e.__class___.__name__}: {e}"
self.logger.error(msg)
self._exception = OpenBBError(msg)
@@ -178,14 +187,6 @@ class WebSocketClient: # pylint: disable=too-many-instance-attributes
if os.path.exists(self.results_file) and not self.save_results: # type: ignore
os.remove(self.results_file) # type: ignore
def _setup_database(self) -> None:
"""Set up the SQLite database and table."""
# pylint: disable=import-outside-toplevel
from openbb_core.provider.utils.helpers import run_async # noqa
from openbb_websockets.helpers import setup_database
run_async(setup_database, self.results_path, self.table_name)
def _log_provider_output(self, output_queue) -> None:
"""Log output from the provider logger, handling exceptions, errors, and messages that are not data."""
# pylint: disable=import-outside-toplevel
@@ -193,7 +194,7 @@ class WebSocketClient: # pylint: disable=too-many-instance-attributes
import queue
import sys
from openbb_core.provider.utils.errors import UnauthorizedError
from openbb_websockets.helpers import clean_message
from openbb_core.provider.utils.websockets.helpers import clean_message
from pydantic import ValidationError
while not self._stop_log_thread_event.is_set():
@@ -274,7 +275,7 @@ class WebSocketClient: # pylint: disable=too-many-instance-attributes
# pylint: disable=import-outside-toplevel
import queue # noqa
import sys
from openbb_websockets.helpers import clean_message
from openbb_core.provider.utils.websockets.helpers import clean_message
while not self._stop_broadcasting_event.is_set():
try:
@@ -321,7 +322,7 @@ class WebSocketClient: # pylint: disable=too-many-instance-attributes
import threading
import time
from openbb_core.app.model.abstract.error import OpenBBError
from openbb_websockets.helpers import decrypt_value
from openbb_core.provider.utils.websockets.helpers import decrypt_value
if self.is_running:
self.logger.info("Provider connection already running.")
@@ -495,59 +496,22 @@ class WebSocketClient: # pylint: disable=too-many-instance-attributes
return False
@property
def results(self) -> Union[list[dict], list["Data"], None]:
def results(self) -> list:
"""Retrieve the deserialized results from the results file."""
# pylint: disable=import-outside-toplevel
import json # noqa
import sqlite3
from openbb_core.app.model.abstract.error import OpenBBError
output: list = []
file_path = self.results_path
if file_path.exists():
with sqlite3.connect(file_path) as conn:
try:
cursor = conn.execute(f"SELECT * FROM {self.table_name}") # noqa
for row in cursor:
_, message = row
if self.data_model:
message = json.loads(message)
if isinstance(message, (str, bytes)):
output.append(
self.data_model.model_validate_json(message)
)
elif isinstance(message, dict):
output.append(self.data_model(**message))
else:
output.append(json.loads(message))
except Exception as e:
raise OpenBBError(f"Error retrieving results: {e}") from e
if output:
return output
self.logger.info("No results found in %s", self.results_file)
return None
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
import sqlite3 # noqa
try:
with sqlite3.connect(self.results_path) as conn:
conn.execute("PRAGMA journal_mode=WAL;")
conn.execute(f"DELETE FROM {self.table_name}") # noqa
conn.commit()
self._setup_database()
self.logger.info(
"Results cleared from table %s in %s",
self.table_name,
self.results_file,
)
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)
@@ -608,7 +572,7 @@ class WebSocketClient: # pylint: disable=too-many-instance-attributes
import psutil
import queue
from openbb_platform_api.utils.api import check_port
from openbb_websockets.helpers import decrypt_value
from openbb_core.provider.utils.websockets.helpers import decrypt_value
if (
self._broadcast_process is not None

5
openbb_platform/extensions/websockets/openbb_websockets/listen.py
View File

@@ -25,7 +25,10 @@ class Listener:
import websockets
from openbb_core.app.model.abstract.error import OpenBBError
from openbb_core.provider.utils.errors import UnauthorizedError
from openbb_websockets.helpers import clean_message, get_logger
from openbb_core.provider.utils.websockets.helpers import (
clean_message,
get_logger,
)
from websockets.exceptions import InvalidStatusCode
kwargs = kwargs or {}

18
openbb_platform/providers/fmp/openbb_fmp/utils/websocket_client.py
View File

@@ -2,18 +2,17 @@
import asyncio
import json
import os
import signal
import sys
import websockets
import websockets.exceptions
from openbb_core.provider.utils.websockets.database import Database
from openbb_core.provider.utils.websockets.helpers import (
get_logger,
handle_termination_signal,
handle_validation_error,
parse_kwargs,
write_to_db,
)
from openbb_core.provider.utils.websockets.message_queue import MessageQueue
from openbb_fmp.models.websocket_connection import FmpWebSocketData
@@ -24,7 +23,13 @@ kwargs = parse_kwargs()
queue = MessageQueue()
command_queue = MessageQueue()
CONNECT_KWARGS = kwargs.pop("connect_kwargs", {})
kwargs["results_file"] = os.path.abspath(kwargs["results_file"])
DATABASE = Database(
results_file=kwargs["results_file"],
table_name=kwargs["table_name"],
limit=kwargs.get("limit"),
logger=logger,
)
async def login(websocket):
@@ -125,12 +130,7 @@ async def process_message(message):
except ValidationError:
raise e from e
if result:
await write_to_db(
result,
kwargs["results_file"],
kwargs["table_name"],
kwargs.get("limit"),
)
await DATABASE._write_to_db(result) # pylint: disable=protected-access
async def connect_and_stream():

23
openbb_platform/providers/intrinio/openbb_intrinio/utils/websocket_client.py
View File

@@ -9,12 +9,12 @@ import sys
from typing import Any
from openbb_core.provider.utils.helpers import run_async
from openbb_core.provider.utils.websockets.database import Database
from openbb_core.provider.utils.websockets.helpers import (
get_logger,
handle_termination_signal,
handle_validation_error,
parse_kwargs,
write_to_db,
)
from openbb_core.provider.utils.websockets.message_queue import MessageQueue
from openbb_intrinio.models.websocket_connection import IntrinioWebSocketData
@@ -26,8 +26,15 @@ kwargs = parse_kwargs()
command_queue = MessageQueue(logger=logger)
CONNECT_KWARGS = kwargs.pop("connect_kwargs", {})
DATABASE = Database(
results_file=kwargs["results_file"],
table_name=kwargs["table_name"],
limit=kwargs.get("limit"),
logger=logger,
)
async def process_message(message):
def process_message(message):
"""Process the message and write to the database."""
result: Any = None
try:
@@ -43,17 +50,12 @@ async def process_message(message):
except ValidationError:
raise e from e
if result:
await write_to_db(
result,
kwargs["results_file"],
kwargs["table_name"],
kwargs["limit"],
)
DATABASE.write_to_db(result)
def on_message(message, backlog):
"""Process the message and write to the database."""
run_async(process_message, message)
process_message(message)
options = {
@@ -73,6 +75,7 @@ client = IntrinioRealtimeClient(
async def subscribe(symbol, event):
"""Subscribe or unsubscribe to a symbol."""
ticker = symbol.split(",") if isinstance(symbol, str) else symbol
ticker = ["lobby"] if "*" in ticker or "LOBBY" in ticker else ticker
try:
if event == "subscribe":
client.join(ticker)
@@ -123,7 +126,7 @@ if __name__ == "__main__":
try:
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
loop.set_exception_handler(lambda loop, context: None)
# loop.set_exception_handler(lambda loop, context: None)
for sig in (signal.SIGINT, signal.SIGTERM):
loop.add_signal_handler(sig, handle_termination_signal, logger)

16
openbb_platform/providers/polygon/openbb_polygon/utils/websocket_client.py
View File

@@ -8,12 +8,12 @@ import sys
import websockets
import websockets.exceptions
from openbb_core.provider.utils.websockets.database import Database
from openbb_core.provider.utils.websockets.helpers import (
get_logger,
handle_termination_signal,
handle_validation_error,
parse_kwargs,
write_to_db,
)
from openbb_core.provider.utils.websockets.message_queue import MessageQueue
from openbb_polygon.models.websocket_connection import (
@@ -32,6 +32,13 @@ FEED = kwargs.pop("feed", None)
ASSET_TYPE = kwargs.pop("asset_type", None)
kwargs["results_file"] = os.path.abspath(kwargs.get("results_file"))
DATABASE = Database(
results_file=kwargs["results_file"],
table_name=kwargs["table_name"],
limit=kwargs.get("limit"),
logger=logger,
)
async def handle_symbol(symbol):
"""Handle the symbol and map it to the correct format."""
@@ -181,12 +188,7 @@ async def process_message(message):
raise e from e
if result:
await write_to_db(
result,
kwargs["results_file"],
kwargs["table_name"],
kwargs.get("limit"),
)
await DATABASE._write_to_db(result)
else:
logger.info("PROVIDER INFO: %s", msg)

16
openbb_platform/providers/tiingo/openbb_tiingo/utils/websocket_client.py
View File

@@ -8,12 +8,12 @@ import sys
import websockets
from openbb_core.provider.utils.errors import UnauthorizedError
from openbb_core.provider.utils.websockets.database import Database
from openbb_core.provider.utils.websockets.helpers import (
get_logger,
handle_termination_signal,
handle_validation_error,
parse_kwargs,
write_to_db,
)
from openbb_core.provider.utils.websockets.message_queue import MessageQueue
from openbb_tiingo.models.websocket_connection import TiingoWebSocketData
@@ -75,6 +75,13 @@ kwargs = parse_kwargs()
CONNECT_KWARGS = kwargs.pop("connect_kwargs", {})
kwargs["results_file"] = os.path.abspath(kwargs["results_file"])
DATABASE = Database(
results_file=kwargs["results_file"],
table_name=kwargs["table_name"],
limit=kwargs.get("limit"),
logger=logger,
)
# Subscribe and unsubscribe events are handled in a separate connection using the subscription_id set by the login event.
async def update_symbols(symbol, event):
@@ -189,12 +196,7 @@ async def process_message(message): # pylint: disable=too-many-branches
raise e from e
if result:
await write_to_db(
result,
kwargs["results_file"],
kwargs["table_name"],
kwargs.get("limit"),
)
await DATABASE._write_to_db(result) # pylint: disable=protected-access
return