superset/tests/integration_tests/dashboards/superset_factory_util.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 typing import Optional

from flask_appbuilder import Model
from flask_appbuilder.security.sqla.models import User

from superset import db
from superset.connectors.sqla.models import SqlaTable, sqlatable_user
from superset.models.core import Database
from superset.models.dashboard import (
    Dashboard,
    dashboard_slices,
    dashboard_user,
    DashboardRoles,
)
from superset.models.slice import Slice, slice_user
from tests.integration_tests.dashboards.dashboard_test_utils import (
    random_slug,
    random_str,
    random_title,
)

logger = logging.getLogger(__name__)

inserted_dashboards_ids = []
inserted_databases_ids = []
inserted_sqltables_ids = []
inserted_slices_ids = []


def create_dashboard_to_db(
    dashboard_title: Optional[str] = None,
    slug: Optional[str] = None,
    published: bool = False,
    owners: Optional[list[User]] = None,
    slices: Optional[list[Slice]] = None,
    css: str = "",
    json_metadata: str = "",
    position_json: str = "",
) -> Dashboard:
    dashboard = create_dashboard(
        dashboard_title,
        slug,
        published,
        owners,
        slices,
        css,
        json_metadata,
        position_json,
    )

    insert_model(dashboard)
    inserted_dashboards_ids.append(dashboard.id)
    return dashboard


def create_dashboard(
    dashboard_title: Optional[str] = None,
    slug: Optional[str] = None,
    published: bool = False,
    owners: Optional[list[User]] = None,
    slices: Optional[list[Slice]] = None,
    css: str = "",
    json_metadata: str = "",
    position_json: str = "",
) -> Dashboard:
    dashboard_title = dashboard_title if dashboard_title is not None else random_title()
    slug = slug if slug is not None else random_slug()
    owners = owners if owners is not None else []
    slices = slices if slices is not None else []
    return Dashboard(
        dashboard_title=dashboard_title,
        slug=slug,
        published=published,
        owners=owners,
        css=css,
        position_json=position_json,
        json_metadata=json_metadata,
        slices=slices,
    )


def insert_model(dashboard: Model) -> None:
    db.session.add(dashboard)
    db.session.commit()
    db.session.refresh(dashboard)


def create_slice_to_db(
    name: Optional[str] = None,
    datasource_id: Optional[int] = None,
    owners: Optional[list[User]] = None,
) -> Slice:
    slice_ = create_slice(datasource_id, name=name, owners=owners)
    insert_model(slice_)
    inserted_slices_ids.append(slice_.id)
    return slice_


def create_slice(
    datasource_id: Optional[int] = None,
    datasource: Optional[SqlaTable] = None,
    name: Optional[str] = None,
    owners: Optional[list[User]] = None,
) -> Slice:
    name = name if name is not None else random_str()
    owners = owners if owners is not None else []
    datasource_type = "table"
    if datasource:
        return Slice(
            slice_name=name,
            table=datasource,
            owners=owners,
            datasource_type=datasource_type,
        )

    datasource_id = (
        datasource_id
        if datasource_id is not None
        else create_datasource_table_to_db(name=name + "_table").id
    )

    return Slice(
        slice_name=name,
        datasource_id=datasource_id,
        owners=owners,
        datasource_type=datasource_type,
    )


def create_datasource_table_to_db(
    name: Optional[str] = None,
    db_id: Optional[int] = None,
    owners: Optional[list[User]] = None,
) -> SqlaTable:
    sqltable = create_datasource_table(name, db_id, owners=owners)
    insert_model(sqltable)
    inserted_sqltables_ids.append(sqltable.id)
    return sqltable


def create_datasource_table(
    name: Optional[str] = None,
    db_id: Optional[int] = None,
    database: Optional[Database] = None,
    owners: Optional[list[User]] = None,
) -> SqlaTable:
    name = name if name is not None else random_str()
    owners = owners if owners is not None else []
    if database:
        return SqlaTable(table_name=name, database=database, owners=owners)
    db_id = db_id if db_id is not None else create_database_to_db(name=name + "_db").id
    return SqlaTable(table_name=name, database_id=db_id, owners=owners)


def create_database_to_db(name: Optional[str] = None) -> Database:
    database = create_database(name)
    insert_model(database)
    inserted_databases_ids.append(database.id)
    return database


def create_database(name: Optional[str] = None) -> Database:
    name = name if name is not None else random_str()
    return Database(database_name=name, sqlalchemy_uri="sqlite:///:memory:")


def delete_all_inserted_objects() -> None:
    delete_all_inserted_dashboards()
    delete_all_inserted_slices()
    delete_all_inserted_tables()
    delete_all_inserted_dbs()


