superset2/tests/integration_tests/charts/commands_tests.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 time
from datetime import datetime
from unittest.mock import patch

import pytest
import yaml
from flask import g  # noqa: F401

from superset import db, security_manager
from superset.commands.chart.create import CreateChartCommand
from superset.commands.chart.exceptions import (
    ChartForbiddenError,
    ChartNotFoundError,
    WarmUpCacheChartNotFoundError,
)
from superset.commands.chart.export import ExportChartsCommand
from superset.commands.chart.fave import AddFavoriteChartCommand
from superset.commands.chart.importers.v1 import ImportChartsCommand
from superset.commands.chart.unfave import DelFavoriteChartCommand
from superset.commands.chart.update import UpdateChartCommand
from superset.commands.chart.warm_up_cache import ChartWarmUpCacheCommand
from superset.commands.exceptions import CommandInvalidError
from superset.commands.importers.exceptions import IncorrectVersionError
from superset.connectors.sqla.models import SqlaTable
from superset.daos.chart import ChartDAO
from superset.models.core import Database
from superset.models.slice import Slice
from superset.utils import json
from superset.utils.core import override_user
from tests.integration_tests.base_tests import SupersetTestCase
from tests.integration_tests.fixtures.birth_names_dashboard import (
    load_birth_names_dashboard_with_slices,  # noqa: F401
    load_birth_names_data,  # noqa: F401
)
from tests.integration_tests.fixtures.energy_dashboard import (
    load_energy_table_data,  # noqa: F401
    load_energy_table_with_slice,  # noqa: F401
)
from tests.integration_tests.fixtures.importexport import (
    chart_config,
    chart_metadata_config,
    database_config,
    database_metadata_config,
    dataset_config,
)


class TestExportChartsCommand(SupersetTestCase):
    @patch("superset.security.manager.g")
    @pytest.mark.usefixtures("load_energy_table_with_slice")
    def test_export_chart_command(self, mock_g):
        mock_g.user = security_manager.find_user("admin")

        example_chart = (
            db.session.query(Slice).filter_by(slice_name="Energy Sankey").one()
        )
        command = ExportChartsCommand([example_chart.id])
        contents = dict(command.run())

        expected = [
            "metadata.yaml",
            f"charts/Energy_Sankey_{example_chart.id}.yaml",
            f"datasets/examples/energy_usage_{example_chart.table.id}.yaml",
            "databases/examples.yaml",
        ]
        assert expected == list(contents.keys())

        metadata = yaml.safe_load(
            contents[f"charts/Energy_Sankey_{example_chart.id}.yaml"]()
        )

        assert metadata == {
            "slice_name": "Energy Sankey",
            "description": None,
            "certified_by": None,
            "certification_details": None,
            "viz_type": "sankey",
            "params": {
                "collapsed_fieldsets": "",
                "groupby": ["source", "target"],
                "metric": "sum__value",
                "row_limit": "5000",
                "slice_name": "Energy Sankey",
                "viz_type": "sankey",
            },
            "cache_timeout": None,
            "dataset_uuid": str(example_chart.table.uuid),
            "uuid": str(example_chart.uuid),
            "version": "1.0.0",
            "query_context": None,
        }

    @patch("superset.utils.core.g")
    @patch("superset.security.manager.g")
    @pytest.mark.usefixtures("load_energy_table_with_slice")
    def test_export_chart_command_no_access(self, utils_mock_g, manager_mock_g):
        """Test that users can't export datasets they don't have access to"""
        manager_mock_g.user = security_manager.find_user("gamma")
        utils_mock_g.user = manager_mock_g.user

        example_chart = db.session.query(Slice).all()[0]
        command = ExportChartsCommand([example_chart.id])
        contents = command.run()
        with self.assertRaises(ChartNotFoundError):  # noqa: PT027
            next(contents)

    @patch("superset.security.manager.g")
    def test_export_chart_command_invalid_dataset(self, mock_g):
        """Test that an error is raised when exporting an invalid dataset"""
        mock_g.user = security_manager.find_user("admin")
        command = ExportChartsCommand([-1])
        contents = command.run()
        with self.assertRaises(ChartNotFoundError):  # noqa: PT027
            next(contents)

    @patch("superset.security.manager.g")
    @pytest.mark.usefixtures("load_energy_table_with_slice")
    def test_export_chart_command_key_order(self, mock_g):
        """Test that they keys in the YAML have the same order as export_fields"""
        mock_g.user = security_manager.find_user("admin")

        example_chart = (
            db.session.query(Slice).filter_by(slice_name="Energy Sankey").one()
        )
        command = ExportChartsCommand([example_chart.id])
        contents = dict(command.run())

        metadata = yaml.safe_load(
            contents[f"charts/Energy_Sankey_{example_chart.id}.yaml"]()
        )
        assert list(metadata.keys()) == [
            "slice_name",
            "description",
            "certified_by",
            "certification_details",
            "viz_type",
            "params",
            "query_context",
            "cache_timeout",
            "uuid",
            "version",
            "dataset_uuid",
        ]

    @patch("superset.security.manager.g")
    @pytest.mark.usefixtures("load_energy_table_with_slice")
    def test_export_chart_command_no_related(self, mock_g):
        """
        Test that only the chart is exported when export_related=False.
        """
        mock_g.user = security_manager.find_user("admin")

        example_chart = (
            db.session.query(Slice).filter_by(slice_name="Energy Sankey").one()
        )
        command = ExportChartsCommand([example_chart.id], export_related=False)
        contents = dict(command.run())

        expected = [
            "metadata.yaml",
            f"charts/Energy_Sankey_{example_chart.id}.yaml",
        ]
        assert expected == list(contents.keys())


