fix(mcp): prevent encoding errors and fix tool bugs on MCP client transports (#38786)

superset/mcp_service/sql_lab/tool/execute_sql.py

@@ -28,6 +28,7 @@ import logging
 from decimal import Decimal
 from typing import Any
 
+import pandas as pd
 from fastmcp import Context
 from superset_core.mcp.decorators import tool, ToolAnnotations
 from superset_core.queries.types import (
@@ -176,6 +177,46 @@ def _sanitize_row_values(rows: list[dict[str, Any]]) -> None:
                 row[key] = str(value)
 
 
+def _data_to_statement_data(data: Any) -> StatementData:
+    """Convert statement data (DataFrame, list, dict, bytes) to StatementData.
+
+    When results come from cache, data may be a dict/list/bytes instead of
+    a pandas DataFrame. This function handles all cases defensively.
+    """
+    from superset.utils import json as json_utils
+
+    if isinstance(data, list):
+        rows_data = data
+    elif isinstance(data, dict):
+        rows_data = data.get("data", [data])
+        if not isinstance(rows_data, list):
+            rows_data = [rows_data]
+    elif isinstance(data, pd.DataFrame):
+        rows_data = data.to_dict(orient="records")
+        _sanitize_row_values(rows_data)
+        return StatementData(
+            rows=rows_data,
+            columns=[
+                ColumnInfo(name=col, type=str(data[col].dtype)) for col in data.columns
+            ],
+        )
+    elif isinstance(data, bytes):
+        try:
+            decoded = json_utils.loads(data)
+            rows_data = decoded if isinstance(decoded, list) else [decoded]
+        except (ValueError, UnicodeDecodeError):
+            rows_data = []
+    else:
+        rows_data = [{"value": str(data)}]
+
+    _sanitize_row_values(rows_data)
+    col_names = list(rows_data[0].keys()) if rows_data else []
+    return StatementData(
+        rows=rows_data,
+        columns=[ColumnInfo(name=col, type="object") for col in col_names],
+    )
+
+
 def _convert_to_response(result: QueryResult) -> ExecuteSqlResponse:
     """Convert QueryResult to ExecuteSqlResponse."""
     if result.status != QueryStatus.SUCCESS:
@@ -193,15 +234,7 @@ def _convert_to_response(result: QueryResult) -> ExecuteSqlResponse:
     for stmt in result.statements:
         stmt_data: StatementData | None = None
         if stmt.data is not None:
-            df = stmt.data
-            rows_data = df.to_dict(orient="records")
-            _sanitize_row_values(rows_data)
-            stmt_data = StatementData(
-                rows=rows_data,
-                columns=[
-                    ColumnInfo(name=col, type=str(df[col].dtype)) for col in df.columns
-                ],
-            )
+            stmt_data = _data_to_statement_data(stmt.data)
             data_bearing_count += 1
 
         statements.append(