diff options
author | montezdesousa <79287829+montezdesousa@users.noreply.github.com> | 2024-01-17 12:22:29 +0000 |
---|---|---|
committer | GitHub <noreply@github.com> | 2024-01-17 12:22:29 +0000 |
commit | a2e3f04711c0eca27cde490a30596b4074894def (patch) | |
tree | 4d574ad723bbea283e790fdfeeb511e86c0d0bf3 | |
parent | 6fb7860a3e33ad5fc942b7261562d2ce7e00f7f6 (diff) |
Log custom headers (#5966)
* add custom headers
* fix: send on startup and defaults
* pylint
* unittests
5 files changed, 57 insertions, 8 deletions
diff --git a/openbb_platform/core/openbb_core/api/router/commands.py b/openbb_platform/core/openbb_core/api/router/commands.py index 99a1f28f0b7..bd627630542 100644 --- a/openbb_platform/core/openbb_core/api/router/commands.py +++ b/openbb_platform/core/openbb_core/api/router/commands.py @@ -1,9 +1,9 @@ import inspect from functools import partial, wraps from inspect import Parameter, Signature, signature -from typing import Any, Callable, Dict, Tuple, TypeVar +from typing import Any, Callable, Dict, Optional, Tuple, TypeVar -from fastapi import APIRouter, Depends +from fastapi import APIRouter, Depends, Header from fastapi.routing import APIRoute from openbb_core.app.charting_service import ChartingService from openbb_core.app.command_runner import CommandRunner @@ -70,6 +70,17 @@ def build_new_signature(path: str, func: Callable) -> Signature: ) ) + if custom_headers := SystemService().system_settings.api_settings.custom_headers: + for name, default in custom_headers.items(): + new_parameter_list.append( + Parameter( + name.replace("-", "_"), + kind=Parameter.POSITIONAL_OR_KEYWORD, + default=default, + annotation=Annotated[Optional[str], Header()], + ) + ) + if Env().API_AUTH: new_parameter_list.append( Parameter( diff --git a/openbb_platform/core/openbb_core/app/command_runner.py b/openbb_platform/core/openbb_core/app/command_runner.py index ae95313b5aa..8e41623f49a 100644 --- a/openbb_platform/core/openbb_core/app/command_runner.py +++ b/openbb_platform/core/openbb_core/app/command_runner.py @@ -305,6 +305,15 @@ class StaticCommandRunner: # If we're on the api we need to remove "chart" here because the parameter is added on # commands.py and the function signature does not expect "chart" kwargs.pop("chart", None) + # We also pop custom headers + + model_headers = ( + SystemService().system_settings.api_settings.custom_headers or {} + ) + custom_headers = { + name: kwargs.pop(name.replace("-", "_"), default) + for name, default in model_headers.items() or {} + } or None try: obbject = await cls._command( @@ -334,6 +343,7 @@ class StaticCommandRunner: func=func, kwargs=kwargs, exec_info=exc_info(), + custom_headers=custom_headers, ) return obbject @@ -414,6 +424,7 @@ class CommandRunner: def user_settings(self, user_settings: UserSettings) -> None: self._user_settings = user_settings + # pylint: disable=W1113 async def run( self, route: str, @@ -434,6 +445,7 @@ class CommandRunner: return await StaticCommandRunner.run(execution_context, *args, **kwargs) + # pylint: disable=W1113 def sync_run( self, route: str, diff --git a/openbb_platform/core/openbb_core/app/logs/logging_service.py b/openbb_platform/core/openbb_core/app/logs/logging_service.py index 854c7f55aac..569411081e1 100644 --- a/openbb_platform/core/openbb_core/app/logs/logging_service.py +++ b/openbb_platform/core/openbb_core/app/logs/logging_service.py @@ -135,7 +135,11 @@ class LoggingService(metaclass=SingletonMeta): return handlers_manager - def _log_startup(self, route: Optional[str] = None) -> None: + def _log_startup( + self, + route: Optional[str] = None, + custom_headers: Optional[Dict[str, Any]] = None, + ) -> None: """ Log startup information. """ @@ -165,6 +169,7 @@ class LoggingService(metaclass=SingletonMeta): else {} ), "SYSTEM": self._system_settings, + "custom_headers": custom_headers, }, default=to_jsonable_python, ), @@ -180,6 +185,7 @@ class LoggingService(metaclass=SingletonMeta): exec_info: Optional[ Tuple[Type[BaseException], BaseException, Optional[TracebackType]] ] = None, + custom_headers: Optional[Dict[str, Any]] = None, ): """ Log command output and relevant information. @@ -208,7 +214,7 @@ class LoggingService(metaclass=SingletonMeta): self._handlers_manager.update_handlers(self._logging_settings) if "login" in route: - self._log_startup(route) + self._log_startup(route, custom_headers) else: logger = logging.getLogger(__name__) @@ -230,6 +236,7 @@ class LoggingService(metaclass=SingletonMeta): "route": route, "input": kwargs, "error": str(openbb_error.original) if openbb_error else None, + "custom_headers": custom_headers, }, default=to_jsonable_python, ) diff --git a/openbb_platform/core/openbb_core/app/model/fast_api_settings.py b/openbb_platform/core/openbb_core/app/model/fast_api_settings.py index 565118dfc57..2795ef310ad 100644 --- a/openbb_platform/core/openbb_core/app/model/fast_api_settings.py +++ b/openbb_platform/core/openbb_core/app/model/fast_api_settings.py @@ -1,5 +1,5 @@ """FastAPI configuration settings model.""" -from typing import List +from typing import Dict, List, Optional from pydantic import BaseModel, ConfigDict, Field, computed_field @@ -41,6 +41,9 @@ class FastAPISettings(BaseModel): ) servers: List[Servers] = Field(default_factory=lambda: [Servers()]) cors: Cors = Field(default_factory=Cors) + custom_headers: Optional[Dict[str, str]] = Field( + default=None, description="Custom headers and respective default value." + ) @computed_field # type: ignore[misc] @property diff --git a/openbb_platform/core/tests/app/logs/test_logging_service.py b/openbb_platform/core/tests/app/logs/test_logging_service.py index 67d90d8ce71..bd7460bc3c2 100644 --- a/openbb_platform/core/tests/app/logs/test_logging_service.py +++ b/openbb_platform/core/tests/app/logs/test_logging_service.py @@ -86,13 +86,16 @@ def test_log_startup(logging_service): ) logging_service._system_settings = "your_system_settings" - logging_service._log_startup(route="test_route") + logging_service._log_startup( + route="test_route", custom_headers={"X-OpenBB-Test": "test"} + ) expected_log_data = { "route": "test_route", "PREFERENCES": "your_preferences", "KEYS": {"username": "defined", "password": "defined"}, "SYSTEM": "your_system_settings", + "custom_headers": {"X-OpenBB-Test": "test"}, } mock_info.assert_called_once_with( "STARTUP: %s ", @@ -102,7 +105,7 @@ def test_log_startup(logging_service): @pytest.mark.parametrize( - "user_settings, system_settings, route, func, kwargs, exec_info", + "user_settings, system_settings, route, func, kwargs, exec_info, custom_headers", [ ( "mock_settings", @@ -111,6 +114,7 @@ def test_log_startup(logging_service): "mock_func", {}, None, + None, ), ( "mock_settings", @@ -119,6 +123,7 @@ def test_log_startup(logging_service): "mock_func", {}, (OpenBBError, OpenBBError("mock_error")), + {"X-OpenBB-Test": "test"}, ), ( "mock_settings", @@ -127,11 +132,19 @@ def test_log_startup(logging_service): "mock_func", {}, None, + {"X-OpenBB-Test1": "test1", "X-OpenBB-Test2": "test2"}, ), ], ) def test_log( - logging_service, user_settings, system_settings, route, func, kwargs, exec_info + logging_service, + user_settings, + system_settings, + route, + func, + kwargs, + exec_info, + custom_headers, ): with patch( "openbb_core.app.logs.logging_service.LoggingSettings", @@ -148,6 +161,7 @@ def test_log( func=func, kwargs=kwargs, exec_info=exec_info, + custom_headers=custom_headers, ) assert mock_log_startup.assert_called_once @@ -165,6 +179,7 @@ def test_log( func=mock_callable, kwargs=kwargs, exec_info=exec_info, + custom_headers=custom_headers, ) message_label = "ERROR" if exec_info else "CMD" @@ -173,6 +188,7 @@ def test_log( "route": route, "input": kwargs, "error": str(exec_info[1]) if exec_info else None, + "custom_headers": custom_headers, } ) log_message = f"{message_label}: {log_message}" |