def _split_query_ast_one_level_recursive(
    query_node: SubQueryNode,
    ast: AstType,
    type_info: TypeInfo,
    edge_to_stitch_fields: Dict[Tuple[str, str], Tuple[str, str]],
    name_assigner: IntermediateOutNameAssigner,
) -> AstType:
    """Return an AST node with which to replace the input AST in the selections that contain it.

    This function examines the selections of the input AST, and recursively calls either
    _split_query_ast_one_level_recursive or _split_query_ast_one_level_recursive_normal_fields
    depending on whether the selections contains a single InlineFragmentNode or a number of normal
    fields.

    Args:
        query_node: SubQueryNode, whose list of child query connections may be modified to
                    include new children.
        ast: The AST that we are trying to split into child components.
             It is not modified by this function.
        type_info: Used to get information about the types of fields while traversing the query AST.
        edge_to_stitch_fields: Mapping (type name, vertex field name) to
                               (source field name, sink field name) used in the @stitch directive
                               for each cross schema edge.
        name_assigner: Object used to generate and keep track of names of newly created
                       @output directives.

    Returns:
        The AST with which to replace the input AST in the selections that contain it.
    """
    if ast.selection_set is None:
        raise AssertionError("AST's selection_set cannot be None.")
    type_info.enter(ast.selection_set)
    selections = ast.selection_set.selections

    type_coercion = try_get_inline_fragment(selections)
    if type_coercion is not None:
        # Case 1: type coercion
        type_info.enter(type_coercion)
        new_type_coercion = _split_query_ast_one_level_recursive(
            query_node, type_coercion, type_info, edge_to_stitch_fields,
            name_assigner)
        type_info.leave(type_coercion)

        if new_type_coercion is type_coercion:
            new_selections: Sequence[SelectionNode] = selections
        else:
            new_selections = [new_type_coercion]
    else:
        # Case 2: normal fields
        new_selections = _split_query_ast_one_level_recursive_normal_fields(
            query_node, selections, type_info, edge_to_stitch_fields,
            name_assigner)
    type_info.leave(ast.selection_set)

    # Return input, or make copy
    if new_selections is not selections:
        new_ast = copy(ast)
        new_ast.selection_set = SelectionSetNode(selections=new_selections)
        return new_ast
    else:
        return ast
def _split_query_ast_one_level_recursive_normal_fields(
    query_node: SubQueryNode,
    selections: List[SelectionNode],
    type_info: TypeInfo,
    edge_to_stitch_fields: Dict[Tuple[str, str], Tuple[str, str]],
    name_assigner: IntermediateOutNameAssigner,
) -> Sequence[SelectionNode]:
    """One case of splitting query, selections contains a number of fields, no inline fragments.

    The input selections will be divided into three sets: property fields, intra-schema vertex
    fields, and cross-schema vertex fields.

    Each cross-schema vertex field will not be included in the output selections. The AST
    branch that each cross-schema vertex field leads to will be made into its own separate query
    AST. The parent and child property fields used in the stitch will be added to the parent and
    child ASTs, if not already present. @output directives will be added to these parent and
    child property fields, if not already present. @filter directives will not be added to
    child property fields in this step. This is because one may choose to rearrange and reroot
    the tree of SubQueryNodes to achieve an execution order with better performance. @filter
    directives should be added only once the tree's structure is fixed.

    _split_query_ast_one_level_recursive will be called recursive on each intra-schema vertex
    field.

    Args:
        query_node: Containing list of child query connections may be modified to
                    include new children.
        selections: Containing a number of property fields and vertex fields.
        type_info: Used to get information about the types of fields while traversing the query AST.
        edge_to_stitch_fields: Mapping (type name, vertex field name) to
                               (source field name, sink field name) used in the @stitch directive
                               for each cross schema edge.
        name_assigner: Used to generate and keep track of names of newly created @output directives.

    Returns:
        List of SelectionNodes to replace the list of selections in the SelectionSet one level
        above. All cross schema edges in the input list will be removed, and in their place,
        property fields added or modified. If no changes were made, the exact input list object
        will be returned.
    """
    parent_type = type_info.get_parent_type()
    if parent_type is None:
        raise AssertionError("parent_type cannot be None.")
    parent_type_name = parent_type.name

    made_changes = False

    # First, collect all property fields, but don't make any changes to them yet
    property_fields_map, vertex_fields = _split_selections_property_and_vertex(
        selections)

    # Second, process cross schema fields. This will modify our record of property fields, and
    # create child SubQueryNodes attached to the input SubQueryNode
    intra_schema_fields, cross_schema_fields = _split_vertex_fields_intra_and_cross_schema(
        vertex_fields, parent_type_name, edge_to_stitch_fields)
    for cross_schema_field in cross_schema_fields:
        type_info.enter(cross_schema_field)
        child_type = type_info.get_type()
        if child_type is not None:
            child_type_name = strip_non_null_and_list_from_type(
                child_type).name
        else:
            raise AssertionError(
                "The query may be invalid against the schema, causing TypeInfo to lose track "
                'of the types of fields. This occurs at the cross schema field "{}", while '
                'splitting the AST "{}"'.format(cross_schema_field,
                                                query_node.query_ast))
        stitch_data_key = (parent_type_name, cross_schema_field.name.value)
        parent_field_name, child_field_name = edge_to_stitch_fields[
            stitch_data_key]
        _process_cross_schema_field(
            query_node,
            cross_schema_field,
            property_fields_map,
            child_type_name,
            parent_field_name,
            child_field_name,
            name_assigner,
        )
        made_changes = True  # Cross schema edges are removed from the output, causing changes
        type_info.leave(cross_schema_field)

    # Third, process intra schema edges by recursing on them
    new_intra_schema_fields: List[SelectionNode] = []
    for intra_schema_field in intra_schema_fields:
        type_info.enter(intra_schema_field)
        new_intra_schema_field = _split_query_ast_one_level_recursive(
            query_node, intra_schema_field, type_info, edge_to_stitch_fields,
            name_assigner)
        if new_intra_schema_field is not intra_schema_field:
            made_changes = True
        new_intra_schema_fields.append(new_intra_schema_field)
        type_info.leave(intra_schema_field)

    # Return input, or make copy
    if made_changes:
        new_selections: Sequence[
            SelectionNode] = _get_selections_from_property_and_vertex_fields(
                property_fields_map, new_intra_schema_fields)
        return new_selections
    else:
        return selections
