From dcfa890608139a7fd64fa422801a23ef0046efd6 Mon Sep 17 00:00:00 2001 From: Danglewood <85772166+deeleeramone@users.noreply.github.com> Date: Mon, 16 Dec 2024 13:57:57 -0800 Subject: [PATCH] move some files around and create unit tests for MessageQueue --- .../provider/utils/websockets/__init__.py | 1 + .../provider/utils/websockets/helpers.py | 278 ++++++++++++++++++ .../utils/websockets/message_queue.py | 62 ++++ .../utils/websockets/test_message_queue.py | 69 +++++ .../fmp/openbb_fmp/utils/websocket_client.py | 6 +- .../openbb_intrinio/utils/websocket_client.py | 8 +- .../openbb_polygon/utils/websocket_client.py | 12 +- .../openbb_tiingo/utils/websocket_client.py | 6 +- 8 files changed, 426 insertions(+), 16 deletions(-) create mode 100644 openbb_platform/core/openbb_core/provider/utils/websockets/__init__.py create mode 100644 openbb_platform/core/openbb_core/provider/utils/websockets/helpers.py create mode 100644 openbb_platform/core/openbb_core/provider/utils/websockets/message_queue.py create mode 100644 openbb_platform/core/tests/provider/utils/websockets/test_message_queue.py diff --git a/openbb_platform/core/openbb_core/provider/utils/websockets/__init__.py b/openbb_platform/core/openbb_core/provider/utils/websockets/__init__.py new file mode 100644 index 00000000000..41e2a62502a --- /dev/null +++ b/openbb_platform/core/openbb_core/provider/utils/websockets/__init__.py @@ -0,0 +1 @@ +"""WebSocket Utilities.""" diff --git a/openbb_platform/core/openbb_core/provider/utils/websockets/helpers.py b/openbb_platform/core/openbb_core/provider/utils/websockets/helpers.py new file mode 100644 index 00000000000..66a74e3f0b1 --- /dev/null +++ b/openbb_platform/core/openbb_core/provider/utils/websockets/helpers.py @@ -0,0 +1,278 @@ +"""WebSocket Helper Functions.""" + +# pylint: disable=protected-access + +import logging +import re +from typing import Any, Optional + +from openbb_core.app.model.abstract.error import OpenBBError +from openbb_core.provider.utils.errors import UnauthorizedError +from pydantic import ValidationError + +AUTH_TOKEN_FILTER = re.compile( + r"(auth_token=)([^&]*)", + re.IGNORECASE | re.MULTILINE, +) + +connected_clients: dict = {} + + +def clean_message(message: str) -> str: + """Clean the message.""" + return AUTH_TOKEN_FILTER.sub(r"\1********", message) + + +def get_logger(name, level=logging.INFO): + """Get a logger instance.""" + # pylint: disable=import-outside-toplevel + import uuid + + logger = logging.getLogger(f"{name}-{uuid.uuid4()}") + handler = logging.StreamHandler() + handler.setLevel(level) + formatter = logging.Formatter("%(message)s\n") + handler.setFormatter(formatter) + logger.addHandler(handler) + logger.setLevel(level) + + return logger + + +def handle_validation_error(logger: logging.Logger, error: ValidationError): + """Log and raise a Pydantic ValidationError from a provider connection.""" + err = f"{error.__class__.__name__} -> {error.title}: {error.json()}" + logger.error(err) + raise error from error + + +async def get_status(name: Optional[str] = None, client: Optional[Any] = None) -> dict: + """Get the status of a client.""" + if name and name not in connected_clients: + raise OpenBBError(f"Client {name} not connected.") + if not name and not client: + raise OpenBBError("Either name or client must be provided.") + client = client if client else connected_clients[name] + provider_pid = client._psutil_process.pid if client.is_running else None + broadcast_pid = ( + client._psutil_broadcast_process.pid if client.is_broadcasting else None + ) + status = { + "name": client.name, + "auth_required": client._auth_token is not None, + "subscribed_symbols": client.symbol, + "is_running": client.is_running, + "provider_pid": provider_pid, + "is_broadcasting": client.is_broadcasting, + "broadcast_address": client.broadcast_address, + "broadcast_pid": broadcast_pid, + "results_file": client.results_file, + "table_name": client.table_name, + "save_results": client.save_results, + } + return status + + +def encrypt_value(key, iv, value): + """Encrypt a value before storing.""" + # pylint: disable=import-outside-toplevel + import base64 # noqa + from cryptography.hazmat.backends import default_backend + from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes + + backend = default_backend() + cipher = Cipher(algorithms.AES(key), modes.CFB(iv), backend=backend) + encryptor = cipher.encryptor() + encrypted_value = encryptor.update(value.encode()) + encryptor.finalize() + return base64.b64encode(encrypted_value).decode() + + +def decrypt_value(key, iv, encrypted_value): + """Decrypt the value for use.""" + # pylint: disable=import-outside-toplevel + import base64 # noqa + from cryptography.hazmat.backends import default_backend + from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes + + backend = default_backend() + cipher = Cipher(algorithms.AES(key), modes.CFB(iv), backend=backend) + decryptor = cipher.decryptor() + decrypted_value = ( + decryptor.update(base64.b64decode(encrypted_value)) + decryptor.finalize() + ) + return decrypted_value.decode() + + +async def check_auth(name: str, auth_token: Optional[str] = None) -> bool: + """Check the auth token.""" + if name not in connected_clients: + raise OpenBBError(f"Client {name} not connected.") + client = connected_clients[name] + if client._auth_token is None: + return True + if auth_token is None: + raise UnauthorizedError(f"Client authorization token is required for {name}.") + if auth_token != client._decrypt_value(client._auth_token): + raise UnauthorizedError(f"Invalid client authorization token for {name}.") + return True + + +def handle_termination_signal(logger): + """Handle termination signals to ensure graceful shutdown.""" + # pylint: disable=import-outside-toplevel + import sys + + logger.info( + "PROVIDER INFO: Termination signal received. WebSocket connection closed." + ) + sys.exit(0) + + +def parse_kwargs(): + """Parse command line keyword arguments.""" + # pylint: disable=import-outside-toplevel + import json + import sys + + args = sys.argv[1:].copy() + _kwargs: dict = {} + for i, arg in enumerate(args): + if arg.startswith("url"): + _kwargs["url"] = arg[4:] + continue + if "=" in arg: + key, value = arg.split("=") + + if key == "connect_kwargs": + value = {} if value == "None" else json.loads(value) + + _kwargs[key] = value + elif arg.startswith("--"): + key = arg[2:] + + if i + 1 < len(args) and not args[i + 1].startswith("--"): + value = args[i + 1] + + if isinstance(value, str) and value.lower() in ["false", "true"]: + value = value.lower() == "true" # type: ignore + elif isinstance(value, str) and value.lower() == "none": + value = None + _kwargs[key] = value + else: + _kwargs[key] = True + + return _kwargs + + +async def setup_database(results_path, table_name): + """Create the SQLite database, if required.""" + # pylint: disable=import-outside-toplevel + import os # noqa + import aiosqlite + + if os.path.exists(results_path): + async with aiosqlite.connect(results_path) as conn: + try: + await conn.execute("SELECT name FROM sqlite_master WHERE type='table';") + except aiosqlite.DatabaseError as e: + raise OpenBBError( + "Unexpected error caused by an invalid SQLite database file." + "Please check the path, and inspect the file if it exists." + + f" -> {e}" + ) from e + + async with aiosqlite.connect(results_path) as conn: + await conn.execute( + f""" + CREATE TABLE IF NOT EXISTS {table_name} ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + message JSON + ) + """ + ) + await conn.commit() + + +async def write_to_db(message, results_path, table_name, limit): + """Write the WebSocket message to the SQLite database.""" + # pylint: disable=import-outside-toplevel + import aiosqlite + + conn = await aiosqlite.connect(results_path) + await conn.execute("PRAGMA journal_mode=WAL;") + try: + # Check if the table exists and create it if it doesn't + await conn.execute( + f""" + CREATE TABLE IF NOT EXISTS {table_name} ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + message JSON + ) + """ + ) + await conn.commit() + + await conn.execute( + f"INSERT INTO {table_name} (message) VALUES (?)", # noqa + (message,), + ) + await conn.commit() + + records = await conn.execute(f"SELECT COUNT(*) FROM {table_name}") # noqa + count = (await records.fetchone())[0] + count = await conn.execute(f"SELECT COUNT(*) FROM {table_name}") # noqa + current_count = int((await count.fetchone())[0]) + limit = 0 if limit is None else int(limit) + + if current_count > limit and limit != 0: + await conn.execute( + f""" + DELETE FROM {table_name} + WHERE id IN ( + SELECT id FROM {table_name} + ORDER BY id DESC + LIMIT -1 OFFSET ? + ) + """, # noqa: S608 + (limit,), + ) + + await conn.commit() + except Exception as e: # pylint: disable=broad-except + raise OpenBBError( + f"Unexpected error encountered while inserting message into the database. -> {e.__class__.__name__}: {e}" + ) from e + + finally: + await conn.close() + + +class StdOutSink: + """Filter stdout for PII.""" + + def write(self, message): + """Write to stdout.""" + # pylint: disable=import-outside-toplevel + import sys + + cleaned_message = AUTH_TOKEN_FILTER.sub(r"\1********", message) + if cleaned_message != message: + cleaned_message = f"{cleaned_message}\n" + sys.__stdout__.write(cleaned_message) # type: ignore + + def flush(self): + """Flush stdout.""" + # pylint: disable=import-outside-toplevel + import sys + + sys.__stdout__.flush() # type: ignore + + +class AuthTokenFilter(logging.Formatter): + """Custom logging formatter to filter auth tokens.""" + + def format(self, record): + """Format the log record.""" + original_message = super().format(record) + cleaned_message = AUTH_TOKEN_FILTER.sub(r"\1********", original_message) + return cleaned_message diff --git a/openbb_platform/core/openbb_core/provider/utils/websockets/message_queue.py b/openbb_platform/core/openbb_core/provider/utils/websockets/message_queue.py new file mode 100644 index 00000000000..f626af8b51b --- /dev/null +++ b/openbb_platform/core/openbb_core/provider/utils/websockets/message_queue.py @@ -0,0 +1,62 @@ +"""Async WebSocket Message Queue.""" + +import logging +from typing import Optional + + +class MessageQueue: + """Async message queue for the WebSocket connection.""" + + def __init__( + self, + max_size: int = 5000, + max_retries=5, + backoff_factor=0.5, + logger: Optional[logging.Logger] = None, + ): + """Initialize the MessageQueue.""" + # pylint: disable=import-outside-toplevel + from asyncio import Queue # noqa + from openbb_core.provider.utils.websockets.helpers import get_logger + + self.queue: Queue = Queue(maxsize=max_size) + self.max_retries = max_retries + self.backoff_factor = backoff_factor + self.logger = ( + logger + if logger + else get_logger("openbb.websocket.queue", level=logging.WARN) + ) + + async def dequeue(self): + """Dequeue a message.""" + return await self.queue.get() + + async def enqueue(self, message): + """Enqueue a message.""" + # pylint: disable=import-outside-toplevel + from asyncio import sleep + + retries = 0 + while retries < self.max_retries: + if self.queue.full(): + retries += 1 + msg = f"Queue is full. Retrying {retries}/{self.max_retries}..." + self.logger.warning(msg) + await sleep(self.backoff_factor * retries) + else: + await self.queue.put(message) + return + + self.logger.warn("Failed to enqueue message after maximum retries.") + + async def process_queue(self, handler): + """Process the message queue.""" + while not self.queue.empty(): + message = await self.queue.get() + await self._process_message(message, handler) + self.queue.task_done() + + async def _process_message(self, message, handler): + """Process the message with the handler coroutine.""" + await handler(message) diff --git a/openbb_platform/core/tests/provider/utils/websockets/test_message_queue.py b/openbb_platform/core/tests/provider/utils/websockets/test_message_queue.py new file mode 100644 index 00000000000..ba0c6102cbf --- /dev/null +++ b/openbb_platform/core/tests/provider/utils/websockets/test_message_queue.py @@ -0,0 +1,69 @@ +"""Unit Tests For MessageQueue Class.""" + +import asyncio + +import pytest +from openbb_core.provider.utils.websockets.message_queue import MessageQueue + +MOCK_MESSAGES = [ + {"message": "test1"}, + {"message": "test2"}, + {"message": "test3"}, +] + + +@pytest.fixture +def message_queue(): + """Return a MessageQueue instance.""" + return MessageQueue(max_size=2, max_retries=2, backoff_factor=0.1) + + +@pytest.mark.asyncio +async def test_enqueue_dequeue(message_queue): + """Test the enqueue and dequeue methods.""" + await message_queue.enqueue(MOCK_MESSAGES[0]) + assert not message_queue.queue.empty() + await message_queue.enqueue(MOCK_MESSAGES[1]) + assert message_queue.queue.qsize() == 2 + assert await message_queue.dequeue() == MOCK_MESSAGES[0] + assert message_queue.queue.qsize() == 1 + assert await message_queue.dequeue() == MOCK_MESSAGES[1] + assert message_queue.queue.empty() + + +@pytest.mark.asyncio +async def test_enqueue_full(message_queue): + """Test the enqueue method when the queue is full.""" + await message_queue.enqueue(MOCK_MESSAGES[0]) + await message_queue.enqueue(MOCK_MESSAGES[1]) + + assert message_queue.queue.full() + + with pytest.warns(Warning): + await message_queue.enqueue(MOCK_MESSAGES[2]) + assert message_queue.queue.qsize() == 2 + + assert await message_queue.dequeue() == MOCK_MESSAGES[0] + assert message_queue.queue.qsize() == 1 + await message_queue.enqueue(MOCK_MESSAGES[2]) + assert await message_queue.dequeue() == MOCK_MESSAGES[1] + assert await message_queue.dequeue() == MOCK_MESSAGES[2] + assert message_queue.queue.empty() + + +@pytest.mark.asyncio +async def test_process_queue(message_queue): + """Test the process_queue method.""" + + NUM_MESSAGES = 0 + + async def handler(message): + """Test handler.""" + nonlocal NUM_MESSAGES + NUM_MESSAGES += 1 + + await message_queue.enqueue(MOCK_MESSAGES[0]) + await message_queue.enqueue(MOCK_MESSAGES[1]) + asyncio.create_task(message_queue.process_queue(handler)) + await asyncio.sleep(0.1) + assert NUM_MESSAGES == 2 diff --git a/openbb_platform/providers/fmp/openbb_fmp/utils/websocket_client.py b/openbb_platform/providers/fmp/openbb_fmp/utils/websocket_client.py index 90a4f53487e..c28715185a9 100644 --- a/openbb_platform/providers/fmp/openbb_fmp/utils/websocket_client.py +++ b/openbb_platform/providers/fmp/openbb_fmp/utils/websocket_client.py @@ -8,15 +8,15 @@ import sys import websockets import websockets.exceptions -from openbb_fmp.models.websocket_connection import FmpWebSocketData -from openbb_websockets.helpers import ( - MessageQueue, +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 from pydantic import ValidationError logger = get_logger("openbb.websocket.fmp") diff --git a/openbb_platform/providers/intrinio/openbb_intrinio/utils/websocket_client.py b/openbb_platform/providers/intrinio/openbb_intrinio/utils/websocket_client.py index 9596f4248bd..470cc375b4b 100644 --- a/openbb_platform/providers/intrinio/openbb_intrinio/utils/websocket_client.py +++ b/openbb_platform/providers/intrinio/openbb_intrinio/utils/websocket_client.py @@ -9,16 +9,16 @@ import sys from typing import Any from openbb_core.provider.utils.helpers import run_async -from openbb_intrinio.models.websocket_connection import IntrinioWebSocketData -from openbb_intrinio.utils.stocks_client import IntrinioRealtimeClient -from openbb_websockets.helpers import ( - MessageQueue, +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 +from openbb_intrinio.utils.stocks_client import IntrinioRealtimeClient from pydantic import ValidationError logger = get_logger("openbb.websocket.intrinio") diff --git a/openbb_platform/providers/polygon/openbb_polygon/utils/websocket_client.py b/openbb_platform/providers/polygon/openbb_polygon/utils/websocket_client.py index 4daab5892a6..a5fe307d682 100644 --- a/openbb_platform/providers/polygon/openbb_polygon/utils/websocket_client.py +++ b/openbb_platform/providers/polygon/openbb_polygon/utils/websocket_client.py @@ -8,18 +8,18 @@ import sys import websockets import websockets.exceptions -from openbb_polygon.models.websocket_connection import ( - FEED_MAP, - PolygonWebSocketData, -) -from openbb_websockets.helpers import ( - MessageQueue, +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 ( + FEED_MAP, + PolygonWebSocketData, +) from pydantic import ValidationError logger = get_logger("openbb.websocket.polygon") diff --git a/openbb_platform/providers/tiingo/openbb_tiingo/utils/websocket_client.py b/openbb_platform/providers/tiingo/openbb_tiingo/utils/websocket_client.py index 463f4ebcf40..e977d2d3110 100644 --- a/openbb_platform/providers/tiingo/openbb_tiingo/utils/websocket_client.py +++ b/openbb_platform/providers/tiingo/openbb_tiingo/utils/websocket_client.py @@ -8,15 +8,15 @@ import sys import websockets from openbb_core.provider.utils.errors import UnauthorizedError -from openbb_tiingo.models.websocket_connection import TiingoWebSocketData -from openbb_websockets.helpers import ( - MessageQueue, +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 from pydantic import ValidationError # These are the data array definitions.