Skip to content

Commit

Permalink
fix(alerts): execute query as report executor (#22167)
Browse files Browse the repository at this point in the history
  • Loading branch information
villebro authored Nov 18, 2022
1 parent 25114a7 commit c3f9f0b
Show file tree
Hide file tree
Showing 2 changed files with 74 additions and 6 deletions.
10 changes: 4 additions & 6 deletions superset/reports/commands/alert.py
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,7 @@
from celery.exceptions import SoftTimeLimitExceeded
from flask_babel import lazy_gettext as _

from superset import app, jinja_context, security_manager
from superset import app, jinja_context
from superset.commands.base import BaseCommand
from superset.reports.commands.exceptions import (
AlertQueryError,
Expand All @@ -36,6 +36,7 @@
AlertValidatorConfigError,
)
from superset.reports.models import ReportSchedule, ReportScheduleValidatorType
from superset.reports.utils import get_executor
from superset.utils.core import override_user
from superset.utils.retries import retry_call

Expand Down Expand Up @@ -148,11 +149,8 @@ def _execute_query(self) -> pd.DataFrame:
rendered_sql, ALERT_SQL_LIMIT
)

with override_user(
security_manager.find_user(
username=app.config["THUMBNAIL_SELENIUM_USER"]
)
):
user = get_executor(self._report_schedule)
with override_user(user):
start = default_timer()
df = self._report_schedule.database.get_df(sql=limited_rendered_sql)
stop = default_timer()
Expand Down
70 changes: 70 additions & 0 deletions tests/integration_tests/reports/alert_tests.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,10 +15,80 @@
# specific language governing permissions and limitations
# under the License.
# pylint: disable=invalid-name, unused-argument, import-outside-toplevel
from contextlib import nullcontext
from typing import List, Optional, Union

import pandas as pd
import pytest
from pytest_mock import MockFixture

from superset.reports.commands.exceptions import AlertQueryError
from superset.reports.models import ReportCreationMethod, ReportScheduleType
from superset.reports.types import ReportScheduleExecutor
from superset.utils.database import get_example_database
from tests.integration_tests.test_app import app


@pytest.mark.parametrize(
"owner_names,creator_name,config,expected_result",
[
(["gamma"], None, [ReportScheduleExecutor.SELENIUM], "admin"),
(["gamma"], None, [ReportScheduleExecutor.OWNER], "gamma"),
(["alpha", "gamma"], "gamma", [ReportScheduleExecutor.CREATOR_OWNER], "gamma"),
(["alpha", "gamma"], "alpha", [ReportScheduleExecutor.CREATOR_OWNER], "alpha"),
(
["alpha", "gamma"],
"admin",
[ReportScheduleExecutor.CREATOR_OWNER],
AlertQueryError(),
),
],
)
def test_execute_query_as_report_executor(
owner_names: List[str],
creator_name: Optional[str],
config: List[ReportScheduleExecutor],
expected_result: Union[str, Exception],
mocker: MockFixture,
app_context: None,
get_user,
) -> None:

from superset.reports.commands.alert import AlertCommand
from superset.reports.models import ReportSchedule

with app.app_context():
original_config = app.config["ALERT_REPORTS_EXECUTE_AS"]
app.config["ALERT_REPORTS_EXECUTE_AS"] = config
owners = [get_user(owner_name) for owner_name in owner_names]
report_schedule = ReportSchedule(
created_by=get_user(creator_name) if creator_name else None,
owners=owners,
type=ReportScheduleType.ALERT,
description="description",
crontab="0 9 * * *",
creation_method=ReportCreationMethod.ALERTS_REPORTS,
sql="SELECT 1",
grace_period=14400,
working_timeout=3600,
database=get_example_database(),
validator_config_json='{"op": "==", "threshold": 1}',
)
command = AlertCommand(report_schedule=report_schedule)
override_user_mock = mocker.patch(
"superset.reports.commands.alert.override_user"
)
cm = (
pytest.raises(type(expected_result))
if isinstance(expected_result, Exception)
else nullcontext()
)
with cm:
command.run()
assert override_user_mock.call_args[0][0].username == expected_result

app.config["ALERT_REPORTS_EXECUTE_AS"] = original_config


def test_execute_query_succeeded_no_retry(
mocker: MockFixture, app_context: None
Expand Down

0 comments on commit c3f9f0b

Please # to comment.