From 3fbfb1cc18c155f1f5f81e8c152cd62d3afdb4b9 Mon Sep 17 00:00:00 2001 From: Elizabeth Thompson Date: Wed, 3 Jan 2024 15:48:09 -0800 Subject: [PATCH 1/4] add logs_context decorator --- superset/utils/decorators.py | 50 +++++++++++++++++++- tests/unit_tests/utils/test_decorators.py | 57 +++++++++++++++++++++++ 2 files changed, 106 insertions(+), 1 deletion(-) diff --git a/superset/utils/decorators.py b/superset/utils/decorators.py index 9c21e3b5ec5a5..35a8136a830e2 100644 --- a/superset/utils/decorators.py +++ b/superset/utils/decorators.py @@ -20,8 +20,9 @@ from collections.abc import Iterator from contextlib import contextmanager from typing import Any, Callable, TYPE_CHECKING +from uuid import UUID -from flask import current_app, Response +from flask import current_app, g, Response from superset.utils import core as utils from superset.utils.dates import now_as_float @@ -61,6 +62,53 @@ def wrapped(*args: Any, **kwargs: Any) -> Any: return decorate +def logs_context(**ctx_kwargs: int | str | UUID | None) -> Callable[..., Any]: + """ + Takes arguments and adds them to the global logs_context. + This is for logging purposes only and values should not be relied on or mutated + """ + + def decorate(f: Callable[..., Any]) -> Callable[..., Any]: + def wrapped(*args: Any, **kwargs: Any) -> Any: + if not hasattr(g, "logs_context"): + g.logs_context = {} + available_logs_context_keys = [ + "slice_id", + "dashboard_id", + "execution_id", + "report_schedule_id", + ] + logs_context_data = { + key: val + for key, val in kwargs.items() + if key in available_logs_context_keys + } + + # if keys are passed in to decorator directly, add them to logs_context + # by overriding values from kwargs + # e.g. @logs_context(slice_id=1, dashboard_id=1) + for key in available_logs_context_keys: + if ctx_kwargs.get(key) is not None: + logs_context_data[key] = ctx_kwargs[key] + try: + # override value from local kwargs from function execution if it exists + # e.g. @logs_context(slice_id=1, dashboard_id=1) + # def my_func(slice_id=None, **kwargs) + # + # my_func(slice_id=2) + logs_context_data[key] = locals()["kwargs"][key] + except KeyError: + # do nothing if the key doesn't exist + pass + + g.logs_context.update(logs_context_data) + return f(*args, **kwargs) + + return wrapped + + return decorate + + @contextmanager def stats_timing(stats_key: str, stats_logger: BaseStatsLogger) -> Iterator[float]: """Provide a transactional scope around a series of operations.""" diff --git a/tests/unit_tests/utils/test_decorators.py b/tests/unit_tests/utils/test_decorators.py index f600334414eeb..32dba078d48a5 100644 --- a/tests/unit_tests/utils/test_decorators.py +++ b/tests/unit_tests/utils/test_decorators.py @@ -16,12 +16,14 @@ # under the License. +import uuid from contextlib import nullcontext from inspect import isclass from typing import Any, Optional from unittest.mock import call, Mock, patch import pytest +from flask import g from superset import app from superset.utils import decorators @@ -85,3 +87,58 @@ def my_func(response: ResponseValues, *args: Any, **kwargs: Any) -> str: with cm: my_func(response_value, 1, 2) mock.assert_called_once_with(expected_result, 1) + + +def test_context_decorator() -> None: + @decorators.logs_context() + def myfunc(*args, **kwargs) -> str: + return "test" + + # should be able to add values to the decorator function directly + @decorators.logs_context(slice_id=1, dashboard_id=1, execution_id=uuid.uuid4()) + def myfunc_with_kwargs(*args, **kwargs) -> str: + return "test" + + # should not add any data to the global.logs_context scope + myfunc(1, 1) + assert g.logs_context == {} + + # should add dashboard_id to the global.logs_context scope + myfunc(1, 1, dashboard_id=1) + assert g.logs_context == {"dashboard_id": 1} + g.logs_context = {} + + # should add slice_id to the global.logs_context scope + myfunc(1, 1, slice_id=1) + assert g.logs_context == {"slice_id": 1} + g.logs_context = {} + + # should add execution_id to the global.logs_context scope + myfunc(1, 1, execution_id=1) + assert g.logs_context == {"execution_id": 1} + g.logs_context = {} + + # should add all three to the global.logs_context scope + myfunc(1, 1, dashboard_id=1, slice_id=1, execution_id=1) + assert g.logs_context == {"dashboard_id": 1, "slice_id": 1, "execution_id": 1} + g.logs_context = {} + + # should overwrite existing values in the global.logs_context scope + g.logs_context = {"dashboard_id": 2, "slice_id": 2, "execution_id": 2} + myfunc(1, 1, dashboard_id=3, slice_id=3, execution_id=3) + assert g.logs_context == {"dashboard_id": 3, "slice_id": 3, "execution_id": 3} + g.logs_context = {} + + # should be able to add values to the decorator function directly + myfunc_with_kwargs() + assert g.logs_context["dashboard_id"] == 1 + assert g.logs_context["slice_id"] == 1 + assert isinstance(g.logs_context["execution_id"], uuid.UUID) + g.logs_context = {} + + # should be able to add values to the decorator function directly + # and overwrite with any kwargs passed in + myfunc_with_kwargs(execution_id=4) + + assert g.logs_context == {"dashboard_id": 1, "slice_id": 1, "execution_id": 4} + g.logs_context = {} From 262744eb8ff8c386eed39d137f60647ecd8090c4 Mon Sep 17 00:00:00 2001 From: Elizabeth Thompson Date: Thu, 4 Jan 2024 17:26:56 -0800 Subject: [PATCH 2/4] add example of context_log --- superset/commands/report/execute.py | 2 + superset/reports/notifications/slack.py | 9 +- superset/utils/decorators.py | 3 +- tests/unit_tests/notifications/slack_tests.py | 87 +++++++++++++++++++ 4 files changed, 99 insertions(+), 2 deletions(-) create mode 100644 tests/unit_tests/notifications/slack_tests.py diff --git a/superset/commands/report/execute.py b/superset/commands/report/execute.py index d4b53e30dd846..349ef0c85b9a7 100644 --- a/superset/commands/report/execute.py +++ b/superset/commands/report/execute.py @@ -71,6 +71,7 @@ from superset.utils.celery import session_scope from superset.utils.core import HeaderDataType, override_user from superset.utils.csv import get_chart_csv_data, get_chart_dataframe +from superset.utils.decorators import logs_context from superset.utils.screenshots import ChartScreenshot, DashboardScreenshot from superset.utils.urls import get_url_path @@ -81,6 +82,7 @@ class BaseReportState: current_states: list[ReportState] = [] initial: bool = False + @logs_context() def __init__( self, session: Session, diff --git a/superset/reports/notifications/slack.py b/superset/reports/notifications/slack.py index fbae398bc50d6..41ed2ceeb37ad 100644 --- a/superset/reports/notifications/slack.py +++ b/superset/reports/notifications/slack.py @@ -22,6 +22,7 @@ import backoff import pandas as pd +from flask import g from flask_babel import gettext as __ from slack_sdk import WebClient from slack_sdk.errors import ( @@ -175,6 +176,7 @@ def send(self) -> None: channel = self._get_channel() body = self._get_body() file_type = "csv" if self._content.csv else "png" + global_logs_context = getattr(g, "logs_context", {}) or {} try: token = app.config["SLACK_API_TOKEN"] if callable(token): @@ -192,7 +194,12 @@ def send(self) -> None: ) else: client.chat_postMessage(channel=channel, text=body) - logger.info("Report sent to slack") + logger.info( + "Report sent to slack", + extra={ + "execution_id": global_logs_context.get("execution_id"), + }, + ) except ( BotUserAccessError, SlackRequestError, diff --git a/superset/utils/decorators.py b/superset/utils/decorators.py index 35a8136a830e2..af67fc5c68086 100644 --- a/superset/utils/decorators.py +++ b/superset/utils/decorators.py @@ -91,7 +91,8 @@ def wrapped(*args: Any, **kwargs: Any) -> Any: if ctx_kwargs.get(key) is not None: logs_context_data[key] = ctx_kwargs[key] try: - # override value from local kwargs from function execution if it exists + # override value from local kwargs from + # function execution if it exists # e.g. @logs_context(slice_id=1, dashboard_id=1) # def my_func(slice_id=None, **kwargs) # diff --git a/tests/unit_tests/notifications/slack_tests.py b/tests/unit_tests/notifications/slack_tests.py new file mode 100644 index 0000000000000..f37998ea8ccb4 --- /dev/null +++ b/tests/unit_tests/notifications/slack_tests.py @@ -0,0 +1,87 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +import uuid +from unittest.mock import MagicMock, patch + +import pandas as pd +from flask import g + + +@patch("superset.reports.notifications.slack.logger") +def test_send_slack( + logger_mock: MagicMock, +) -> None: + # `superset.models.helpers`, a dependency of following imports, + # requires app context + from superset.reports.models import ReportRecipients, ReportRecipientType + from superset.reports.notifications.base import NotificationContent + from superset.reports.notifications.slack import SlackNotification, WebClient + + execution_id = uuid.uuid4() + g.logs_context = {"execution_id": execution_id} + content = NotificationContent( + name="test alert", + header_data={ + "notification_format": "PNG", + "notification_type": "Alert", + "owners": [1], + "notification_source": None, + "chart_id": None, + "dashboard_id": None, + }, + embedded_data=pd.DataFrame( + { + "A": [1, 2, 3], + "B": [4, 5, 6], + "C": ["111", "222", '333'], + } + ), + description='

