superset2/superset/mcp_service/common/cache_schemas.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.

"""
Cache control schemas for MCP tools.

These schemas provide cache control parameters that leverage Superset's
existing cache infrastructure including query result cache, metadata cache,
form data cache, and dashboard cache.
"""

from typing import Annotated, Any

from pydantic import BaseModel, Field, model_validator


class CacheControlMixin(BaseModel):
    """
    Mixin for cache control parameters that can be added to any request schema.

    Leverages Superset's existing cache layers:
    - Query Result Cache: Caches actual query results from customer databases
    - Metadata Cache: Caches table schemas, column info, etc.
    - Form Data Cache: Caches chart configurations for explore URLs
    - Dashboard Cache: Caches rendered dashboard components
    """

    use_cache: bool = Field(default=True, description="Use cache if available")

    force_refresh: bool = Field(
        default=False, description="Invalidate cache and fetch fresh data"
    )


class QueryCacheControl(CacheControlMixin):
    """
    Cache control specifically for data queries.

    Used by tools that execute SQL queries against customer databases
    like get_chart_data, chart previews, and chart creation.
    """

    cache_timeout: int | None = Field(
        default=None, description="Cache timeout override in seconds (0 to disable)"
    )


class MetadataCacheControl(CacheControlMixin):
    """
    Cache control for metadata operations.

    Used by tools that fetch database metadata like table schemas,
    column information, metrics, and dataset listings.
    """

    refresh_metadata: bool = Field(
        default=False, description="Refresh metadata cache for schema changes"
    )


class FormDataCacheControl(CacheControlMixin):
    """
    Cache control for form data and chart configurations.

    Used by tools that work with chart configurations and explore URLs
    like generate_explore_link and chart preview updates.
    """

    cache_form_data: bool = Field(
        default=True, description="Cache form data for future use"
    )


class CreatedByMeMixin(BaseModel):
    """Mixin that adds a created_by_me filter flag to list request schemas.

    Provides a clean caller-facing alternative to exposing foreign key IDs.
    The server translates the flag into the appropriate FK filter and injects
    the current user's ID automatically.
    """

    created_by_me: Annotated[
        bool,
        Field(
            default=False,
            description=(
                "When true, return only items created by the current user. "
                "Can be combined with 'filters' but not with 'search'."
            ),
        ),
    ]

    @model_validator(mode="after")
    def _validate_created_by_me_with_search(self) -> Any:
        if getattr(self, "search", None) and self.created_by_me:
            raise ValueError(
                "'created_by_me' cannot be combined with 'search'. "
                "Use 'created_by_me' alone or with 'filters'."
            )
        return self


class OwnedByMeMixin(BaseModel):
    """Mixin that adds an owned_by_me filter flag to list request schemas.

    Provides a clean caller-facing alternative to exposing M2M owner IDs.
    The server translates the flag into the appropriate owner filter and injects
    the current user's ID automatically.

    When combined with created_by_me, returns items where the current user is
    either the creator OR an owner (union, not intersection).
    """

    owned_by_me: Annotated[
        bool,
        Field(
            default=False,
            description=(
                "When true, return only items where the current user is listed as "
                "an owner. Can be combined with 'filters' but not with 'search'. "
                "Can be combined with 'created_by_me' to return items where the "
                "current user is either the creator or an owner."
            ),
        ),
    ]

    @model_validator(mode="after")
    def _validate_owned_by_me(self) -> Any:
        if getattr(self, "search", None) and self.owned_by_me:
            raise ValueError(
                "'owned_by_me' cannot be combined with 'search'. "
                "Use 'owned_by_me' alone or with 'filters'."
            )
        return self


class CacheStatus(BaseModel):
    """
    Information about cache usage in tool responses.

    Provides transparency about whether data was served from cache
    or freshly fetched, helping users understand data freshness.
    """

    cache_hit: bool = Field(
        description="Whether the data was served from cache (True) or "
        "freshly fetched (False)"
    )

    cache_type: str | None = Field(
        default=None,
        description=(
            "Type of cache used: 'query', 'metadata', 'form_data', 'dashboard', "
            "or 'none' if no cache was used"
        ),
    )

    cache_age_seconds: int | None = Field(
        default=None, description="Age of cached data in seconds, if served from cache"
    )

    cache_key: str | None = Field(
        default=None,
        description="Cache key used (for debugging), truncated if too long",
    )

    refreshed: bool = Field(
        default=False, description="Whether cache was refreshed as part of this request"
    )