superset2/superset/mcp_service/utils/schema_utils.py

# 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.

"""
Generic utilities for flexible schema input handling in MCP tools.

This module provides utilities to accept both JSON string and object formats
for input parameters, making MCP tools more flexible for different clients.
"""

from __future__ import annotations

import logging
from typing import Any, Callable, List, Type, TypeVar

from pydantic import BaseModel, ValidationError

logger = logging.getLogger(__name__)

T = TypeVar("T")


class JSONParseError(ValueError):
    """Raised when JSON parsing fails with helpful context."""

    def __init__(self, value: Any, error: Exception, param_name: str = "parameter"):
        self.value = value
        self.original_error = error
        self.param_name = param_name
        super().__init__(
            f"Failed to parse {param_name} from JSON string: {error}. "
            f"Received value: {value!r}"
        )


def parse_json_or_passthrough(
    value: Any, param_name: str = "parameter", strict: bool = False
) -> Any:
    """
    Parse a value that can be either a JSON string or a native Python object.

    This function handles the common pattern where API parameters can be provided
    as either:
    - A JSON string (e.g., from CLI tools or tests): '{"key": "value"}'
    - A native Python object (e.g., from LLM clients): {"key": "value"}

    Args:
        value: The input value to parse. Can be a string, list, dict, or any JSON-
            serializable type.
        param_name: Name of the parameter for error messages (default: "parameter")
        strict: If True, raises JSONParseError on parse failures. If False, logs
            warning and returns original value (default: False)

    Returns:
        Parsed Python object if value was a JSON string, otherwise returns value
        unchanged.

    Raises:
        JSONParseError: If strict=True and JSON parsing fails.

    Examples:
        >>> parse_json_or_passthrough('[1, 2, 3]', 'numbers')
        [1, 2, 3]

        >>> parse_json_or_passthrough([1, 2, 3], 'numbers')
        [1, 2, 3]

        >>> parse_json_or_passthrough('{"key": "value"}', 'config')
        {'key': 'value'}

        >>> parse_json_or_passthrough({'key': 'value'}, 'config')
        {'key': 'value'}
    """
    # If not a string, return as-is (already in object form)
    if not isinstance(value, str):
        return value

    # Try to parse as JSON
    try:
        from superset.utils import json

        parsed = json.loads(value)
        logger.debug("Successfully parsed %s from JSON string", param_name)
        return parsed
    except (ValueError, TypeError) as e:
        error_msg = (
            f"Failed to parse {param_name} from JSON string: {e}. Received: {value!r}"
        )

        if strict:
            raise JSONParseError(value, e, param_name) from None

        logger.warning("%s. Returning original value.", error_msg)
        return value


def parse_json_or_list(
    value: Any, param_name: str = "parameter", item_separator: str = ","
) -> List[Any]:
    """
    Parse a value into a list, accepting JSON string, list, or comma-separated string.

    This function provides maximum flexibility for list parameters by accepting:
    - JSON array string: '["item1", "item2"]'
    - Python list: ["item1", "item2"]
    - Comma-separated string: "item1, item2, item3"
    - Empty/None: returns empty list

    Args:
        value: Input value to parse into a list
        param_name: Name of the parameter for error messages
        item_separator: Separator for comma-separated strings (default: ",")

    Returns:
        List of items. Returns empty list if value is None or empty.

    Examples:
        >>> parse_json_or_list('["a", "b"]', 'items')
        ['a', 'b']

        >>> parse_json_or_list(['a', 'b'], 'items')
        ['a', 'b']

        >>> parse_json_or_list('a, b, c', 'items')
        ['a', 'b', 'c']

        >>> parse_json_or_list(None, 'items')
        []
    """
    # Handle None and empty values
    if value is None or value == "":
        return []

    # Already a list, return as-is
    if isinstance(value, list):
        return value

    # Try to parse as JSON
    if isinstance(value, str):
        try:
            from superset.utils import json

            parsed = json.loads(value)
            # If successfully parsed and it's a list, return it
            if isinstance(parsed, list):
                logger.debug("Successfully parsed %s from JSON string", param_name)
                return parsed

            # If parsed to non-list (e.g., single value), wrap in list
            logger.debug(
                "Parsed %s from JSON to non-list, wrapping in list", param_name
            )
            return [parsed]
        except (ValueError, TypeError):
            # Not valid JSON, try comma-separated parsing
            logger.debug(
                "Could not parse %s as JSON, trying comma-separated", param_name
            )
            items = [
                item.strip() for item in value.split(item_separator) if item.strip()
            ]
            return items

    # For any other type, wrap in a list
    logger.debug("Wrapping %s value in list", param_name)
    return [value]


