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

move client to core.provider

Browse Source
This commit is contained in:
Danglewood
2024-12-17 23:05:48 -08:00
parent 70caac38e9
commit 88978b35bb
14 changed files with 238 additions and 124 deletions
Download Patch File Download Diff File Expand all files Collapse all files

258
openbb_platform/extensions/websockets/openbb_websockets/client.py → openbb_platform/core/openbb_core/provider/utils/websockets/client.py
View File

@@ -1,48 +1,60 @@
"""WebSocket Client module for interacting with a provider websocket in a non-blocking pattern."""
"""Module for running OpenBB Provider websocket connection scripts."""
# 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
import logging
from pydantic import BaseModel
class WebSocketClient: # pylint: disable=too-many-instance-attributes
"""Client for interacting with a websocket server in a non-blocking pattern.
"""
Client for interacting with a websocket server in a non-blocking pattern, and handling the subprocesses.
Parameters
----------
name : str
Name to assign the WebSocket connection. Used to identify and manage multiple instances.
Name to assign the WebSocket connection. Used to identify and manage multiple instances from the API.
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.
The Python module for the provider websocket_client module.
Runs in a separate thread, and is an equivalent to 'python -m module'.
Example: 'openbb_fmp.utils.websocket_client'.
Pass additional keyword arguments to the script by including **kwargs.
symbol : Optional[str]
The symbol(s) requested to subscribe. Enter multiple symbols separated by commas without spaces.
The symbol(s) requested to subscribe on start. Enter multiple symbols separated by commas, without spaces.
Where supported by the provider, * represents all symbols within the feed.
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.
Default is 5000. 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.
File is discarded when the Python session ends unless 'save_results' is set to True.
The connection can be re-established with the same results file to continue writing.
EACH NEW CONNECTION SHOULD HAVE A UNIQUE RESULTS FILE.
If the intention is to permanently store the results for historical records,
save the current session to a new file and copy new records at periodic intervals into the master.
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]
data_model : Optional[BaseModel]
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.
Used to limit access to the broadcast stream. When provided, listeners should supply as a URL parameter.
Example: 'ws://127.0.0.1:6666/?auth_token=SOME_TOKEN'.
When provided, the auth_token is required to interact with the instance of this class
from the REST 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.
A pre-configured logger to use for this instance. 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'}}.
{'api_key': 'MY_API_KEY', 'connect_kwargs': {'key': 'value'}}.
Properties
----------
@@ -55,10 +67,9 @@ class WebSocketClient: # pylint: disable=too-many-instance-attributes
is_broadcasting : bool
Check if the broadcast server process is running.
broadcast_address : str
URI address for the results broadcast server.
URI address for connecting to the broadcast stream.
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.
All stored results from the provider connection.
Clear the results by deleting the property. e.g., del client.results
Methods
@@ -68,15 +79,17 @@ class WebSocketClient: # pylint: disable=too-many-instance-attributes
disconnect
Disconnect from the provider WebSocket.
subscribe
Subscribe to a new symbol or list of symbols.
Send a subscribe message to the provider connection.
unsubscribe
Unsubscribe from a symbol or list of symbols.
Send an unsubscribe message to the provider connection.
start_broadcasting
Start the broadcast server to stream results over a network connection.
Start the broadcast server to stream results over a network.
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.
query_database
Run a SELECT query to the database. Returns a list of deserialized results, or validated models.
"""
def __init__( # noqa: PLR0913 # pylint: disable=too-many-arguments,too-many-positional-arguments,too-many-locals
@@ -84,13 +97,13 @@ class WebSocketClient: # pylint: disable=too-many-instance-attributes
name: str,
module: str,
symbol: Optional[str] = None,
limit: Optional[int] = 1000,
limit: Optional[int] = 5000,
results_file: Optional[str] = None,
table_name: Optional[str] = None,
save_results: bool = False,
data_model: Optional["Data"] = None,
data_model: Optional["BaseModel"] = None,
auth_token: Optional[str] = None,
logger: Optional[logging.Logger] = None,
logger: Optional["logging.Logger"] = None,
**kwargs,
) -> None:
"""Initialize the WebSocketClient class."""
@@ -162,20 +175,21 @@ class WebSocketClient: # pylint: disable=too-many-instance-attributes
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}"
f" {e.__class__.__name__ if hasattr(e, '__class__') else e.__name__} -> {e}"
)
self.logger.error(msg)
self._exception = OpenBBError(msg)
self._atexit()
raise OpenBBError(msg) from e
def _atexit(self) -> None:
"""Clean up the WebSocket client processes at exit."""
"""Clean up the running processes at exit."""
# pylint: disable=import-outside-toplevel
import os
@@ -310,13 +324,23 @@ class WebSocketClient: # pylint: disable=too-many-instance-attributes
output = output.replace("INFO:", "BROADCAST INFO:") + "\n"
output = output[0] if isinstance(output, tuple) else output
output = clean_message(output)
if (
output.startswith("BROADCAST ERROR:")
or "unexpected error" in output.lower()
):
self._psutil_broadcast_process.kill()
self._broadcast_process.wait()
self._broadcast_thread.join()
sys.stdout.write(output + "\n")
sys.stdout.flush()
break
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."""
"""Connect to the provider client connection."""
# pylint: disable=import-outside-toplevel
import os # noqa
import psutil
@@ -367,35 +391,44 @@ class WebSocketClient: # pylint: disable=too-many-instance-attributes
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)
try:
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()
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()
self._log_thread = threading.Thread(
target=self._log_provider_output,
args=(log_output_queue,),
)
self._log_thread.daemon = True
self._log_thread.start()
except Exception as e: # pylint: disable=broad-except
msg = f"Unexpected error -> {e.__class__.__name__ if hasattr(e, '__class__') else e.__name__} -> {e}"
self.logger.error(msg)
self._atexit()
raise OpenBBError(msg) from e
# Give it some startup time to allow the connection to be established and for exceptions to populate.
time.sleep(2)
@@ -406,12 +439,14 @@ class WebSocketClient: # pylint: disable=too-many-instance-attributes
raise OpenBBError(exc)
if not self.is_running:
self.logger.error("The provider server failed to start.")
self.logger.error(
"Unexpected error -> Provider connection process failed to start."
)
def send_message(
self, message, target: Literal["provider", "broadcast"] = "provider"
) -> None:
"""Send a message to the WebSocket process."""
"""Write to the provider, or broadcast, process stdin."""
if target == "provider":
self._provider_message_queue.put(message)
read_message_queue(self, self._provider_message_queue)
@@ -420,10 +455,10 @@ class WebSocketClient: # pylint: disable=too-many-instance-attributes
read_message_queue(self, self._broadcast_message_queue, target="broadcast")
def disconnect(self) -> None:
"""Disconnect from the provider WebSocket."""
"""Disconnect from the provider connection."""
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.")
self.logger.info("Provider client connection is not running.")
return
if (
self._psutil_process is not None
@@ -435,13 +470,18 @@ class WebSocketClient: # pylint: disable=too-many-instance-attributes
self._thread.join()
self._log_thread.join()
self._stop_log_thread_event.clear()
self.logger.info("Disconnected from the provider WebSocket.")
self.logger.info("Disconnected from the provider server.")
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."""
"""
Send a subscribe message to the active provider connection.
Messages are sent as JSON strings formatted as:
{"event": "subscribe", "symbol": "AAPL,MSFT"}
"""
# pylint: disable=import-outside-toplevel
import json # noqa
import time
@@ -463,7 +503,12 @@ class WebSocketClient: # pylint: disable=too-many-instance-attributes
self._symbol = ",".join(new_symbols)
def unsubscribe(self, symbol) -> None:
"""Unsubscribe from a symbol or list of symbols."""
"""
Unsubscribe from a symbol or list of symbols.
Messages are sent as JSON strings formatted as:
{"event": "unsubscribe", "symbol": "AAPL,MSFT"}
"""
# pylint: disable=import-outside-toplevel
import json # noqa
import time
@@ -500,23 +545,33 @@ class WebSocketClient: # pylint: disable=too-many-instance-attributes
@property
def results(self) -> list:
"""Retrieve the deserialized results from the results file."""
"""
Retrieve the deserialized results from the active Database.
Clear the results by deleting the property. e.g., del client.results
"""
# 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
msg = (
"Error retrieving results:"
f" {e.__class__.__name__ if hasattr(e,'__class__') else e.__name__} -> {e}"
)
raise OpenBBError(msg) from e
@results.deleter
def results(self):
"""Clear results stored from the WebSocket stream."""
# pylint: disable=import-outside-toplevel
"""Clear results stored by the active WebSocket stream."""
try:
self.database.clear_results()
except Exception as e: # pylint: disable=broad-except
msg = f"Error clearing results: {e.__class__.__name__}: {e}"
msg = (
"Error clearing results:"
f" {e.__class__.__name__ if hasattr(e, '__class__') else e.__name__} -> {e}"
)
self.logger.error(msg)
@property
@@ -560,13 +615,63 @@ class WebSocketClient: # pylint: disable=too-many-instance-attributes
else None
)
def query_database(
self,
sql: Optional[str] = None,
limit: Optional[int] = 25,
) -> list:
"""
Make a SELECT query to the database for results.
The database always contains two columns:
"id" - an auto-incrementing ID
"message" - a JSON serialized row of data
Parameters
----------
sql : Optional[str]
SQL query to execute. Default is None.
limit : Optional[int]
Limit the number of records returned, by most recent. Default is 25, set to None to return all records.
Returns
-------
list
A list of deserialized results from the database.
If a 'data_model' was supplied at initialization, it will be a list of validated models.
"""
if not sql:
query = f"SELECT message FROM {self.table_name} ORDER BY id DESC" # noqa
if limit is not None:
query += f" LIMIT {limit};"
else:
query = (
sql.replace(";", "") + f" LIMIT {limit}"
if limit is not None and "LIMIT" not in sql.upper()
else sql
)
return self.database.query(query)
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."""
"""
Broadcast results over a network connection.
Parameters
----------
host : str
The host address to broadcast results to. Default is 127.0.0.1
port : int
The port to broadcast results to. Default is 6666
If the port is already in use, the next available port is used.
**kwargs: dict
Additional keyword arguments to pass to the `uvicorn.run`.
"""
# pylint: disable=import-outside-toplevel
import os # noqa
import subprocess
@@ -692,9 +797,10 @@ def non_blocking_websocket(client, output_queue, provider_message_queue) -> None
if output:
output_queue.put(output.strip())
except Exception as e:
except Exception as e: # pylint: disable=broad-except
msg = (
f"Unexpected error in non_blocking_websocket: {e.__class__.__name__} -> {e}"
"Unexpected error in non_blocking_websocket:"
f" {e.__class__.__name__ if hasattr(e, '__class__') else e.__name__} -> {e}"
)
client.logger.error(msg)
raise e from e
@@ -710,7 +816,7 @@ def send_message(
# pylint: disable=import-outside-toplevel
import json
if isinstance(message, (dict, list)):
if not isinstance(message, str):
message = json.dumps(message)
try:
if target == "provider":
@@ -726,7 +832,10 @@ def send_message(
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}"
msg = (
f"Error sending message to the {target} process:"
f" {e.__class__.__name__ if hasattr(e, '__class__') else e.__name__} -> {e}"
)
client.logger.error(msg)
@@ -747,8 +856,8 @@ def read_message_queue(
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}"
"Error while attempting to transmit from the outgoing message queue:"
f"{e.__class__.__name__ if hasattr(e, '__class__') else e.__name__} -> {e} -> {message}"
)
client.logger.error(err)
@@ -767,7 +876,8 @@ def non_blocking_broadcast(client, output_queue, broadcast_message_queue) -> Non
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}"
f"Unexpected error in non_blocking_broadcast:"
f" {e.__class__.__name__ if hasattr(e, '__class__') else e.__name__} -> {e}"
)
client.logger.error(err)
finally:

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

