betodealmeida commented on code in PR #37183:
URL: https://github.com/apache/superset/pull/37183#discussion_r2701266150
##########
superset/mcp_service/chart/tool/get_chart_data.py:
##########
@@ -123,16 +146,87 @@ async def get_chart_data( # noqa: C901
start_time = time.time()
+ # Track whether we're using unsaved state
+ using_unsaved_state = False
+ cached_form_data_dict = None
+
try:
await ctx.report_progress(2, 4, "Preparing data query")
from superset.charts.schemas import ChartDataQueryContextSchema
from superset.commands.chart.data.get_data_command import
ChartDataCommand
+ # Check if form_data_key is provided - use cached form_data instead
+ if request.form_data_key:
+ await ctx.info(
+ "Retrieving unsaved chart state from cache:
form_data_key=%s"
+ % (request.form_data_key,)
+ )
+ cached_form_data = _get_cached_form_data(request.form_data_key)
+
+ if cached_form_data:
Review Comment:
```suggestion
if cached_form_data :=
_get_cached_form_data(request.form_data_key):
```
##########
superset/mcp_service/chart/tool/get_chart_data.py:
##########
@@ -123,16 +146,87 @@ async def get_chart_data( # noqa: C901
start_time = time.time()
+ # Track whether we're using unsaved state
+ using_unsaved_state = False
+ cached_form_data_dict = None
+
try:
await ctx.report_progress(2, 4, "Preparing data query")
from superset.charts.schemas import ChartDataQueryContextSchema
from superset.commands.chart.data.get_data_command import
ChartDataCommand
+ # Check if form_data_key is provided - use cached form_data instead
+ if request.form_data_key:
+ await ctx.info(
+ "Retrieving unsaved chart state from cache:
form_data_key=%s"
+ % (request.form_data_key,)
+ )
+ cached_form_data = _get_cached_form_data(request.form_data_key)
+
+ if cached_form_data:
+ try:
+ cached_form_data_dict =
utils_json.loads(cached_form_data)
+ using_unsaved_state = True
+ await ctx.info(
+ "Using cached form_data from form_data_key for
data query"
+ )
+ except (TypeError, ValueError) as e:
+ await ctx.warning(
+ "Failed to parse cached form_data: %s. "
+ "Falling back to saved chart configuration." %
str(e)
+ )
+ else:
+ await ctx.warning(
+ "form_data_key provided but no cached data found. "
+ "The cache may have expired. Using saved chart
configuration."
+ )
+
# Use the chart's saved query_context - this is the key!
# The query_context contains all the information needed to
reproduce
# the chart's data exactly as shown in the visualization
query_context_json = None
- if chart.query_context:
+
+ # If using cached form_data, we need to build query_context from it
+ if using_unsaved_state and cached_form_data_dict:
+ # Build query context from cached form_data (unsaved state)
+ from superset.common.query_context_factory import
QueryContextFactory
Review Comment:
Let's move imports to top-level, unless they're here to prevent circular
imports.
##########
superset/mcp_service/chart/tool/get_chart_info.py:
##########
@@ -36,6 +36,22 @@
logger = logging.getLogger(__name__)
+def _get_cached_form_data(form_data_key: str) -> str | None:
+ """Retrieve form_data from cache using form_data_key.
+
+ Returns the JSON string of form_data if found, None otherwise.
+ """
+ from superset.commands.explore.form_data.get import GetFormDataCommand
+ from superset.commands.explore.form_data.parameters import
CommandParameters
Review Comment:
Move to top-level if possible.
##########
superset/mcp_service/dashboard/tool/get_dashboard_info.py:
##########
@@ -40,6 +41,23 @@
logger = logging.getLogger(__name__)
+def _get_permalink_state(permalink_key: str) -> Dict[str, Any] | None:
+ """Retrieve dashboard filter state from permalink.
+
+ Returns the permalink value containing dashboardId and state if found,
+ None otherwise.
+ """
+ from superset.commands.dashboard.permalink.get import
GetDashboardPermalinkCommand
+
+ try:
+ result = GetDashboardPermalinkCommand(permalink_key).run()
+ # Convert TypedDict to regular dict for type compatibility
+ return dict(result) if result else None
Review Comment:
`result` is already a dict, all you're doing here is losing type information.
It's better to make:
```python
def _get_permalink_state(permalink_key: str) -> DashboardPermalinkValue |
None:
```
And fix whatever is calling `_get_permalink_state()`.
##########
superset/mcp_service/dashboard/tool/get_dashboard_info.py:
##########
@@ -40,6 +41,23 @@
logger = logging.getLogger(__name__)
+def _get_permalink_state(permalink_key: str) -> Dict[str, Any] | None:
+ """Retrieve dashboard filter state from permalink.
+
+ Returns the permalink value containing dashboardId and state if found,
+ None otherwise.
+ """
+ from superset.commands.dashboard.permalink.get import
GetDashboardPermalinkCommand
+
+ try:
+ result = GetDashboardPermalinkCommand(permalink_key).run()
+ # Convert TypedDict to regular dict for type compatibility
+ return dict(result) if result else None
+ except Exception as e:
Review Comment:
+1
##########
superset/mcp_service/chart/tool/get_chart_data.py:
##########
@@ -42,6 +42,22 @@
logger = logging.getLogger(__name__)
+def _get_cached_form_data(form_data_key: str) -> str | None:
+ """Retrieve form_data from cache using form_data_key.
+
+ Returns the JSON string of form_data if found, None otherwise.
+ """
+ from superset.commands.explore.form_data.get import GetFormDataCommand
+ from superset.commands.explore.form_data.parameters import
CommandParameters
+
+ try:
+ cmd_params = CommandParameters(key=form_data_key)
+ return GetFormDataCommand(cmd_params).run()
+ except Exception as e:
Review Comment:
^ this
##########
superset/mcp_service/chart/tool/get_chart_info.py:
##########
@@ -36,6 +36,22 @@
logger = logging.getLogger(__name__)
+def _get_cached_form_data(form_data_key: str) -> str | None:
+ """Retrieve form_data from cache using form_data_key.
+
+ Returns the JSON string of form_data if found, None otherwise.
+ """
+ from superset.commands.explore.form_data.get import GetFormDataCommand
+ from superset.commands.explore.form_data.parameters import
CommandParameters
+
+ try:
+ cmd_params = CommandParameters(key=form_data_key)
+ return GetFormDataCommand(cmd_params).run()
+ except Exception as e:
+ logger.warning("Failed to retrieve form_data from cache: %s", e)
+ return None
Review Comment:
+1
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]