This is a test alert


', + ) + with patch.object( + WebClient, "chat_postMessage", return_value=True + ) as chat_post_message_mock: + SlackNotification( + recipient=ReportRecipients( + type=ReportRecipientType.SLACK, + recipient_config_json='{"target": "some_channel"}', + ), + content=content, + ).send() + logger_mock.info.assert_called_with( + "Report sent to slack", extra={"execution_id": execution_id} + ) + chat_post_message_mock.assert_called_with( + channel="some_channel", + text="""*test alert* + +

This is a test alert


+ + + +``` +| | A | B | C | +|---:|----:|----:|:-----------------------------------------| +| 0 | 1 | 4 | 111 | +| 1 | 2 | 5 | 222 | +| 2 | 3 | 6 | 333 | +``` +""", + ) + + # reset g.logs_context + g.logs_context = None From 519ea973c26cf011b8de1cc6da01c82eb7c414c1 Mon Sep 17 00:00:00 2001 From: Elizabeth Thompson Date: Fri, 12 Jan 2024 12:23:17 -0800 Subject: [PATCH 3/4] improve tests and decorator --- superset/utils/decorators.py | 64 ++++++++++---- tests/unit_tests/notifications/slack_tests.py | 7 +- tests/unit_tests/utils/test_decorators.py | 83 ++++++++++++++----- 3 files changed, 111 insertions(+), 43 deletions(-) diff --git a/superset/utils/decorators.py b/superset/utils/decorators.py index af67fc5c68086..fe8e292cffc97 100644 --- a/superset/utils/decorators.py +++ b/superset/utils/decorators.py @@ -62,7 +62,10 @@ def wrapped(*args: Any, **kwargs: Any) -> Any: return decorate -def logs_context(**ctx_kwargs: int | str | UUID | None) -> Callable[..., Any]: +def logs_context( + context_func: Callable[..., dict[Any, Any]] | None = None, + **ctx_kwargs: int | str | UUID | None, +) -> Callable[..., Any]: """ Takes arguments and adds them to the global logs_context. This is for logging purposes only and values should not be relied on or mutated @@ -72,35 +75,60 @@ def decorate(f: Callable[..., Any]) -> Callable[..., Any]: def wrapped(*args: Any, **kwargs: Any) -> Any: if not hasattr(g, "logs_context"): g.logs_context = {} + + # limit data that can be saved to logs_context + # in order to prevent antipatterns available_logs_context_keys = [ "slice_id", "dashboard_id", + "dataset_id", "execution_id", "report_schedule_id", ] + # set value from kwargs from + # wrapper function if it exists + # e.g. @logs_context() + # def my_func(slice_id=None, **kwargs) + # + # my_func(slice_id=2) logs_context_data = { key: val for key, val in kwargs.items() if key in available_logs_context_keys + if val is not None } - # if keys are passed in to decorator directly, add them to logs_context - # by overriding values from kwargs - # e.g. @logs_context(slice_id=1, dashboard_id=1) - for key in available_logs_context_keys: - if ctx_kwargs.get(key) is not None: - logs_context_data[key] = ctx_kwargs[key] - try: - # override value from local kwargs from - # function execution if it exists - # e.g. @logs_context(slice_id=1, dashboard_id=1) - # def my_func(slice_id=None, **kwargs) - # - # my_func(slice_id=2) - logs_context_data[key] = locals()["kwargs"][key] - except KeyError: - # do nothing if the key doesn't exist - pass + try: + # if keys are passed in to decorator directly, add them to logs_context + # by overriding values from kwargs + # e.g. @logs_context(slice_id=1, dashboard_id=1) + logs_context_data.update( + { + key: ctx_kwargs.get(key) + for key in available_logs_context_keys + if ctx_kwargs.get(key) is not None + } + ) + + if context_func is not None: + # if a context function is passed in, call it and add the + # returned values to logs_context + # context_func=lambda *args, **kwargs: { + # "slice_id": 1, "dashboard_id": 1 + # } + logs_context_data.update( + { + key: value + for key, value in context_func(*args, **kwargs).items() + if key in available_logs_context_keys + if value is not None + } + ) + + except (TypeError, KeyError): + # do nothing if the key doesn't exist + # or context is not callable + pass g.logs_context.update(logs_context_data) return f(*args, **kwargs) diff --git a/tests/unit_tests/notifications/slack_tests.py b/tests/unit_tests/notifications/slack_tests.py index f37998ea8ccb4..7e6cc3afc1ff4 100644 --- a/tests/unit_tests/notifications/slack_tests.py +++ b/tests/unit_tests/notifications/slack_tests.py @@ -21,9 +21,11 @@ from flask import g +@patch("superset.reports.notifications.slack.g") @patch("superset.reports.notifications.slack.logger") def test_send_slack( logger_mock: MagicMock, + flask_global_mock: MagicMock, ) -> None: # `superset.models.helpers`, a dependency of following imports, # requires app context @@ -32,7 +34,7 @@ def test_send_slack( from superset.reports.notifications.slack import SlackNotification, WebClient execution_id = uuid.uuid4() - g.logs_context = {"execution_id": execution_id} + flask_global_mock.logs_context = {"execution_id": execution_id} content = NotificationContent( name="test alert", header_data={ @@ -82,6 +84,3 @@ def test_send_slack( ``` """, ) - - # reset g.logs_context - g.logs_context = None diff --git a/tests/unit_tests/utils/test_decorators.py b/tests/unit_tests/utils/test_decorators.py index 32dba078d48a5..8e73d1f428d1c 100644 --- a/tests/unit_tests/utils/test_decorators.py +++ b/tests/unit_tests/utils/test_decorators.py @@ -23,7 +23,6 @@ from unittest.mock import call, Mock, patch import pytest -from flask import g from superset import app from superset.utils import decorators @@ -89,7 +88,10 @@ def my_func(response: ResponseValues, *args: Any, **kwargs: Any) -> str: mock.assert_called_once_with(expected_result, 1) -def test_context_decorator() -> None: +@patch("superset.utils.decorators.g") +def test_context_decorator(flask_g_mock) -> None: + flask_g_mock.logs_context = {} + @decorators.logs_context() def myfunc(*args, **kwargs) -> str: return "test" @@ -99,46 +101,85 @@ def myfunc(*args, **kwargs) -> str: def myfunc_with_kwargs(*args, **kwargs) -> str: return "test" + @decorators.logs_context(bad_context=1) + def myfunc_with_dissallowed_kwargs(*args, **kwargs) -> str: + return "test" + + @decorators.logs_context( + context_func=lambda *args, **kwargs: {"slice_id": kwargs["chart_id"]} + ) + def myfunc_with_context(*args, **kwargs) -> str: + return "test" + # should not add any data to the global.logs_context scope myfunc(1, 1) - assert g.logs_context == {} + assert flask_g_mock.logs_context == {} # should add dashboard_id to the global.logs_context scope myfunc(1, 1, dashboard_id=1) - assert g.logs_context == {"dashboard_id": 1} - g.logs_context = {} + assert flask_g_mock.logs_context == {"dashboard_id": 1} + # reset logs_context + flask_g_mock.logs_context = {} # should add slice_id to the global.logs_context scope myfunc(1, 1, slice_id=1) - assert g.logs_context == {"slice_id": 1} - g.logs_context = {} + assert flask_g_mock.logs_context == {"slice_id": 1} + # reset logs_context + flask_g_mock.logs_context = {} # should add execution_id to the global.logs_context scope myfunc(1, 1, execution_id=1) - assert g.logs_context == {"execution_id": 1} - g.logs_context = {} + assert flask_g_mock.logs_context == {"execution_id": 1} + # reset logs_context + flask_g_mock.logs_context = {} # should add all three to the global.logs_context scope myfunc(1, 1, dashboard_id=1, slice_id=1, execution_id=1) - assert g.logs_context == {"dashboard_id": 1, "slice_id": 1, "execution_id": 1} - g.logs_context = {} + assert flask_g_mock.logs_context == { + "dashboard_id": 1, + "slice_id": 1, + "execution_id": 1, + } + # reset logs_context + flask_g_mock.logs_context = {} # should overwrite existing values in the global.logs_context scope - g.logs_context = {"dashboard_id": 2, "slice_id": 2, "execution_id": 2} + flask_g_mock.logs_context = {"dashboard_id": 2, "slice_id": 2, "execution_id": 2} myfunc(1, 1, dashboard_id=3, slice_id=3, execution_id=3) - assert g.logs_context == {"dashboard_id": 3, "slice_id": 3, "execution_id": 3} - g.logs_context = {} + assert flask_g_mock.logs_context == { + "dashboard_id": 3, + "slice_id": 3, + "execution_id": 3, + } + # reset logs_context + flask_g_mock.logs_context = {} + + # should not add a value that does not exist in the global.logs_context scope + myfunc_with_dissallowed_kwargs() + assert flask_g_mock.logs_context == {} + # reset logs_context + flask_g_mock.logs_context = {} # should be able to add values to the decorator function directly myfunc_with_kwargs() - assert g.logs_context["dashboard_id"] == 1 - assert g.logs_context["slice_id"] == 1 - assert isinstance(g.logs_context["execution_id"], uuid.UUID) - g.logs_context = {} + assert flask_g_mock.logs_context["dashboard_id"] == 1 + assert flask_g_mock.logs_context["slice_id"] == 1 + assert isinstance(flask_g_mock.logs_context["execution_id"], uuid.UUID) + # reset logs_context + flask_g_mock.logs_context = {} # should be able to add values to the decorator function directly - # and overwrite with any kwargs passed in + # and it will overwrite any kwargs passed into the decorated function myfunc_with_kwargs(execution_id=4) - assert g.logs_context == {"dashboard_id": 1, "slice_id": 1, "execution_id": 4} - g.logs_context = {} + assert flask_g_mock.logs_context["dashboard_id"] == 1 + assert flask_g_mock.logs_context["slice_id"] == 1 + assert isinstance(flask_g_mock.logs_context["execution_id"], uuid.UUID) + # reset logs_context + flask_g_mock.logs_context = {} + + # should be able to pass a callable context to the decorator + myfunc_with_context(chart_id=1) + assert flask_g_mock.logs_context == {"slice_id": 1} + # reset logs_context + flask_g_mock.logs_context = {} From 329f68446f867d070ae9a882c293fcfd39658f49 Mon Sep 17 00:00:00 2001 From: Elizabeth Thompson Date: Fri, 12 Jan 2024 14:50:35 -0800 Subject: [PATCH 4/4] add more tests and logging --- superset/utils/decorators.py | 7 +- tests/unit_tests/utils/test_decorators.py | 124 +++++++++++++++++----- 2 files changed, 100 insertions(+), 31 deletions(-) diff --git a/superset/utils/decorators.py b/superset/utils/decorators.py index fe8e292cffc97..43b6909caf3fb 100644 --- a/superset/utils/decorators.py +++ b/superset/utils/decorators.py @@ -16,6 +16,7 @@ # under the License. from __future__ import annotations +import logging import time from collections.abc import Iterator from contextlib import contextmanager @@ -27,6 +28,8 @@ from superset.utils import core as utils from superset.utils.dates import now_as_float +logger = logging.getLogger(__name__) + if TYPE_CHECKING: from superset.stats_logger import BaseStatsLogger @@ -125,10 +128,10 @@ def wrapped(*args: Any, **kwargs: Any) -> Any: } ) - except (TypeError, KeyError): + except (TypeError, KeyError, AttributeError): # do nothing if the key doesn't exist # or context is not callable - pass + logger.warning("Invalid data was passed to the logs context decorator") g.logs_context.update(logs_context_data) return f(*args, **kwargs) diff --git a/tests/unit_tests/utils/test_decorators.py b/tests/unit_tests/utils/test_decorators.py index 8e73d1f428d1c..6cfe55b6411d9 100644 --- a/tests/unit_tests/utils/test_decorators.py +++ b/tests/unit_tests/utils/test_decorators.py @@ -90,13 +90,10 @@ def my_func(response: ResponseValues, *args: Any, **kwargs: Any) -> str: @patch("superset.utils.decorators.g") def test_context_decorator(flask_g_mock) -> None: - flask_g_mock.logs_context = {} - @decorators.logs_context() def myfunc(*args, **kwargs) -> str: return "test" - # should be able to add values to the decorator function directly @decorators.logs_context(slice_id=1, dashboard_id=1, execution_id=uuid.uuid4()) def myfunc_with_kwargs(*args, **kwargs) -> str: return "test" @@ -111,39 +108,36 @@ def myfunc_with_dissallowed_kwargs(*args, **kwargs) -> str: def myfunc_with_context(*args, **kwargs) -> str: return "test" - # should not add any data to the global.logs_context scope + ### should not add any data to the global.logs_context scope + flask_g_mock.logs_context = {} myfunc(1, 1) assert flask_g_mock.logs_context == {} - # should add dashboard_id to the global.logs_context scope + ### should add dashboard_id to the global.logs_context scope + flask_g_mock.logs_context = {} myfunc(1, 1, dashboard_id=1) assert flask_g_mock.logs_context == {"dashboard_id": 1} - # reset logs_context - flask_g_mock.logs_context = {} - # should add slice_id to the global.logs_context scope + ### should add slice_id to the global.logs_context scope + flask_g_mock.logs_context = {} myfunc(1, 1, slice_id=1) assert flask_g_mock.logs_context == {"slice_id": 1} - # reset logs_context - flask_g_mock.logs_context = {} - # should add execution_id to the global.logs_context scope + ### should add execution_id to the global.logs_context scope + flask_g_mock.logs_context = {} myfunc(1, 1, execution_id=1) assert flask_g_mock.logs_context == {"execution_id": 1} - # reset logs_context - flask_g_mock.logs_context = {} - # should add all three to the global.logs_context scope + ### should add all three to the global.logs_context scope + flask_g_mock.logs_context = {} myfunc(1, 1, dashboard_id=1, slice_id=1, execution_id=1) assert flask_g_mock.logs_context == { "dashboard_id": 1, "slice_id": 1, "execution_id": 1, } - # reset logs_context - flask_g_mock.logs_context = {} - # should overwrite existing values in the global.logs_context scope + ### should overwrite existing values in the global.logs_context scope flask_g_mock.logs_context = {"dashboard_id": 2, "slice_id": 2, "execution_id": 2} myfunc(1, 1, dashboard_id=3, slice_id=3, execution_id=3) assert flask_g_mock.logs_context == { @@ -151,35 +145,107 @@ def myfunc_with_context(*args, **kwargs) -> str: "slice_id": 3, "execution_id": 3, } - # reset logs_context + + ### Test when g.logs_context already exists + flask_g_mock.logs_context = {"slice_id": 2, "dashboard_id": 2} + args = (3, 4) + kwargs = {"slice_id": 3, "dashboard_id": 3} + myfunc(*args, **kwargs) + assert flask_g_mock.logs_context == {"slice_id": 3, "dashboard_id": 3} + + ### Test when kwargs contain additional keys flask_g_mock.logs_context = {} + args = (1, 2) + kwargs = { + "slice_id": 1, + "dashboard_id": 1, + "dataset_id": 1, + "execution_id": 1, + "report_schedule_id": 1, + "extra_key": 1, + } + myfunc(*args, **kwargs) + assert flask_g_mock.logs_context == { + "slice_id": 1, + "dashboard_id": 1, + "dataset_id": 1, + "execution_id": 1, + "report_schedule_id": 1, + } - # should not add a value that does not exist in the global.logs_context scope + ### should not add a value that does not exist in the global.logs_context scope + flask_g_mock.logs_context = {} myfunc_with_dissallowed_kwargs() assert flask_g_mock.logs_context == {} - # reset logs_context - flask_g_mock.logs_context = {} - # should be able to add values to the decorator function directly + ### should be able to add values to the decorator function directly + flask_g_mock.logs_context = {} myfunc_with_kwargs() assert flask_g_mock.logs_context["dashboard_id"] == 1 assert flask_g_mock.logs_context["slice_id"] == 1 assert isinstance(flask_g_mock.logs_context["execution_id"], uuid.UUID) - # reset logs_context - flask_g_mock.logs_context = {} - # should be able to add values to the decorator function directly + ### should be able to add values to the decorator function directly # and it will overwrite any kwargs passed into the decorated function + flask_g_mock.logs_context = {} myfunc_with_kwargs(execution_id=4) assert flask_g_mock.logs_context["dashboard_id"] == 1 assert flask_g_mock.logs_context["slice_id"] == 1 assert isinstance(flask_g_mock.logs_context["execution_id"], uuid.UUID) - # reset logs_context - flask_g_mock.logs_context = {} - # should be able to pass a callable context to the decorator + ### should be able to pass a callable context to the decorator + flask_g_mock.logs_context = {} myfunc_with_context(chart_id=1) assert flask_g_mock.logs_context == {"slice_id": 1} - # reset logs_context + + ### Test when context_func returns additional keys + # it should use the context_func values flask_g_mock.logs_context = {} + args = (1, 2) + kwargs = {"slice_id": 1, "dashboard_id": 1} + + @decorators.logs_context( + context_func=lambda *args, **kwargs: { + "slice_id": 2, + "dashboard_id": 2, + "dataset_id": 2, + "execution_id": 2, + "report_schedule_id": 2, + "extra_key": 2, + } + ) + def myfunc_with_extra_keys_context(*args, **kwargs) -> str: + return "test" + + myfunc_with_extra_keys_context( + *args, + **kwargs, + ) + assert flask_g_mock.logs_context == { + "slice_id": 2, + "dashboard_id": 2, + "dataset_id": 2, + "execution_id": 2, + "report_schedule_id": 2, + } + + ### Test when context_func does not return a dictionary + flask_g_mock.logs_context = {} + + @decorators.logs_context(context_func=lambda: "foo") # type: ignore + def myfunc_with_bad_return_value() -> str: + return "test" + + myfunc_with_bad_return_value() + assert flask_g_mock.logs_context == {} + + ### Test when context_func is not callable + flask_g_mock.logs_context = {} + + @decorators.logs_context(context_func="foo") # type: ignore + def context_func_not_callable() -> str: + return "test" + + context_func_not_callable() + assert flask_g_mock.logs_context == {}