def test_execute_query_succeeded_with_retries(mocker: MockFixture,
app_context: None) -> None:
from superset.reports.commands.alert import AlertCommand, AlertQueryError
execute_query_mock = mocker.patch(
"superset.reports.commands.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())
command.validate()
assert execute_query_mock.call_count == expected_max_retries
def next(self) -> None:
self.set_state_and_log(ReportState.WORKING)
try:
# If it's an alert check if the alert is triggered
if self._report_schedule.type == ReportScheduleType.ALERT:
if not AlertCommand(self._report_schedule).run():
self.set_state_and_log(ReportState.NOOP)
return
self.send()
self.set_state_and_log(ReportState.SUCCESS)
except CommandException as first_ex:
self.set_state_and_log(ReportState.ERROR,
error_message=str(first_ex))
# TODO (dpgaspar) convert this logic to a new state eg: ERROR_ON_GRACE
if not self.is_in_error_grace_period():
try:
self.send_error(
f"Error occurred for {self._report_schedule.type}:"
f" {self._report_schedule.name}",
str(first_ex),
)
self.set_state_and_log(
ReportState.ERROR,
error_message=REPORT_SCHEDULE_ERROR_NOTIFICATION_MARKER,
)
except CommandException as second_ex:
self.set_state_and_log(ReportState.ERROR,
error_message=str(second_ex))
raise first_ex
def next(self) -> None:
if self._report_schedule.type == ReportScheduleType.ALERT:
if self.is_in_grace_period():
self.set_state_and_log(
ReportState.GRACE,
error_message=str(ReportScheduleAlertGracePeriodError()),
)
return
self.set_state_and_log(ReportState.WORKING)
try:
if not AlertCommand(self._report_schedule).run():
self.set_state_and_log(ReportState.NOOP)
return
except CommandException as ex:
self.send_error(
f"Error occurred for {self._report_schedule.type}:"
f" {self._report_schedule.name}",
str(ex),
)
self.set_state_and_log(
ReportState.ERROR,
error_message=REPORT_SCHEDULE_ERROR_NOTIFICATION_MARKER,
)
raise ex
try:
self.send()
self.set_state_and_log(ReportState.SUCCESS)
except CommandException as ex:
self.set_state_and_log(ReportState.ERROR, error_message=str(ex))
def test_execute_query_succeeded_no_retry(mocker: MockFixture,
app_context: None) -> None:
from superset.reports.commands.alert import AlertCommand
execute_query_mock = mocker.patch(
"superset.reports.commands.alert.AlertCommand._execute_query",
side_effect=lambda: pd.DataFrame([{
"sample_col": 0
}]),
)
command = AlertCommand(report_schedule=mocker.Mock())
command.validate()
assert execute_query_mock.call_count == 1
def test_execute_query_failed_no_retry(mocker: MockFixture,
app_context: None) -> None:
from superset.reports.commands.alert import AlertCommand, AlertQueryTimeout
execute_query_mock = mocker.patch(
"superset.reports.commands.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())
try:
command.validate()
except AlertQueryTimeout:
pass
assert execute_query_mock.call_count == 1
def test_execute_query_failed_max_retries(mocker: MockFixture,
app_context: None) -> None:
from superset.reports.commands.alert import AlertCommand, AlertQueryError
execute_query_mock = mocker.patch(
"superset.reports.commands.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())
try:
command.validate()
except AlertQueryError:
pass
# Should match the value defined in superset_test_config.py
assert execute_query_mock.call_count == 3
def next(self) -> None:
self.set_state_and_log(ReportState.WORKING)
try:
# If it's an alert check if the alert is triggered
if self._report_schedule.type == ReportScheduleType.ALERT:
if not AlertCommand(self._report_schedule).run():
self.set_state_and_log(ReportState.NOOP)
return
self.send()
self.set_state_and_log(ReportState.SUCCESS)
except CommandException as ex:
self.set_state_and_log(ReportState.ERROR, error_message=str(ex))
raise ex
def run(self) -> None:
with session_scope(nullpool=True) as session:
try:
start_dttm = datetime.utcnow()
self.validate(session=session)
if not self._model:
raise ReportScheduleExecuteUnexpectedError()
self.set_state_and_log(session, start_dttm, ReportLogState.WORKING)
# If it's an alert check if the alert is triggered
if self._model.type == ReportScheduleType.ALERT:
if not AlertCommand(self._model).run():
self.set_state_and_log(session, start_dttm, ReportLogState.NOOP)
return
self._send(self._model)
# Log, state and TS
self.set_state_and_log(session, start_dttm, ReportLogState.SUCCESS)
except ReportScheduleAlertGracePeriodError as ex:
self.set_state_and_log(
session, start_dttm, ReportLogState.NOOP, error_message=str(ex)
)
except ReportSchedulePreviousWorkingError as ex:
self.create_log(
session,
start_dttm,
datetime.utcnow(),
state=ReportLogState.ERROR,
error_message=str(ex),
)
session.commit()
raise
except CommandException as ex:
self.set_state_and_log(
session, start_dttm, ReportLogState.ERROR, error_message=str(ex)
)
# We want to actually commit the state and log inside the scope
session.commit()
raise