Repository: apache/superset
File: superset/mcp_service/chart/tool/get_chart_data.py (Python)

Code before:
            )
        )
        logger.info("Getting data for chart %s: %s", chart.id, chart.slice_name)

        import time

        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 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)

Review comment (performance):
```suggestion
                if cached_form_data := _get_cached_form_data(request.form_data_key):
```

Code after:
            "Chart found successfully: chart_id=%s, chart_name=%s, viz_type=%s"
            % (
                chart.id,
                chart.slice_name,
                chart.viz_type,
            )
        )
        logger.info("Getting data for chart %s: %s", chart.id, chart.slice_name)

        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,)
                )
                if cached_form_data := _get_cached_form_data(request.form_data_key):
                    try:
                        parsed_form_data = utils_json.loads(cached_form_data)
                        # Only use if it's actually a dict (not null, list, etc.)
                        if isinstance(parsed_form_data, dict):
                            cached_form_data_dict = parsed_form_data
                            using_unsaved_state = True
                            await ctx.info(
                                "Using cached form_data from form_data_key "
                                "for data query"
                            )
                        else:
                            await ctx.warning(
                                "Cached form_data is not a JSON object. "
                                "Falling back to saved chart configuration."
                            )
                    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."