Log custom headers (#5966)

* add custom headers

* fix: send on startup and defaults

* pylint

* unittests
This commit is contained in:
montezdesousa
2024-01-17 12:22:29 +00:00
committed by GitHub
parent 6fb7860a3e
commit a2e3f04711
5 changed files with 57 additions and 8 deletions

View File

@@ -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(

View File

@@ -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,

View File

@@ -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,
)

View File

@@ -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

View File

@@ -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}"