def parse_json_or_model(
    value: Any, model_class: Type[BaseModel], param_name: str = "parameter"
) -> BaseModel:
    """
    Parse a value into a Pydantic model, accepting JSON string or dict.

    Args:
        value: Input value to parse (JSON string, dict, or model instance)
        model_class: Pydantic model class to validate against
        param_name: Name of the parameter for error messages

    Returns:
        Validated Pydantic model instance

    Raises:
        ValidationError: If the value cannot be parsed or validated

    Examples:
        >>> class MyModel(BaseModel):
        ...     name: str
        ...     value: int

        >>> parse_json_or_model('{"name": "test", "value": 42}', MyModel)
        MyModel(name='test', value=42)

        >>> parse_json_or_model({"name": "test", "value": 42}, MyModel)
        MyModel(name='test', value=42)
    """
    # If already an instance of the model, return as-is
    if isinstance(value, model_class):
        return value

    # Parse JSON string if needed
    parsed_value = parse_json_or_passthrough(value, param_name, strict=True)

    # Validate and construct the model
    try:
        return model_class.model_validate(parsed_value)
    except ValidationError:
        raise


def parse_json_or_model_list(
    value: Any,
    model_class: Type[BaseModel],
    param_name: str = "parameter",
) -> List[BaseModel]:
    """
    Parse a value into a list of Pydantic models, accepting JSON string or list.

    Args:
        value: Input value to parse (JSON string, list of dicts, or list of models)
        model_class: Pydantic model class for list items
        param_name: Name of the parameter for error messages

    Returns:
        List of validated Pydantic model instances

    Raises:
        ValidationError: If any item cannot be parsed or validated

    Examples:
        >>> class Item(BaseModel):
        ...     name: str

        >>> parse_json_or_model_list('[{"name": "a"}, {"name": "b"}]', Item)
        [Item(name='a'), Item(name='b')]

        >>> parse_json_or_model_list([{"name": "a"}, {"name": "b"}], Item)
        [Item(name='a'), Item(name='b')]
    """
    # Handle None and empty
    if value is None or value == "":
        return []

    # Parse to list first
    items = parse_json_or_list(value, param_name)

    # Validate each item against the model
    validated_items = []
    for i, item in enumerate(items):
        try:
            if isinstance(item, model_class):
                validated_items.append(item)
            else:
                validated_items.append(model_class.model_validate(item))
        except ValidationError:
            logger.debug(
                "Failed to validate %s[%s] against %s",
                param_name,
                i,
                model_class.__name__,
            )
            # Re-raise original validation error
            raise

    return validated_items


# Pydantic validator decorators for common use cases
def json_or_passthrough_validator(
    param_name: str | None = None, strict: bool = False
) -> Callable[[Type[BaseModel], Any, Any], Any]:
    """
    Decorator factory for Pydantic field validators that accept JSON or objects.

    This creates a validator that can be used with Pydantic's @field_validator
    decorator to automatically parse JSON strings.

    Args:
        param_name: Parameter name for error messages (uses field name if None)
        strict: Whether to raise errors on parse failures

    Returns:
        Validator function compatible with @field_validator

    Example:
        >>> class MySchema(BaseModel):
        ...     config: dict
        ...
        ...     @field_validator('config', mode='before')
        ...     @classmethod
        ...     def parse_config(cls, v):
        ...         return parse_json_or_passthrough(v, 'config')
    """

    def validator(cls: Type[BaseModel], v: Any, info: Any = None) -> Any:
        # Use field name from validation info if param_name not provided
        field_name = param_name or (info.field_name if info else "field")
        return parse_json_or_passthrough(v, field_name, strict)

    return validator


def json_or_list_validator(
    param_name: str | None = None, item_separator: str = ","
) -> Callable[[Type[BaseModel], Any, Any], List[Any]]:
    """
    Decorator factory for Pydantic validators that parse values into lists.

    Args:
        param_name: Parameter name for error messages
        item_separator: Separator for comma-separated strings

    Returns:
        Validator function compatible with @field_validator

    Example:
        >>> class MySchema(BaseModel):
        ...     items: List[str]
        ...
        ...     @field_validator('items', mode='before')
        ...     @classmethod
        ...     def parse_items(cls, v):
        ...         return parse_json_or_list(v, 'items')
    """

    def validator(cls: Type[BaseModel], v: Any, info: Any = None) -> List[Any]:
        field_name = param_name or (info.field_name if info else "field")
        return parse_json_or_list(v, field_name, item_separator)

    return validator


def json_or_model_list_validator(
    model_class: Type[BaseModel], param_name: str | None = None
) -> Callable[[Type[BaseModel], Any, Any], List[BaseModel]]:
    """
    Decorator factory for Pydantic validators that parse lists of models.

    Args:
        model_class: Pydantic model class for list items
        param_name: Parameter name for error messages

    Returns:
        Validator function compatible with @field_validator

    Example:
        >>> class FilterModel(BaseModel):
        ...     col: str
        ...     value: str
        ...
        >>> class MySchema(BaseModel):
        ...     filters: List[FilterModel]
        ...
        ...     @field_validator('filters', mode='before')
        ...     @classmethod
        ...     def parse_filters(cls, v):
        ...         return parse_json_or_model_list(v, FilterModel, 'filters')
    """

    def validator(cls: Type[BaseModel], v: Any, info: Any = None) -> List[BaseModel]:
        field_name = param_name or (info.field_name if info else "field")
        return parse_json_or_model_list(v, model_class, field_name)

    return validator