class TestImportChartsCommand(SupersetTestCase):
    @patch("superset.utils.core.g")
    @patch("superset.security.manager.g")
    @patch("superset.commands.database.importers.v1.utils.add_permissions")
    def test_import_v1_chart(self, mock_add_permissions, sm_g, utils_g) -> None:
        """Test that we can import a chart"""
        admin = sm_g.user = utils_g.user = security_manager.find_user("admin")
        contents = {
            "metadata.yaml": yaml.safe_dump(chart_metadata_config),
            "databases/imported_database.yaml": yaml.safe_dump(database_config),
            "datasets/imported_dataset.yaml": yaml.safe_dump(dataset_config),
            "charts/imported_chart.yaml": yaml.safe_dump(chart_config),
        }
        command = ImportChartsCommand(contents)
        command.run()

        chart: Slice = (
            db.session.query(Slice).filter_by(uuid=chart_config["uuid"]).one()
        )
        dataset = chart.datasource
        assert json.loads(chart.params) == {
            "annotation_layers": [],
            "color_picker": {"a": 1, "b": 135, "g": 122, "r": 0},
            "datasource": dataset.uid if dataset else None,
            "js_columns": ["color"],
            "js_data_mutator": "data => data.map(d => ({\\n    ...d,\\n    color: colors.hexToRGB(d.extraProps.color)\\n}));",  # noqa: E501
            "js_onclick_href": "",
            "js_tooltip": "",
            "line_column": "path_json",
            "line_type": "json",
            "line_width": 150,
            "mapbox_style": "mapbox://styles/mapbox/light-v9",
            "reverse_long_lat": False,
            "row_limit": 5000,
            "slice_id": 43,
            "time_grain_sqla": None,
            "time_range": " : ",
            "viewport": {
                "altitude": 1.5,
                "bearing": 0,
                "height": 1094,
                "latitude": 37.73671752604488,
                "longitude": -122.18885402582598,
                "maxLatitude": 85.05113,
                "maxPitch": 60,
                "maxZoom": 20,
                "minLatitude": -85.05113,
                "minPitch": 0,
                "minZoom": 0,
                "pitch": 0,
                "width": 669,
                "zoom": 9.51847667620428,
            },
            "viz_type": "deck_path",
        }

        dataset = (
            db.session.query(SqlaTable).filter_by(uuid=dataset_config["uuid"]).one()
        )
        table_name = dataset.table_name if dataset else None
        assert table_name == "imported_dataset"
        assert chart.table == dataset

        database = (
            db.session.query(Database).filter_by(uuid=database_config["uuid"]).one()
        )
        assert database.database_name == "imported_database"
        assert chart.table.database == database

        assert chart.owners == [admin]

        db.session.delete(chart)
        db.session.delete(dataset)
        db.session.delete(database)
        db.session.commit()

    @patch("superset.security.manager.g")
    @patch("superset.commands.database.importers.v1.utils.add_permissions")
    def test_import_v1_chart_multiple(self, mock_add_permissions, sm_g):
        """Test that a chart can be imported multiple times"""
        sm_g.user = security_manager.find_user("admin")
        contents = {
            "metadata.yaml": yaml.safe_dump(chart_metadata_config),
            "databases/imported_database.yaml": yaml.safe_dump(database_config),
            "datasets/imported_dataset.yaml": yaml.safe_dump(dataset_config),
            "charts/imported_chart.yaml": yaml.safe_dump(chart_config),
        }
        command = ImportChartsCommand(contents, overwrite=True)
        command.run()
        command.run()

        dataset = (
            db.session.query(SqlaTable).filter_by(uuid=dataset_config["uuid"]).one()
        )
        charts = db.session.query(Slice).filter_by(datasource_id=dataset.id).all()
        assert len(charts) == 1

        database = dataset.database

        db.session.delete(charts[0])
        db.session.delete(dataset)
        db.session.delete(database)
        db.session.commit()

    @patch("superset.commands.database.importers.v1.utils.add_permissions")
    def test_import_v1_chart_validation(self, mock_add_permissions):
        """Test different validations applied when importing a chart"""
        # metadata.yaml must be present
        contents = {
            "databases/imported_database.yaml": yaml.safe_dump(database_config),
            "datasets/imported_dataset.yaml": yaml.safe_dump(dataset_config),
            "charts/imported_chart.yaml": yaml.safe_dump(chart_config),
        }
        command = ImportChartsCommand(contents)
        with pytest.raises(IncorrectVersionError) as excinfo:
            command.run()
        assert str(excinfo.value) == "Missing metadata.yaml"

        # version should be 1.0.0
        contents["metadata.yaml"] = yaml.safe_dump(
            {
                "version": "2.0.0",
                "type": "SqlaTable",
                "timestamp": "2020-11-04T21:27:44.423819+00:00",
            }
        )
        command = ImportChartsCommand(contents)
        with pytest.raises(IncorrectVersionError) as excinfo:
            command.run()
        assert str(excinfo.value) == "Must be equal to 1.0.0."

        # type should be Slice
        contents["metadata.yaml"] = yaml.safe_dump(database_metadata_config)
        command = ImportChartsCommand(contents)
        with pytest.raises(CommandInvalidError) as excinfo:
            command.run()
        assert str(excinfo.value).startswith("Error importing chart")
        assert excinfo.value.normalized_messages() == {
            "metadata.yaml": {"type": ["Must be equal to Slice."]}
        }

        # must also validate datasets and databases
        broken_config = database_config.copy()
        del broken_config["database_name"]
        contents["metadata.yaml"] = yaml.safe_dump(chart_metadata_config)
        contents["databases/imported_database.yaml"] = yaml.safe_dump(broken_config)
        command = ImportChartsCommand(contents)
        with pytest.raises(CommandInvalidError) as excinfo:
            command.run()
        assert str(excinfo.value).startswith("Error importing chart")
        assert excinfo.value.normalized_messages() == {
            "databases/imported_database.yaml": {
                "database_name": ["Missing data for required field."],
            }
        }


