Support for multiple keys

graphene_federation/main.py

@@ -20,9 +20,11 @@ def build_schema(query, mutation=None, **kwargs):
     return graphene.Schema(query=_get_query(schema, query), mutation=mutation, **kwargs)
 
 
-def key(fields: str):
+def key(fields: str, *args: str):
     def decorator(Type):
         register_entity(Type.__name__, Type)
-        setattr(Type, '_sdl', '@key(fields: "%s")' % fields)
+        keys = [fields] + list(args)
+        sdl = " ".join(['@key(fields: "%s")' % key for key in keys])
+        setattr(Type, '_sdl', sdl)
         return Type
     return decorator

integration_tests/service_a/src/schema.py

@@ -12,11 +12,17 @@ class FunnyText(ObjectType):
     id = external(Int(required=True))
 
 
+@extend(fields='email')
+class User(ObjectType):
+    email = external(String())
+
+
 class Post(ObjectType):
     id = Int(required=True)
     title = String(required=True)
     text = Field(lambda: FunnyText)
     files = List(NonNull(FileNode))
+    author = Field(lambda: User)
 
 
 class Query(ObjectType):
@@ -28,6 +34,7 @@ def resolve_posts(root, info):
             Post(id=1, title='title1', text=FunnyText(id=1), files=[FileNode(id=1)]),
             Post(id=2, title='title2', text=FunnyText(id=2), files=[FileNode(id=2), FileNode(id=3)]),
             Post(id=3, title='title3', text=FunnyText(id=3)),
+            Post(id=4, title='title4', text=FunnyText(id=4), author=User(email="frank@frank.com")),
         ]
 
     def resolve_goodbye(root, info):

integration_tests/service_b/src/schema.py

@@ -26,6 +26,20 @@ def __resolve_reference(self, info, **kwargs):
         return FileNode(id=self.id, name=f'file_{self.id}')
 
 
+@key('id', 'email')
+class User(ObjectType):
+    id = Int(required=True)
+    email = String()
+
+    def __resolve_reference(self, info, **kwargs):
+        if hasattr(info, 'id'):
+            return User(id=self.id, email=f'name_{self.id}')
+
+        user_id = 1001 if self.email == "frank@frank.com" else hash(self.email) % 10000000
+
+        return User(id=user_id, email=self.email)
+
+
 # to test that @key applied only to FileNode, but not to FileNodeAnother
 class FileNodeAnother(ObjectType):
     id = Int(required=True)
@@ -36,6 +50,13 @@ class Query(ObjectType):
     file = Field(lambda: FileNode)
 
 
+types = [
+    FileNode,
+    FunnyText,
+    FileNodeAnother,
+    User
+]
+
 class FunnyMutation(Mutation):
     result = String(required=True)
 
@@ -48,4 +69,15 @@ class Mutation(ObjectType):
     funny_mutation = FunnyMutation.Field()
 
 
-schema = build_schema(Query, Mutation, types=[FileNode, FunnyText, FileNodeAnother])
+class Query(ObjectType):
+    file = Field(lambda: FileNode)
+
+
+types = [
+    FileNode,
+    FunnyText,
+    FileNodeAnother,
+    User
+]
+
+schema = build_schema(Query, Mutation, types=types)

integration_tests/tests/tests/test_main.py

@@ -31,6 +31,10 @@ def test_external_types():
                        id
                        name
                     }
+                    author {
+                        id
+                        email
+                    }
                 }
             }
         """,
@@ -42,16 +46,19 @@ def test_external_types():
     )
     assert response.status_code == 200
     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': 2, 'name': 'file_2'}, {'id': 3, 'name': 'file_3'}] == posts[1]['files']
-    assert {'id': 2, 'body': 'funny_text_2'} == posts[1]['text']
+    assert 4 == len(posts)
+    assert [{"id": 1, "name": "file_1"}] == posts[0]['files']
+    assert {"id": 1, "body": "funny_text_1"} == 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 posts[2]['files'] is None
-    assert {'id': 3, 'body': 'funny_text_3'} == posts[2]['text']
+    assert {"id": 3, "body": "funny_text_3"} == posts[2]['text']
+    assert posts[3]["author"] is not None
+    author = posts[3]["author"]
+    assert {"id": 1001, "email": "frank@frank.com", } == author
 
 
-def test_key_decorator_applied_by_exact_match_only():
+def fetch_sdl():
     query = {
         'query': """
             query {
@@ -64,7 +71,11 @@ def test_key_decorator_applied_by_exact_match_only():
     }
     response = requests.post('http://service_b:5000/graphql', json=query)
     assert response.status_code == 200
-    sdl = response.json()['data']['_service']['sdl']
+    return response.json()['data']['_service']['sdl']
+
+
+def test_key_decorator_applied_by_exact_match_only():
+    sdl = fetch_sdl()
     assert 'type FileNode  @key(fields: "id")' in sdl
     assert 'type FileNodeAnother  @key(fields: "id")' not in sdl
 
@@ -85,3 +96,7 @@ def test_mutation_is_accessible_in_federation():
     assert 'errors' not in response.json()
     assert response.json()['data']['funnyMutation']['result'] == 'Funny'
 
+
+def test_multiple_key_decorators_apply_multiple_key_annotations():
+    sdl = fetch_sdl()
+    assert 'type User  @key(fields: "id") @key(fields: "email")' in sdl