def delete_all_inserted_dashboards():
    try:
        dashboards_to_delete: list[Dashboard] = (
            db.session.query(Dashboard)
            .filter(Dashboard.id.in_(inserted_dashboards_ids))
            .all()
        )
        for dashboard in dashboards_to_delete:
            try:
                delete_dashboard(dashboard, False)
            except Exception as ex:
                logger.error(f"failed to delete {dashboard.id}", exc_info=True)
                raise ex
        if len(inserted_dashboards_ids) > 0:
            db.session.commit()
            inserted_dashboards_ids.clear()
    except Exception as ex2:
        logger.error("delete_all_inserted_dashboards failed", exc_info=True)
        raise ex2


def delete_dashboard(dashboard: Dashboard, do_commit: bool = False) -> None:
    logger.info(f"deleting dashboard{dashboard.id}")
    delete_dashboard_roles_associations(dashboard)
    delete_dashboard_users_associations(dashboard)
    delete_dashboard_slices_associations(dashboard)
    db.session.delete(dashboard)
    if do_commit:
        db.session.commit()


def delete_dashboard_users_associations(dashboard: Dashboard) -> None:
    db.session.execute(
        dashboard_user.delete().where(dashboard_user.c.dashboard_id == dashboard.id)
    )


def delete_dashboard_roles_associations(dashboard: Dashboard) -> None:
    db.session.execute(
        DashboardRoles.delete().where(DashboardRoles.c.dashboard_id == dashboard.id)
    )


def delete_dashboard_slices_associations(dashboard: Dashboard) -> None:
    db.session.execute(
        dashboard_slices.delete().where(dashboard_slices.c.dashboard_id == dashboard.id)
    )


def delete_all_inserted_slices():
    try:
        slices_to_delete: list[Slice] = (
            db.session.query(Slice).filter(Slice.id.in_(inserted_slices_ids)).all()
        )
        for slice in slices_to_delete:
            try:
                delete_slice(slice, False)
            except Exception as ex:
                logger.error(f"failed to delete {slice.id}", exc_info=True)
                raise ex
        if len(inserted_slices_ids) > 0:
            db.session.commit()
            inserted_slices_ids.clear()
    except Exception as ex2:
        logger.error("delete_all_inserted_slices failed", exc_info=True)
        raise ex2


def delete_slice(slice_: Slice, do_commit: bool = False) -> None:
    logger.info(f"deleting slice{slice_.id}")
    delete_slice_users_associations(slice_)
    db.session.delete(slice_)
    if do_commit:
        db.session.commit()


def delete_slice_users_associations(slice_: Slice) -> None:
    db.session.execute(slice_user.delete().where(slice_user.c.slice_id == slice_.id))


def delete_all_inserted_tables():
    try:
        tables_to_delete: list[SqlaTable] = (
            db.session.query(SqlaTable)
            .filter(SqlaTable.id.in_(inserted_sqltables_ids))
            .all()
        )
        for table in tables_to_delete:
            try:
                delete_sqltable(table, False)
            except Exception as ex:
                logger.error(f"failed to delete {table.id}", exc_info=True)
                raise ex
        if len(inserted_sqltables_ids) > 0:
            db.session.commit()
            inserted_sqltables_ids.clear()
    except Exception as ex2:
        logger.error("delete_all_inserted_tables failed", exc_info=True)
        raise ex2


def delete_sqltable(table: SqlaTable, do_commit: bool = False) -> None:
    logger.info(f"deleting table{table.id}")
    delete_table_users_associations(table)
    db.session.delete(table)
    if do_commit:
        db.session.commit()


def delete_table_users_associations(table: SqlaTable) -> None:
    db.session.execute(
        sqlatable_user.delete().where(sqlatable_user.c.table_id == table.id)
    )


def delete_all_inserted_dbs():
    try:
        databases_to_delete: list[Database] = (
            db.session.query(Database)
            .filter(Database.id.in_(inserted_databases_ids))
            .all()
        )
        for database in databases_to_delete:
            try:
                delete_database(database, False)
            except Exception as ex:
                logger.error(f"failed to delete {database.id}", exc_info=True)
                raise ex
        if len(inserted_databases_ids) > 0:
            db.session.commit()
            inserted_databases_ids.clear()
    except Exception as ex2:
        logger.error("delete_all_inserted_databases failed", exc_info=True)
        raise ex2


def delete_database(database: Database, do_commit: bool = False) -> None:
    logger.info(f"deleting database{database.id}")
    db.session.delete(database)
    if do_commit:
        db.session.commit()