Repository: strawberry-graphql/strawberry

Issue:
Wrong info type passed to resolve_reference classmethod
<!-- Provide a general summary of the bug in the title above. -->

<!--- This template is entirely optional and can be removed, but is here to help both you and us. -->
<!--- Anything on lines wrapped in comments like these will not show up in the final text. -->

## Describe the Bug

<!-- A clear and concise description of what the bug is. -->

Take the following example:

```py
class User:
    id: strawberry.ID
    email: str

    @classmethod
    def resolve_reference(cls, id: strawberry.ID, info: strawberry.Info) -> Self:
      print(info.selected_fields)
```

This will result in an error saying the attribute doesn't exist, even though it should. The problem appears to stem from the fact that `GraphQLResolveInfo` is passed through instead in this instance rather than `Info` which provides the attribute.

(This is a federated example.)

## System Information

 - Operating system: MacOS (14.4)
 - Strawberry version (if applicable): 0.232.0

## Additional Context

<!-- Add any other relevant information about the problem here. -->

https://discord.com/channels/689806334337482765/1244954381326090331


Discussion:

Resolving patch:
diff --git a/strawberry/federation/schema.py b/strawberry/federation/schema.py
--- a/strawberry/federation/schema.py
+++ b/strawberry/federation/schema.py
@@ -1,5 +1,4 @@
 from collections import defaultdict
-from copy import copy
 from functools import cached_property, partial
 from itertools import chain
 from typing import (
@@ -10,6 +9,7 @@
     Iterable,
     List,
     Mapping,
+    NewType,
     Optional,
     Set,
     Type,
@@ -17,36 +17,36 @@
     cast,
 )
 
-from graphql import (
-    GraphQLError,
-    GraphQLField,
-    GraphQLInterfaceType,
-    GraphQLList,
-    GraphQLNonNull,
-    GraphQLScalarType,
-    GraphQLUnionType,
-)
-from graphql.type.definition import GraphQLArgument
+from graphql import GraphQLError
 
+from strawberry.annotation import StrawberryAnnotation
+from strawberry.custom_scalar import scalar
 from strawberry.printer import print_schema
 from strawberry.schema import Schema as BaseSchema
+from strawberry.type import (
+    StrawberryContainer,
+    WithStrawberryObjectDefinition,
+    get_object_definition,
+)
+from strawberry.types.info import Info
 from strawberry.types.types import StrawberryObjectDefinition
+from strawberry.union import StrawberryUnion
 from strawberry.utils.inspect import get_func_args
 
 from .schema_directive import StrawberryFederationSchemaDirective
 
 if TYPE_CHECKING:
     from graphql import ExecutionContext as GraphQLExecutionContext
-    from graphql import GraphQLObjectType
 
     from strawberry.custom_scalar import ScalarDefinition, ScalarWrapper
     from strawberry.enum import EnumDefinition
     from strawberry.extensions import SchemaExtension
     from strawberry.federation.schema_directives import ComposeDirective
     from strawberry.schema.config import StrawberryConfig
-    from strawberry.schema.types.concrete_type import TypeMap
     from strawberry.schema_directive import StrawberrySchemaDirective
-    from strawberry.union import StrawberryUnion
+
+
+FederationAny = scalar(NewType("_Any", object), name="_Any")  # type: ignore
 
 
 class Schema(BaseSchema):
@@ -67,7 +67,8 @@ def __init__(
         schema_directives: Iterable[object] = (),
         enable_federation_2: bool = False,
     ) -> None:
