Source code for operators.grapi.update_application_answer

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 UpdateApplicationAnswerOperator(GrapiBaseOperator):

    ui_color = "#b4ffb3"

    GRAPHQL_QUERY = """
    mutation UpdateApplicationAnswerProtected($data: ApplicationAnswerProtectedInput!, $where: ApplicationAnswerProtectedWhereInput!) {
        updateApplicationAnswerProtected(data: $data, where: $where) {
            id
        }
    }
    """
    data = None
    expense_line_item = None
    application_id = None

    @apply_defaults
    def __init__(
        self, expense_line_item: str, application_id: str, data: dict, *args, **kwargs
    ):
        super().__init__(*args, **kwargs)
        self.data = data
        self.expense_line_item = expense_line_item
        self.application_id = application_id



    def get_request_payload(self) -> Dict[str, Any]:
        """Prepare the GraphQL request payload."""

        variables = {
            "data": self.data,
            "where": {
                "id": self.expense_line_item,
                "applicationId": self.application_id,
            },
        }

        print({
            "query": self.GRAPHQL_QUERY,
            "variables": variables,
        })
        return {
            "query": self.GRAPHQL_QUERY,
            "variables": variables,
        }





    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")

        return result["data"]




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

            self.log.info(f"Updating application resource")
            results = super().execute(context)
            self.log.info(
                f"Updating application resource completed successfully: {results}"
            )

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