Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: global logs context #26418

Merged
merged 4 commits into from
Jan 16, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions superset/commands/report/execute.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -81,6 +82,7 @@ class BaseReportState:
current_states: list[ReportState] = []
initial: bool = False

@logs_context()
def __init__(
self,
session: Session,
Expand Down
9 changes: 8 additions & 1 deletion superset/reports/notifications/slack.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 (
Expand Down Expand Up @@ -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):
Expand All @@ -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,
Expand Down
82 changes: 81 additions & 1 deletion superset/utils/decorators.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,16 +16,20 @@
# under the License.
from __future__ import annotations

import logging
import time
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

logger = logging.getLogger(__name__)

if TYPE_CHECKING:
from superset.stats_logger import BaseStatsLogger

Expand Down Expand Up @@ -61,6 +65,82 @@
return decorate


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

def decorate(f: Callable[..., Any]) -> Callable[..., Any]:
def wrapped(*args: Any, **kwargs: Any) -> Any:
if not hasattr(g, "logs_context"):
g.logs_context = {}

Check warning on line 80 in superset/utils/decorators.py

View check run for this annotation

Codecov / codecov/patch

superset/utils/decorators.py#L80

Added line #L80 was not covered by tests

# 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
}

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, AttributeError):
# do nothing if the key doesn't exist
# or context is not callable
logger.warning("Invalid data was passed to the logs context decorator")

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."""
Expand Down
86 changes: 86 additions & 0 deletions tests/unit_tests/notifications/slack_tests.py
Original file line number Diff line number Diff line change
@@ -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> |
```
""",
)
Loading
Loading