operation_name_and_type.py

import logging
from typing import Optional, Literal, List, cast

from graphql import DocumentNode as GraphQLDocumentNode
from graphql.language import OperationDefinitionNode, DefinitionNode

from strawberry.extensions import Extension


def get_first_operation(
    definitions: List[DefinitionNode],
) -> Optional[OperationDefinitionNode]:
    definition = next(
        (
            cast(OperationDefinitionNode, node)
            for node in definitions
            if isinstance(node, OperationDefinitionNode)
        ),
        None,
    )
    return definition


def get_operation_name(
    operation_name: Optional[str], graphql_document: GraphQLDocumentNode
) -> str:
    # If an operation_name has been specified then use that
    if operation_name:
        return operation_name

    definition = get_first_operation(graphql_document.definitions)
    if not definition:
        raise RuntimeError("Can't get GraphQL operation")

    if not definition.name:
        return "unnamed"

    return definition.name.value


graphql_operation_types = Literal["QUERY", "MUTATION"]


def get_operation_type(
    operation_name: Optional[str], graphql_document: GraphQLDocumentNode
) -> graphql_operation_types:
    definition: Optional[OperationDefinitionNode] = None

    # If no operation_name has been specified then use the first
    # OperationDefinitionNode
    if not operation_name:
        definition = get_first_operation(graphql_document.definitions)
    else:
        for d in graphql_document.definitions:
            d = cast(OperationDefinitionNode, d)
            if d.name and d.name.value == operation_name:
                definition = d
                break

    if not definition:
        raise RuntimeError("Can't get GraphQL operation type")

    return cast(graphql_operation_types, definition.operation.name)


class MyExtension(Extension):
    def on_request_end(self):
        execution_context = self.execution_context

        if not execution_context.graphql_document:
            logger.warning("Execution context doesn't have a graphql_document")
            return

        operation_name = get_operation_name(
            execution_context.operation_name, execution_context.graphql_document
        )
        operation_type = get_operation_type(
            execution_context.operation_name, execution_context.graphql_document
        )