feat: global logs context (#26418)

tests/unit_tests/notifications/slack_tests.py

@@ -0,0 +1,86 @@
+# 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.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
+    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()
+    flask_global_mock.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", '<a href="http://www.example.com">333</a>'],
+            }
+        ),
+        description='<p>This is <a href="#">a test</a> alert</p><br />',
+    )
+    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*
+
+<p>This is <a href="#">a test</a> alert</p><br />
+
+<None|Explore in Superset>
+
+```
+|    |   A |   B | C                                        |
+|---:|----:|----:|:-----------------------------------------|
+|  0 |   1 |   4 | 111                                      |
+|  1 |   2 |   5 | 222                                      |
+|  2 |   3 |   6 | <a href="http://www.example.com">333</a> |
+```
+""",
+        )

tests/unit_tests/utils/test_decorators.py

@@ -16,6 +16,7 @@
 # under the License.
 
 
+import uuid
 from contextlib import nullcontext
 from inspect import isclass
 from typing import Any, Optional
@@ -85,3 +86,166 @@ def test_statsd_gauge(
         with cm:
             my_func(response_value, 1, 2)
             mock.assert_called_once_with(expected_result, 1)
+
+
+@patch("superset.utils.decorators.g")
+def test_context_decorator(flask_g_mock) -> None:
+    @decorators.logs_context()
+    def myfunc(*args, **kwargs) -> str:
+        return "test"
+
+    @decorators.logs_context(slice_id=1, dashboard_id=1, execution_id=uuid.uuid4())
+    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
+    flask_g_mock.logs_context = {}
+    myfunc(1, 1)
+    assert flask_g_mock.logs_context == {}
+
+    ### 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}
+
+    ### 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}
+
+    ### 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}
+
+    ### 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,
+    }
+
+    ### 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 == {
+        "dashboard_id": 3,
+        "slice_id": 3,
+        "execution_id": 3,
+    }
+
+    ### 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
+    flask_g_mock.logs_context = {}
+    myfunc_with_dissallowed_kwargs()
+    assert flask_g_mock.logs_context == {}
+
+    ### 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)
+
+    ### 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)
+
+    ### 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}
+
+    ### 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 == {}