def _split_query_one_level(
    query_node: SubQueryNode,
    merged_schema_descriptor: MergedSchemaDescriptor,
    edge_to_stitch_fields: Dict[Tuple[str, str], Tuple[str, str]],
    name_assigner: IntermediateOutNameAssigner,
) -> None:
    """Split the query node, creating children out of all branches across cross schema edges.

    The input query_node will be modified. Its query_ast will be replaced by a new AST with
    branches leading out of cross schema edges removed, and new property fields and @output
    directives added as necessary. Its child_query_connections will be modified by tacking
    on SubQueryNodes created from these cut-off branches.

    Args:
        query_node: Query to be split into its child components. Its query_ast
                    will be replaced (but the original AST will not be modified) and its
                    child_query_connections will be modified.
        merged_schema_descriptor: The schema that the query AST contained in the input
                                  query_node targets.
        edge_to_stitch_fields: Mapping (type name, vertex field name) to
                               (source field name, sink field name) used in the @stitch directive
                               for each cross schema edge.
        name_assigner: Object used to generate and keep track of names of newly created
                       @output directive.

    Raises:
        - GraphQLValidationError if the query AST contained in the input query_node is invalid,
          for example, having an @output directive on a cross schema edge
        - SchemaStructureError if the merged_schema_descriptor provided appears to be invalid
          or inconsistent
    """
    type_info = TypeInfo(merged_schema_descriptor.schema)

    operation_definition = get_only_query_definition(query_node.query_ast,
                                                     GraphQLValidationError)
    if not isinstance(operation_definition, OperationDefinitionNode):
        raise AssertionError(
            f"Expected operation_definition to be an OperationDefinitionNode, but it was of"
            f"type {type(operation_definition)}. This should be impossible.")

    type_info.enter(operation_definition)
    new_operation_definition = _split_query_ast_one_level_recursive(
        query_node, operation_definition, type_info, edge_to_stitch_fields,
        name_assigner)
    type_info.leave(operation_definition)

    if new_operation_definition is not operation_definition:
        query_node.query_ast = DocumentNode(
            definitions=[new_operation_definition])

    # Check resulting AST is valid
    validation_errors = validate(merged_schema_descriptor.schema,
                                 query_node.query_ast)
    if len(validation_errors) > 0:
        raise AssertionError(
            'The resulting split query "{}" is invalid, with the following error messages: {}'
            "".format(query_node.query_ast, validation_errors))

    # Set schema id, check for consistency
    visitor = TypeInfoVisitor(
        type_info,
        SchemaIdSetterVisitor(type_info, query_node,
                              merged_schema_descriptor.type_name_to_schema_id),
    )
    visit(query_node.query_ast, visitor)

    if query_node.schema_id is None:
        raise AssertionError(
            'Unreachable code reached. The schema id of query piece "{}" has not been '
            "determined.".format(query_node.query_ast))