superset/commands/report/execute.py
# 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 logging
from datetime import datetime, timedelta
from typing import Any, Optional, Union
from uuid import UUID
import pandas as pd
from celery.exceptions import SoftTimeLimitExceeded
from superset import app, db, security_manager
from superset.commands.base import BaseCommand
from superset.commands.dashboard.permalink.create import CreateDashboardPermalinkCommand
from superset.commands.exceptions import CommandException
from superset.commands.report.alert import AlertCommand
from superset.commands.report.exceptions import (
ReportScheduleAlertGracePeriodError,
ReportScheduleClientErrorsException,
ReportScheduleCsvFailedError,
ReportScheduleCsvTimeout,
ReportScheduleDataFrameFailedError,
ReportScheduleDataFrameTimeout,
ReportScheduleExecuteUnexpectedError,
ReportScheduleNotFoundError,
ReportSchedulePreviousWorkingError,
ReportScheduleScreenshotFailedError,
ReportScheduleScreenshotTimeout,
ReportScheduleStateNotFoundError,
ReportScheduleSystemErrorsException,
ReportScheduleUnexpectedError,
ReportScheduleWorkingTimeoutError,
)
from superset.common.chart_data import ChartDataResultFormat, ChartDataResultType
from superset.daos.report import (
REPORT_SCHEDULE_ERROR_NOTIFICATION_MARKER,
ReportScheduleDAO,
)
from superset.errors import ErrorLevel, SupersetError, SupersetErrorType
from superset.exceptions import SupersetErrorsException, SupersetException
from superset.extensions import feature_flag_manager, machine_auth_provider_factory
from superset.reports.models import (
ReportDataFormat,
ReportExecutionLog,
ReportRecipients,
ReportRecipientType,
ReportSchedule,
ReportScheduleType,
ReportSourceFormat,
ReportState,
)
from superset.reports.notifications import create_notification
from superset.reports.notifications.base import NotificationContent
from superset.reports.notifications.exceptions import NotificationError
from superset.tasks.utils import get_executor
from superset.utils import json
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.pdf import build_pdf_from_screenshots
from superset.utils.screenshots import ChartScreenshot, DashboardScreenshot
from superset.utils.urls import get_url_path
logger = logging.getLogger(__name__)
class BaseReportState:
current_states: list[ReportState] = []
initial: bool = False
@logs_context()
def __init__(
self,
report_schedule: ReportSchedule,
scheduled_dttm: datetime,
execution_id: UUID,
) -> None:
self._report_schedule = report_schedule
self._scheduled_dttm = scheduled_dttm
self._start_dttm = datetime.utcnow()
self._execution_id = execution_id
def update_report_schedule_and_log(
self,
state: ReportState,
error_message: Optional[str] = None,
) -> None:
"""
Update the report schedule state et al. and reflect the change in the execution
log.
"""
self.update_report_schedule(state)
self.create_log(error_message)
def update_report_schedule(self, state: ReportState) -> None:
"""
Update the report schedule state et al.
When the report state is WORKING we must ensure that the values from the last
execution run are cleared to ensure that they are not propagated to the
execution log.
"""
if state == ReportState.WORKING:
self._report_schedule.last_value = None
self._report_schedule.last_value_row_json = None
self._report_schedule.last_state = state
self._report_schedule.last_eval_dttm = datetime.utcnow()
db.session.commit()
def create_log(self, error_message: Optional[str] = None) -> None:
"""
Creates a Report execution log, uses the current computed last_value for Alerts
"""
log = ReportExecutionLog(
scheduled_dttm=self._scheduled_dttm,
start_dttm=self._start_dttm,
end_dttm=datetime.utcnow(),
value=self._report_schedule.last_value,
value_row_json=self._report_schedule.last_value_row_json,
state=self._report_schedule.last_state,
error_message=error_message,
report_schedule=self._report_schedule,
uuid=self._execution_id,
)
db.session.add(log)
db.session.commit()
def _get_url(
self,
user_friendly: bool = False,
result_format: Optional[ChartDataResultFormat] = None,
**kwargs: Any,
) -> str:
"""
Get the url for this report schedule: chart or dashboard
"""
force = "true" if self._report_schedule.force_screenshot else "false"
if self._report_schedule.chart:
if result_format in {
ChartDataResultFormat.CSV,
ChartDataResultFormat.JSON,
}:
return get_url_path(
"ChartDataRestApi.get_data",
pk=self._report_schedule.chart_id,
format=result_format.value,
type=ChartDataResultType.POST_PROCESSED.value,
force=force,
)
return get_url_path(
"ExploreView.root",
user_friendly=user_friendly,
form_data=json.dumps({"slice_id": self._report_schedule.chart_id}),
force=force,
**kwargs,
)
# If we need to render dashboard in a specific state, use stateful permalink
if dashboard_state := self._report_schedule.extra.get("dashboard"):
permalink_key = CreateDashboardPermalinkCommand(
dashboard_id=str(self._report_schedule.dashboard.uuid),
state=dashboard_state,
).run()
return get_url_path("Superset.dashboard_permalink", key=permalink_key)
dashboard = self._report_schedule.dashboard
dashboard_id_or_slug = (
dashboard.uuid if dashboard and dashboard.uuid else dashboard.id
)
return get_url_path(
"Superset.dashboard",
user_friendly=user_friendly,
dashboard_id_or_slug=dashboard_id_or_slug,
force=force,
**kwargs,
)
def _get_screenshots(self) -> list[bytes]:
"""
Get chart or dashboard screenshots
:raises: ReportScheduleScreenshotFailedError
"""
url = self._get_url()
_, username = get_executor(
executor_types=app.config["ALERT_REPORTS_EXECUTE_AS"],
model=self._report_schedule,
)
user = security_manager.find_user(username)
if self._report_schedule.chart:
window_width, window_height = app.config["WEBDRIVER_WINDOW"]["slice"]
window_size = (
self._report_schedule.custom_width or window_width,
self._report_schedule.custom_height or window_height,
)
screenshot: Union[ChartScreenshot, DashboardScreenshot] = ChartScreenshot(
url,
self._report_schedule.chart.digest,
window_size=window_size,
thumb_size=app.config["WEBDRIVER_WINDOW"]["slice"],
)
else:
window_width, window_height = app.config["WEBDRIVER_WINDOW"]["dashboard"]
window_size = (
self._report_schedule.custom_width or window_width,
self._report_schedule.custom_height or window_height,
)
screenshot = DashboardScreenshot(
url,
self._report_schedule.dashboard.digest,
window_size=window_size,
thumb_size=app.config["WEBDRIVER_WINDOW"]["dashboard"],
)
try:
image = screenshot.get_screenshot(user=user)
except SoftTimeLimitExceeded as ex:
logger.warning("A timeout occurred while taking a screenshot.")
raise ReportScheduleScreenshotTimeout() from ex
except Exception as ex:
raise ReportScheduleScreenshotFailedError(
f"Failed taking a screenshot {str(ex)}"
) from ex
if not image:
raise ReportScheduleScreenshotFailedError()
return [image]
def _get_pdf(self) -> bytes:
"""
Get chart or dashboard pdf
:raises: ReportSchedulePdfFailedError
"""
screenshots = self._get_screenshots()
pdf = build_pdf_from_screenshots(screenshots)
return pdf
def _get_csv_data(self) -> bytes:
url = self._get_url(result_format=ChartDataResultFormat.CSV)
_, username = get_executor(
executor_types=app.config["ALERT_REPORTS_EXECUTE_AS"],
model=self._report_schedule,
)
user = security_manager.find_user(username)
auth_cookies = machine_auth_provider_factory.instance.get_auth_cookies(user)
if self._report_schedule.chart.query_context is None:
logger.warning("No query context found, taking a screenshot to generate it")
self._update_query_context()
try:
logger.info("Getting chart from %s as user %s", url, user.username)
csv_data = get_chart_csv_data(chart_url=url, auth_cookies=auth_cookies)
except SoftTimeLimitExceeded as ex:
raise ReportScheduleCsvTimeout() from ex
except Exception as ex:
raise ReportScheduleCsvFailedError(
f"Failed generating csv {str(ex)}"
) from ex
if not csv_data:
raise ReportScheduleCsvFailedError()
return csv_data
def _get_embedded_data(self) -> pd.DataFrame:
"""
Return data as a Pandas dataframe, to embed in notifications as a table.
"""
url = self._get_url(result_format=ChartDataResultFormat.JSON)
_, username = get_executor(
executor_types=app.config["ALERT_REPORTS_EXECUTE_AS"],
model=self._report_schedule,
)
user = security_manager.find_user(username)
auth_cookies = machine_auth_provider_factory.instance.get_auth_cookies(user)
if self._report_schedule.chart.query_context is None:
logger.warning("No query context found, taking a screenshot to generate it")
self._update_query_context()
try:
logger.info("Getting chart from %s as user %s", url, user.username)
dataframe = get_chart_dataframe(url, auth_cookies)
except SoftTimeLimitExceeded as ex:
raise ReportScheduleDataFrameTimeout() from ex
except Exception as ex:
raise ReportScheduleDataFrameFailedError(
f"Failed generating dataframe {str(ex)}"
) from ex
if dataframe is None:
raise ReportScheduleCsvFailedError()
return dataframe
def _update_query_context(self) -> None:
"""
Update chart query context.
To load CSV data from the endpoint the chart must have been saved
with its query context. For charts without saved query context we
get a screenshot to force the chart to produce and save the query
context.
"""
try:
self._get_screenshots()
except (
ReportScheduleScreenshotFailedError,
ReportScheduleScreenshotTimeout,
) as ex:
raise ReportScheduleCsvFailedError(
"Unable to fetch data because the chart has no query context "
"saved, and an error occurred when fetching it via a screenshot. "
"Please try loading the chart and saving it again."
) from ex
def _get_log_data(self) -> HeaderDataType:
chart_id = None
dashboard_id = None
report_source = None
if self._report_schedule.chart:
report_source = ReportSourceFormat.CHART
chart_id = self._report_schedule.chart_id
else:
report_source = ReportSourceFormat.DASHBOARD
dashboard_id = self._report_schedule.dashboard_id
log_data: HeaderDataType = {
"notification_type": self._report_schedule.type,
"notification_source": report_source,
"notification_format": self._report_schedule.report_format,
"chart_id": chart_id,
"dashboard_id": dashboard_id,
"owners": self._report_schedule.owners,
}
return log_data
def _get_notification_content(self) -> NotificationContent:
"""
Gets a notification content, this is composed by a title and a screenshot
:raises: ReportScheduleScreenshotFailedError
"""
csv_data = None
screenshot_data = []
pdf_data = None
embedded_data = None
error_text = None
header_data = self._get_log_data()
url = self._get_url(user_friendly=True)
if (
feature_flag_manager.is_feature_enabled("ALERTS_ATTACH_REPORTS")
or self._report_schedule.type == ReportScheduleType.REPORT
):
if self._report_schedule.report_format == ReportDataFormat.PNG:
screenshot_data = self._get_screenshots()
if not screenshot_data:
error_text = "Unexpected missing screenshot"
elif self._report_schedule.report_format == ReportDataFormat.PDF:
pdf_data = self._get_pdf()
if not pdf_data:
error_text = "Unexpected missing pdf"
elif (
self._report_schedule.chart
and self._report_schedule.report_format == ReportDataFormat.CSV
):
csv_data = self._get_csv_data()
if not csv_data:
error_text = "Unexpected missing csv file"
if error_text:
return NotificationContent(
name=self._report_schedule.name,
text=error_text,
header_data=header_data,
)
if (
self._report_schedule.chart
and self._report_schedule.report_format == ReportDataFormat.TEXT
):
embedded_data = self._get_embedded_data()
if self._report_schedule.email_subject:
name = self._report_schedule.email_subject
else:
if self._report_schedule.chart:
name = (
f"{self._report_schedule.name}: "
f"{self._report_schedule.chart.slice_name}"
)
else:
name = (
f"{self._report_schedule.name}: "
f"{self._report_schedule.dashboard.dashboard_title}"
)
return NotificationContent(
name=name,
url=url,
screenshots=screenshot_data,
pdf=pdf_data,
description=self._report_schedule.description,
csv=csv_data,
embedded_data=embedded_data,
header_data=header_data,
)
def _send(
self,
notification_content: NotificationContent,
recipients: list[ReportRecipients],
) -> None:
"""
Sends a notification to all recipients
:raises: CommandException
"""
notification_errors: list[SupersetError] = []
for recipient in recipients:
notification = create_notification(recipient, notification_content)
try:
if app.config["ALERT_REPORTS_NOTIFICATION_DRY_RUN"]:
logger.info(
"Would send notification for alert %s, to %s. "
"ALERT_REPORTS_NOTIFICATION_DRY_RUN is enabled, "
"set it to False to send notifications.",
self._report_schedule.name,
recipient.recipient_config_json,
)
else:
notification.send()
except (NotificationError, SupersetException) as ex:
# collect errors but keep processing them
notification_errors.append(
SupersetError(
message=ex.message,
error_type=SupersetErrorType.REPORT_NOTIFICATION_ERROR,
level=ErrorLevel.ERROR
if ex.status >= 500
else ErrorLevel.WARNING,
)
)
if notification_errors:
# log all errors but raise based on the most severe
for error in notification_errors:
logger.warning(str(error))
if any(error.level == ErrorLevel.ERROR for error in notification_errors):
raise ReportScheduleSystemErrorsException(errors=notification_errors)
if any(error.level == ErrorLevel.WARNING for error in notification_errors):
raise ReportScheduleClientErrorsException(errors=notification_errors)
def send(self) -> None:
"""
Creates the notification content and sends them to all recipients
:raises: CommandException
"""
notification_content = self._get_notification_content()
self._send(notification_content, self._report_schedule.recipients)
def send_error(self, name: str, message: str) -> None:
"""
Creates and sends a notification for an error, to all recipients
:raises: CommandException
"""
header_data = self._get_log_data()
logger.info(
"header_data in notifications for alerts and reports %s, taskid, %s",
header_data,
self._execution_id,
)
notification_content = NotificationContent(
name=name, text=message, header_data=header_data
)
# filter recipients to recipients who are also owners
owner_recipients = [
ReportRecipients(
type=ReportRecipientType.EMAIL,
recipient_config_json=json.dumps({"target": owner.email}),
)
for owner in self._report_schedule.owners
]
self._send(notification_content, owner_recipients)
def is_in_grace_period(self) -> bool:
"""
Checks if an alert is in it's grace period
"""
last_success = ReportScheduleDAO.find_last_success_log(self._report_schedule)
return (
last_success is not None
and self._report_schedule.grace_period
and datetime.utcnow()
- timedelta(seconds=self._report_schedule.grace_period)
< last_success.end_dttm
)
def is_in_error_grace_period(self) -> bool:
"""
Checks if an alert/report on error is in it's notification grace period
"""
last_success = ReportScheduleDAO.find_last_error_notification(
self._report_schedule
)
if not last_success:
return False
return (
last_success is not None
and self._report_schedule.grace_period
and datetime.utcnow()
- timedelta(seconds=self._report_schedule.grace_period)
< last_success.end_dttm
)
def is_on_working_timeout(self) -> bool:
"""
Checks if an alert is in a working timeout
"""
last_working = ReportScheduleDAO.find_last_entered_working_log(
self._report_schedule
)
if not last_working:
return False
return (
self._report_schedule.working_timeout is not None
and self._report_schedule.last_eval_dttm is not None
and datetime.utcnow()
- timedelta(seconds=self._report_schedule.working_timeout)
> last_working.end_dttm
)
def next(self) -> None:
raise NotImplementedError()
class ReportNotTriggeredErrorState(BaseReportState):
"""
Handle Not triggered and Error state
next final states:
- Not Triggered
- Success
- Error
"""
current_states = [ReportState.NOOP, ReportState.ERROR]
initial = True
def next(self) -> None:
self.update_report_schedule_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.update_report_schedule_and_log(ReportState.NOOP)
return
self.send()
self.update_report_schedule_and_log(ReportState.SUCCESS)
except (SupersetErrorsException, Exception) as first_ex:
error_message = str(first_ex)
if isinstance(first_ex, SupersetErrorsException):
error_message = ";".join([error.message for error in first_ex.errors])
self.update_report_schedule_and_log(
ReportState.ERROR, error_message=error_message
)
# TODO (dpgaspar) convert this logic to a new state eg: ERROR_ON_GRACE
if not self.is_in_error_grace_period():
second_error_message = REPORT_SCHEDULE_ERROR_NOTIFICATION_MARKER
try:
self.send_error(
f"Error occurred for {self._report_schedule.type}:"
f" {self._report_schedule.name}",
str(first_ex),
)
except SupersetErrorsException as second_ex:
second_error_message = ";".join(
[error.message for error in second_ex.errors]
)
except Exception as second_ex: # pylint: disable=broad-except
second_error_message = str(second_ex)
finally:
self.update_report_schedule_and_log(
ReportState.ERROR, error_message=second_error_message
)
raise first_ex
class ReportWorkingState(BaseReportState):
"""
Handle Working state
next states:
- Error
- Working
"""
current_states = [ReportState.WORKING]
def next(self) -> None:
if self.is_on_working_timeout():
exception_timeout = ReportScheduleWorkingTimeoutError()
self.update_report_schedule_and_log(
ReportState.ERROR,
error_message=str(exception_timeout),
)
raise exception_timeout
exception_working = ReportSchedulePreviousWorkingError()
self.update_report_schedule_and_log(
ReportState.WORKING,
error_message=str(exception_working),
)
raise exception_working
class ReportSuccessState(BaseReportState):
"""
Handle Success, Grace state
next states:
- Grace
- Not triggered
- Success
"""
current_states = [ReportState.SUCCESS, ReportState.GRACE]
def next(self) -> None:
if self._report_schedule.type == ReportScheduleType.ALERT:
if self.is_in_grace_period():
self.update_report_schedule_and_log(
ReportState.GRACE,
error_message=str(ReportScheduleAlertGracePeriodError()),
)
return
self.update_report_schedule_and_log(ReportState.WORKING)
try:
if not AlertCommand(self._report_schedule).run():
self.update_report_schedule_and_log(ReportState.NOOP)
return
except Exception as ex:
self.send_error(
f"Error occurred for {self._report_schedule.type}:"
f" {self._report_schedule.name}",
str(ex),
)
self.update_report_schedule_and_log(
ReportState.ERROR,
error_message=REPORT_SCHEDULE_ERROR_NOTIFICATION_MARKER,
)
raise ex
try:
self.send()
self.update_report_schedule_and_log(ReportState.SUCCESS)
except Exception as ex: # pylint: disable=broad-except
self.update_report_schedule_and_log(
ReportState.ERROR, error_message=str(ex)
)
class ReportScheduleStateMachine: # pylint: disable=too-few-public-methods
"""
Simple state machine for Alerts/Reports states
"""
states_cls = [ReportWorkingState, ReportNotTriggeredErrorState, ReportSuccessState]
def __init__(
self,
task_uuid: UUID,
report_schedule: ReportSchedule,
scheduled_dttm: datetime,
):
self._execution_id = task_uuid
self._report_schedule = report_schedule
self._scheduled_dttm = scheduled_dttm
def run(self) -> None:
for state_cls in self.states_cls:
if (self._report_schedule.last_state is None and state_cls.initial) or (
self._report_schedule.last_state in state_cls.current_states
):
state_cls(
self._report_schedule,
self._scheduled_dttm,
self._execution_id,
).next()
break
else:
raise ReportScheduleStateNotFoundError()
class AsyncExecuteReportScheduleCommand(BaseCommand):
"""
Execute all types of report schedules.
- On reports takes chart or dashboard screenshots and sends configured notifications
- On Alerts uses related Command AlertCommand and sends configured notifications
"""
def __init__(self, task_id: str, model_id: int, scheduled_dttm: datetime):
self._model_id = model_id
self._model: Optional[ReportSchedule] = None
self._scheduled_dttm = scheduled_dttm
self._execution_id = UUID(task_id)
def run(self) -> None:
try:
self.validate()
if not self._model:
raise ReportScheduleExecuteUnexpectedError()
_, username = get_executor(
executor_types=app.config["ALERT_REPORTS_EXECUTE_AS"],
model=self._model,
)
user = security_manager.find_user(username)
with override_user(user):
logger.info(
"Running report schedule %s as user %s",
self._execution_id,
username,
)
ReportScheduleStateMachine(
self._execution_id, self._model, self._scheduled_dttm
).run()
except CommandException as ex:
raise ex
except Exception as ex:
raise ReportScheduleUnexpectedError(str(ex)) from ex
def validate(self) -> None:
# Validate/populate model exists
logger.info(
"session is validated: id %s, executionid: %s",
self._model_id,
self._execution_id,
)
self._model = (
db.session.query(ReportSchedule).filter_by(id=self._model_id).one_or_none()
)
if not self._model:
raise ReportScheduleNotFoundError()