chore(mcp): extract shared chart helpers and ASCII rendering into separate modules (#39438)

superset/mcp_service/chart/tool/get_chart_data.py

@@ -25,15 +25,19 @@ from typing import Any, Dict, List, TYPE_CHECKING
 
 from fastmcp import Context
 from flask import current_app
+from sqlalchemy.exc import SQLAlchemyError
 from superset_core.mcp.decorators import tool, ToolAnnotations
 
 if TYPE_CHECKING:
     from superset.models.slice import Slice
 
 from superset.commands.exceptions import CommandException
-from superset.commands.explore.form_data.parameters import CommandParameters
 from superset.exceptions import OAuth2Error, OAuth2RedirectError, SupersetException
 from superset.extensions import event_logger
+from superset.mcp_service.chart.chart_helpers import (
+    find_chart_by_identifier,
+    get_cached_form_data,
+)
 from superset.mcp_service.chart.chart_utils import validate_chart_dataset
 from superset.mcp_service.chart.schemas import (
     ChartData,
@@ -62,21 +66,6 @@ def _apply_extra_form_data(
     merge_extra_filters(form_data)
 
 
-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
-
-    try:
-        cmd_params = CommandParameters(key=form_data_key)
-        return GetFormDataCommand(cmd_params).run()
-    except (KeyError, ValueError, CommandException) as e:
-        logger.warning("Failed to retrieve form_data from cache: %s", e)
-        return None
-
-
 @tool(
     tags=["data"],
     class_permission_name="Chart",
@@ -127,7 +116,6 @@ async def get_chart_data(  # noqa: C901
 
     try:
         await ctx.report_progress(1, 4, "Looking up chart")
-        from superset.daos.chart import ChartDAO
         from superset.utils import json as utils_json
 
         chart = None
@@ -141,7 +129,7 @@ async def get_chart_data(  # noqa: C901
                     "No chart identifier - querying data from unsaved chart cache: "
                     "form_data_key=%s" % (request.form_data_key,)
                 )
-                cached_form_data = _get_cached_form_data(request.form_data_key)
+                cached_form_data = get_cached_form_data(request.form_data_key)
                 if not cached_form_data:
                     return ChartError(
                         error="No cached chart data found for form_data_key. "
@@ -166,25 +154,13 @@ async def get_chart_data(  # noqa: C901
 
         # Find the chart by identifier
         with event_logger.log_context(action="mcp.get_chart_data.chart_lookup"):
-            if isinstance(request.identifier, int) or (
-                isinstance(request.identifier, str) and request.identifier.isdigit()
-            ):
-                chart_id = (
-                    int(request.identifier)
-                    if isinstance(request.identifier, str)
-                    else request.identifier
+            await ctx.debug("Looking up chart: identifier=%s" % (request.identifier,))
+            if request.identifier is None:
+                return ChartError(
+                    error="Chart identifier is required",
+                    error_type="ValidationError",
                 )
-                await ctx.debug(
-                    "Performing ID-based chart lookup: chart_id=%s" % (chart_id,)
-                )
-                chart = ChartDAO.find_by_id(chart_id)
-            elif isinstance(request.identifier, str):
-                await ctx.debug(
-                    "Performing UUID-based chart lookup: uuid=%s"
-                    % (request.identifier,)
-                )
-                # Try UUID lookup using DAO flexible method
-                chart = ChartDAO.find_by_id(request.identifier, id_column="uuid")
+            chart = find_chart_by_identifier(request.identifier)
 
         if not chart:
             await ctx.error("Chart not found: identifier=%s" % (request.identifier,))
@@ -239,7 +215,7 @@ async def get_chart_data(  # noqa: C901
                         "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):
+                    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.)
@@ -818,7 +794,15 @@ async def get_chart_data(  # noqa: C901
             error=OAUTH2_CONFIG_ERROR_MESSAGE,
             error_type="OAUTH2_REDIRECT_ERROR",
         )
-    except Exception as e:
+    except (
+        SupersetException,
+        CommandException,
+        SQLAlchemyError,
+        KeyError,
+        ValueError,
+        TypeError,
+        AttributeError,
+    ) as e:
         await ctx.error(
             "Chart data retrieval failed: identifier=%s, error=%s, error_type=%s"
             % (