tests/integration_tests/reports/alert_tests.py - superset - Git at Google

 # 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.
 # pylint: disable=invalid-name, unused-argument, import-outside-toplevel
 import uuid
 from contextlib import nullcontext, suppress
 from typing import Optional, Union

 import pandas as pd
 import pytest
 from flask.ctx import AppContext
 from pytest_mock import MockerFixture

 from superset.commands.report.exceptions import AlertQueryError
 from superset.reports.models import ReportCreationMethod, ReportScheduleType
 from superset.tasks.types import ExecutorType, FixedExecutor
 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, [FixedExecutor("admin")], "admin"),
         (["gamma"], None, [ExecutorType.OWNER], "gamma"),
         (
             ["alpha", "gamma"],
             "gamma",
             [ExecutorType.CREATOR_OWNER],
             "gamma",
         ),
         (
             ["alpha", "gamma"],
             "alpha",
             [ExecutorType.CREATOR_OWNER],
             "alpha",
         ),
         (
             ["alpha", "gamma"],
             "admin",
             [ExecutorType.CREATOR_OWNER],
             AlertQueryError(),
         ),
         (["gamma"], None, [ExecutorType.CURRENT_USER], AlertQueryError()),
     ],
 )
 def test_execute_query_as_report_executor(
     owner_names: list[str],
     creator_name: Optional[str],
     config: list[ExecutorType],
     expected_result: Union[tuple[ExecutorType, str], Exception],
     mocker: MockerFixture,
     app_context: AppContext,
     get_user,
 ) -> None:
     from superset.commands.report.alert import AlertCommand
     from superset.reports.models import ReportSchedule

     original_config = app.config["ALERT_REPORTS_EXECUTORS"]
     app.config["ALERT_REPORTS_EXECUTORS"] = 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, execution_id=uuid.uuid4())
     override_user_mock = mocker.patch("superset.commands.report.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_EXECUTORS"] = original_config


 def test_execute_query_mutate_query_enabled(
     mocker: MockerFixture,
     app_context: AppContext,
     get_user,
 ) -> None:
     from superset.commands.report.alert import AlertCommand
     from superset.reports.models import ReportSchedule

     default_alert_mutate_ff = app.config["MUTATE_ALERT_QUERY"]

     app.config["MUTATE_ALERT_QUERY"] = True
     mocker.patch("superset.commands.report.alert.override_user")
     mock_df = mocker.MagicMock(spec=pd.DataFrame)
     mock_df.empty = True
     mock_database = get_example_database()
     mock_get_df = mocker.patch.object(mock_database, "get_df", return_value=mock_df)
     mock_limited_sql = mocker.patch.object(mock_database, "apply_limit_to_sql")
     mock_mutate_call = mocker.patch.object(mock_database, "mutate_sql_based_on_config")

     report_schedule = ReportSchedule(
         created_by=get_user("admin"),
         owners=[get_user("admin")],
         type=ReportScheduleType.ALERT,
         description="description",
         crontab="0 9 * * *",
         creation_method=ReportCreationMethod.ALERTS_REPORTS,
         sql="SELECT 1",
         grace_period=14400,
         working_timeout=3600,
         database=mock_database,
         validator_config_json='{"op": "==", "threshold": 1}',
     )
     AlertCommand(report_schedule=report_schedule, execution_id=uuid.uuid4()).run()

     mock_mutate_call.assert_called_once_with(mock_limited_sql.return_value)
     mock_get_df.assert_called_once_with(sql=mock_mutate_call.return_value)

     app.config["MUTATE_ALERT_QUERY"] = default_alert_mutate_ff


 def test_execute_query_mutate_query_disabled(
     mocker: MockerFixture,
     app_context: AppContext,
     get_user,
 ) -> None:
     from superset.commands.report.alert import AlertCommand
     from superset.reports.models import ReportSchedule

     default_alert_mutate_ff = app.config["MUTATE_ALERT_QUERY"]

     app.config["MUTATE_ALERT_QUERY"] = False
     mocker.patch("superset.commands.report.alert.override_user")
     mock_database = mocker.MagicMock()

     report_schedule = ReportSchedule(
         created_by=get_user("admin"),
         owners=[get_user("admin")],
         type=ReportScheduleType.ALERT,
         description="description",
         crontab="0 9 * * *",
         creation_method=ReportCreationMethod.ALERTS_REPORTS,
         sql="SELECT 1",
         grace_period=14400,
         working_timeout=3600,
         database=mock_database,
         validator_config_json='{"op": "==", "threshold": 1}',
     )
     AlertCommand(report_schedule=report_schedule, execution_id=uuid.uuid4()).run()

     mock_database.mutate_sql_based_on_config.assert_not_called()
     mock_database.get_df.assert_called_once_with(
         sql=mock_database.apply_limit_to_sql.return_value
     )

     app.config["MUTATE_ALERT_QUERY"] = default_alert_mutate_ff


 def test_execute_query_succeeded_no_retry(
     mocker: MockerFixture, app_context: None
 ) -> None:
     from superset.commands.report.alert import AlertCommand

     execute_query_mock = mocker.patch(
         "superset.commands.report.alert.AlertCommand._execute_query",
         side_effect=lambda: pd.DataFrame([{"sample_col": 0}]),
     )

     command = AlertCommand(report_schedule=mocker.Mock(), execution_id=uuid.uuid4())

     command.validate()

     assert execute_query_mock.call_count == 1


 def test_execute_query_succeeded_with_retries(
     mocker: MockerFixture, app_context: None
 ) -> None:
     from superset.commands.report.alert import AlertCommand, AlertQueryError

     execute_query_mock = mocker.patch(
         "superset.commands.report.alert.AlertCommand._execute_query"
     )

     query_executed_count = 0
     # Should match the value defined in superset_test_config.py
     expected_max_retries = 3

     def _mocked_execute_query() -> pd.DataFrame:
         nonlocal query_executed_count
         query_executed_count += 1

         if query_executed_count < expected_max_retries:
             raise AlertQueryError()
         else:
             return pd.DataFrame([{"sample_col": 0}])

     execute_query_mock.side_effect = _mocked_execute_query
     execute_query_mock.__name__ = "mocked_execute_query"

     command = AlertCommand(report_schedule=mocker.Mock(), execution_id=uuid.uuid4())

     command.validate()

     assert execute_query_mock.call_count == expected_max_retries


 def test_execute_query_failed_no_retry(
     mocker: MockerFixture, app_context: None
 ) -> None:
     from superset.commands.report.alert import AlertCommand, AlertQueryTimeout

     execute_query_mock = mocker.patch(
         "superset.commands.report.alert.AlertCommand._execute_query"
     )

     def _mocked_execute_query() -> None:
         raise AlertQueryTimeout

     execute_query_mock.side_effect = _mocked_execute_query
     execute_query_mock.__name__ = "mocked_execute_query"

     command = AlertCommand(report_schedule=mocker.Mock(), execution_id=uuid.uuid4())

     with suppress(AlertQueryTimeout):
         command.validate()
     assert execute_query_mock.call_count == 1


 def test_execute_query_failed_max_retries(
     mocker: MockerFixture, app_context: None
 ) -> None:
     from superset.commands.report.alert import AlertCommand, AlertQueryError

     execute_query_mock = mocker.patch(
         "superset.commands.report.alert.AlertCommand._execute_query"
     )

     def _mocked_execute_query() -> None:
         raise AlertQueryError

     execute_query_mock.side_effect = _mocked_execute_query
     execute_query_mock.__name__ = "mocked_execute_query"

     command = AlertCommand(report_schedule=mocker.Mock(), execution_id=uuid.uuid4())

     with suppress(AlertQueryError):
         command.validate()
     # Should match the value defined in superset_test_config.py
     assert execute_query_mock.call_count == 3


 def test_get_alert_metadata_from_object(
     mocker: MockerFixture,
     app_context: AppContext,
     get_user,
 ) -> None:
     from superset.commands.report.alert import AlertCommand
     from superset.reports.models import ReportSchedule

     app.config["ALERT_REPORTS_EXECUTORS"] = [ExecutorType.OWNER]

     mock_database = mocker.MagicMock()
     mock_exec_id = uuid.uuid4()
     report_schedule = ReportSchedule(
         created_by=get_user("admin"),
         owners=[get_user("admin")],
         type=ReportScheduleType.ALERT,
         description="description",
         crontab="0 9 * * *",
         creation_method=ReportCreationMethod.ALERTS_REPORTS,
         sql="SELECT 1",
         grace_period=14400,
         working_timeout=3600,
         database=mock_database,
         validator_config_json='{"op": "==", "threshold": 1}',
     )

     cm = AlertCommand(report_schedule=report_schedule, execution_id=mock_exec_id)
     assert cm._get_alert_metadata_from_object() == {
         "report_schedule_id": report_schedule.id,
         "execution_id": mock_exec_id,
     }
	# 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.
	# pylint: disable=invalid-name, unused-argument, import-outside-toplevel
	import uuid
	from contextlib import nullcontext, suppress
	from typing import Optional, Union

	import pandas as pd
	import pytest
	from flask.ctx import AppContext
	from pytest_mock import MockerFixture

	from superset.commands.report.exceptions import AlertQueryError
	from superset.reports.models import ReportCreationMethod, ReportScheduleType
	from superset.tasks.types import ExecutorType, FixedExecutor
	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, [FixedExecutor("admin")], "admin"),
	(["gamma"], None, [ExecutorType.OWNER], "gamma"),
	(
	["alpha", "gamma"],
	"gamma",
	[ExecutorType.CREATOR_OWNER],
	"gamma",
	),
	(
	["alpha", "gamma"],
	"alpha",
	[ExecutorType.CREATOR_OWNER],
	"alpha",
	),
	(
	["alpha", "gamma"],
	"admin",
	[ExecutorType.CREATOR_OWNER],
	AlertQueryError(),
	),
	(["gamma"], None, [ExecutorType.CURRENT_USER], AlertQueryError()),
	],
	)
	def test_execute_query_as_report_executor(
	owner_names: list[str],
	creator_name: Optional[str],
	config: list[ExecutorType],
	expected_result: Union[tuple[ExecutorType, str], Exception],
	mocker: MockerFixture,
	app_context: AppContext,
	get_user,
	) -> None:
	from superset.commands.report.alert import AlertCommand
	from superset.reports.models import ReportSchedule

	original_config = app.config["ALERT_REPORTS_EXECUTORS"]
	app.config["ALERT_REPORTS_EXECUTORS"] = 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, execution_id=uuid.uuid4())
	override_user_mock = mocker.patch("superset.commands.report.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_EXECUTORS"] = original_config


	def test_execute_query_mutate_query_enabled(
	mocker: MockerFixture,
	app_context: AppContext,
	get_user,
	) -> None:
	from superset.commands.report.alert import AlertCommand
	from superset.reports.models import ReportSchedule

	default_alert_mutate_ff = app.config["MUTATE_ALERT_QUERY"]

	app.config["MUTATE_ALERT_QUERY"] = True
	mocker.patch("superset.commands.report.alert.override_user")
	mock_df = mocker.MagicMock(spec=pd.DataFrame)
	mock_df.empty = True
	mock_database = get_example_database()
	mock_get_df = mocker.patch.object(mock_database, "get_df", return_value=mock_df)
	mock_limited_sql = mocker.patch.object(mock_database, "apply_limit_to_sql")
	mock_mutate_call = mocker.patch.object(mock_database, "mutate_sql_based_on_config")

	report_schedule = ReportSchedule(
	created_by=get_user("admin"),
	owners=[get_user("admin")],
	type=ReportScheduleType.ALERT,
	description="description",
	crontab="0 9 * * *",
	creation_method=ReportCreationMethod.ALERTS_REPORTS,
	sql="SELECT 1",
	grace_period=14400,
	working_timeout=3600,
	database=mock_database,
	validator_config_json='{"op": "==", "threshold": 1}',
	)
	AlertCommand(report_schedule=report_schedule, execution_id=uuid.uuid4()).run()

	mock_mutate_call.assert_called_once_with(mock_limited_sql.return_value)
	mock_get_df.assert_called_once_with(sql=mock_mutate_call.return_value)

	app.config["MUTATE_ALERT_QUERY"] = default_alert_mutate_ff


	def test_execute_query_mutate_query_disabled(
	mocker: MockerFixture,
	app_context: AppContext,
	get_user,
	) -> None:
	from superset.commands.report.alert import AlertCommand
	from superset.reports.models import ReportSchedule

	default_alert_mutate_ff = app.config["MUTATE_ALERT_QUERY"]

	app.config["MUTATE_ALERT_QUERY"] = False
	mocker.patch("superset.commands.report.alert.override_user")
	mock_database = mocker.MagicMock()

	report_schedule = ReportSchedule(
	created_by=get_user("admin"),
	owners=[get_user("admin")],
	type=ReportScheduleType.ALERT,
	description="description",
	crontab="0 9 * * *",
	creation_method=ReportCreationMethod.ALERTS_REPORTS,
	sql="SELECT 1",
	grace_period=14400,
	working_timeout=3600,
	database=mock_database,
	validator_config_json='{"op": "==", "threshold": 1}',
	)
	AlertCommand(report_schedule=report_schedule, execution_id=uuid.uuid4()).run()

	mock_database.mutate_sql_based_on_config.assert_not_called()
	mock_database.get_df.assert_called_once_with(
	sql=mock_database.apply_limit_to_sql.return_value
	)

	app.config["MUTATE_ALERT_QUERY"] = default_alert_mutate_ff


	def test_execute_query_succeeded_no_retry(
	mocker: MockerFixture, app_context: None
	) -> None:
	from superset.commands.report.alert import AlertCommand

	execute_query_mock = mocker.patch(
	"superset.commands.report.alert.AlertCommand._execute_query",
	side_effect=lambda: pd.DataFrame([{"sample_col": 0}]),
	)

	command = AlertCommand(report_schedule=mocker.Mock(), execution_id=uuid.uuid4())

	command.validate()

	assert execute_query_mock.call_count == 1


	def test_execute_query_succeeded_with_retries(
	mocker: MockerFixture, app_context: None
	) -> None:
	from superset.commands.report.alert import AlertCommand, AlertQueryError

	execute_query_mock = mocker.patch(
	"superset.commands.report.alert.AlertCommand._execute_query"
	)

	query_executed_count = 0
	# Should match the value defined in superset_test_config.py
	expected_max_retries = 3

	def _mocked_execute_query() -> pd.DataFrame:
	nonlocal query_executed_count
	query_executed_count += 1

	if query_executed_count < expected_max_retries:
	raise AlertQueryError()
	else:
	return pd.DataFrame([{"sample_col": 0}])

	execute_query_mock.side_effect = _mocked_execute_query
	execute_query_mock.__name__ = "mocked_execute_query"

	command = AlertCommand(report_schedule=mocker.Mock(), execution_id=uuid.uuid4())

	command.validate()

	assert execute_query_mock.call_count == expected_max_retries


	def test_execute_query_failed_no_retry(
	mocker: MockerFixture, app_context: None
	) -> None:
	from superset.commands.report.alert import AlertCommand, AlertQueryTimeout

	execute_query_mock = mocker.patch(
	"superset.commands.report.alert.AlertCommand._execute_query"
	)

	def _mocked_execute_query() -> None:
	raise AlertQueryTimeout

	execute_query_mock.side_effect = _mocked_execute_query
	execute_query_mock.__name__ = "mocked_execute_query"

	command = AlertCommand(report_schedule=mocker.Mock(), execution_id=uuid.uuid4())

	with suppress(AlertQueryTimeout):
	command.validate()
	assert execute_query_mock.call_count == 1


	def test_execute_query_failed_max_retries(
	mocker: MockerFixture, app_context: None
	) -> None:
	from superset.commands.report.alert import AlertCommand, AlertQueryError

	execute_query_mock = mocker.patch(
	"superset.commands.report.alert.AlertCommand._execute_query"
	)

	def _mocked_execute_query() -> None:
	raise AlertQueryError

	execute_query_mock.side_effect = _mocked_execute_query
	execute_query_mock.__name__ = "mocked_execute_query"

	command = AlertCommand(report_schedule=mocker.Mock(), execution_id=uuid.uuid4())

	with suppress(AlertQueryError):
	command.validate()
	# Should match the value defined in superset_test_config.py
	assert execute_query_mock.call_count == 3


	def test_get_alert_metadata_from_object(
	mocker: MockerFixture,
	app_context: AppContext,
	get_user,
	) -> None:
	from superset.commands.report.alert import AlertCommand
	from superset.reports.models import ReportSchedule

	app.config["ALERT_REPORTS_EXECUTORS"] = [ExecutorType.OWNER]

	mock_database = mocker.MagicMock()
	mock_exec_id = uuid.uuid4()
	report_schedule = ReportSchedule(
	created_by=get_user("admin"),
	owners=[get_user("admin")],
	type=ReportScheduleType.ALERT,
	description="description",
	crontab="0 9 * * *",
	creation_method=ReportCreationMethod.ALERTS_REPORTS,
	sql="SELECT 1",
	grace_period=14400,
	working_timeout=3600,
	database=mock_database,
	validator_config_json='{"op": "==", "threshold": 1}',
	)

	cm = AlertCommand(report_schedule=report_schedule, execution_id=mock_exec_id)
	assert cm._get_alert_metadata_from_object() == {
	"report_schedule_id": report_schedule.id,
	"execution_id": mock_exec_id,
	}