This is an automated email from the ASF dual-hosted git repository.
aminghadersohi pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/superset.git
The following commit(s) were added to refs/heads/master by this push:
new 7a9fe06762b fix(mcp): stop data tool responses crashing on binary and
NaT values (#44495)
7a9fe06762b is described below
commit 7a9fe06762b45d104ad098779116afdfee50bf30
Author: Amin Ghadersohi <[email protected]>
AuthorDate: Wed Sep 23 10:39:01 2026 +1000
fix(mcp): stop data tool responses crashing on binary and NaT values
(#44495)
---
docs/admin_docs/configuration/mcp-server.mdx | 15 +
superset/mcp_service/chart/schemas.py | 33 +--
superset/mcp_service/chart/tool/get_chart_data.py | 31 +-
superset/mcp_service/dashboard/schemas.py | 13 +-
superset/mcp_service/dataset/schemas.py | 13 +-
superset/mcp_service/semantic_layer/schemas.py | 11 +-
superset/mcp_service/sql_lab/schemas.py | 18 +-
superset/mcp_service/sql_lab/tool/execute_sql.py | 19 --
superset/mcp_service/utils/response_utils.py | 8 +-
superset/mcp_service/utils/serialization.py | 225 +++++++++++++++
.../mcp_service/chart/tool/test_get_chart_data.py | 25 +-
.../mcp_service/sql_lab/tool/test_execute_sql.py | 104 +++----
.../mcp_service/test_response_serialization.py | 321 +++++++++++++++++++++
.../mcp_service/utils/test_serialization.py | 294 +++++++++++++++++++
14 files changed, 1002 insertions(+), 128 deletions(-)
diff --git a/docs/admin_docs/configuration/mcp-server.mdx
b/docs/admin_docs/configuration/mcp-server.mdx
index d07c06e54c2..319db21ccc4 100644
--- a/docs/admin_docs/configuration/mcp-server.mdx
+++ b/docs/admin_docs/configuration/mcp-server.mdx
@@ -1175,6 +1175,21 @@ answer. Dashboard access checks and row-level security
remain independent of thi
display-context projection. Integrations must ensure the lookup actually
happens
and the assistant uses the result; tool guidance alone cannot guarantee this.
+## Result serialization
+
+Data-bearing tool responses sanitize row values, column samples, and statistics
+even when Python producers supply tuples, sets, generators, or read-only
mappings.
+Missing and non-finite values become JSON `null`; non-UTF-8 binary values
become
+base64 strings prefixed with `base64:`. Boolean row counts are treated as
+unavailable, producing `null` for nullable counts and `0` for non-nullable
counts,
+rather than being interpreted as one or zero rows.
+
+## Chart data column counts
+
+For saved and unsaved charts, `get_chart_data` column metadata excludes missing
+values from `unique_count`. Values serialized as JSON `null`, including NaN,
+NaT, and non-finite numbers, contribute to `null_count` instead.
+
## Dataset query dimensions and nested fields
The `query_dataset` tool accepts exact column names registered on the selected
diff --git a/superset/mcp_service/chart/schemas.py
b/superset/mcp_service/chart/schemas.py
index 4113cf02408..077bb756c9d 100644
--- a/superset/mcp_service/chart/schemas.py
+++ b/superset/mcp_service/chart/schemas.py
@@ -75,6 +75,13 @@ from superset.mcp_service.utils.sanitization import (
sanitize_user_input,
sanitize_user_input_with_changes,
)
+from superset.mcp_service.utils.serialization import (
+ JsonSafeMapping,
+ JsonSafeRows,
+ JsonSafeValues,
+ OptionalRowCount,
+ RowCount,
+)
logger = logging.getLogger(__name__)
@@ -4026,7 +4033,7 @@ class DataColumn(BaseModel):
name: str = Field(..., description="Column name")
display_name: str = Field(..., description="Human-readable column name")
data_type: str = Field(..., description="Inferred data type")
- sample_values: List[Any] = Field(description="Representative sample
values")
+ sample_values: JsonSafeValues = Field(description="Representative sample
values")
null_count: int = Field(
description="Number of null values. Approximate — see
'statistics.sampled_rows'"
" if the source result set was larger than the row cap used to compute
it."
@@ -4036,7 +4043,7 @@ class DataColumn(BaseModel):
"'statistics.sampled_rows' if the source result set was larger than
the "
"row cap used to compute it."
)
- statistics: Dict[str, Any] | None = Field(
+ statistics: JsonSafeMapping | None = Field(
None,
description="Additional column statistics, when available. May include
"
"'sampled_rows' (any column type) when null_count/unique_count were "
@@ -4052,9 +4059,9 @@ class ChartQueryResult(BaseModel):
query_index: int = Field(description="Zero-based query position")
columns: list[str] = Field(description="Column names returned by the
query")
- data: list[dict[str, Any]] = Field(description="Actual data rows")
- row_count: int = Field(description="Rows returned")
- total_rows: int | None = Field(None, description="Total available rows")
+ data: JsonSafeRows = Field(description="Actual data rows")
+ row_count: RowCount = Field(description="Rows returned")
+ total_rows: OptionalRowCount = Field(None, description="Total available
rows")
class ChartData(BaseModel):
@@ -4067,7 +4074,7 @@ class ChartData(BaseModel):
# Enhanced data description
columns: List[DataColumn] = Field(description="Rich column metadata")
- data: List[Dict[str, Any]] = Field(description="Actual data rows")
+ data: JsonSafeRows = Field(description="Actual data rows")
query_results: list[ChartQueryResult] | None = Field(
None,
description=(
@@ -4077,18 +4084,8 @@ class ChartData(BaseModel):
)
# Data insights
- row_count: int = Field(description="Rows returned")
- total_rows: int | None = Field(description="Total available rows")
-
- @field_validator("total_rows", mode="before")
- @classmethod
- def _coerce_total_rows(cls, v: Any) -> int | None:
- if v is None:
- return None
- try:
- return int(v)
- except (TypeError, ValueError):
- return None
+ row_count: RowCount = Field(description="Rows returned")
+ total_rows: OptionalRowCount = Field(description="Total available rows")
data_freshness: datetime | None = Field(description="When data was last
updated")
diff --git a/superset/mcp_service/chart/tool/get_chart_data.py
b/superset/mcp_service/chart/tool/get_chart_data.py
index dff7cf87b41..5369a436c3c 100644
--- a/superset/mcp_service/chart/tool/get_chart_data.py
+++ b/superset/mcp_service/chart/tool/get_chart_data.py
@@ -60,6 +60,7 @@ from superset.mcp_service.utils.oauth2_utils import (
build_oauth2_redirect_message,
OAUTH2_CONFIG_ERROR_MESSAGE,
)
+from superset.mcp_service.utils.serialization import is_missing_value
from superset.utils.core import GenericDataType
logger = logging.getLogger(__name__)
@@ -858,7 +859,7 @@ async def execute_chart_data( # noqa: C901
sample_values = [
row.get(col_name)
for row in data[:3]
- if row.get(col_name) is not None
+ if not is_missing_value(row.get(col_name))
]
# Use SQL-derived GenericDataType when available,
@@ -878,8 +879,16 @@ async def execute_chart_data( # noqa: C901
display_name=col_name.replace("_", " ").title(),
data_type=data_type,
sample_values=sample_values[:3],
- null_count=sum(1 for row in data if row.get(col_name)
is None),
- unique_count=len({str(row.get(col_name)) for row in
data}),
+ null_count=sum(
+ 1 for row in data if
is_missing_value(row.get(col_name))
+ ),
+ unique_count=len(
+ {
+ str(row.get(col_name))
+ for row in data
+ if not is_missing_value(row.get(col_name))
+ }
+ ),
)
)
@@ -1199,7 +1208,9 @@ async def _query_from_form_data( # noqa: C901
columns = []
for col_name in raw_columns:
sample_values = [
- row.get(col_name) for row in data[:3] if row.get(col_name) is
not None
+ row.get(col_name)
+ for row in data[:3]
+ if not is_missing_value(row.get(col_name))
]
data_type = "string"
if sample_values and all(
@@ -1212,8 +1223,16 @@ async def _query_from_form_data( # noqa: C901
display_name=col_name.replace("_", " ").title(),
data_type=data_type,
sample_values=sample_values[:3],
- null_count=sum(1 for row in data if row.get(col_name) is
None),
- unique_count=len({str(row.get(col_name)) for row in data}),
+ null_count=sum(
+ 1 for row in data if
is_missing_value(row.get(col_name))
+ ),
+ unique_count=len(
+ {
+ str(row.get(col_name))
+ for row in data
+ if not is_missing_value(row.get(col_name))
+ }
+ ),
)
)
diff --git a/superset/mcp_service/dashboard/schemas.py
b/superset/mcp_service/dashboard/schemas.py
index 498f3c9a215..935b268ea80 100644
--- a/superset/mcp_service/dashboard/schemas.py
+++ b/superset/mcp_service/dashboard/schemas.py
@@ -117,6 +117,7 @@ from superset.mcp_service.utils.sanitization import (
sanitize_user_input,
sanitize_user_input_with_changes,
)
+from superset.mcp_service.utils.serialization import JsonSafeRows,
OptionalRowCount
from superset.mcp_service.utils.url_utils import get_superset_base_url
from superset.utils.core import DatasourceType
from superset.utils.json import loads as json_loads
@@ -2907,11 +2908,11 @@ class DashboardChartQueryData(BaseModel):
query_index: int = Field(..., description="Zero-based query position")
columns: list[str] = Field(default_factory=list, description="Result
column names")
- sample_data: list[dict[str, Any]] = Field(
+ sample_data: JsonSafeRows = Field(
default_factory=list, description="A few example data rows"
)
- row_count: int | None = Field(None, description="Rows returned by this
query")
- total_rows: int | None = Field(
+ row_count: OptionalRowCount = Field(None, description="Rows returned by
this query")
+ total_rows: OptionalRowCount = Field(
None, description="Total rows available for this query when known"
)
truncated: bool = Field(
@@ -2926,11 +2927,11 @@ class DashboardChartData(BaseModel):
chart_name: str = Field(..., description="Chart name")
chart_type: str = Field(..., description="Chart viz type")
columns: list[str] = Field(default_factory=list, description="Result
column names")
- sample_data: list[dict[str, Any]] = Field(
+ sample_data: JsonSafeRows = Field(
default_factory=list, description="A few example data rows"
)
- row_count: int | None = Field(None, description="Rows returned by the
query")
- total_rows: int | None = Field(
+ row_count: OptionalRowCount = Field(None, description="Rows returned by
the query")
+ total_rows: OptionalRowCount = Field(
None,
description=(
"Total rows available when known; null when the fetch was capped
with "
diff --git a/superset/mcp_service/dataset/schemas.py
b/superset/mcp_service/dataset/schemas.py
index 4e4de483a8d..3cc6b594dc4 100644
--- a/superset/mcp_service/dataset/schemas.py
+++ b/superset/mcp_service/dataset/schemas.py
@@ -60,6 +60,11 @@ from superset.mcp_service.system.schemas import (
TagInfo,
)
from superset.mcp_service.utils.response_utils import humanize_timestamp
+from superset.mcp_service.utils.serialization import (
+ JsonSafeRows,
+ OptionalRowCount,
+ RowCount,
+)
from superset.sql.parse import has_aggregate
from superset.utils import json
@@ -886,11 +891,9 @@ class QueryDatasetResponse(BaseModel):
columns: List[DataColumn] = Field(
default_factory=list, description="Column metadata for returned data"
)
- data: List[Dict[str, Any]] = Field(
- default_factory=list, description="Query result rows"
- )
- row_count: int = Field(0, description="Number of rows returned")
- total_rows: int | None = Field(
+ data: JsonSafeRows = Field(default_factory=list, description="Query result
rows")
+ row_count: RowCount = Field(0, description="Number of rows returned")
+ total_rows: OptionalRowCount = Field(
None, description="Total row count from the query engine"
)
from_dttm: datetime | None = Field(
diff --git a/superset/mcp_service/semantic_layer/schemas.py
b/superset/mcp_service/semantic_layer/schemas.py
index 7339afbe5fd..11ea6f7b743 100644
--- a/superset/mcp_service/semantic_layer/schemas.py
+++ b/superset/mcp_service/semantic_layer/schemas.py
@@ -28,6 +28,11 @@ from superset.mcp_service.chart.schemas import DataColumn,
PerformanceMetadata
from superset.mcp_service.common.cache_schemas import CacheStatus
from superset.mcp_service.common.error_schemas import MCPBaseError
from superset.mcp_service.common.time_range_validation import
validate_time_range
+from superset.mcp_service.utils.serialization import (
+ JsonSafeRows,
+ OptionalRowCount,
+ RowCount,
+)
# ---------------------------------------------------------------------------
# Shared error schema
@@ -302,9 +307,9 @@ class GetTableResponse(BaseModel):
"""Response schema for get_table."""
columns: list[DataColumn]
- data: list[dict[str, Any]]
- row_count: int
- total_rows: int | None = None
+ data: JsonSafeRows
+ row_count: RowCount
+ total_rows: OptionalRowCount = None
from_dttm: datetime | None = Field(
None,
description=(
diff --git a/superset/mcp_service/sql_lab/schemas.py
b/superset/mcp_service/sql_lab/schemas.py
index 7fc03bb708f..11cc8942aa0 100644
--- a/superset/mcp_service/sql_lab/schemas.py
+++ b/superset/mcp_service/sql_lab/schemas.py
@@ -28,6 +28,12 @@ from pydantic import (
model_serializer,
)
+from superset.mcp_service.utils.serialization import (
+ JsonSafeRows,
+ OptionalRowCount,
+ RowCount,
+)
+
class _SchemaFieldNormalizer(BaseModel):
"""Mixin that renames schema_name → schema in JSON output.
@@ -133,9 +139,7 @@ class ColumnInfo(BaseModel):
class StatementData(BaseModel):
"""Row data and column metadata for a single SQL statement."""
- rows: list[dict[str, Any]] = Field(
- ..., description="Result rows as list of dictionaries"
- )
+ rows: JsonSafeRows = Field(..., description="Result rows as list of
dictionaries")
columns: list[ColumnInfo] = Field(..., description="Column metadata
information")
@@ -146,7 +150,7 @@ class StatementInfo(BaseModel):
executed_sql: str = Field(
..., description="SQL after transformations (RLS, mutations, limits)"
)
- row_count: int = Field(..., description="Number of rows returned/affected")
+ row_count: RowCount = Field(..., description="Number of rows
returned/affected")
execution_time_ms: float | None = Field(
None, description="Statement execution time in milliseconds"
)
@@ -164,14 +168,14 @@ class ExecuteSqlResponse(BaseModel):
"""Response schema for SQL execution results."""
success: bool = Field(..., description="Whether query executed
successfully")
- rows: list[dict[str, Any]] | None = Field(
+ rows: JsonSafeRows | None = Field(
None, description="Query result rows as list of dictionaries"
)
columns: list[ColumnInfo] | None = Field(
None, description="Column metadata information"
)
- row_count: int | None = Field(None, description="Number of rows returned")
- affected_rows: int | None = Field(
+ row_count: OptionalRowCount = Field(None, description="Number of rows
returned")
+ affected_rows: OptionalRowCount = Field(
None, description="Number of rows affected (for DML queries)"
)
execution_time: float | None = Field(
diff --git a/superset/mcp_service/sql_lab/tool/execute_sql.py
b/superset/mcp_service/sql_lab/tool/execute_sql.py
index 4e5d980c765..fa4d13d9e6a 100644
--- a/superset/mcp_service/sql_lab/tool/execute_sql.py
+++ b/superset/mcp_service/sql_lab/tool/execute_sql.py
@@ -23,7 +23,6 @@ Database.execute() API with RLS, template rendering, and
security validation.
"""
import logging
-from decimal import Decimal
from typing import Any
import pandas as pd
@@ -311,22 +310,6 @@ async def _log_execution_result(
)
-def _sanitize_row_values(rows: list[dict[str, Any]]) -> None:
- """Sanitize non-serializable values in rows for JSON serialization."""
- for row in rows:
- for key, value in row.items():
- if isinstance(value, (bytes, memoryview)):
- raw = bytes(value) if isinstance(value, memoryview) else value
- try:
- row[key] = raw.decode("utf-8")
- except (UnicodeDecodeError, AttributeError):
- row[key] = raw.hex()
- elif isinstance(value, Decimal):
- row[key] = float(value)
- elif not isinstance(value, (str, int, float, bool, type(None),
list, dict)):
- row[key] = str(value)
-
-
def _data_to_statement_data(data: Any) -> StatementData:
"""Convert statement data (DataFrame, list, dict, bytes) to StatementData.
@@ -343,7 +326,6 @@ def _data_to_statement_data(data: Any) -> StatementData:
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=[
@@ -359,7 +341,6 @@ def _data_to_statement_data(data: Any) -> StatementData:
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,
diff --git a/superset/mcp_service/utils/response_utils.py
b/superset/mcp_service/utils/response_utils.py
index b50c4b89223..a3ba54f33b8 100644
--- a/superset/mcp_service/utils/response_utils.py
+++ b/superset/mcp_service/utils/response_utils.py
@@ -64,6 +64,8 @@ if TYPE_CHECKING:
import humanize
+from superset.mcp_service.utils.serialization import is_missing_value
+
def humanize_timestamp(dt: datetime | None) -> str | None:
"""Convert a datetime to a humanized string like '2 hours ago'."""
@@ -192,7 +194,9 @@ def format_data_columns(
columns_meta: list[DataColumn] = []
for col_name in raw_columns:
sample_values = [
- row.get(col_name) for row in data[:3] if row.get(col_name) is not
None
+ row.get(col_name)
+ for row in data[:3]
+ if not is_missing_value(row.get(col_name))
]
data_type: str = "string"
if sample_values:
@@ -207,7 +211,7 @@ def format_data_columns(
unique_vals: set[str] = set()
for row in stats_rows:
val = row.get(col_name)
- if val is None:
+ if is_missing_value(val):
null_count += 1
else:
unique_vals.add(str(val))
diff --git a/superset/mcp_service/utils/serialization.py
b/superset/mcp_service/utils/serialization.py
new file mode 100644
index 00000000000..e39893bd93b
--- /dev/null
+++ b/superset/mcp_service/utils/serialization.py
@@ -0,0 +1,225 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""JSON-safety helpers for MCP response schemas.
+
+FastMCP serialises every tool return value with
+``pydantic_core.to_json(value, fallback=str)``. The ``fallback`` only fires for
+types pydantic does not recognise, so values that *are* recognised but cannot
+be rendered as JSON raise ``PydanticSerializationError`` and abort the whole
+tool call — after the warehouse query already succeeded. Warehouse result sets
+routinely carry such values:
+
+* ``bytes`` from ``BLOB``/``BYTEA``/``VARBINARY`` columns that are not valid
+ UTF-8 (``invalid utf-8 sequence of 1 bytes from index N``);
+* ``pandas.NaT`` for a null timestamp — a ``datetime`` subclass backed by a
+ float, which pydantic tries to render as a datetime
+ (``TypeError: 'float' object cannot be interpreted as an integer``);
+* ``float`` NaN/Infinity, which serialise to the non-standard JSON literals
+ ``NaN``/``Infinity`` that strict JSON clients reject;
+* ``str`` holding lone surrogates, produced by drivers decoding with
+ ``surrogatepass`` (``surrogates not allowed``).
+
+Row counts are a second failure source: a DBAPI cursor may report ``rowcount``
+as a ``float``, which fails validation of an ``int`` field before serialisation
+is even reached.
+
+The annotated types exported here (:data:`JsonSafeRows`,
+:data:`JsonSafeValues`, :data:`JsonSafeMapping`, :data:`RowCount`,
+:data:`OptionalRowCount`) attach the coercion as a ``BeforeValidator``, so any
+response schema that uses them is protected at construction time. Prefer them
+over per-field validators so new data-bearing tools inherit the protection.
+"""
+
+from __future__ import annotations
+
+import base64
+import math
+from collections.abc import Mapping
+from datetime import date, datetime, time, timedelta
+from typing import Annotated, Any
+from uuid import UUID
+
+import numpy as np
+import pandas as pd
+from pydantic import BeforeValidator
+
+from superset.utils.json import base_json_conv
+
+# Undecodable binary is returned base64-encoded behind this marker. Base64 is
+# lossless and always valid UTF-8; decoding with a lenient codec instead (e.g.
+# UTF-16) can emit lone surrogates, which is a second way to break the
+# serializer.
+BINARY_PREFIX = "base64:"
+
+# Values nested deeper than this (a cyclic or pathologically nested JSON
+# column) are rendered as text rather than walked.
+MAX_DEPTH = 20
+
+# Types pydantic renders natively and that never need rewriting.
+_PASSTHROUGH_TYPES = (datetime, date, time, timedelta, UUID)
+
+# Marks a value the scalar pass did not handle; ``None`` is a real result.
+_UNHANDLED = object()
+
+
+def _decode_binary(value: bytes | bytearray | memoryview) -> str:
+ """Render binary column data as a JSON-safe string."""
+ raw = bytes(value)
+ try:
+ return _sanitize_str(raw.decode("utf-8"))
+ except UnicodeDecodeError:
+ return BINARY_PREFIX + base64.b64encode(raw).decode("ascii")
+
+
+def _sanitize_str(value: str) -> str:
+ """Replace lone surrogates, which cannot be encoded as UTF-8."""
+ if value.isascii():
+ return value
+ try:
+ value.encode("utf-8")
+ except UnicodeEncodeError:
+ return value.encode("utf-8", "replace").decode("utf-8")
+ return value
+
+
+def is_missing_value(value: Any) -> bool:
+ """True when a result value is null-like: ``None``, NaN, ``pandas.NaT``,
+ ``pandas.NA`` or ``numpy`` NaT.
+
+ Use this instead of ``value is None`` when summarising result data, so a
+ column full of NaN is not reported as having no nulls. This is exactly the
+ set of values :func:`sanitize_json_value` renders as JSON ``null``, so a
+ null count derived from it agrees with the payload.
+
+ ``pandas.isna`` returns an array for array-likes, so only a scalar boolean
+ result counts as missing.
+ """
+ if value is None:
+ return True
+ if isinstance(value, (float, np.floating)):
+ # NaN and the infinities, at every numpy width.
+ return not math.isfinite(value)
+ if isinstance(value, str):
+ return False
+ try:
+ missing = pd.isna(value)
+ except (TypeError, ValueError):
+ return False
+ return missing is True or (isinstance(missing, np.bool_) and bool(missing))
+
+
+def _sanitize_scalar(value: Any) -> Any:
+ """Sanitize a leaf value, or return :data:`_UNHANDLED` when ``value`` needs
+ the slower container/conversion path."""
+ if value is None or value is True or value is False:
+ return value
+ if isinstance(value, str):
+ return _sanitize_str(value)
+ if isinstance(value, int):
+ return value
+ if isinstance(value, float):
+ # Covers numpy.float64 (a float subclass) and NaN/Infinity.
+ return float(value) if math.isfinite(value) else None
+ if isinstance(value, (bytes, bytearray, memoryview)):
+ return _decode_binary(value)
+ return _UNHANDLED
+
+
+def _sanitize_other(value: Any, depth: int) -> Any:
+ """Sanitize containers and the long tail of warehouse value types."""
+ if isinstance(value, Mapping):
+ return {
+ _sanitize_str(key if isinstance(key, str) else str(key)): (
+ sanitize_json_value(item, depth + 1)
+ )
+ for key, item in value.items()
+ }
+ if isinstance(value, (list, tuple, set, frozenset)):
+ return [sanitize_json_value(item, depth + 1) for item in value]
+ # ``pandas.NaT`` is a ``datetime`` subclass, so the missing-value check has
+ # to run before the passthrough types.
+ if is_missing_value(value):
+ return None
+ if isinstance(value, np.generic):
+ # Every numpy scalar width, not just the few base_json_conv knows.
+ return sanitize_json_value(value.item(), depth + 1)
+ if isinstance(value, _PASSTHROUGH_TYPES):
+ return value
+ try:
+ # Normalises ndarrays, sets, Decimal and similar; the result may itself
+ # be a container or string that needs another pass. Recursion is
bounded
+ # by MAX_DEPTH.
+ converted = base_json_conv(value)
+ except TypeError:
+ # Unknown to Superset too. Render it as text rather than leaving it for
+ # FastMCP's ``fallback=str``, which the structured-content path
+ # (``to_jsonable_python``) does not apply.
+ return _sanitize_str(str(value))
+ return sanitize_json_value(converted, depth + 1)
+
+
+def sanitize_json_value(value: Any, depth: int = 0) -> Any:
+ """Return ``value`` with anything ``pydantic_core.to_json`` cannot encode
+ replaced by a JSON-safe equivalent."""
+ if (sanitized := _sanitize_scalar(value)) is not _UNHANDLED:
+ return sanitized
+ if depth >= MAX_DEPTH:
+ return _sanitize_str(str(value))
+ return _sanitize_other(value, depth)
+
+
+def sanitize_mapping(value: Any) -> Any:
+ """Sanitize a mapping, leaving non-mapping input to pydantic."""
+ if not isinstance(value, Mapping):
+ return value
+ return sanitize_json_value(value)
+
+
+def coerce_optional_int(value: Any) -> Any:
+ """Coerce an engine-reported count to ``int``, or ``None`` when it is not
+ a usable count (e.g. a boolean or ``NaN``)."""
+ if isinstance(value, (bool, np.bool_)):
+ return None
+ if value is None or isinstance(value, int):
+ return value
+ try:
+ coerced = int(value)
+ except (TypeError, ValueError, OverflowError):
+ return None
+ return coerced
+
+
+def coerce_int(value: Any) -> Any:
+ """Like :func:`coerce_optional_int` but yields ``0`` instead of ``None``,
+ for count fields that are not nullable. A count the engine did not report
+ therefore reads as ``0``; callers needing the true size of a result should
+ measure the rows themselves."""
+ coerced = coerce_optional_int(value)
+ return 0 if coerced is None else coerced
+
+
+#: A mapping of arbitrary values drawn from result data (e.g. column stats).
+JsonSafeMapping = Annotated[dict[str, Any], BeforeValidator(sanitize_mapping)]
+#: Result rows, sanitized element by element after pydantic coerces the
iterable.
+JsonSafeRows = list[JsonSafeMapping]
+#: Column samples, sanitized after pydantic coerces the iterable to a list.
+JsonSafeValues = list[Annotated[Any, BeforeValidator(sanitize_json_value)]]
+#: A non-nullable row count reported by the query engine.
+RowCount = Annotated[int, BeforeValidator(coerce_int)]
+#: A nullable row count reported by the query engine.
+OptionalRowCount = Annotated[int | None, BeforeValidator(coerce_optional_int)]
diff --git a/tests/unit_tests/mcp_service/chart/tool/test_get_chart_data.py
b/tests/unit_tests/mcp_service/chart/tool/test_get_chart_data.py
index 3436551a332..792e7a3fdac 100644
--- a/tests/unit_tests/mcp_service/chart/tool/test_get_chart_data.py
+++ b/tests/unit_tests/mcp_service/chart/tool/test_get_chart_data.py
@@ -25,6 +25,7 @@ from types import SimpleNamespace
from typing import Any
from unittest.mock import MagicMock, patch
+import pandas as pd
import pytest
from fastmcp import Client
@@ -588,7 +589,10 @@ class TestUnsavedChartDataQueryConstruction:
assert not isinstance(result, ChartError)
assert [row["team"] for row in result.data] == ["Empty", "Blue", "NaN"]
assert result.row_count == result.total_rows == 3
- assert result.data_quality["completeness"] == pytest.approx(5 / 6)
+ # Both the None and the NaN metric serialize as null, so both count
+ # toward incompleteness: 2 missing cells out of 3 rows x 2 columns.
+ assert result.data_quality["completeness"] == pytest.approx(4 / 6)
+ assert [row["saved_sla"] for row in result.data] == [None, 98.5, None]
query = captured[0]["queries"][0]
assert query["metrics"] == ["saved_sla"]
assert query["orderby"] == [("saved_sla", False)]
@@ -1815,6 +1819,14 @@ class TestSavedChartExtraFormDataFilters:
{"team": "Infinity", "saved_sla": float("inf")},
{"team": "Negative infinity", "saved_sla": -float("inf")},
]
+ rows.extend(
+ [
+ {"team": "NaT", "saved_sla": pd.NaT},
+ {"team": "NA", "saved_sla": pd.NA},
+ ]
+ if export_format == "json"
+ else []
+ )
if has_finite:
rows.append({"team": "Blue", "saved_sla": 42})
source_rowcount = len(rows) + 7
@@ -1887,9 +1899,18 @@ class TestSavedChartExtraFormDataFilters:
expected_groups = [row["team"] for row in rows]
if export_format == "json":
assert [row["team"] for row in data["data"]] == expected_groups
+ # NaN and the infinities are not valid JSON, so they serialize as
+ # null and count as missing alongside None, NaT and NA: six
+ # missing cells out of len(rows) x 2 columns.
assert data["data_quality"]["completeness"] == pytest.approx(
- 1 - 1 / (len(rows) * 2)
+ 1 - 6 / (len(rows) * 2)
+ )
+ assert [row["saved_sla"] for row in data["data"][:6]] == [None] * 6
+ metric_column = next(
+ column for column in data["columns"] if column["name"] ==
"saved_sla"
)
+ assert metric_column["null_count"] == 6
+ assert metric_column["unique_count"] == int(has_finite)
assert data.get("query_results") is None
elif export_format == "csv":
import csv
diff --git a/tests/unit_tests/mcp_service/sql_lab/tool/test_execute_sql.py
b/tests/unit_tests/mcp_service/sql_lab/tool/test_execute_sql.py
index b88ff947139..909bf61486e 100644
--- a/tests/unit_tests/mcp_service/sql_lab/tool/test_execute_sql.py
+++ b/tests/unit_tests/mcp_service/sql_lab/tool/test_execute_sql.py
@@ -22,6 +22,7 @@ These tests mock Database.execute() to test the MCP tool's
parameter mapping
and response conversion logic.
"""
+import base64
import logging
from decimal import Decimal
from typing import Any
@@ -42,6 +43,7 @@ from superset.exceptions import (
)
from superset.mcp_service.app import mcp
from superset.mcp_service.sql_lab.schemas import ColumnInfo
+from superset.mcp_service.utils.serialization import BINARY_PREFIX
logging.basicConfig(level=logging.DEBUG)
logger = logging.getLogger(__name__)
@@ -1154,8 +1156,11 @@ class TestExecuteSql:
row = data["rows"][0]
# UTF-8 decodable bytes should become string
assert row["utf8_data"] == "hello world"
- # Non-UTF-8 bytes should become hex
- assert row["binary_data"] == "000102ff"
+ # Non-UTF-8 bytes are base64-encoded rather than failing the call
+ assert (
+ row["binary_data"]
+ == BINARY_PREFIX +
base64.b64encode(b"\x00\x01\x02\xff").decode()
+ )
@patch("superset.security_manager")
@patch("superset.db")
@@ -1215,53 +1220,55 @@ class TestExecuteSql:
assert isinstance(row["price"], float)
-class TestSanitizeRowValues:
- """Unit tests for _sanitize_row_values helper function."""
-
- def test_sanitize_utf8_bytes(self):
- from superset.mcp_service.sql_lab.tool.execute_sql import
_sanitize_row_values
+class TestStatementRowSanitization:
+ """Row values are made JSON-safe by the response schema itself, so every
+ tool that returns rows gets the same treatment."""
- rows = [{"data": b"hello"}]
- _sanitize_row_values(rows)
- assert rows[0]["data"] == "hello"
+ @staticmethod
+ def _rows(rows: list[dict[str, Any]]) -> list[dict[str, Any]]:
+ from superset.mcp_service.sql_lab.tool.execute_sql import (
+ _data_to_statement_data,
+ )
- def test_sanitize_non_utf8_bytes(self):
- from superset.mcp_service.sql_lab.tool.execute_sql import
_sanitize_row_values
+ return _data_to_statement_data(rows).rows
- rows = [{"data": b"\x00\xff"}]
- _sanitize_row_values(rows)
- assert rows[0]["data"] == "00ff"
+ def test_utf8_bytes_decode_to_text(self):
+ assert self._rows([{"data": b"hello"}])[0]["data"] == "hello"
- def test_sanitize_memoryview(self):
- from superset.mcp_service.sql_lab.tool.execute_sql import
_sanitize_row_values
+ def test_non_utf8_bytes_are_base64_encoded(self):
+ value = self._rows([{"data": b"\x00\xff"}])[0]["data"]
+ assert value == BINARY_PREFIX + base64.b64encode(b"\x00\xff").decode()
- rows = [{"data": memoryview(b"test")}]
- _sanitize_row_values(rows)
- assert rows[0]["data"] == "test"
+ def test_memoryview_is_handled(self):
+ assert self._rows([{"data": memoryview(b"test")}])[0]["data"] == "test"
- def test_sanitize_decimal(self):
- from superset.mcp_service.sql_lab.tool.execute_sql import
_sanitize_row_values
+ def test_decimal_becomes_float(self):
+ row = self._rows([{"price": Decimal("19.99"), "count":
Decimal("42")}])[0]
+ assert row["price"] == 19.99
+ assert isinstance(row["price"], float)
+ assert row["count"] == 42.0
- rows = [{"price": Decimal("19.99"), "count": Decimal("42")}]
- _sanitize_row_values(rows)
- assert rows[0]["price"] == 19.99
- assert isinstance(rows[0]["price"], float)
- assert rows[0]["count"] == 42.0
+ def test_missing_values_become_none(self):
+ """The old tool-level helper stringified NaT to "NaT"."""
+ row = self._rows([{"ts": pd.NaT, "ratio": float("nan")}])[0]
+ assert row["ts"] is None
+ assert row["ratio"] is None
- def test_sanitize_custom_type_uses_str(self):
- from superset.mcp_service.sql_lab.tool.execute_sql import
_sanitize_row_values
+ def test_nested_containers_are_sanitized(self):
+ """The old tool-level helper skipped list and dict values, which is how
+ binary data still reached the serializer."""
+ row = self._rows([{"payload": {"blob": b"\xff"}, "items":
[b"\xff"]}])[0]
+ assert row["payload"]["blob"].startswith(BINARY_PREFIX)
+ assert row["items"][0].startswith(BINARY_PREFIX)
+ def test_custom_type_uses_str(self):
class CustomType:
def __str__(self):
return "custom_value"
- rows = [{"data": CustomType()}]
- _sanitize_row_values(rows)
- assert rows[0]["data"] == "custom_value"
+ assert self._rows([{"data": CustomType()}])[0]["data"] ==
"custom_value"
def test_preserves_json_serializable_types(self):
- from superset.mcp_service.sql_lab.tool.execute_sql import
_sanitize_row_values
-
rows = [
{
"str_val": "hello",
@@ -1273,33 +1280,10 @@ class TestSanitizeRowValues:
"dict_val": {"a": 1},
}
]
- original = [dict(row) for row in rows]
- _sanitize_row_values(rows)
- assert rows == original
-
- def test_sanitize_empty_rows(self):
- from superset.mcp_service.sql_lab.tool.execute_sql import
_sanitize_row_values
-
- rows: list[dict[str, Any]] = []
- _sanitize_row_values(rows)
- assert rows == []
+ assert self._rows(rows) == rows
- def test_sanitize_mixed_types_in_single_row(self):
- from superset.mcp_service.sql_lab.tool.execute_sql import
_sanitize_row_values
-
- rows = [
- {
- "id": 1,
- "name": "test",
- "price": Decimal("9.99"),
- "blob": b"\x00\x01\x02\xff",
- }
- ]
- _sanitize_row_values(rows)
- assert rows[0]["id"] == 1
- assert rows[0]["name"] == "test"
- assert rows[0]["price"] == 9.99
- assert rows[0]["blob"] == "000102ff"
+ def test_empty_rows(self):
+ assert self._rows([]) == []
class TestExecuteSqlOAuth2:
diff --git a/tests/unit_tests/mcp_service/test_response_serialization.py
b/tests/unit_tests/mcp_service/test_response_serialization.py
new file mode 100644
index 00000000000..300b5d0ef9d
--- /dev/null
+++ b/tests/unit_tests/mcp_service/test_response_serialization.py
@@ -0,0 +1,321 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""Serialization contract for the data-bearing MCP tool responses.
+
+FastMCP renders every tool return value through
+``pydantic_core.to_json(value, fallback=str)``. ``fallback`` only fires for
+types pydantic does not recognise, so a recognised-but-unrenderable value
+(non-UTF-8 ``bytes``, ``pandas.NaT``) raises ``PydanticSerializationError`` and
+destroys the whole tool call after the warehouse query already succeeded.
+These tests drive that real serializer — asserting on field values alone is
+what let the failure recur after it was first patched.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Callable
+
+import pandas as pd
+import pytest
+from fastmcp.tools.base import default_serializer
+from pydantic import BaseModel
+from pydantic_core import to_jsonable_python
+
+from superset.mcp_service.chart.schemas import (
+ ChartData,
+ ChartQueryResult,
+ DataColumn,
+ PerformanceMetadata,
+)
+from superset.mcp_service.common.cache_schemas import CacheStatus
+from superset.mcp_service.dashboard.schemas import (
+ DashboardChartData,
+ DashboardChartQueryData,
+)
+from superset.mcp_service.dataset.schemas import QueryDatasetResponse
+from superset.mcp_service.semantic_layer.schemas import GetTableResponse
+from superset.mcp_service.sql_lab.schemas import (
+ ExecuteSqlResponse,
+ StatementData,
+ StatementInfo,
+)
+from superset.mcp_service.utils.serialization import BINARY_PREFIX
+from superset.utils.json import loads as json_loads
+
+#: A row as a driver can really hand it back: an undecodable BLOB and a null
+#: timestamp materialised as ``pandas.NaT``.
+HOSTILE_ROW: dict[str, Any] = {
+ "blob": b"\xff\xfe\x00\x01",
+ "ts": pd.NaT,
+ "ratio": float("nan"),
+ "nested": {"inner": b"\x80"},
+}
+
+#: Some engines report ``rowcount`` as a float. A fractional value is used
+#: deliberately: pydantic accepts an integral float for an ``int`` field on its
+#: own, so only a fractional one proves the coercion is doing the work.
+FLOAT_COUNT = 5.7
+
+
+def _serialize(response: BaseModel) -> dict[str, Any]:
+ """Serialize exactly as FastMCP does and parse back as strict JSON.
+
+ ``loads`` defaults to ``allow_nan=False``, so a payload carrying the
+ non-standard ``NaN``/``Infinity`` literals fails here too."""
+ payload = json_loads(default_serializer(response))
+ # Structured content takes the other pydantic_core entry point.
+ to_jsonable_python(value=response)
+ return payload
+
+
+def _assert_row_is_safe(row: dict[str, Any]) -> None:
+ """Assert that binary, missing, and nested result values are JSON-safe."""
+ assert row["blob"].startswith(BINARY_PREFIX)
+ assert row["ts"] is None
+ assert row["ratio"] is None
+ assert row["nested"]["inner"].startswith(BINARY_PREFIX)
+
+
+def _execute_sql_response() -> ExecuteSqlResponse:
+ """Build a SQL Lab response with hostile rows and statement data."""
+ return ExecuteSqlResponse(
+ success=True,
+ rows=[HOSTILE_ROW],
+ columns=[],
+ row_count=FLOAT_COUNT,
+ affected_rows=FLOAT_COUNT,
+ statements=[
+ StatementInfo(
+ original_sql="SELECT 1",
+ executed_sql="SELECT 1",
+ row_count=FLOAT_COUNT,
+ data=StatementData(rows=[HOSTILE_ROW], columns=[]),
+ )
+ ],
+ )
+
+
+def _data_column() -> DataColumn:
+ """Column metadata derives its samples and stats from the rows, so it
+ carries the same hostile values."""
+ return DataColumn(
+ name="blob",
+ display_name="blob",
+ data_type="bytes",
+ sample_values=[b"\xff\xfe", pd.NaT],
+ null_count=1,
+ unique_count=2,
+ statistics={"min": b"\xff\xfe"},
+ )
+
+
+def _chart_data() -> ChartData:
+ """Build a chart response with hostile rows, metadata, and query
results."""
+ return ChartData(
+ chart_id=1,
+ chart_name="chart",
+ chart_type="table",
+ columns=[_data_column()],
+ data=[HOSTILE_ROW],
+ query_results=[
+ ChartQueryResult(
+ query_index=0,
+ columns=["blob"],
+ data=[HOSTILE_ROW],
+ row_count=FLOAT_COUNT,
+ total_rows=FLOAT_COUNT,
+ )
+ ],
+ row_count=FLOAT_COUNT,
+ total_rows=FLOAT_COUNT,
+ data_freshness=None,
+ summary="summary",
+ insights=[],
+ data_quality={},
+ recommended_visualizations=[],
+ performance=PerformanceMetadata(
+ query_duration_ms=42, cache_status="fresh_query"
+ ),
+ cache_status=CacheStatus(cache_hit=False),
+ )
+
+
+def _query_dataset_response() -> QueryDatasetResponse:
+ """Build a dataset response with hostile rows and column metadata."""
+ return QueryDatasetResponse(
+ dataset_id=1,
+ dataset_name="dataset",
+ columns=[_data_column()],
+ data=[HOSTILE_ROW],
+ row_count=FLOAT_COUNT,
+ total_rows=FLOAT_COUNT,
+ )
+
+
+def _get_table_response() -> GetTableResponse:
+ """Build a semantic-layer table response with hostile warehouse values."""
+ return GetTableResponse(
+ columns=[_data_column()],
+ data=[HOSTILE_ROW],
+ row_count=FLOAT_COUNT,
+ total_rows=FLOAT_COUNT,
+ summary="summary",
+ source="builtin",
+ )
+
+
+def _dashboard_chart_data() -> DashboardChartData:
+ """Build a dashboard chart response with hostile samples and query data."""
+ return DashboardChartData(
+ chart_id=1,
+ chart_name="chart",
+ chart_type="table",
+ columns=["blob"],
+ sample_data=[HOSTILE_ROW],
+ row_count=FLOAT_COUNT,
+ total_rows=FLOAT_COUNT,
+ queries=[
+ DashboardChartQueryData(
+ query_index=0,
+ columns=["blob"],
+ sample_data=[HOSTILE_ROW],
+ row_count=FLOAT_COUNT,
+ total_rows=FLOAT_COUNT,
+ )
+ ],
+ )
+
+
[email protected](
+ "build,rows_field,count_fields",
+ [
+ pytest.param(
+ _execute_sql_response,
+ "rows",
+ ("row_count", "affected_rows"),
+ id="execute_sql",
+ ),
+ pytest.param(
+ _chart_data, "data", ("row_count", "total_rows"),
id="get_chart_data"
+ ),
+ pytest.param(
+ _query_dataset_response,
+ "data",
+ ("row_count", "total_rows"),
+ id="query_dataset",
+ ),
+ pytest.param(
+ _get_table_response, "data", ("row_count", "total_rows"),
id="get_table"
+ ),
+ pytest.param(
+ _dashboard_chart_data,
+ "sample_data",
+ ("row_count", "total_rows"),
+ id="get_dashboard_data",
+ ),
+ ],
+)
+def test_hostile_result_set_serializes(
+ build: Callable[[], BaseModel],
+ rows_field: str,
+ count_fields: tuple[str, ...],
+) -> None:
+ """Float row counts and a non-UTF-8 binary column must not abort the
call."""
+ payload = _serialize(build())
+
+ for field in count_fields:
+ assert payload[field] == 5, field
+ _assert_row_is_safe(payload[rows_field][0])
+
+
+def test_execute_sql_statements_are_sanitized() -> None:
+ payload = _serialize(_execute_sql_response())
+ statement = payload["statements"][0]
+
+ assert statement["row_count"] == 5
+ _assert_row_is_safe(statement["data"]["rows"][0])
+
+
+def test_chart_query_results_are_sanitized() -> None:
+ payload = _serialize(_chart_data())
+ result = payload["query_results"][0]
+
+ assert result["row_count"] == result["total_rows"] == 5
+ _assert_row_is_safe(result["data"][0])
+
+
[email protected](
+ "build",
+ [_chart_data, _query_dataset_response, _get_table_response],
+ ids=["get_chart_data", "query_dataset", "get_table"],
+)
+def test_column_metadata_is_sanitized(build: Callable[[], BaseModel]) -> None:
+ """Column samples and statistics are lifted straight out of the rows."""
+ column = _serialize(build())["columns"][0]
+
+ assert column["sample_values"][0].startswith(BINARY_PREFIX)
+ assert column["sample_values"][1] is None
+ assert column["statistics"]["min"].startswith(BINARY_PREFIX)
+ assert column["null_count"] == 1
+ assert column["unique_count"] == 2
+
+
[email protected](
+ "count,expected",
+ [
+ (5.7, 5),
+ (float("nan"), None),
+ ("unknown", None),
+ (None, None),
+ (True, None),
+ (False, None),
+ ],
+ ids=["fractional-float", "nan", "non-numeric", "missing", "true", "false"],
+)
+def test_unusable_nullable_count_becomes_null(count: Any, expected: int |
None) -> None:
+ """A count the engine reports in an unusable form must not fail the
call."""
+ response = QueryDatasetResponse(
+ dataset_id=1, dataset_name="dataset", data=[], row_count=0,
total_rows=count
+ )
+
+ assert _serialize(response)["total_rows"] == expected
+
+
+def test_unusable_required_count_becomes_zero() -> None:
+ """A non-nullable count falls back to 0 rather than failing validation."""
+ response = StatementInfo(
+ original_sql="SET x", executed_sql="SET x", row_count="unknown"
+ )
+
+ assert _serialize(response)["row_count"] == 0
+
+
+def test_clean_results_are_untouched() -> None:
+ """The coercion must not disturb ordinary payloads."""
+ response = QueryDatasetResponse(
+ dataset_id=1,
+ dataset_name="dataset",
+ data=[{"count": 3, "label": "a", "ratio": 1.5}],
+ row_count=1,
+ total_rows=None,
+ )
+ payload = _serialize(response)
+
+ assert payload["data"] == [{"count": 3, "label": "a", "ratio": 1.5}]
+ assert payload["row_count"] == 1
+ assert payload["total_rows"] is None
diff --git a/tests/unit_tests/mcp_service/utils/test_serialization.py
b/tests/unit_tests/mcp_service/utils/test_serialization.py
new file mode 100644
index 00000000000..ac1df122d71
--- /dev/null
+++ b/tests/unit_tests/mcp_service/utils/test_serialization.py
@@ -0,0 +1,294 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""Unit tests for the MCP JSON-safety helpers."""
+
+from __future__ import annotations
+
+import base64
+from collections import deque, UserDict
+from collections.abc import Callable
+from decimal import Decimal
+from types import MappingProxyType
+from typing import Any
+from unittest.mock import patch
+
+import numpy as np
+import pandas as pd
+import pytest
+from fastmcp.tools.base import default_serializer
+from pydantic import TypeAdapter, ValidationError
+from pydantic_core import to_jsonable_python
+
+from superset.mcp_service.utils.serialization import (
+ BINARY_PREFIX,
+ coerce_int,
+ coerce_optional_int,
+ is_missing_value,
+ JsonSafeMapping,
+ JsonSafeRows,
+ JsonSafeValues,
+ MAX_DEPTH,
+ OptionalRowCount,
+ RowCount,
+ sanitize_json_value,
+ sanitize_mapping,
+)
+from superset.utils.json import loads as json_loads
+
+
+def test_utf8_bytes_decode_to_text() -> None:
+ assert sanitize_json_value(b"hello") == "hello"
+
+
+def test_non_utf8_bytes_become_base64() -> None:
+ raw = b"\xff\xfe\x00\x01"
+ assert sanitize_json_value(raw) == BINARY_PREFIX +
base64.b64encode(raw).decode()
+
+
[email protected]("value", [bytearray(b"\x80\x81"),
memoryview(b"\x80\x81")])
+def test_bytearray_and_memoryview_are_handled(value: Any) -> None:
+ """``bytes`` is not the only binary type a DBAPI driver returns."""
+ assert (
+ sanitize_json_value(value)
+ == BINARY_PREFIX + base64.b64encode(b"\x80\x81").decode()
+ )
+
+
[email protected](
+ "value", [pd.NaT, np.datetime64("NaT"), float("nan"), np.nan, pd.NA]
+)
+def test_missing_values_become_null(value: Any) -> None:
+ """``pandas.NaT`` is a ``datetime`` subclass backed by a float; pydantic
+ tries to render it as a datetime and raises ``TypeError: 'float' object
+ cannot be interpreted as an integer``."""
+ assert sanitize_json_value(value) is None
+
+
[email protected]("value", [float("inf"), float("-inf")])
+def test_infinities_become_null(value: float) -> None:
+ """``Infinity`` is not valid JSON, so strict clients reject it."""
+ assert sanitize_json_value(value) is None
+
+
+def test_lone_surrogates_are_replaced() -> None:
+ """Drivers decoding with ``surrogatepass`` produce strings that cannot be
+ encoded back to UTF-8."""
+ value = b"\xed\xa0\xbd".decode("utf-8", "surrogatepass")
+ sanitized = sanitize_json_value(value)
+ assert sanitized.encode("utf-8") # no UnicodeEncodeError
+ assert "\ud83d" not in sanitized
+
+
+def test_ordinary_scalars_pass_through() -> None:
+ assert sanitize_json_value("abc") == "abc"
+ assert sanitize_json_value(7) == 7
+ assert sanitize_json_value(True) is True
+ assert sanitize_json_value(1.5) == 1.5
+ assert sanitize_json_value(None) is None
+
+
[email protected](
+ "value,expected",
+ [
+ (np.int64(3), 3),
+ (np.int32(3), 3),
+ (np.int8(3), 3),
+ (np.uint16(3), 3),
+ (np.float64(1.5), 1.5),
+ (np.float32(1.5), 1.5),
+ (np.bool_(True), True),
+ ],
+)
+def test_numpy_scalars_of_every_width_are_normalised(value: Any, expected:
Any) -> None:
+ """Rendering a numpy scalar as text would silently turn numbers into
+ strings for the caller."""
+ sanitized = sanitize_json_value(value)
+ assert sanitized == expected
+ assert isinstance(sanitized, type(expected))
+
+
[email protected]("value", [np.float32("inf"), np.float64("-inf")])
+def test_non_finite_numpy_floats_become_null(value: Any) -> None:
+ assert sanitize_json_value(value) is None
+
+
+def test_numpy_arrays_become_lists() -> None:
+ assert sanitize_json_value(np.array([1, 2])) == [1, 2]
+
+
+def test_decimal_becomes_float() -> None:
+ """A JSON number is more useful to the caller than a quoted string."""
+ sanitized = sanitize_json_value(Decimal("19.99"))
+ assert sanitized == 19.99
+ assert isinstance(sanitized, float)
+
+
+def test_unknown_types_are_stringified() -> None:
+ """The structured-content path (``to_jsonable_python``) has no
+ ``fallback=str``, so unknown values cannot be left as-is."""
+
+ class Custom:
+ def __str__(self) -> str:
+ return "custom_value"
+
+ assert sanitize_json_value(Custom()) == "custom_value"
+
+
[email protected](
+ "value,expected",
+ [
+ (None, True),
+ (pd.NaT, True),
+ (pd.NA, True),
+ (float("nan"), True),
+ (np.float32("nan"), True),
+ (0, False),
+ ("", False),
+ ("NaN", False),
+ ("NaT", False),
+ ("None", False),
+ (b"\xff", False),
+ ([], False),
+ ([1, 2], False),
+ ],
+)
+def test_is_missing_value(value: Any, expected: bool) -> None:
+ assert is_missing_value(value) is expected
+
+
+def test_nested_containers_are_walked() -> None:
+ value = {"a": [{"b": b"\xff"}], "c": (pd.NaT, float("nan"))}
+ assert sanitize_json_value(value) == {
+ "a": [{"b": BINARY_PREFIX + base64.b64encode(b"\xff").decode()}],
+ "c": [None, None],
+ }
+
+
+def test_non_string_mapping_keys_are_stringified() -> None:
+ assert sanitize_json_value({1: "a"}) == {"1": "a"}
+
+
+def test_deep_nesting_is_capped_instead_of_recursing() -> None:
+ value: Any = b"\xff"
+ for _ in range(MAX_DEPTH + 5):
+ value = [value]
+ # The cap renders the tail as text rather than recursing forever; the
+ # point is that it returns something JSON-encodable at all.
+ assert json_loads(default_serializer(sanitize_json_value(value)))
+
+
+def test_sanitize_mapping_leaves_non_mappings_to_pydantic() -> None:
+ assert sanitize_mapping(5) == 5
+ assert sanitize_mapping({"a": b"\xff"}) == {
+ "a": BINARY_PREFIX + base64.b64encode(b"\xff").decode()
+ }
+
+
[email protected](
+ "value,expected",
+ [
+ (None, None),
+ (5, 5),
+ (5.0, 5),
+ (5.9, 5),
+ (Decimal("5.9"), 5),
+ ("nonsense", None),
+ (float("nan"), None),
+ (object(), None),
+ ],
+)
+def test_coerce_optional_int(value: Any, expected: int | None) -> None:
+ assert coerce_optional_int(value) == expected
+
+
+def test_coerce_int_falls_back_to_zero() -> None:
+ assert coerce_int(None) == 0
+ assert coerce_int("nonsense") == 0
+ assert coerce_int(5.9) == 5
+
+
[email protected]("container", [list, tuple, set, frozenset, iter,
deque])
[email protected]("value", [b"\xff", pd.NaT, float("inf"), "\ud800"])
+def test_json_safe_values_sanitize_coerced_iterables(
+ container: Callable[[list[Any]], Any], value: Any
+) -> None:
+ """Every accepted iterable must sanitize its elements before
serialization."""
+ values = TypeAdapter(JsonSafeValues).validate_python(container([value]))
+ expected = [sanitize_json_value(value)]
+ assert json_loads(default_serializer(values)) == expected
+ assert to_jsonable_python(values) == expected
+
+
+def test_json_safe_values_sanitize_generator() -> None:
+ """One-shot generators must be consumed once and sanitized element by
element."""
+ values = TypeAdapter(JsonSafeValues).validate_python(
+ value for value in [b"\xff", pd.NaT]
+ )
+ assert json_loads(default_serializer(values)) == ["base64:/w==", None]
+ assert to_jsonable_python(values) == ["base64:/w==", None]
+
+
[email protected]("container", [dict, MappingProxyType, UserDict])
+def test_json_safe_mapping_sanitizes_general_mappings(
+ container: Callable[[dict[str, Any]], Any],
+) -> None:
+ """Mapping implementations and nested mappings retain their JSON object
shape."""
+ mapping = TypeAdapter(JsonSafeMapping).validate_python(
+ container({"blob": b"\xff", "nested": container({"ts": pd.NaT})})
+ )
+ expected = {"blob": "base64:/w==", "nested": {"ts": None}}
+ assert json_loads(default_serializer(mapping)) == expected
+ assert to_jsonable_python(mapping) == expected
+
+
[email protected]("container", [list, tuple, iter, deque])
+def test_json_safe_rows_sanitize_coerced_iterables(
+ container: Callable[[list[Any]], Any],
+) -> None:
+ """Row mappings must be sanitized after the outer iterable is coerced."""
+ rows = TypeAdapter(JsonSafeRows).validate_python(
+ container([MappingProxyType({"blob": b"\xff", "ts": pd.NaT})])
+ )
+ expected = [{"blob": "base64:/w==", "ts": None}]
+ assert json_loads(default_serializer(rows)) == expected
+ assert to_jsonable_python(rows) == expected
+
+
[email protected]("value", ["text", b"bytes", bytearray(b"bytes"), {},
5, None])
[email protected]("alias", [JsonSafeRows, JsonSafeValues])
+def test_json_safe_sequences_reject_invalid_containers(value: Any, alias: Any)
-> None:
+ """Sanitization must not broaden the list schema's accepted input types."""
+ with pytest.raises(ValidationError):
+ TypeAdapter(alias).validate_python(value)
+
+
[email protected]("value", [True, False, np.bool_(True),
np.bool_(False)])
+def test_boolean_counts_are_unavailable(value: Any) -> None:
+ """Booleans are not warehouse row counts, including numpy booleans."""
+ assert coerce_optional_int(value) is None
+ assert coerce_int(value) == 0
+ assert TypeAdapter(OptionalRowCount).validate_python(value) is None
+ assert TypeAdapter(RowCount).validate_python(value) == 0
+
+
+def test_string_missing_check_does_not_call_pandas() -> None:
+ """Literal strings cannot represent missing scalars and need no pandas
call."""
+ with patch("superset.mcp_service.utils.serialization.pd.isna") as isna:
+ for value in ("", "NaN", "NaT", "None", "ordinary"):
+ assert is_missing_value(value) is False
+ isna.assert_not_called()
