superset/tests/integration_tests/reports/utils.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 json
from contextlib import contextmanager
from typing import Any, Optional
from uuid import uuid4

from flask_appbuilder.security.sqla.models import User

from superset import db, security_manager
from superset.models.core import Database
from superset.models.dashboard import Dashboard
from superset.models.slice import Slice
from superset.reports.models import (
    ReportDataFormat,
    ReportExecutionLog,
    ReportRecipients,
    ReportRecipientType,
    ReportSchedule,
    ReportScheduleType,
    ReportState,
)
from superset.utils.core import override_user
from tests.integration_tests.test_app import app
from tests.integration_tests.utils import read_fixture

TEST_ID = str(uuid4())
CSV_FILE = read_fixture("trends.csv")
SCREENSHOT_FILE = read_fixture("sample.png")
DEFAULT_OWNER_EMAIL = "admin@fab.org"


def insert_report_schedule(
    type: str,
    name: str,
    crontab: str,
    owners: list[User],
    timezone: Optional[str] = None,
    sql: Optional[str] = None,
    description: Optional[str] = None,
    chart: Optional[Slice] = None,
    dashboard: Optional[Dashboard] = None,
    database: Optional[Database] = None,
    validator_type: Optional[str] = None,
    validator_config_json: Optional[str] = None,
    log_retention: Optional[int] = None,
    last_state: Optional[ReportState] = None,
    grace_period: Optional[int] = None,
    recipients: Optional[list[ReportRecipients]] = None,
    report_format: Optional[ReportDataFormat] = None,
    logs: Optional[list[ReportExecutionLog]] = None,
    extra: Optional[dict[Any, Any]] = None,
    force_screenshot: bool = False,
) -> ReportSchedule:
    owners = owners or []
    recipients = recipients or []
    logs = logs or []
    last_state = last_state or ReportState.NOOP

    with override_user(owners[0]):
        report_schedule = ReportSchedule(
            type=type,
            name=name,
            crontab=crontab,
            timezone=timezone,
            sql=sql,
            description=description,
            chart=chart,
            dashboard=dashboard,
            database=database,
            owners=owners,
            validator_type=validator_type,
            validator_config_json=validator_config_json,
            log_retention=log_retention,
            grace_period=grace_period,
            recipients=recipients,
            logs=logs,
            last_state=last_state,
            report_format=report_format,
            extra=extra,
            force_screenshot=force_screenshot,
        )
    db.session.add(report_schedule)
    db.session.commit()
    return report_schedule


def create_report_notification(
    email_target: Optional[str] = None,
    slack_channel: Optional[str] = None,
    chart: Optional[Slice] = None,
    dashboard: Optional[Dashboard] = None,
    database: Optional[Database] = None,
    sql: Optional[str] = None,
    report_type: ReportScheduleType = ReportScheduleType.REPORT,
    validator_type: Optional[str] = None,
    validator_config_json: Optional[str] = None,
    grace_period: Optional[int] = None,
    report_format: Optional[ReportDataFormat] = None,
    name: Optional[str] = None,
    extra: Optional[dict[str, Any]] = None,
    force_screenshot: bool = False,
    owners: Optional[list[User]] = None,
) -> ReportSchedule:
    if not owners:
        owners = [
            (
                db.session.query(security_manager.user_model)
                .filter_by(email=DEFAULT_OWNER_EMAIL)
                .one_or_none()
            )
        ]

    if slack_channel:
        recipient = ReportRecipients(
            type=ReportRecipientType.SLACK,
            recipient_config_json=json.dumps(
                {
                    "target": slack_channel,
                }
            ),
        )
    else:
        recipient = ReportRecipients(
            type=ReportRecipientType.EMAIL,
            recipient_config_json=json.dumps({"target": email_target}),
        )

    if name is None:
        name = "report_with_csv" if report_format else "report"

    report_schedule = insert_report_schedule(
        report_type,
        name=name,
        crontab="0 9 * * *",
        description="Daily report",
        sql=sql,
        chart=chart,
        dashboard=dashboard,
        database=database,
        recipients=[recipient],
        owners=owners,
        validator_type=validator_type,
        validator_config_json=validator_config_json,
        grace_period=grace_period,
        report_format=report_format or ReportDataFormat.PNG,
        extra=extra,
        force_screenshot=force_screenshot,
    )
    return report_schedule


def cleanup_report_schedule(report_schedule: ReportSchedule) -> None:
    db.session.query(ReportExecutionLog).filter(
        ReportExecutionLog.report_schedule == report_schedule
    ).delete()
    db.session.query(ReportRecipients).filter(
        ReportRecipients.report_schedule == report_schedule
    ).delete()

    db.session.delete(report_schedule)
    db.session.commit()


@contextmanager
def create_dashboard_report(dashboard, extra, **kwargs):
    report_schedule = create_report_notification(
        email_target="target@example.com",
        dashboard=dashboard,
        extra={
            "dashboard": extra,
        },
        **kwargs
    )
    error = None

    try:
        yield report_schedule
    except Exception as ex:  # pylint: disable=broad-except
        error = ex

    # make sure to clean up in case of yield exceptions
    cleanup_report_schedule(report_schedule)

    if error:
        raise error