Extended types should support interfaces and type extension

README.md

@@ -67,6 +67,16 @@ class Message(graphene.ObjectType):
         return 1
 
 ```
+
+### __resolve_reference
+* Each type which is decorated with `@key` or `@extend` is added to `_Entity` union
+* `__resolve_reference` method can be defined for each type that is an entity. This method is called whenever an entity is requested as part of the fulfilling a query plan.
+If not explicitly defined, default resolver is used. Default resolver just creates instance of type with passed fieldset as kwargs, see [`entity.get_entity_query`](graphene_federation/entity.py) for more details
+* You should define `__resolve_reference`, if you need to extract object before passing it to fields resolvers (example: [FileNode](integration_tests/service_b/src/schema.py))
+* You should not define `__resolve_reference`, if fileds resolvers need only data passed in fieldset (example: [FunnyText](integration_tests/service_a/src/schema.py))
+* read more in [official documentation](https://www.apollographql.com/docs/apollo-server/api/apollo-federation/#__resolvereference)
+------------------------
+
 For more details see [examples](examples/)
 
 Or better check [integration_tests](integration_tests/)

graphene_federation/entity.py

@@ -30,13 +30,19 @@ def resolve_entities(parent, info, representations):
             entities = []
             for representation in representations:
                 model = custom_entities[representation["__typename"]]
-                resolver = getattr(model, "_%s__resolve_reference" % representation["__typename"])
                 model_aguments = representation.copy()
                 model_aguments.pop("__typename")
+                model_instance = model(**model_aguments)
 
-                entities.append(
-                    resolver(model(**model_aguments), info)
-                )
+                try:
+                    resolver = getattr(
+                        model, "_%s__resolve_reference" % representation["__typename"])
+                except AttributeError:
+                    pass
+                else:
+                    model_instance = resolver(model_instance, info)
+
+                entities.append(model_instance)
             return entities
 
     return EntityQuery

graphene_federation/extend.py

@@ -1,3 +1,4 @@
+from .entity import register_entity
 extended_types = {}
 
 
@@ -7,6 +8,7 @@ def register_extend_type(typename, Type):
 
 def extend(fields: str):
     def decorator(Type):
+        register_entity(Type.__name__, Type)
         register_extend_type(Type.__name__, Type)
         setattr(Type, '_sdl', '@key(fields: "%s")' % fields)
         return Type

graphene_federation/service.py

@@ -11,8 +11,9 @@ def _mark_external(entity_name, entity, schema):
         field = getattr(entity, field_name, None)
         if field is not None and getattr(field, '_external', False):
             # todo write tests on regexp
-            pattern = re.compile("(\s%s\s*\{[^\}]*\s%s[\s]*:[\s]*[^\s]+)(\s)" % (entity_name, field_name))
-            schema = pattern.sub('\g<1> @external ', schema)
+            pattern = re.compile(
+                r"(\s%s\s[^\{]*\{[^\}]*\s%s[\s]*:[\s]*[^\s]+)(\s)" % (entity_name, field_name))
+            schema = pattern.sub(r'\g<1> @external ', schema)
 
     return schema
 
@@ -34,9 +35,11 @@ def get_sdl(schema, custom_entities):
     for entity_name, entity in extended_types.items():
         string_schema = _mark_external(entity_name, entity, string_schema)
 
-        type_def = "type %s " % entity_name
-        repl_str = "extend %s %s " % (type_def, entity._sdl)
-        pattern = re.compile(type_def)
+        type_def_re = r"type %s ([^\{]*)" % entity_name
+        type_def = r"type %s " % entity_name
+        repl_str = r"extend %s \1 %s " % (type_def, entity._sdl)
+        pattern = re.compile(type_def_re)
+
         string_schema = pattern.sub(repl_str, string_schema)
 
     return string_schema

integration_tests/service_a/src/schema.py

@@ -1,16 +1,37 @@
-from graphene import ObjectType, String, Int, List, NonNull, Field
+from graphene import ObjectType, String, Int, List, NonNull, Field, Interface
 from graphene_federation import build_schema, extend, external
 
 
+class DecoratedText(Interface):
+    color = Int(required=True)
+
+
 @extend(fields='id')
 class FileNode(ObjectType):
     id = external(Int(required=True))
 
 
 @extend(fields='id')
 class FunnyText(ObjectType):
+    class Meta:
+        interfaces = (DecoratedText,)
     id = external(Int(required=True))
 
+    def resolve_color(self, info, **kwargs):
+        return self.id + 2
+
+
+class FunnyTextAnother(ObjectType):
+    """
+    To test @extend on types with same prefix
+    """
+    class Meta:
+        interfaces = (DecoratedText,)
+    id = Int(required=True)
+
+    def resolve_color(self, info, **kwargs):
+        return self.id + 2
+
 
 class Post(ObjectType):
     id = Int(required=True)
@@ -34,4 +55,4 @@ def resolve_goodbye(root, info):
         return 'See ya!'
 
 
-schema = build_schema(query=Query)
+schema = build_schema(query=Query, types=[FunnyTextAnother])

integration_tests/tests/tests/test_main.py

@@ -26,6 +26,7 @@ def test_external_types():
                     text {
                        id
                        body
+                       color
                     }
                     files {
                        id
@@ -44,11 +45,11 @@ def test_external_types():
     posts = json.loads(response.content)['data']['posts']
     assert 3 == len(posts)
     assert [{'id': 1, 'name': 'file_1'}] == posts[0]['files']
-    assert {'id': 1, 'body': 'funny_text_1'} == posts[0]['text']
+    assert {'id': 1, 'body': 'funny_text_1', 'color': 3} == posts[0]['text']
     assert [{'id': 2, 'name': 'file_2'}, {'id': 3, 'name': 'file_3'}] == posts[1]['files']
-    assert {'id': 2, 'body': 'funny_text_2'} == posts[1]['text']
+    assert {'id': 2, 'body': 'funny_text_2', 'color': 4} == posts[1]['text']
     assert posts[2]['files'] is None
-    assert {'id': 3, 'body': 'funny_text_3'} == posts[2]['text']
+    assert {'id': 3, 'body': 'funny_text_3', 'color': 5} == posts[2]['text']
 
 
 def test_key_decorator_applied_by_exact_match_only():