class TestChartsCreateCommand(SupersetTestCase):
    @patch("superset.utils.core.g")
    @patch("superset.commands.chart.create.g")
    @patch("superset.security.manager.g")
    @pytest.mark.usefixtures("load_energy_table_with_slice")
    def test_create_v1_response(self, mock_sm_g, mock_c_g, mock_u_g):
        """Test that the create chart command creates a chart"""
        user = security_manager.find_user(username="admin")
        mock_u_g.user = mock_c_g.user = mock_sm_g.user = user
        chart_data = {
            "slice_name": "new chart",
            "description": "new description",
            "owners": [user.id],
            "viz_type": "new_viz_type",
            "params": json.dumps({"viz_type": "new_viz_type"}),
            "cache_timeout": 1000,
            "datasource_id": 1,
            "datasource_type": "table",
        }
        command = CreateChartCommand(chart_data)
        chart = command.run()
        chart = db.session.query(Slice).get(chart.id)
        assert chart.viz_type == "new_viz_type"
        json_params = json.loads(chart.params)
        assert json_params == {"viz_type": "new_viz_type"}
        assert chart.slice_name == "new chart"
        assert chart.owners == [user]
        db.session.delete(chart)
        db.session.commit()


class TestChartsUpdateCommand(SupersetTestCase):
    @patch("superset.commands.chart.update.g")
    @patch("superset.utils.core.g")
    @patch("superset.security.manager.g")
    @pytest.mark.usefixtures("load_energy_table_with_slice")
    def test_update_sets_last_saved_at(self, mock_sm_g, mock_c_g, mock_u_g):
        """Test that update sets last_saved_at when previously unset"""
        pk = db.session.query(Slice).all()[0].id
        user = security_manager.find_user(username="admin")
        mock_u_g.user = mock_c_g.user = mock_sm_g.user = user

        # Explicitly set last_saved_at to None to test None -> datetime transition
        chart_to_update = db.session.query(Slice).get(pk)
        chart_to_update.last_saved_at = None
        db.session.commit()

        command = UpdateChartCommand(
            pk,
            {"description": "test", "owners": [user.id]},
        )
        command.run()

        chart = db.session.query(Slice).get(pk)
        assert chart.last_saved_at is not None
        assert chart.last_saved_by == user

    @patch("superset.commands.chart.update.g")
    @patch("superset.utils.core.g")
    @patch("superset.security.manager.g")
    @pytest.mark.usefixtures("load_energy_table_with_slice")
    def test_update_changes_last_saved_at(self, mock_sm_g, mock_c_g, mock_u_g):
        """Test that update changes last_saved_at when it already has a value"""
        pk = db.session.query(Slice).all()[0].id
        user = security_manager.find_user(username="admin")
        mock_u_g.user = mock_c_g.user = mock_sm_g.user = user

        chart_to_update = db.session.query(Slice).get(pk)
        chart_to_update.last_saved_at = datetime.now()
        db.session.commit()
        # Refresh to get the database value with MySQL's truncated microseconds
        db.session.refresh(chart_to_update)
        last_saved_before = chart_to_update.last_saved_at

        command = UpdateChartCommand(
            pk,
            {"description": "test", "owners": [user.id]},
        )
        # Sleep to ensure timestamp differs at MySQL's second precision (DATETIME(0))
        time.sleep(1)
        command.run()

        chart = db.session.query(Slice).get(pk)
        assert chart.last_saved_at.replace(microsecond=0) != last_saved_before.replace(
            microsecond=0
        )
        assert chart.last_saved_by == user

    @patch("superset.utils.core.g")
    @patch("superset.security.manager.g")
    @pytest.mark.usefixtures("load_energy_table_with_slice")
    @pytest.mark.skip(reason="This test will be changed to use the api/v1/data")
    def test_query_context_update_command(self, mock_sm_g, mock_g):
        """
        Test that a user can generate the chart query context
        payload without affecting owners
        """
        chart = db.session.query(Slice).all()[0]
        pk = chart.id
        admin = security_manager.find_user(username="admin")
        chart.owners = [admin]
        db.session.commit()

        user = security_manager.find_user(username="alpha")
        mock_g.user = mock_sm_g.user = user
        query_context = json.dumps({"foo": "bar"})
        json_obj = {
            "query_context_generation": True,
            "query_context": query_context,
        }
        command = UpdateChartCommand(pk, json_obj)
        command.run()
        chart = db.session.query(Slice).get(pk)
        assert chart.query_context == query_context
        assert len(chart.owners) == 1
        assert chart.owners[0] == admin

    @patch("superset.commands.chart.update.g")
    @patch("superset.utils.core.g")
    @patch("superset.security.manager.g")
    @pytest.mark.usefixtures("load_energy_table_with_slice")
    def test_update_chart_dashboard_security_existing_relationship(
        self, mock_sm_g, mock_u_g, mock_c_g
    ):
        """Test that chart owners can update charts linked to inaccessible
        dashboards (existing relationships)"""
        from superset.models.dashboard import Dashboard

        # Create a chart owned by alpha
        admin = security_manager.find_user(username="admin")
        alpha = security_manager.find_user(username="alpha")

        # Set user context for dashboard creation
        mock_u_g.user = mock_c_g.user = mock_sm_g.user = admin

        chart = db.session.query(Slice).first()
        chart.owners = [alpha]

        # Create a dashboard owned by admin (not accessible to alpha)
        admin_dashboard = Dashboard(
            dashboard_title="Admin Dashboard",
            slug="admin-dashboard",
            owners=[admin],
            published=False,
        )
        db.session.add(admin_dashboard)

        # Link chart to admin's dashboard (alpha owns chart, admin owns dashboard)
        chart.dashboards.append(admin_dashboard)
        db.session.commit()

        # Alpha should still be able to update their chart
        # even though it's linked to admin's dashboard
        mock_u_g.user = mock_c_g.user = mock_sm_g.user = alpha

        json_obj = {
            "description": "Updated description",
            "dashboards": [
                d.id for d in chart.dashboards
            ],  # Keep existing relationships
        }
        command = UpdateChartCommand(chart.id, json_obj)
        command.run()

        # Should succeed - alpha can update their chart
        updated_chart = db.session.query(Slice).get(chart.id)
        assert updated_chart.description == "Updated description"

        # Clean up
        db.session.delete(admin_dashboard)
        db.session.commit()

    @patch("superset.commands.chart.update.g")
    @patch("superset.utils.core.g")
    @patch("superset.security.manager.g")
    @pytest.mark.usefixtures("load_energy_table_with_slice")
    def test_update_chart_dashboard_security_new_unauthorized_relationship(
        self, mock_sm_g, mock_u_g, mock_c_g
    ):
        """Test that users cannot add charts to dashboards they don't have access to"""
        from superset.commands.chart.exceptions import ChartInvalidError
        from superset.models.dashboard import Dashboard

        admin = security_manager.find_user(username="admin")
        alpha = security_manager.find_user(username="alpha")

        # Set user context for dashboard creation
        mock_u_g.user = mock_c_g.user = mock_sm_g.user = admin

        # Create chart owned by alpha
        chart = db.session.query(Slice).first()
        chart.owners = [alpha]

        # Create private dashboard owned by admin (not accessible to alpha)
        admin_dashboard = Dashboard(
            dashboard_title="Admin Private Dashboard",
            slug="admin-private-dashboard",
            owners=[admin],
            published=False,  # Private dashboard
        )
        db.session.add(admin_dashboard)
        db.session.commit()

        # Alpha tries to add their chart to admin's private dashboard
        mock_u_g.user = mock_c_g.user = mock_sm_g.user = alpha

        json_obj = {
            "description": "Trying to add to unauthorized dashboard",
            "dashboards": [admin_dashboard.id],  # NEW unauthorized relationship
        }
        command = UpdateChartCommand(chart.id, json_obj)

        # Should fail - alpha cannot access admin's private dashboard
        with self.assertRaises(ChartInvalidError):  # noqa: PT027
            command.run()

        # Clean up
        db.session.delete(admin_dashboard)
        db.session.commit()

    @patch("superset.commands.chart.update.g")
    @patch("superset.utils.core.g")
    @patch("superset.security.manager.g")
    @pytest.mark.usefixtures("load_energy_table_with_slice")
    def test_update_chart_dashboard_security_admin_bypass(
        self, mock_sm_g, mock_u_g, mock_c_g
    ):
        """Test that admins can add charts to any dashboard"""
        from superset.models.dashboard import Dashboard

        admin = security_manager.find_user(username="admin")
        alpha = security_manager.find_user(username="alpha")

        # Set user context for dashboard creation
        mock_u_g.user = mock_c_g.user = mock_sm_g.user = alpha

        # Create chart owned by admin
        chart = db.session.query(Slice).first()
        chart.owners = [admin]

        # Create private dashboard owned by alpha
        alpha_dashboard = Dashboard(
            dashboard_title="Alpha Private Dashboard",
            slug="alpha-private-dashboard",
            owners=[alpha],
            published=False,
        )
        db.session.add(alpha_dashboard)
        db.session.commit()

        # Admin should be able to add chart to any dashboard
        mock_u_g.user = mock_c_g.user = mock_sm_g.user = admin

        json_obj = {
            "description": "Admin adding to any dashboard",
            "dashboards": [alpha_dashboard.id],
        }
        command = UpdateChartCommand(chart.id, json_obj)
        command.run()

        # Should succeed - admin has access to all dashboards
        updated_chart = db.session.query(Slice).get(chart.id)
        assert alpha_dashboard in updated_chart.dashboards

        # Clean up
        db.session.delete(alpha_dashboard)
        db.session.commit()