-        query = self._get_federation_query_type(query)
+        query = self._get_federation_query_type(query, mutation, subscription, types)
+        types = [*types, FederationAny]
 
         super().__init__(
             query=query,
@@ -84,16 +85,19 @@ def __init__(
 
         self.schema_directives = list(schema_directives)
 
-        self._add_scalars()
-        self._add_entities_to_query()
-
         if enable_federation_2:
             composed_directives = self._add_compose_directives()
             self._add_link_directives(composed_directives)  # type: ignore
         else:
             self._remove_resolvable_field()
 
-    def _get_federation_query_type(self, query: Optional[Type]) -> Type:
+    def _get_federation_query_type(
+        self,
+        query: Optional[Type[WithStrawberryObjectDefinition]],
+        mutation: Optional[Type[WithStrawberryObjectDefinition]],
+        subscription: Optional[Type[WithStrawberryObjectDefinition]],
+        additional_types: Iterable[Type[WithStrawberryObjectDefinition]],
+    ) -> Type:
         """Returns a new query type that includes the _service field.
 
         If the query type is provided, it will be used as the base for the new
@@ -109,13 +113,6 @@ def _get_federation_query_type(self, query: Optional[Type]) -> Type:
         added if the schema contains an entity type.
         """
 
-        # note we don't add the _entities field here, as we need to know if the
-        # schema contains an entity type first and we do that by leveraging
-        # the schema converter type map, so we don't have to do that twice
-        # TODO: ideally we should be able to do this without using the schema
-        # converter, but for now this is the easiest way to do it
-        # see `_add_entities_to_query`
-
         import strawberry
         from strawberry.tools.create_type import create_type
         from strawberry.tools.merge_types import merge_types
@@ -132,6 +129,19 @@ def service() -> Service:
 
         fields = [service]
 
+        entity_type = _get_entity_type(query, mutation, subscription, additional_types)
+
+        if entity_type:
+            self.entities_resolver.__annotations__["return"] = List[
+                Optional[entity_type]  # type: ignore
+            ]
+
+            entities_field = strawberry.field(
+                name="_entities", resolver=self.entities_resolver
+            )
+
+            fields.insert(0, entities_field)
+
         FederationQuery = create_type(name="Query", fields=fields)
 
         if query is None:
@@ -139,10 +149,7 @@ def service() -> Service:
 
         query_type = merge_types(
             "Query",
-            (
-                FederationQuery,
-                query,
-            ),
+            (FederationQuery, query),
         )
 
         # TODO: this should be probably done in merge_types
@@ -151,31 +158,9 @@ def service() -> Service:
 
         return query_type
 
-    def _add_entities_to_query(self) -> None:
-        entity_type = _get_entity_type(self.schema_converter.type_map)
-
-        if not entity_type:
-            return
-
-        self._schema.type_map[entity_type.name] = entity_type
-        fields = {"_entities": self._get_entities_field(entity_type)}
-
-        # Copy the query type, update it to use the modified fields
-        query_type = cast("GraphQLObjectType", self._schema.query_type)
-        fields.update(query_type.fields)
-
-        query_type = copy(query_type)
-        query_type.fields = fields
-
-        self._schema.query_type = query_type
-        self._schema.type_map[query_type.name] = query_type
-
     def entities_resolver(
-        self,
-        root,  # noqa: ANN001
-        info,  # noqa: ANN001
-        representations,  # noqa: ANN001
-    ) -> List[object]:
+        self, info: Info, representations: List[FederationAny]
+    ) -> List[FederationAny]:
         results = []
 
         for representation in representations:
@@ -198,9 +183,8 @@ def entities_resolver(
             else:
                 from strawberry.arguments import convert_argument
 
-                strawberry_schema = info.schema.extensions["strawberry-definition"]
-                config = strawberry_schema.config
-                scalar_registry = strawberry_schema.schema_converter.scalar_registry
+                config = info.schema.config
+                scalar_registry = info.schema.schema_converter.scalar_registry
 
                 get_result = partial(
                     convert_argument,
@@ -222,11 +206,6 @@ def entities_resolver(
 
         return results
 
-    def _add_scalars(self) -> None:
-        self.Any = GraphQLScalarType("_Any")
-
-        self._schema.type_map["_Any"] = self.Any
-
     def _remove_resolvable_field(self) -> None:
         # this might be removed when we remove support for federation 1
         # or when we improve how we print the directives
@@ -335,17 +314,6 @@ def _add_compose_directives(self) -> List["ComposeDirective"]:
 
         return compose_directives
 
-    def _get_entities_field(self, entity_type: GraphQLUnionType) -> GraphQLField:
-        return GraphQLField(
-            GraphQLNonNull(GraphQLList(entity_type)),
-            args={
-                "representations": GraphQLArgument(
-                    GraphQLNonNull(GraphQLList(GraphQLNonNull(self.Any)))
-                )
-            },
-            resolve=self.entities_resolver,
-        )
-
     def _warn_for_federation_directives(self) -> None:
         # this is used in the main schema to raise if there's a directive
         # that's for federation, but in this class we don't want to warn,
@@ -354,33 +322,56 @@ def _warn_for_federation_directives(self) -> None:
         pass
 
 
-def _get_entity_type(type_map: "TypeMap") -> Optional[GraphQLUnionType]:
-    # https://www.apollographql.com/docs/apollo-server/federation/federation-spec/#resolve-requests-for-entities
+def _get_entity_type(
+    query: Optional[Type[WithStrawberryObjectDefinition]],
+    mutation: Optional[Type[WithStrawberryObjectDefinition]],
+    subscription: Optional[Type[WithStrawberryObjectDefinition]],
+    additional_types: Iterable[Type[WithStrawberryObjectDefinition]],
+) -> Optional[StrawberryUnion]:
+    # recursively iterate over the schema to find all types annotated with @key
+    # if no types are annotated with @key, then the _Entity union and Query._entities
+    # field should not be added to the schema
 
-    # To implement the _Entity union, each type annotated with @key
-    # should be added to the _Entity union.
+    entity_types = set()
 
-    federation_key_types = [
-        type.implementation
-        for type in type_map.values()
-        if _has_federation_keys(type.definition)
-        # TODO: check this
-        and not isinstance(type.implementation, GraphQLInterfaceType)
-    ]
+    # need a stack to keep track of the types we need to visit
+    stack: List[Any] = [query, mutation, subscription, *additional_types]
 
-    # If no types are annotated with the key directive, then the _Entity
-    # union and Query._entities field should be removed from the schema.
-    if not federation_key_types:
-        return None
+    seen = set()
+
+    while stack:
+        type_ = stack.pop()
 
-    entity_type = GraphQLUnionType("_Entity", federation_key_types)  # type: ignore
+        if type_ is None:
+            continue
 
-    def _resolve_type(self: Any, value: Any, info: Any) -> str:
-        return self.__strawberry_definition__.name
+        while isinstance(type_, StrawberryContainer):
+            type_ = type_.of_type
+
+        type_definition = get_object_definition(type_, strict=False)
+
+        if type_definition is None:
+            continue
+
+        if type_definition.is_object_type and _has_federation_keys(type_definition):
+            entity_types.add(type_)
+
+        for field in type_definition.fields:
+            if field.type and field.type in seen:
+                continue
+
+            seen.add(field.type)
+            stack.append(field.type)
+
+    if not entity_types:
+        return None
 
-    entity_type.resolve_type = _resolve_type
+    sorted_types = sorted(entity_types, key=lambda t: t.__strawberry_definition__.name)
 
-    return entity_type
+    return StrawberryUnion(
+        "_Entity",
+        type_annotations=tuple(StrawberryAnnotation(type_) for type_ in sorted_types),
+    )
 
 
 def _is_key(directive: Any) -> bool:
diff --git a/strawberry/types/types.py b/strawberry/types/types.py
--- a/strawberry/types/types.py
+++ b/strawberry/types/types.py
@@ -146,6 +146,10 @@ def is_specialized_generic(self) -> bool:
     def specialized_type_var_map(self) -> Optional[Dict[str, type]]:
         return get_specialized_type_var_map(self.origin)
 
+    @property
+    def is_object_type(self) -> bool:
+        return not self.is_input and not self.is_interface
+
     @property
     def type_params(self) -> List[TypeVar]:
         type_params: List[TypeVar] = []


Test patch:
diff --git a/tests/federation/test_entities.py b/tests/federation/test_entities.py
--- a/tests/federation/test_entities.py
+++ b/tests/federation/test_entities.py
@@ -46,11 +46,11 @@ def test_info_param_in_resolve_reference():
     @strawberry.federation.type(keys=["upc"])
     class Product:
         upc: str
-        info: str
+        debug_field_name: str
 
         @classmethod
-        def resolve_reference(cls, info, upc) -> "Product":
-            return Product(upc=upc, info=info)
+        def resolve_reference(cls, info: strawberry.Info, upc: str) -> "Product":
+            return Product(upc=upc, debug_field_name=info.field_name)
 
     @strawberry.federation.type(extend=True)
     class Query:
@@ -65,7 +65,7 @@ def top_products(self, first: int) -> typing.List[Product]:
             _entities(representations: $representations) {
                 ... on Product {
                     upc
-                    info
+                    debugFieldName
                 }
             }
         }
@@ -80,10 +80,15 @@ def top_products(self, first: int) -> typing.List[Product]:
 
     assert not result.errors
 
-    assert (
-        "GraphQLResolveInfo(field_name='_entities', field_nodes=[FieldNode"
-        in result.data["_entities"][0]["info"]
-    )
+    assert result.data == {
+        "_entities": [
+            {
+                "upc": "B00005N5PF",
+                # _entities is the field that's called by federation
+                "debugFieldName": "_entities",
+            }
+        ]
+    }
 
 
 def test_does_not_need_custom_resolve_reference_for_basic_things():
diff --git a/tests/federation/test_schema.py b/tests/federation/test_schema.py
--- a/tests/federation/test_schema.py
+++ b/tests/federation/test_schema.py
@@ -23,6 +23,28 @@ def top_products(self, first: int) -> List[Product]:
 
     schema = strawberry.federation.Schema(query=Query, enable_federation_2=True)
 
+    expected_sdl = textwrap.dedent("""
+        type Product {
+          upc: String!
+          name: String
+          price: Int
+          weight: Int
+        }
+
+        extend type Query {
+          _service: _Service!
+          topProducts(first: Int!): [Product!]!
+        }
+
+        scalar _Any
+
+        type _Service {
+          sdl: String!
+        }
+    """).strip()
+
+    assert str(schema) == expected_sdl
+
     query = """
         query {
             __type(name: "_Entity") {
@@ -57,6 +79,35 @@ def top_products(self, first: int) -> List[Product]:
 
     schema = strawberry.federation.Schema(query=Query, enable_federation_2=True)
 
+    expected_sdl = textwrap.dedent("""
+        schema @link(url: "https://specs.apollo.dev/federation/v2.7", import: ["@key"]) {
+          query: Query
+        }
+
+        type Product @key(fields: "upc") {
+          upc: String!
+          name: String
+          price: Int
+          weight: Int
+        }
+
+        extend type Query {
+          _entities(representations: [_Any!]!): [_Entity]!
+          _service: _Service!
+          topProducts(first: Int!): [Product!]!
+        }
+
+        scalar _Any
+
+        union _Entity = Product
+
+        type _Service {
+          sdl: String!
+        }
+    """).strip()
+
+    assert str(schema) == expected_sdl
+
     query = """
         query {
             __type(name: "_Entity") {