From 35864748f2d7dde53ebbe707bb7a79ef5cf8765d Mon Sep 17 00:00:00 2001
From: Ville Brofeldt <33317356+villebro@users.noreply.github.com>
Date: Tue, 24 Aug 2021 18:37:34 +0300
Subject: [PATCH] fix(explore): retain chart ownership on query context update
 (#16419)

---
 superset/charts/commands/update.py            | 12 ++++----
 .../charts/commands_tests.py                  | 30 ++++++++++++++++++-
 2 files changed, 34 insertions(+), 8 deletions(-)

diff --git a/superset/charts/commands/update.py b/superset/charts/commands/update.py
index b09845512c0c4..c63d0bccca2e0 100644
--- a/superset/charts/commands/update.py
+++ b/superset/charts/commands/update.py
@@ -84,13 +84,17 @@ def validate(self) -> None:
         if not self._model:
             raise ChartNotFoundError()
 
-        # Check ownership; when only updating query context we ignore
+        # Check and update ownership; when only updating query context we ignore
         # ownership so the update can be performed by report workers
         if not is_query_context_update(self._properties):
             try:
                 check_ownership(self._model)
+                owners = self.populate_owners(self._actor, owner_ids)
+                self._properties["owners"] = owners
             except SupersetSecurityException as ex:
                 raise ChartForbiddenError() from ex
+            except ValidationError as ex:
+                exceptions.append(ex)
 
         # Validate/Populate datasource
         if datasource_id is not None:
@@ -107,12 +111,6 @@ def validate(self) -> None:
                 exceptions.append(DashboardsNotFoundValidationError())
             self._properties["dashboards"] = dashboards
 
-        # Validate/Populate owner
-        try:
-            owners = self.populate_owners(self._actor, owner_ids)
-            self._properties["owners"] = owners
-        except ValidationError as ex:
-            exceptions.append(ex)
         if exceptions:
             exception = ChartInvalidError()
             exception.add_list(exceptions)
diff --git a/tests/integration_tests/charts/commands_tests.py b/tests/integration_tests/charts/commands_tests.py
index 238a54e29c7f4..7e718f8ebd07e 100644
--- a/tests/integration_tests/charts/commands_tests.py
+++ b/tests/integration_tests/charts/commands_tests.py
@@ -321,7 +321,7 @@ def test_update_v1_response(self, mock_sm_g, mock_g):
         json_obj = {
             "description": "test for update",
             "cache_timeout": None,
-            "owners": [1],
+            "owners": [actor.id],
         }
         command = UpdateChartCommand(actor, model_id, json_obj)
         last_saved_before = db.session.query(Slice).get(pk).last_saved_at
@@ -329,3 +329,31 @@ def test_update_v1_response(self, mock_sm_g, mock_g):
         chart = db.session.query(Slice).get(pk)
         assert chart.last_saved_at != last_saved_before
         assert chart.last_saved_by == actor
+
+    @patch("superset.views.base.g")
+    @patch("superset.security.manager.g")
+    @pytest.mark.usefixtures("load_energy_table_with_slice")
+    def test_query_context_update_command(self, mock_sm_g, mock_g):
+        """"
+        Test that a user can generate the chart query context
+        payloadwithout 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()
+
+        actor = security_manager.find_user(username="alpha")
+        mock_g.user = mock_sm_g.user = actor
+        query_context = json.dumps({"foo": "bar"})
+        json_obj = {
+            "query_context_generation": True,
+            "query_context": query_context,
+        }
+        command = UpdateChartCommand(actor, 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