class TestChartWarmUpCacheCommand(SupersetTestCase):
    def test_warm_up_cache_command_chart_not_found(self):
        with self.assertRaises(WarmUpCacheChartNotFoundError):  # noqa: PT027
            ChartWarmUpCacheCommand(99999, None, None).run()

    @pytest.mark.usefixtures("load_birth_names_dashboard_with_slices")
    @pytest.mark.skip(reason="This test will be changed to use the api/v1/data")
    def test_warm_up_cache(self):
        slc = self.get_slice("Top 10 Girl Name Share")
        result = ChartWarmUpCacheCommand(slc.id, None, None).run()
        assert result == {
            "chart_id": slc.id,
            "viz_error": None,
            "viz_status": "success",
        }

        # can just pass in chart as well
        result = ChartWarmUpCacheCommand(slc, None, None).run()
        assert result == {
            "chart_id": slc.id,
            "viz_error": None,
            "viz_status": "success",
        }


class TestFavoriteChartCommand(SupersetTestCase):
    @pytest.mark.usefixtures("load_energy_table_with_slice")
    def test_fave_unfave_chart_command(self):
        """Test that a user can fave/unfave a chart"""
        with self.client.application.test_request_context():
            example_chart = db.session.query(Slice).all()[0]

            # Assert that the chart exists
            assert example_chart is not None

            with override_user(security_manager.find_user("admin")):
                AddFavoriteChartCommand(example_chart.id).run()

                # Assert that the dashboard was faved
                ids = ChartDAO.favorited_ids([example_chart])
                assert example_chart.id in ids

                DelFavoriteChartCommand(example_chart.id).run()

                # Assert that the chart was unfaved
                ids = ChartDAO.favorited_ids([example_chart])
                assert example_chart.id not in ids

    @pytest.mark.usefixtures("load_energy_table_with_slice")
    def test_fave_unfave_chart_command_not_found(self):
        """Test that faving / unfaving a non-existing chart raises an exception"""
        with self.client.application.test_request_context():
            example_chart_id = 1234

            with override_user(security_manager.find_user("admin")):
                with self.assertRaises(ChartNotFoundError):  # noqa: PT027
                    AddFavoriteChartCommand(example_chart_id).run()

                with self.assertRaises(ChartNotFoundError):  # noqa: PT027
                    DelFavoriteChartCommand(example_chart_id).run()

    @pytest.mark.usefixtures("load_energy_table_with_slice")
    @patch("superset.daos.base.BaseDAO.find_by_id")
    def test_fave_unfave_chart_command_forbidden(self, mock_find_by_id):
        """Test that faving / unfaving raises an exception for a chart the user doesn't own"""  # noqa: E501
        with self.client.application.test_request_context():
            example_chart = db.session.query(Slice).all()[0]
            mock_find_by_id.return_value = example_chart

            # Assert that the chart exists
            assert example_chart is not None

            with override_user(security_manager.find_user("gamma")):
                with self.assertRaises(ChartForbiddenError):  # noqa: PT027
                    AddFavoriteChartCommand(example_chart.id).run()

                with self.assertRaises(ChartForbiddenError):  # noqa: PT027
                    DelFavoriteChartCommand(example_chart.id).run()