fix(charts): add query_context validation to prevent incomplete metadata

superset/charts/schemas.py

@@ -213,7 +213,7 @@ class ChartPostSchema(Schema):
     query_context = fields.String(
         metadata={"description": query_context_description},
         allow_none=True,
-        validate=utils.validate_json,
+        validate=utils.validate_query_context_metadata,
     )
     query_context_generation = fields.Boolean(
         metadata={"description": query_context_generation_description}, allow_none=True
@@ -274,7 +274,9 @@ class ChartPutSchema(Schema):
         validate=utils.validate_json,
     )
     query_context = fields.String(
-        metadata={"description": query_context_description}, allow_none=True
+        metadata={"description": query_context_description},
+        allow_none=True,
+        validate=utils.validate_query_context_metadata,
     )
     query_context_generation = fields.Boolean(
         metadata={"description": query_context_generation_description}, allow_none=True

superset/utils/schema.py

@@ -50,4 +50,37 @@ def validate_json(value: Union[bytes, bytearray, str]) -> None:
     try:
         json.validate_json(value)
     except json.JSONDecodeError as ex:
-        raise ValidationError("JSON not valid") from ex
+        error_msg = "JSON not valid"
+        raise ValidationError(error_msg) from ex
+
+
+def validate_query_context_metadata(value: Union[bytes, bytearray, str, None]) -> None:
+    """
+    Validator for query_context field to ensure it contains required metadata.
+
+    Validates that the query_context JSON contains the required 'datasource' and
+    'queries' fields needed for chart data retrieval.
+
+    :raises ValidationError: if value is not valid JSON or missing required fields
+    :param value: a JSON string that should contain datasource and queries metadata
+    """
+    if value is None or value == "":
+        return  # Allow None values and empty strings
+
+    # Reuse existing JSON validation logic
+    validate_json(value)
+
+    # Parse and validate the structure
+    parsed_data = json.loads(value)
+
+    # Validate required fields exist in the query_context
+    if not isinstance(parsed_data, dict):
+        error_msg = "Query context must be a valid JSON object"
+        raise ValidationError(error_msg)
+
+    # When query_context is provided (not None), validate it has required fields
+    required_fields = {"datasource", "queries"}
+    missing_fields: set[str] = required_fields - parsed_data.keys()
+    if missing_fields:
+        fields_str = ", ".join(sorted(missing_fields))
+        raise ValidationError(f"Query context is missing required fields: {fields_str}")

tests/unit_tests/charts/test_schemas.py

@@ -22,8 +22,11 @@
 from superset.charts.schemas import (
     ChartDataProphetOptionsSchema,
     ChartDataQueryObjectSchema,
+    ChartPostSchema,
+    ChartPutSchema,
     get_time_grain_choices,
 )
+from superset.utils import json
 
 
 def test_get_time_grain_choices(app_context: None) -> None:
@@ -152,3 +155,120 @@ def test_time_grain_validation_with_config_addons(app_context: None) -> None:
     }
     result = schema.load(custom_data)
     assert result["time_grain"] == "PT10M"
+
+
+def test_chart_post_schema_query_context_validation(app_context: None) -> None:
+    """Test that ChartPostSchema validates query_context contains required metadata"""
+    schema = ChartPostSchema()
+
+    # Valid query_context with datasource and queries should pass
+    valid_query_context = json.dumps(
+        {
+            "datasource": {"type": "table", "id": 1},
+            "queries": [{"metrics": ["count"], "columns": []}],
+        }
+    )
+    valid_data = {
+        "slice_name": "Test Chart",
+        "datasource_id": 1,
+        "datasource_type": "table",
+        "query_context": valid_query_context,
+    }
+    result = schema.load(valid_data)
+    assert result["query_context"] == valid_query_context
+
+    # None query_context should be allowed (allow_none=True)
+    none_data = {
+        "slice_name": "Test Chart",
+        "datasource_id": 1,
+        "datasource_type": "table",
+        "query_context": None,
+    }
+    result = schema.load(none_data)
+    assert result["query_context"] is None
+
+    # Query context missing 'datasource' field should fail
+    missing_datasource = json.dumps(
+        {"queries": [{"metrics": ["count"], "columns": []}]}
+    )
+    invalid_data_1 = {
+        "slice_name": "Test Chart",
+        "datasource_id": 1,
+        "datasource_type": "table",
+        "query_context": missing_datasource,
+    }
+    with pytest.raises(ValidationError) as exc_info:
+        schema.load(invalid_data_1)
+    assert "query_context" in exc_info.value.messages
+    assert "datasource" in str(exc_info.value.messages["query_context"])
+
+    # Query context missing 'queries' field should fail
+    missing_queries = json.dumps({"datasource": {"type": "table", "id": 1}})
+    invalid_data_2 = {
+        "slice_name": "Test Chart",
+        "datasource_id": 1,
+        "datasource_type": "table",
+        "query_context": missing_queries,
+    }
+    with pytest.raises(ValidationError) as exc_info:
+        schema.load(invalid_data_2)
+    assert "query_context" in exc_info.value.messages
+    assert "queries" in str(exc_info.value.messages["query_context"])
+
+    # Query context missing both 'datasource' and 'queries' should fail
+    empty_query_context = json.dumps({})
+    invalid_data_3 = {
+        "slice_name": "Test Chart",
+        "datasource_id": 1,
+        "datasource_type": "table",
+        "query_context": empty_query_context,
+    }
+    with pytest.raises(ValidationError) as exc_info:
+        schema.load(invalid_data_3)
+    assert "query_context" in exc_info.value.messages
+    assert "datasource" in str(exc_info.value.messages["query_context"])
+    assert "queries" in str(exc_info.value.messages["query_context"])
+
+    # Invalid JSON should fail
+    invalid_json = "not valid json"
+    invalid_data_4 = {
+        "slice_name": "Test Chart",
+        "datasource_id": 1,
+        "datasource_type": "table",
+        "query_context": invalid_json,
+    }
+    with pytest.raises(ValidationError) as exc_info:
+        schema.load(invalid_data_4)
+    assert "query_context" in exc_info.value.messages
+
+
+def test_chart_put_schema_query_context_validation(app_context: None) -> None:
+    """Test that ChartPutSchema validates query_context contains required metadata"""
+    schema = ChartPutSchema()
+
+    # Valid query_context with datasource and queries should pass
+    valid_query_context = json.dumps(
+        {
+            "datasource": {"type": "table", "id": 1},
+            "queries": [{"metrics": ["count"], "columns": []}],
+        }
+    )
+    valid_data = {
+        "slice_name": "Updated Chart",
+        "query_context": valid_query_context,
+    }
+    result = schema.load(valid_data)
+    assert result["query_context"] == valid_query_context
+
+    # Query context missing required fields should fail
+    missing_datasource = json.dumps(
+        {"queries": [{"metrics": ["count"], "columns": []}]}
+    )
+    invalid_data = {
+        "slice_name": "Updated Chart",
+        "query_context": missing_datasource,
+    }
+    with pytest.raises(ValidationError) as exc_info:
+        schema.load(invalid_data)
+    assert "query_context" in exc_info.value.messages
+    assert "datasource" in str(exc_info.value.messages["query_context"])