mirror of
https://github.com/OpenBB-finance/OpenBB.git
synced 2026-10-28 21:12:16 +00:00
playing with some settings and update database test
This commit is contained in:
Danglewood
@@ -852,7 +852,7 @@ class WebSocketClient: # pylint: disable=too-many-instance-attributes
|
||||
def __repr__(self):
|
||||
"""Return the WebSocketClient representation."""
|
||||
return (
|
||||
f"WebSocketClient(module={[d for d in self.module if "api_key" not in d]}, symbol={self.symbol}, "
|
||||
f"WebSocketClient(module={[d for d in self.module if 'api_key' not in d]}, 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}, "
|
||||
|
||||
@@ -8,8 +8,12 @@ from contextlib import asynccontextmanager
|
||||
from datetime import timedelta
|
||||
from typing import TYPE_CHECKING, Any, Iterable, Optional, Union
|
||||
|
||||
import uvloop
|
||||
from openbb_core.app.model.abstract.error import OpenBBError
|
||||
from openbb_core.provider.utils.helpers import run_async
|
||||
from openbb_core.provider.utils.websockets.helpers import kill_thread
|
||||
|
||||
asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import logging
|
||||
@@ -65,30 +69,6 @@ CHECK_FOR = (
|
||||
)
|
||||
|
||||
|
||||
def kill_thread(thread: threading.Thread) -> None:
|
||||
"""Kill thread by setting a stop flag."""
|
||||
# pylint: disable=import-outside-toplevel
|
||||
import ctypes
|
||||
|
||||
if hasattr(thread, "loop") and thread.loop:
|
||||
for task in asyncio.all_tasks(thread.loop):
|
||||
task.cancel()
|
||||
|
||||
if not thread.is_alive():
|
||||
return
|
||||
|
||||
thread_id = thread.ident
|
||||
if thread_id is None:
|
||||
return
|
||||
|
||||
res = ctypes.pythonapi.PyThreadState_SetAsyncExc(
|
||||
ctypes.c_long(thread_id), ctypes.py_object(SystemExit)
|
||||
)
|
||||
if res > 1:
|
||||
ctypes.pythonapi.PyThreadState_SetAsyncExc(ctypes.c_long(thread_id), None)
|
||||
raise SystemError("PyThreadState_SetAsyncExc failed")
|
||||
|
||||
|
||||
class Database:
|
||||
"""
|
||||
Class to read from, and write to, the SQL file using aiosqlite.
|
||||
@@ -160,7 +140,7 @@ class Database:
|
||||
data_model: Optional["BaseModel"] = None,
|
||||
limit: Optional[int] = None,
|
||||
logger: Optional["logging.Logger"] = None,
|
||||
loop: Optional["asyncio.AbstractEventLoop"] = None,
|
||||
loop: Optional[asyncio.AbstractEventLoop] = None,
|
||||
**kwargs,
|
||||
):
|
||||
"""Initialize the ResultsDB class."""
|
||||
@@ -212,7 +192,7 @@ class Database:
|
||||
|
||||
try:
|
||||
if self.results_file is not None and os.path.exists(self.results_file): # type: ignore
|
||||
async with self.get_connection("read") as conn:
|
||||
async with self.get_connection("write") as conn:
|
||||
try:
|
||||
cursor = await conn.execute(
|
||||
"SELECT name FROM sqlite_master WHERE type='table';"
|
||||
@@ -238,7 +218,8 @@ class Database:
|
||||
)
|
||||
pragmas = [
|
||||
"PRAGMA journal_mode=WAL",
|
||||
"PRAGMA synchronous=OFF",
|
||||
"PRAGMA synchronous=off",
|
||||
"PRAGMA temp_store=memory",
|
||||
]
|
||||
|
||||
for pragma in pragmas:
|
||||
@@ -288,6 +269,15 @@ class Database:
|
||||
|
||||
if name not in self._connections:
|
||||
conn = await aiosqlite.connect(results_file, **conn_kwargs)
|
||||
pragmas = [
|
||||
"PRAGMA journal_mode=WAL",
|
||||
"PRAGMA synchronous=off",
|
||||
"PRAGMA temp_store=memory",
|
||||
]
|
||||
for pragma in pragmas:
|
||||
await conn.execute(pragma)
|
||||
|
||||
await conn.commit()
|
||||
self._connections[name] = conn
|
||||
|
||||
yield self._connections[name]
|
||||
@@ -373,30 +363,39 @@ class Database:
|
||||
params = ()
|
||||
async with conn.execute(query, params) as cursor:
|
||||
async for row in cursor:
|
||||
rows.append(await self._deserialize_row(row))
|
||||
rows.append(await self._deserialize_row(row, cursor))
|
||||
|
||||
return rows
|
||||
|
||||
except Exception as e: # pylint: disable=broad-except
|
||||
raise OpenBBError(e) from e
|
||||
|
||||
async def _deserialize_row(self, row) -> dict:
|
||||
"""Deserialize a row from the SQLite database."""
|
||||
async def _deserialize_row(self, row, cursor) -> dict:
|
||||
"""Deserialize a row from the SQLite database.
|
||||
Handles both single message column and multiple extracted fields."""
|
||||
# pylint: disable=import-outside-toplevel
|
||||
import json
|
||||
|
||||
try:
|
||||
return (
|
||||
json.loads(row[0])
|
||||
if (
|
||||
(
|
||||
if len(row) == 1:
|
||||
# Single column case (full message)
|
||||
return (
|
||||
json.loads(row[0])
|
||||
if (
|
||||
isinstance(row[0], str)
|
||||
and (row[0].startswith("{") or row[0].startswith("["))
|
||||
)
|
||||
or isinstance(row[0], bytes)
|
||||
else row[0]
|
||||
)
|
||||
else row[0]
|
||||
)
|
||||
else:
|
||||
# Multiple column case (extracted fields)
|
||||
return {cursor.description[i][0]: row[i] for i in range(len(row))}
|
||||
|
||||
except (json.JSONDecodeError, AttributeError) as e:
|
||||
self.logger.error(f"Failed to deserialize row: {e}")
|
||||
return row[0] if len(row) == 1 else dict(enumerate(row))
|
||||
|
||||
except Exception as e: # pylint: disable=broad-except
|
||||
msg = (
|
||||
"Unexpected error while deserializing row -> "
|
||||
@@ -453,14 +452,12 @@ class Database:
|
||||
|
||||
async def _query_db(self, sql, parameters: Optional[Iterable[Any]] = None) -> list:
|
||||
"""Query the SQLite database."""
|
||||
# pylint: disable=import-outside-toplevel
|
||||
import json
|
||||
|
||||
if not sql or sql in ("", "''"):
|
||||
raise OpenBBError("Empty query not allowed.")
|
||||
query = (
|
||||
sql
|
||||
if sql.startswith("SELECT")
|
||||
if sql.strip().startswith("SELECT")
|
||||
else f"SELECT message FROM {self.table_name} WHERE {sql}" # noqa
|
||||
)
|
||||
if not query.endswith(";"):
|
||||
@@ -479,17 +476,7 @@ class Database:
|
||||
query, parameters
|
||||
) as cursor:
|
||||
async for row in cursor:
|
||||
rows.append(
|
||||
json.loads(row[0])
|
||||
if (
|
||||
(
|
||||
isinstance(row[0], str)
|
||||
and (row[0].startswith("{") or row[0].startswith("["))
|
||||
)
|
||||
or isinstance(row[0], bytes)
|
||||
)
|
||||
else row[0]
|
||||
)
|
||||
rows.append(await self._deserialize_row(row, cursor))
|
||||
return rows
|
||||
except Exception as e: # pylint: disable=broad-except
|
||||
raise OpenBBError(e) from e
|
||||
@@ -563,7 +550,7 @@ class Database:
|
||||
def create_writer(
|
||||
self,
|
||||
queue: Optional["MessageQueue"] = None,
|
||||
batch_size: int = 25000,
|
||||
batch_size: int = 200,
|
||||
prune_interval: Optional[int] = None,
|
||||
export_directory: Optional[Union[str, "Path"]] = None,
|
||||
export_interval: Optional[int] = None,
|
||||
@@ -624,7 +611,7 @@ class DatabaseWriter:
|
||||
self,
|
||||
database: Database,
|
||||
queue: Optional["MessageQueue"] = None,
|
||||
batch_size: int = 25000,
|
||||
batch_size: int = 200,
|
||||
prune_interval: Optional[int] = None,
|
||||
export_directory: Optional[Union[str, "Path"]] = None,
|
||||
export_interval: Optional[int] = None,
|
||||
@@ -687,7 +674,8 @@ class DatabaseWriter:
|
||||
async def start_writer(self):
|
||||
"""Start writing tasks."""
|
||||
if not self.writer_running:
|
||||
asyncio.get_event_loop().run_in_executor(None, self.batch_processor.start())
|
||||
if not self.batch_processor.is_alive():
|
||||
self.batch_processor.start()
|
||||
await self._create_connection()
|
||||
for _ in range(self.num_workers):
|
||||
task = asyncio.create_task(self._process_queue())
|
||||
@@ -700,7 +688,7 @@ class DatabaseWriter:
|
||||
await asyncio.gather(*self.write_tasks, return_exceptions=True)
|
||||
if self._conn:
|
||||
await self._conn.close()
|
||||
self.batch_processor.running = False
|
||||
self.batch_processor.stop()
|
||||
kill_thread(self.batch_processor)
|
||||
|
||||
async def _process_queue(self):
|
||||
@@ -718,7 +706,7 @@ class DatabaseWriter:
|
||||
except asyncio.TimeoutError:
|
||||
break
|
||||
if batch:
|
||||
self._write_batch(batch)
|
||||
await self._write_batch(batch)
|
||||
batch = []
|
||||
else:
|
||||
await asyncio.sleep(0.001)
|
||||
@@ -734,9 +722,9 @@ class DatabaseWriter:
|
||||
while not self.queue.queue.empty() and len(batch) < self.batch_size:
|
||||
batch.append(await self.queue.dequeue())
|
||||
|
||||
self._write_batch(batch)
|
||||
await self._write_batch(batch)
|
||||
|
||||
def _write_batch(self, batch):
|
||||
async def _write_batch(self, batch):
|
||||
"""Write the batch of messages to the database."""
|
||||
if not batch:
|
||||
return
|
||||
@@ -757,6 +745,7 @@ class DatabaseWriter:
|
||||
chunk_size = 20000
|
||||
minutes = self.export_interval or 5
|
||||
latest_date = None
|
||||
earliest_date = None
|
||||
if not self._export_running or not self.export_thread:
|
||||
return
|
||||
try:
|
||||
@@ -767,6 +756,12 @@ class DatabaseWriter:
|
||||
ORDER BY json_extract(message, '$.date') DESC LIMIT 1
|
||||
""" # noqa
|
||||
|
||||
earliest_query = f"""
|
||||
SELECT json_extract(message, '$.date')
|
||||
FROM {self.database.table_name}
|
||||
ORDER BY json_extract(message, '$.date') ASC LIMIT 1
|
||||
""" # noqa
|
||||
|
||||
async with self.database.get_connection("read") as conn:
|
||||
try:
|
||||
async with conn.execute(latest_query) as cursor:
|
||||
@@ -774,19 +769,29 @@ class DatabaseWriter:
|
||||
if not latest_date:
|
||||
return
|
||||
|
||||
async with conn.execute(earliest_query) as cursor:
|
||||
earliest_date = await cursor.fetchone()
|
||||
if not latest_date:
|
||||
return
|
||||
|
||||
except asyncio.InvalidStateError as e:
|
||||
self.database.logger.error(f"Database connection state error: {e}")
|
||||
self._export_running = False
|
||||
return
|
||||
|
||||
latest_timestamp = to_datetime(latest_date[0])
|
||||
|
||||
earliest_timestamp = to_datetime(earliest_date[0])
|
||||
# Round down to nearest interval
|
||||
cutoff_time = latest_timestamp - timedelta(
|
||||
minutes=latest_timestamp.minute % minutes,
|
||||
seconds=latest_timestamp.second,
|
||||
microseconds=latest_timestamp.microsecond,
|
||||
)
|
||||
earliest_time = earliest_timestamp - timedelta(
|
||||
minutes=latest_timestamp.minute % minutes,
|
||||
seconds=latest_timestamp.second,
|
||||
microseconds=latest_timestamp.microsecond,
|
||||
)
|
||||
|
||||
# If we have processed data before, use that as reference
|
||||
if self._last_processed_timestamp:
|
||||
@@ -794,8 +799,9 @@ class DatabaseWriter:
|
||||
else:
|
||||
start_time = cutoff_time - timedelta(minutes=minutes)
|
||||
|
||||
cutoff_time = self._last_processed_timestamp + timedelta(minutes=minutes)
|
||||
start_time = self._last_processed_timestamp
|
||||
if start_time < earliest_timestamp:
|
||||
start_time = earliest_time
|
||||
|
||||
results_file = (
|
||||
self.export_directory
|
||||
+ "/"
|
||||
@@ -1042,9 +1048,7 @@ class DatabaseWriter:
|
||||
count = await self.database._query_db(query, (cutoff_timestamp,))
|
||||
|
||||
if count[0] > 0:
|
||||
# task = asyncio.create_task(self._export_database())
|
||||
task = await asyncio.to_thread(self._export_database)
|
||||
await asyncio.gather(task)
|
||||
await self._export_database()
|
||||
else:
|
||||
await asyncio.sleep(minutes * 60)
|
||||
# Stagger slightly so things don't happen exactly on the minute.
|
||||
@@ -1058,13 +1062,11 @@ class DatabaseWriter:
|
||||
|
||||
class BatchProcessor(threading.Thread):
|
||||
"""
|
||||
Batch processor for writing messages to the SQLite database in a new thread.
|
||||
|
||||
This class is a thread intended for use as a subprocess and is called by `DatabaseWriter.start_writer()`.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self, database_writer: DatabaseWriter, num_workers=120, collection_time=0.35
|
||||
self, database_writer: DatabaseWriter, num_workers=120, collection_time=0.25
|
||||
):
|
||||
"""Initialize the BatchProcessor class."""
|
||||
# pylint: disable=import-outside-toplevel
|
||||
@@ -1078,20 +1080,42 @@ class BatchProcessor(threading.Thread):
|
||||
self.num_workers = num_workers
|
||||
self.workers: list = []
|
||||
self.collection_time = collection_time
|
||||
self._shutdown = threading.Event()
|
||||
|
||||
def run(self):
|
||||
"""Run the batch processor as tasks."""
|
||||
|
||||
self.loop = asyncio.new_event_loop()
|
||||
asyncio.set_event_loop(self.loop)
|
||||
try:
|
||||
self.loop = asyncio.new_event_loop()
|
||||
asyncio.set_event_loop(self.loop)
|
||||
# Create worker tasks
|
||||
for _ in range(self.num_workers):
|
||||
worker = self.loop.create_task(self._worker())
|
||||
self.workers.append(worker)
|
||||
# Run workers
|
||||
self.loop.run_until_complete(asyncio.gather(*self.workers))
|
||||
while self.running and not self._shutdown.is_set():
|
||||
try:
|
||||
self.loop.run_until_complete(self._worker())
|
||||
except (SystemExit, KeyboardInterrupt):
|
||||
self.running = False
|
||||
break
|
||||
except Exception as e:
|
||||
self.writer.database.logger.error(f"Batch processing error: {e}")
|
||||
break
|
||||
finally:
|
||||
self._cleanup()
|
||||
|
||||
def stop(self):
|
||||
"""Signal thread to stop gracefully."""
|
||||
self.running = False
|
||||
self._shutdown.set()
|
||||
if self.loop and self.loop.is_running():
|
||||
self.loop.call_soon_threadsafe(self.loop.stop)
|
||||
|
||||
def _cleanup(self):
|
||||
"""Clean up resources on shutdown"""
|
||||
if self.loop:
|
||||
pending = asyncio.all_tasks(self.loop)
|
||||
for task in pending:
|
||||
task.cancel()
|
||||
self.loop.run_until_complete(
|
||||
asyncio.gather(*pending, return_exceptions=True)
|
||||
)
|
||||
self.loop.close()
|
||||
|
||||
async def _worker(self):
|
||||
@@ -1101,7 +1125,7 @@ class BatchProcessor(threading.Thread):
|
||||
batch_size = 200
|
||||
while self.running:
|
||||
try:
|
||||
batch = []
|
||||
batch: list = []
|
||||
total = 0
|
||||
collection_start = time.time()
|
||||
|
||||
@@ -1136,6 +1160,7 @@ class BatchProcessor(threading.Thread):
|
||||
for b in batch:
|
||||
values.extend([(msg,) for msg in b])
|
||||
async with self.writer.database.get_connection("write") as conn:
|
||||
await conn.execute("PRAGMA temp_store=memory")
|
||||
async with conn.cursor() as cursor:
|
||||
await cursor.executemany(query, values)
|
||||
await conn.commit()
|
||||
|
||||
@@ -4,9 +4,13 @@
|
||||
|
||||
import logging
|
||||
import re
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from pydantic import ValidationError
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import threading
|
||||
|
||||
AUTH_TOKEN_FILTER = re.compile(
|
||||
r"(auth_token=)([^&]*)",
|
||||
re.IGNORECASE | re.MULTILINE,
|
||||
@@ -73,13 +77,10 @@ def decrypt_value(key, iv, encrypted_value):
|
||||
|
||||
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)
|
||||
raise SystemExit("Termination signal received.")
|
||||
|
||||
|
||||
def parse_kwargs() -> dict:
|
||||
@@ -127,3 +128,28 @@ def parse_kwargs() -> dict:
|
||||
_kwargs[key] = True
|
||||
|
||||
return _kwargs
|
||||
|
||||
|
||||
def kill_thread(thread: "threading.Thread") -> None:
|
||||
"""Kill thread by setting a stop flag."""
|
||||
# pylint: disable=import-outside-toplevel
|
||||
import asyncio
|
||||
import ctypes
|
||||
|
||||
if hasattr(thread, "loop") and thread.loop:
|
||||
for task in asyncio.all_tasks(thread.loop):
|
||||
task.cancel()
|
||||
|
||||
if not thread.is_alive():
|
||||
return
|
||||
|
||||
thread_id = thread.ident
|
||||
if thread_id is None:
|
||||
return
|
||||
|
||||
res = ctypes.pythonapi.PyThreadState_SetAsyncExc(
|
||||
ctypes.c_long(thread_id), ctypes.py_object(SystemExit)
|
||||
)
|
||||
if res > 1:
|
||||
ctypes.pythonapi.PyThreadState_SetAsyncExc(ctypes.c_long(thread_id), None)
|
||||
raise SystemError("PyThreadState_SetAsyncExc failed")
|
||||
|
||||
@@ -1,7 +1,11 @@
|
||||
"""Unit Tests For Database Operations."""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
|
||||
import pytest
|
||||
from openbb_core.provider.utils.websockets.database import Database
|
||||
from openbb_core.provider.utils.websockets.database import Database, DatabaseWriter
|
||||
from openbb_core.provider.utils.websockets.message_queue import MessageQueue
|
||||
|
||||
MOCK_MESSAGES = [
|
||||
{"type": "trade", "symbol": "test1", "price": 100},
|
||||
@@ -16,6 +20,19 @@ def database():
|
||||
return Database(table_name="test")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def database_writer():
|
||||
"""Return a MessageQueue instance."""
|
||||
writer = DatabaseWriter(database=Database(table_name="test"), queue=MessageQueue())
|
||||
yield writer
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def message_queue():
|
||||
"""Return a MessageQueue instance."""
|
||||
return MessageQueue()
|
||||
|
||||
|
||||
def test_setup_database(database):
|
||||
"""Test if the database was setup."""
|
||||
assert database
|
||||
@@ -114,3 +131,36 @@ def test_limit():
|
||||
database.write_to_db(MOCK_MESSAGES[0])
|
||||
query = "SELECT id FROM test_limit"
|
||||
assert database.query(query)[0] > 3
|
||||
|
||||
|
||||
def test_batch_process_thread(database_writer):
|
||||
"""Test if the batch process thread starts and stops."""
|
||||
database_writer.batch_processor.start()
|
||||
assert database_writer.batch_processor.is_alive()
|
||||
database_writer.batch_processor.stop()
|
||||
database_writer.batch_processor.join(timeout=1)
|
||||
assert not database_writer.batch_processor.is_alive()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_database_writer(database_writer):
|
||||
"""Test if the database writer is working."""
|
||||
assert database_writer
|
||||
writer = database_writer
|
||||
assert isinstance(writer, DatabaseWriter)
|
||||
assert len(writer.database.fetch_all()) == 0
|
||||
await writer.start_writer()
|
||||
assert not writer._export_running
|
||||
assert not writer._prune_running
|
||||
|
||||
for message in MOCK_MESSAGES:
|
||||
await writer.queue.enqueue(json.dumps(message))
|
||||
|
||||
await asyncio.sleep(0.5)
|
||||
|
||||
await writer.stop_writer()
|
||||
|
||||
messages = writer.database.fetch_all()
|
||||
assert len(messages) == 3
|
||||
|
||||
assert not writer.batch_processor.is_alive()
|
||||
|
||||
@@ -66,8 +66,10 @@ class StdOutSink:
|
||||
def write(self, message):
|
||||
"""Write to stdout."""
|
||||
# pylint: disable=import-outside-toplevel
|
||||
import os
|
||||
import sys
|
||||
|
||||
os.set_blocking(sys.__stdout__.fileno(), False) # type: ignore
|
||||
cleaned_message = AUTH_TOKEN_FILTER.sub(r"\1********", message)
|
||||
if cleaned_message != message:
|
||||
cleaned_message = f"{cleaned_message}\n"
|
||||
|
||||
@@ -36,9 +36,9 @@ import os
|
||||
import signal
|
||||
import sys
|
||||
import time
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
|
||||
import orjson
|
||||
import uvloop
|
||||
import websockets
|
||||
from openbb_core.provider.utils.websockets.database import Database, DatabaseWriter
|
||||
from openbb_core.provider.utils.websockets.helpers import (
|
||||
@@ -54,6 +54,9 @@ from openbb_polygon.models.websocket_connection import (
|
||||
)
|
||||
from pydantic import ValidationError
|
||||
from websockets.asyncio.client import connect
|
||||
from websockets.extensions import permessage_deflate
|
||||
|
||||
asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
|
||||
|
||||
URL_MAP = {
|
||||
"stock": "wss://socket.polygon.io/stocks",
|
||||
@@ -78,7 +81,6 @@ ASSET_TYPE = kwargs.pop("asset_type", "crypto")
|
||||
kwargs["results_file"] = os.path.abspath(kwargs.get("results_file"))
|
||||
URL = URL_MAP.get(ASSET_TYPE)
|
||||
SUBSCRIBED_SYMBOLS: set = set()
|
||||
message_thread_pool = ThreadPoolExecutor(max_workers=16)
|
||||
|
||||
if not kwargs.get("api_key"):
|
||||
raise ValueError("No API key provided.")
|
||||
@@ -86,6 +88,7 @@ if not kwargs.get("api_key"):
|
||||
if not URL:
|
||||
raise ValueError("Invalid asset type provided.")
|
||||
|
||||
|
||||
DATABASE = DatabaseWriter(
|
||||
database=Database(
|
||||
results_file=kwargs["results_file"],
|
||||
@@ -158,8 +161,8 @@ async def login(websocket):
|
||||
login_event = f'{{"action":"auth","params":"{kwargs["api_key"]}"}}'
|
||||
try:
|
||||
await websocket.send(login_event)
|
||||
res = await websocket.recv()
|
||||
response = json.loads(res)
|
||||
res = await websocket.recv(decode=False)
|
||||
response = orjson.loads(res)
|
||||
messages = response if isinstance(response, list) else [response]
|
||||
for msg in messages:
|
||||
if msg.get("status") == "connected":
|
||||
@@ -213,7 +216,7 @@ async def subscribe(websocket, symbol, event):
|
||||
async def read_stdin():
|
||||
"""Read from stdin and queue commands."""
|
||||
while True:
|
||||
line = await asyncio.get_event_loop().run_in_executor(None, sys.stdin.readline)
|
||||
line = await asyncio.to_thread(sys.stdin.readline)
|
||||
sys.stdin.flush()
|
||||
if line:
|
||||
try:
|
||||
@@ -280,22 +283,21 @@ async def connect_and_stream():
|
||||
|
||||
conn_kwargs.update(
|
||||
{
|
||||
"ping_interval": None,
|
||||
"ping_timeout": None,
|
||||
"close_timeout": 1,
|
||||
"max_size": 2**63,
|
||||
"max_size": 32768,
|
||||
"max_queue": None,
|
||||
}
|
||||
)
|
||||
|
||||
async def message_receiver(websocket):
|
||||
"""Message receiver."""
|
||||
try:
|
||||
async for message in websocket:
|
||||
while True:
|
||||
try:
|
||||
message = await websocket.recv(decode=False)
|
||||
input_queue.queue.put_nowait(message)
|
||||
|
||||
except Exception as e:
|
||||
raise e from e
|
||||
except Exception as e:
|
||||
raise e from e
|
||||
|
||||
async def process_input_messages(message):
|
||||
"""Process the messages offloaded from the websocket."""
|
||||
@@ -325,7 +327,8 @@ async def connect_and_stream():
|
||||
|
||||
# Run processing in thread
|
||||
process = asyncio.to_thread(_process_in_thread)
|
||||
asyncio.create_task(process)
|
||||
process_task = asyncio.create_task(process)
|
||||
tasks.add(process_task)
|
||||
|
||||
try:
|
||||
handler_task = asyncio.create_task(
|
||||
@@ -337,15 +340,25 @@ async def connect_and_stream():
|
||||
|
||||
await DATABASE.start_writer()
|
||||
|
||||
for i in range(16):
|
||||
processor_task = asyncio.create_task(
|
||||
input_queue.process_queue(
|
||||
lambda message: process_input_messages(message)
|
||||
)
|
||||
)
|
||||
tasks.add(processor_task)
|
||||
processor_task = asyncio.create_task(
|
||||
input_queue.process_queue(lambda message: process_input_messages(message))
|
||||
)
|
||||
tasks.add(processor_task)
|
||||
|
||||
async for websocket in connect(URL, **conn_kwargs):
|
||||
conn_kwargs["extensions"] = [
|
||||
permessage_deflate.ClientPerMessageDeflateFactory(
|
||||
server_max_window_bits=11,
|
||||
client_max_window_bits=11,
|
||||
compress_settings={
|
||||
"memLevel": 1,
|
||||
"level": 1,
|
||||
},
|
||||
),
|
||||
]
|
||||
async for websocket in connect(
|
||||
URL,
|
||||
**conn_kwargs,
|
||||
):
|
||||
try:
|
||||
if not any(
|
||||
task.name == "cmd_task" for task in tasks if hasattr(task, "name")
|
||||
@@ -357,10 +370,9 @@ async def connect_and_stream():
|
||||
|
||||
await login(websocket)
|
||||
|
||||
response = await websocket.recv()
|
||||
messages = orjson.loads(response)
|
||||
response = await websocket.recv(decode=False)
|
||||
|
||||
await process_message(messages)
|
||||
await process_message(orjson.loads(response))
|
||||
|
||||
await subscribe(websocket, kwargs["symbol"], "subscribe")
|
||||
|
||||
@@ -406,7 +418,7 @@ async def connect_and_stream():
|
||||
|
||||
if __name__ == "__main__":
|
||||
try:
|
||||
loop = asyncio.new_event_loop()
|
||||
loop = uvloop.new_event_loop()
|
||||
asyncio.set_event_loop(loop)
|
||||
loop.set_exception_handler(lambda loop, context: None)
|
||||
|
||||
|
||||
@@ -1,12 +1,13 @@
|
||||
"""FMP WebSocket server."""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
import signal
|
||||
import sys
|
||||
from datetime import datetime, UTC
|
||||
from datetime import UTC, datetime
|
||||
|
||||
import orjson
|
||||
import uvloop
|
||||
import websockets
|
||||
from openbb_core.provider.utils.errors import UnauthorizedError
|
||||
from openbb_core.provider.utils.websockets.database import Database, DatabaseWriter
|
||||
@@ -21,6 +22,8 @@ from openbb_tiingo.models.websocket_connection import TiingoWebSocketData
|
||||
from pandas import to_datetime
|
||||
from pydantic import ValidationError
|
||||
|
||||
asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
|
||||
|
||||
URL_MAP = {
|
||||
"stock": "wss://api.tiingo.com/iex",
|
||||
"fx": "wss://api.tiingo.com/fx",
|
||||
@@ -97,7 +100,6 @@ DATABASE = DatabaseWriter(
|
||||
logger=logger,
|
||||
),
|
||||
queue=db_queue,
|
||||
batch_size=100,
|
||||
)
|
||||
|
||||
|
||||
@@ -120,9 +122,9 @@ async def update_symbols(symbol, event):
|
||||
}
|
||||
|
||||
async with websockets.connect(URL) as websocket:
|
||||
await websocket.send(json.dumps(update_event))
|
||||
response = await websocket.recv()
|
||||
message = json.loads(response)
|
||||
await websocket.send(orjson.dumps(update_event))
|
||||
response = await websocket.recv(decode=False)
|
||||
message = orjson.loads(response)
|
||||
if "tickers" in message.get("data", {}):
|
||||
tickers = message["data"]["tickers"]
|
||||
threshold_level = message["data"].get("thresholdLevel")
|
||||
@@ -155,7 +157,7 @@ async def read_stdin_and_update_symbols():
|
||||
f" Database Queue : {db_queue.queue.qsize()}"
|
||||
)
|
||||
else:
|
||||
line = json.loads(line.strip())
|
||||
line = orjson.loads(line.strip())
|
||||
|
||||
if line:
|
||||
symbol = line.get("symbol")
|
||||
@@ -167,7 +169,7 @@ async def process_message(message): # pylint: disable=too-many-branches
|
||||
"""Process the message and write to the database."""
|
||||
result: dict = {}
|
||||
data_message: dict = {}
|
||||
message = message if isinstance(message, (dict, list)) else json.loads(message)
|
||||
message = message if isinstance(message, (dict, list)) else orjson.loads(message)
|
||||
msg: str = ""
|
||||
if message.get("messageType") == "E":
|
||||
response = message.get("response", {})
|
||||
@@ -232,7 +234,7 @@ async def process_message(message): # pylint: disable=too-many-branches
|
||||
raise e from e
|
||||
|
||||
if result:
|
||||
await db_queue.enqueue(result)
|
||||
db_queue.queue.put_nowait(result)
|
||||
|
||||
|
||||
async def connect_and_stream():
|
||||
@@ -247,9 +249,7 @@ async def connect_and_stream():
|
||||
{
|
||||
"ping_interval": 8,
|
||||
"ping_timeout": 8,
|
||||
"read_limit": 2**256,
|
||||
"close_timeout": 1,
|
||||
"max_queue": 10000,
|
||||
}
|
||||
)
|
||||
|
||||
@@ -268,38 +268,50 @@ async def connect_and_stream():
|
||||
async def message_receiver(websocket):
|
||||
"""Receive messages from the WebSocket."""
|
||||
while True:
|
||||
message = await websocket.recv()
|
||||
await input_queue.enqueue(message)
|
||||
message = await websocket.recv(decode=False)
|
||||
input_queue.queue.put_nowait(orjson.loads(message))
|
||||
|
||||
stdin_task = asyncio.create_task(read_stdin_and_update_symbols())
|
||||
tasks.add(stdin_task)
|
||||
|
||||
try:
|
||||
await DATABASE.start_writer()
|
||||
websocket = await websockets.connect(URL, **conn_kwargs)
|
||||
receiver_task = asyncio.create_task(message_receiver(websocket))
|
||||
tasks.add(receiver_task)
|
||||
await websocket.send(json.dumps(subscribe_event))
|
||||
logger.info("PROVIDER INFO: WebSocket connection established.")
|
||||
for _ in range(9):
|
||||
process_task = asyncio.create_task(
|
||||
input_queue.process_queue(lambda message: process_message(message))
|
||||
)
|
||||
tasks.add(process_task)
|
||||
|
||||
await asyncio.gather(*tasks, return_exceptions=True)
|
||||
async for websocket in websockets.connect(URL, **conn_kwargs):
|
||||
try:
|
||||
if not any(
|
||||
task.name == "receiver_task"
|
||||
for task in tasks
|
||||
if hasattr(task, "name")
|
||||
):
|
||||
receiver_task = asyncio.create_task(
|
||||
message_receiver(websocket), name="receiver_task"
|
||||
)
|
||||
tasks.add(receiver_task)
|
||||
|
||||
except UnauthorizedError as e:
|
||||
logger.error(str(e))
|
||||
sys.exit(1)
|
||||
await websocket.send(orjson.dumps(subscribe_event))
|
||||
logger.info("PROVIDER INFO: WebSocket connection established.")
|
||||
for _ in range(9):
|
||||
process_task = asyncio.create_task(
|
||||
input_queue.process_queue(
|
||||
lambda message: process_message(message)
|
||||
)
|
||||
)
|
||||
tasks.add(process_task)
|
||||
|
||||
except websockets.ConnectionClosed as e:
|
||||
msg = f"PROVIDER INFO: The WebSocket connection was closed -> {e}"
|
||||
logger.info(msg)
|
||||
# Attempt to reopen the connection
|
||||
logger.info("PROVIDER INFO: Attempting to reconnect...")
|
||||
await asyncio.sleep(1)
|
||||
await connect_and_stream()
|
||||
await asyncio.gather(*tasks, return_exceptions=True)
|
||||
|
||||
except UnauthorizedError as e:
|
||||
logger.error(str(e))
|
||||
sys.exit(1)
|
||||
|
||||
except websockets.ConnectionClosed as e:
|
||||
msg = f"PROVIDER INFO: The WebSocket connection was closed -> {e}"
|
||||
logger.info(msg)
|
||||
# Attempt to reopen the connection
|
||||
logger.info("PROVIDER INFO: Attempting to reconnect...")
|
||||
await asyncio.sleep(1)
|
||||
continue
|
||||
|
||||
except websockets.WebSocketException as e:
|
||||
logger.info(str(e))
|
||||
@@ -311,7 +323,6 @@ async def connect_and_stream():
|
||||
sys.exit(1)
|
||||
|
||||
finally:
|
||||
await websocket.close()
|
||||
for task in tasks:
|
||||
task.cancel()
|
||||
await task
|
||||
@@ -322,9 +333,10 @@ async def connect_and_stream():
|
||||
|
||||
if __name__ == "__main__":
|
||||
try:
|
||||
loop = asyncio.new_event_loop()
|
||||
loop = uvloop.new_event_loop()
|
||||
asyncio.set_event_loop(loop)
|
||||
loop.set_exception_handler(lambda loop, context: None)
|
||||
|
||||
for sig in (signal.SIGINT, signal.SIGTERM):
|
||||
loop.add_signal_handler(signal.SIGTERM, handle_termination_signal, logger)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user