@@ -66,7 +66,6 @@ class Database:
table_name: Optional[str] = None,
data_model: Optional["BaseModel"] = None,
limit: Optional[int] = None,
keep_results: Optional[bool] = False,
logger: Optional["logging.Logger"] = None,
loop: Optional["asyncio.AbstractEventLoop"] = None,
**kwargs,
@@ -87,6 +86,10 @@ class Database:
temp_file_path = temp_file.name
self.results_path = Path(temp_file_path).absolute()
self.results_file = temp_file_path
elif results_file and "://" in results_file:
self.results_file = results_file
self.results_path = results_file
kwargs["uri"] = True
else:
self.results_path = Path(results_file).absolute()
self.results_file = results_file

3
openbb_platform/extensions/websockets/openbb_websockets/models.py
View File

@@ -9,10 +9,9 @@ from openbb_core.provider.abstract.query_params import QueryParams
from openbb_core.provider.utils.descriptions import (
DATA_DESCRIPTIONS,
)
from openbb_core.provider.utils.websockets.client import WebSocketClient
from pydantic import ConfigDict, Field, field_validator, model_validator
from openbb_websockets.client import WebSocketClient
class WebSocketQueryParams(QueryParams):
"""Query parameters for WebSocket connection creation."""

12
openbb_platform/providers/fmp/openbb_fmp/models/websocket_connection.py
View File

