mirror of
https://github.com/OpenBB-finance/OpenBB.git
synced 2026-11-02 15:32:04 +00:00
move some files around and create unit tests for MessageQueue
This commit is contained in:
Danglewood
@@ -0,0 +1 @@
|
||||
"""WebSocket Utilities."""
|
||||
@@ -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
|
||||
@@ -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)
|
||||
@@ -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
|
||||
@@ -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")
|
||||
|
||||
@@ -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")
|
||||
|
||||
@@ -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")
|
||||
|
||||
@@ -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.
|
||||
|
||||
Reference in New Issue
Block a user