Source code for operators.grapi.invalidate_cache

from typing import Any, Dict, List, Optional
from airflow.exceptions import AirflowException
from airflow.utils.decorators import apply_defaults
import requests
from operators.grapi.grapi_base import GrapiBaseOperator
from config import Config



class InvalidateCacheOperator(GrapiBaseOperator):
    """
    Operator to invalidate cache using the invalidateCache GraphQL mutation.

    :param tenant: The tenant identifier
    :param modules: List of modules to invalidate cache for. For the data team we can always set this to ["FabricLocationStatistics"].
    :param stage: API stage to use ("main" or "stage")
    :param timeout: Request timeout in seconds
    :param retry_count: Number of retry attempts
    :param retry_delay: Delay between retries in seconds
    """

    template_fields = ["tenant", "modules"]
    ui_color = "#ffecb3"

    GRAPHQL_QUERY = """
    mutation InvalidateCache($data: InvalidateCacheInput!) {
        invalidateCache(data: $data)
    }
    """

    @apply_defaults
    def __init__(self, modules: List[str]=["FabricLocationStatistics"], *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.modules = modules



    def get_request_payload(self) -> Dict[str, Any]:
        """Prepare the GraphQL request payload."""
        if not self.modules:
            raise Exception("Missing required field: modules")

        if not isinstance(self.modules, list):
            raise Exception("modules must be a list of strings")

        return {
            "query": self.GRAPHQL_QUERY,
            "variables": {"data": {"modules": self.modules}},
            "operationName": "InvalidateCache",
        }




    def post_process_response(
        self, response: requests.Response, context: Optional[Dict[str, Any]] = None
    ) -> Dict[str, Any]:
        """Parse and validate GraphQL response."""
        import json

        try:
            result = response.json()
        except json.JSONDecodeError as e:
            raise Exception(f"Failed to parse response JSON: {str(e)}")

        # Check for GraphQL errors
        if "errors" in result:
            error_msg = json.dumps(result["errors"])
            self.log.error(f"GraphQL errors: {error_msg}")
            raise Exception(f"GraphQL request failed with errors: {error_msg}")

        # Check if data exists
        if "data" not in result:
            raise Exception("Response missing 'data' field")

        # Log cache invalidation result
        invalidation_result = result.get("data", {}).get("invalidateCache")
        self.log.info(f"Cache invalidation result: {invalidation_result}")

        return result["data"]




    def execute(self, context):
        """Execute the GraphQL mutation."""
        try:

            self.log.info(f"Invalidating cache for modules: {self.modules}")
            results=super().execute(context)
            self.log.info(f"Cache invalidation completed successfully: {results}")

        except Exception as e:
            self.log.error(f"InvalidateCache failed: {str(e)}", exc_info=True)
            raise AirflowException(f"Execution failed: {str(e)}")