@@ -7,7 +7,7 @@ from typing import Any, Literal, Optional
from openbb_core.app.model.abstract.error import OpenBBError
from openbb_core.provider.abstract.fetcher import Fetcher
from openbb_websockets.client import WebSocketClient
from openbb_core.provider.utils.websockets.client import WebSocketClient
from openbb_websockets.models import (
WebSocketConnection,
WebSocketData,
@@ -15,12 +15,6 @@ from openbb_websockets.models import (
)
from pydantic import Field, field_validator
URL_MAP = {
"stock": "wss://websockets.financialmodelingprep.com",
"fx": "wss://forex.financialmodelingprep.com",
"crypto": "wss://crypto.financialmodelingprep.com",
}
class FmpWebSocketQueryParams(WebSocketQueryParams):
"""FMP WebSocket query parameters."""
@@ -151,12 +145,10 @@ class FmpWebSocketFetcher(Fetcher[FmpWebSocketQueryParams, FmpWebSocketConnectio
import asyncio
api_key = credentials.get("fmp_api_key") if credentials else ""
url = URL_MAP[query.asset_type]
symbol = query.symbol.lower()
kwargs = {
"url": url,
"asset_type": query.asset_type,
"api_key": api_key,
"connect_kwargs": query.connect_kwargs,
}

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

@@ -18,11 +18,20 @@ from openbb_core.provider.utils.websockets.message_queue import MessageQueue
from openbb_fmp.models.websocket_connection import FmpWebSocketData
from pydantic import ValidationError
URL_MAP = {
"stock": "wss://websockets.financialmodelingprep.com",
"fx": "wss://forex.financialmodelingprep.com",
"crypto": "wss://crypto.financialmodelingprep.com",
}
logger = get_logger("openbb.websocket.fmp")
kwargs = parse_kwargs()
queue = MessageQueue()
command_queue = MessageQueue()
CONNECT_KWARGS = kwargs.pop("connect_kwargs", {})
URL = URL_MAP.get(kwargs.pop("asset_type"), None)
if not URL:
raise ValueError("Invalid asset type provided.")
DATABASE = Database(
results_file=kwargs["results_file"],
@@ -143,7 +152,7 @@ async def connect_and_stream():
stdin_task = asyncio.create_task(read_stdin_and_queue_commands())
try:
websocket = await websockets.connect(kwargs["url"], **CONNECT_KWARGS)
websocket = await websockets.connect(URL, **CONNECT_KWARGS)
await login(websocket)
await subscribe(websocket, kwargs["symbol"], "subscribe")

2
openbb_platform/providers/fmp/tests/test_fmp_fetchers.py
View File

@@ -7,6 +7,7 @@ from unittest.mock import MagicMock, patch
import pytest
from openbb_core.app.service.user_service import UserService
from openbb_core.provider.utils.websockets.client import WebSocketClient
from openbb_fmp.models.analyst_estimates import FMPAnalystEstimatesFetcher
from openbb_fmp.models.available_indices import FMPAvailableIndicesFetcher
from openbb_fmp.models.balance_sheet import FMPBalanceSheetFetcher
@@ -79,7 +80,6 @@ from openbb_fmp.models.websocket_connection import (
)
from openbb_fmp.models.world_news import FMPWorldNewsFetcher
from openbb_fmp.models.yield_curve import FMPYieldCurveFetcher
from openbb_websockets.client import WebSocketClient
test_credentials = UserService().default_user_settings.credentials.model_dump(
mode="json"

2
openbb_platform/providers/intrinio/openbb_intrinio/models/websocket_connection.py
View File

@@ -7,8 +7,8 @@ from typing import Any, Literal, Optional
from openbb_core.app.model.abstract.error import OpenBBError
from openbb_core.provider.abstract.fetcher import Fetcher
from openbb_core.provider.utils.websockets.client import WebSocketClient
from openbb_intrinio.utils.references import TRADE_CONDITIONS, VENUES
from openbb_websockets.client import WebSocketClient
from openbb_websockets.models import (
WebSocketConnection,
WebSocketData,

2
openbb_platform/providers/intrinio/tests/test_intrinio_fetchers.py
View File

@@ -7,6 +7,7 @@ from unittest.mock import MagicMock, patch
import pytest
from openbb_core.app.service.user_service import UserService
from openbb_core.provider.utils.websockets.client import WebSocketClient
from openbb_intrinio.models.balance_sheet import IntrinioBalanceSheetFetcher
from openbb_intrinio.models.calendar_ipo import IntrinioCalendarIpoFetcher
from openbb_intrinio.models.cash_flow import IntrinioCashFlowStatementFetcher
@@ -73,7 +74,6 @@ from openbb_intrinio.models.websocket_connection import (
IntrinioWebSocketFetcher,
)
from openbb_intrinio.models.world_news import IntrinioWorldNewsFetcher
from openbb_websockets.client import WebSocketClient
test_credentials = UserService().default_user_settings.credentials.model_dump(
mode="json"

16
openbb_platform/providers/polygon/openbb_polygon/models/websocket_connection.py
View File

@@ -9,6 +9,7 @@ from openbb_core.app.model.abstract.error import OpenBBError
from openbb_core.provider.abstract.data import Data
from openbb_core.provider.abstract.fetcher import Fetcher
from openbb_core.provider.utils.descriptions import DATA_DESCRIPTIONS
from openbb_core.provider.utils.websockets.client import WebSocketClient
from openbb_polygon.utils.constants import (
CRYPTO_EXCHANGE_MAP,
FX_EXCHANGE_MAP,
@@ -20,7 +21,6 @@ from openbb_polygon.utils.constants import (
STOCK_TRADE_CONDITIONS,
)
from openbb_polygon.utils.helpers import map_tape
from openbb_websockets.client import WebSocketClient
from openbb_websockets.models import (
WebSocketConnection,
WebSocketData,
@@ -28,16 +28,6 @@ from openbb_websockets.models import (
)
from pydantic import Field, field_validator, model_validator
URL_MAP = {
"stock": "wss://socket.polygon.io/stocks",
"stock_delayed": "wss://delayed.polygon.io/stocks",
"options": "wss://socket.polygon.io/options",
"options_delayed": "wss://delayed.polygon.io/options",
"fx": "wss://socket.polygon.io/forex",
"crypto": "wss://socket.polygon.io/crypto",
"index": "wss://socket.polygon.io/indices",
"index_delayed": "wss://delayed.polygon.io/indices",
}
ASSET_CHOICES = [
"stock",
@@ -1227,12 +1217,8 @@ class PolygonWebSocketFetcher(
) -> dict:
"""Extract data from the WebSocket."""
api_key = credentials.get("polygon_api_key") if credentials else ""
url = URL_MAP[query.asset_type]
symbol = query.symbol.upper()
kwargs = {
"url": url,
"asset_type": query.asset_type,
"feed": query.feed,
"api_key": api_key,

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

@@ -22,6 +22,17 @@ from openbb_polygon.models.websocket_connection import (
)
from pydantic import ValidationError
URL_MAP = {
"stock": "wss://socket.polygon.io/stocks",
"stock_delayed": "wss://delayed.polygon.io/stocks",
"options": "wss://socket.polygon.io/options",
"options_delayed": "wss://delayed.polygon.io/options",
"fx": "wss://socket.polygon.io/forex",
"crypto": "wss://socket.polygon.io/crypto",
"index": "wss://socket.polygon.io/indices",
"index_delayed": "wss://delayed.polygon.io/indices",
}
logger = get_logger("openbb.websocket.polygon")
queue = MessageQueue(logger=logger)
command_queue = MessageQueue(logger=logger)
@@ -29,8 +40,12 @@ command_queue = MessageQueue(logger=logger)
kwargs = parse_kwargs()
CONNECT_KWARGS = kwargs.pop("connect_kwargs", {})
FEED = kwargs.pop("feed", None)
ASSET_TYPE = kwargs.pop("asset_type", None)
ASSET_TYPE = kwargs.pop("asset_type", "crypto")
kwargs["results_file"] = os.path.abspath(kwargs.get("results_file"))
URL = URL_MAP.get(ASSET_TYPE)
if not URL:
raise ValueError("Invalid asset type provided.")
DATABASE = Database(
results_file=kwargs["results_file"],
@@ -208,7 +223,7 @@ async def connect_and_stream():
connect_kwargs["close_timeout"] = None
try:
async with websockets.connect(kwargs["url"], **connect_kwargs) as websocket:
async with websockets.connect(URL, **connect_kwargs) as websocket:
await login(websocket)
response = await websocket.recv()

2
openbb_platform/providers/polygon/tests/test_polygon_fetchers.py
View File

@@ -6,6 +6,7 @@ from unittest.mock import MagicMock, patch
import pytest
from openbb_core.app.service.user_service import UserService
from openbb_core.provider.utils.websockets.client import WebSocketClient
from openbb_polygon.models.balance_sheet import PolygonBalanceSheetFetcher
from openbb_polygon.models.cash_flow import PolygonCashFlowStatementFetcher
from openbb_polygon.models.company_news import PolygonCompanyNewsFetcher
@@ -25,7 +26,6 @@ from openbb_polygon.models.websocket_connection import (
PolygonWebSocketData,
PolygonWebSocketFetcher,
)
from openbb_websockets.client import WebSocketClient
test_credentials = UserService().default_user_settings.credentials.model_dump(
mode="json"

12
openbb_platform/providers/tiingo/openbb_tiingo/models/websocket_connection.py
View File

@@ -10,7 +10,7 @@ from openbb_core.provider.abstract.fetcher import Fetcher
from openbb_core.provider.utils.descriptions import (
QUERY_DESCRIPTIONS,
)
from openbb_websockets.client import WebSocketClient
from openbb_core.provider.utils.websockets.client import WebSocketClient
from openbb_websockets.models import (
WebSocketConnection,
WebSocketData,
@@ -18,12 +18,6 @@ from openbb_websockets.models import (
)
from pydantic import Field, field_validator, model_validator
URL_MAP = {
"stock": "wss://api.tiingo.com/iex",
"fx": "wss://api.tiingo.com/fx",
"crypto": "wss://api.tiingo.com/crypto",
}
# These are the data array order of definitions.
IEX_FIELDS = [
"type",
@@ -244,7 +238,6 @@ class TiingoWebSocketFetcher(
from asyncio import sleep
api_key = credentials.get("tiingo_token") if credentials else ""
url = URL_MAP[query.asset_type]
threshold_level = (
5
if query.asset_type == "fx" or query.feed == "trade"
@@ -254,12 +247,11 @@ class TiingoWebSocketFetcher(
else 0
)
)
symbol = query.symbol.lower()
kwargs = {
"url": url,
"api_key": api_key,
"asset_type": query.asset_type,
"threshold_level": threshold_level,
"connect_kwargs": query.connect_kwargs,
}

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

@@ -19,6 +19,12 @@ from openbb_core.provider.utils.websockets.message_queue import MessageQueue
from openbb_tiingo.models.websocket_connection import TiingoWebSocketData
from pydantic import ValidationError
URL_MAP = {
"stock": "wss://api.tiingo.com/iex",
"fx": "wss://api.tiingo.com/fx",
"crypto": "wss://api.tiingo.com/crypto",
}
# These are the data array definitions.
IEX_FIELDS = [
"type",
@@ -74,6 +80,10 @@ queue = MessageQueue(logger=logger)
kwargs = parse_kwargs()
CONNECT_KWARGS = kwargs.pop("connect_kwargs", {})
kwargs["results_file"] = os.path.abspath(kwargs["results_file"])
URL = URL_MAP.get(kwargs.pop("asset_type", "crypto"))
if not URL:
raise ValueError("Invalid asset type provided.")
DATABASE = Database(
results_file=kwargs["results_file"],
@@ -86,8 +96,6 @@ DATABASE = Database(
# 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):
"""Update the symbols to subscribe to."""
url = kwargs["url"]
if not SUBSCRIPTION_ID:
logger.error(
"PROVIDER ERROR: Must be assigned a subscription ID to update symbols. Try logging in."
@@ -103,7 +111,7 @@ async def update_symbols(symbol, event):
},
}
async with websockets.connect(url) as websocket:
async with websockets.connect(URL) as websocket:
await websocket.send(json.dumps(update_event))
response = await websocket.recv()
message = json.loads(response)
@@ -228,7 +236,7 @@ async def connect_and_stream():
try:
try:
async with websockets.connect(kwargs["url"], **connect_kwargs) as websocket:
async with websockets.connect(URL, **connect_kwargs) as websocket:
logger.info("PROVIDER INFO: WebSocket connection established.")
await websocket.send(json.dumps(subscribe_event))
while True:

2
openbb_platform/providers/tiingo/tests/test_tiingo_fetchers.py
View File

@@ -6,6 +6,7 @@ from unittest.mock import MagicMock, patch
import pytest
from openbb_core.app.service.user_service import UserService
from openbb_provider.utils.websockets.client import WebSocketClient
from openbb_tiingo.models.company_news import TiingoCompanyNewsFetcher
from openbb_tiingo.models.crypto_historical import TiingoCryptoHistoricalFetcher
from openbb_tiingo.models.currency_historical import TiingoCurrencyHistoricalFetcher
@@ -17,7 +18,6 @@ from openbb_tiingo.models.websocket_connection import (
TiingoWebSocketFetcher,
)
from openbb_tiingo.models.world_news import TiingoWorldNewsFetcher
from openbb_websockets.client import WebSocketClient
test_credentials = UserService().default_user_settings.credentials.model_dump(
mode="json"