feat: Improve state key generation for dashboards and charts (#18576)

* feat: Improve state key generation for dashboards and charts

superset/explore/form_data/commands/create.py

@@ -15,8 +15,8 @@
 # specific language governing permissions and limitations
 # under the License.
 import logging
-from secrets import token_urlsafe
 
+from flask import session
 from sqlalchemy.exc import SQLAlchemyError
 
 from superset.commands.base import BaseCommand
@@ -25,6 +25,7 @@ from superset.explore.form_data.commands.state import TemporaryExploreState
 from superset.explore.form_data.utils import check_access
 from superset.extensions import cache_manager
 from superset.key_value.commands.exceptions import KeyValueCreateFailedError
+from superset.key_value.utils import cache_key, random_key
 
 logger = logging.getLogger(__name__)
 
@@ -37,10 +38,14 @@ class CreateFormDataCommand(BaseCommand):
         try:
             dataset_id = self._cmd_params.dataset_id
             chart_id = self._cmd_params.chart_id
+            tab_id = self._cmd_params.tab_id
             actor = self._cmd_params.actor
             form_data = self._cmd_params.form_data
             check_access(dataset_id, chart_id, actor)
-            key = token_urlsafe(48)
+            contextual_key = cache_key(session.get("_id"), tab_id, dataset_id, chart_id)
+            key = cache_manager.explore_form_data_cache.get(contextual_key)
+            if not key or not tab_id:
+                key = random_key()
             if form_data:
                 state: TemporaryExploreState = {
                     "owner": actor.get_user_id(),
@@ -49,6 +54,7 @@ class CreateFormDataCommand(BaseCommand):
                     "form_data": form_data,
                 }
                 cache_manager.explore_form_data_cache.set(key, state)
+                cache_manager.explore_form_data_cache.set(contextual_key, key)
             return key
         except SQLAlchemyError as ex:
             logger.exception("Error running create command")

superset/explore/form_data/commands/delete.py

@@ -17,6 +17,7 @@
 import logging
 from abc import ABC
 
+from flask import session
 from sqlalchemy.exc import SQLAlchemyError
 
 from superset.commands.base import BaseCommand
@@ -28,6 +29,7 @@ from superset.key_value.commands.exceptions import (
     KeyValueAccessDeniedError,
     KeyValueDeleteFailedError,
 )
+from superset.key_value.utils import cache_key
 
 logger = logging.getLogger(__name__)
 
@@ -44,9 +46,16 @@ class DeleteFormDataCommand(BaseCommand, ABC):
                 key
             )
             if state:
-                check_access(state["dataset_id"], state["chart_id"], actor)
+                dataset_id = state["dataset_id"]
+                chart_id = state["chart_id"]
+                check_access(dataset_id, chart_id, actor)
                 if state["owner"] != actor.get_user_id():
                     raise KeyValueAccessDeniedError()
+                tab_id = self._cmd_params.tab_id
+                contextual_key = cache_key(
+                    session.get("_id"), tab_id, dataset_id, chart_id
+                )
+                cache_manager.explore_form_data_cache.delete(contextual_key)
                 return cache_manager.explore_form_data_cache.delete(key)
             return False
         except SQLAlchemyError as ex:

superset/explore/form_data/commands/parameters.py

@@ -25,5 +25,6 @@ class CommandParameters:
     actor: User
     dataset_id: int = 0
     chart_id: int = 0
+    tab_id: Optional[int] = None
     key: Optional[str] = None
     form_data: Optional[str] = None

superset/explore/form_data/commands/update.py

@@ -16,7 +16,9 @@
 # under the License.
 import logging
 from abc import ABC
+from typing import Optional
 
+from flask import session
 from sqlalchemy.exc import SQLAlchemyError
 
 from superset.commands.base import BaseCommand
@@ -28,6 +30,7 @@ from superset.key_value.commands.exceptions import (
     KeyValueAccessDeniedError,
     KeyValueUpdateFailedError,
 )
+from superset.key_value.utils import cache_key, random_key
 
 logger = logging.getLogger(__name__)
 
@@ -38,7 +41,7 @@ class UpdateFormDataCommand(BaseCommand, ABC):
     ):
         self._cmd_params = cmd_params
 
-    def run(self) -> bool:
+    def run(self) -> Optional[str]:
         try:
             dataset_id = self._cmd_params.dataset_id
             chart_id = self._cmd_params.chart_id
@@ -53,14 +56,25 @@ class UpdateFormDataCommand(BaseCommand, ABC):
                 user_id = actor.get_user_id()
                 if state["owner"] != user_id:
                     raise KeyValueAccessDeniedError()
+
+                # Generate a new key if tab_id changes or equals 0
+                tab_id = self._cmd_params.tab_id
+                contextual_key = cache_key(
+                    session.get("_id"), tab_id, dataset_id, chart_id
+                )
+                key = cache_manager.explore_form_data_cache.get(contextual_key)
+                if not key or not tab_id:
+                    key = random_key()
+                    cache_manager.explore_form_data_cache.set(contextual_key, key)
+
                 new_state: TemporaryExploreState = {
                     "owner": actor.get_user_id(),
                     "dataset_id": dataset_id,
                     "chart_id": chart_id,
                     "form_data": form_data,
                 }
-                return cache_manager.explore_form_data_cache.set(key, new_state)
-            return False
+                cache_manager.explore_form_data_cache.set(key, new_state)
+            return key
         except SQLAlchemyError as ex:
             logger.exception("Error running update command")
             raise KeyValueUpdateFailedError() from ex