Federation v2 support added

README.md

@@ -171,7 +171,6 @@ There is also a cool [example](https://github.com/preply/graphene-federation/iss
 ## Known issues
 
 1. decorators will not work properly on fields with custom names for example `some_field = String(name='another_name')`
-1. `@key` decorator will not work on [compound primary key](https://www.apollographql.com/docs/federation/entities-advanced/#compound-keys)
 
 ------------------------
 

examples/entities.py

@@ -39,7 +39,7 @@ def resolve_file(self, **kwargs):
 '''
 result = schema.execute(query)
 print(result.data)
-# OrderedDict([('_service', OrderedDict([('sdl', '   type File @key(fields: "id") {   id: Int!   name: String }  extend type Query {   hello: String   file: File } ')]))])
+# {'_service': {'sdl': 'type Query {\n  file: File\n}\n\ntype File @key(fields: "id") {\n  id: Int!\n  name: String\n}'}}
 
 query ='''
     query entities($_representations: [_Any!]!) {
@@ -62,4 +62,4 @@ def resolve_file(self, **kwargs):
     ]
 })
 print(result.data)
-# OrderedDict([('_entities', [OrderedDict([('id', 1), ('name', 'test_name')])])])
+# {'_entities': [{'id': 1, 'name': 'test_name'}]}

examples/extend.py

@@ -28,4 +28,4 @@ def resolve_file(self, **kwargs):
 '''
 result = schema.execute(query)
 print(result.data)
-# OrderedDict([('_service', OrderedDict([('sdl', '   extend type Message @key(fields: "id") {   id: Int! @external }  type Query {   message: Message } ')]))])
+# {'sdl': 'type Query {\n  message: Message\n}\n\nextend type Message @key(fields: "id") {\n  id: Int! @external\n}'}}

examples/inaccessible.py

@@ -0,0 +1,70 @@
+import graphene
+
+from graphene_federation import inaccessible, external, provides, key, override, shareable
+
+from graphene_federation import build_schema
+
+
+@key(fields="x")
+class Position(graphene.ObjectType):
+    x = graphene.Int(required=True)
+    y = external(graphene.Int(required=True))
+    z = inaccessible(graphene.Int(required=True))
+    a = provides(graphene.Int(), fields="x")
+    b = override(graphene.Int(required=True), from_="h")
+
+
+@inaccessible
+class ReviewInterface(graphene.Interface):
+    interfaced_body = graphene.String(required=True)
+
+
+@inaccessible
+class Review(graphene.ObjectType):
+    class Meta:
+        interfaces = (ReviewInterface,)
+
+    id = graphene.Int(required=True)
+    body = graphene.String(required=True)
+
+
+@inaccessible
+class Human(graphene.ObjectType):
+    name = graphene.String()
+    born_in = graphene.String()
+
+
+@inaccessible
+class Droid(graphene.ObjectType):
+    name = graphene.String()
+    primary_function = graphene.String()
+
+
+@inaccessible
+class Starship(graphene.ObjectType):
+    name = graphene.String()
+    length = graphene.Int()
+
+
+@inaccessible
+class SearchResult(graphene.Union):
+    class Meta:
+        types = (Human, Droid, Starship)
+
+
+class Query(graphene.ObjectType):
+    position = graphene.Field(Position)
+
+
+schema = build_schema(Query, enable_federation_2=True, types=(ReviewInterface, SearchResult, Review))
+
+query = '''
+    query getSDL {
+      _service {
+         sdl
+      }
+    }
+'''
+result = schema.execute(query)
+print(result.data)
+# {'_service': {'sdl': 'extend schema @link(url: "https://specs.apollo.dev/federation/v2.0", import: ["@external", "@key", "@override", "@provides", "@inaccessible"])\ntype Query {\n  position: Position\n}\n\ntype Position @key(fields: "x") {\n  x: Int!\n  y: Int! @external\n  z: Int! @inaccessible\n  a: Int @provides(fields: "x")\n  b: Int! @override(from: "h")\n}'}}

examples/override.py

@@ -0,0 +1,28 @@
+import graphene
+
+from graphene_federation import build_schema, shareable, external, key, override, inaccessible
+
+
+@key(fields="id")
+class Product(graphene.ObjectType):
+    id = graphene.ID(required=True)
+    in_stock = override(graphene.Boolean(required=True), "Products")
+    out_stock = inaccessible(graphene.Boolean(required=True))
+
+
+class Query(graphene.ObjectType):
+    position = graphene.Field(Product)
+
+
+schema = build_schema(Query, enable_federation_2=True)
+
+query = '''
+    query getSDL {
+      _service {
+         sdl
+      }
+    }
+'''
+result = schema.execute(query)
+print(result.data)
+# {'_service': {'sdl': 'extend schema @link(url: "https://specs.apollo.dev/federation/v2.0", import: ["@key", "@override", "@inaccessible"])\ntype Query {\n  position: Product\n}\n\ntype Product @key(fields: "id") {\n  id: ID!\n  inStock: Boolean! @override(from: "Products")\n  outStock: Boolean! @inaccessible\n}'}}

examples/shareable.py

@@ -0,0 +1,54 @@
+import graphene
+from graphene import Interface
+
+from graphene_federation.shareable import shareable
+
+from graphene_federation import build_schema
+
+
+@shareable
+class Position(graphene.ObjectType):
+    x = graphene.Int(required=True)
+    y = shareable(graphene.Int(required=True))
+
+
+@shareable
+class Human(graphene.ObjectType):
+    name = graphene.String()
+    born_in = graphene.String()
+
+
+@shareable
+class Droid(graphene.ObjectType):
+    name = shareable(graphene.String())
+    primary_function = graphene.String()
+
+
+@shareable
+class Starship(graphene.ObjectType):
+    name = graphene.String()
+    length = shareable(graphene.Int())
+
+
+@shareable
+class SearchResult(graphene.Union):
+    class Meta:
+        types = (Human, Droid, Starship)
+
+
+class Query(graphene.ObjectType):
+    position = graphene.Field(Position)
+
+
+schema = build_schema(Query, enable_federation_2=True, types=(SearchResult,))
+
+query = '''
+    query getSDL {
+      _service {
+         sdl
+      }
+    }
+'''
+result = schema.execute(query)
+print(result.data)
+# {'_service': {'sdl': 'extend schema @link(url: "https://specs.apollo.dev/federation/v2.0", import: ["@shareable"])\ntype Query {\n  position: Position\n}\n\ntype Position  @shareable {\n  x: Int!\n  y: Int! @shareable\n}'}}

examples/tag.py

@@ -0,0 +1,29 @@
+import graphene
+
+from graphene_federation import build_schema, key, inaccessible, shareable
+from graphene_federation.tag import tag
+
+
+class Product(graphene.ObjectType):
+    id = graphene.ID(required=True)
+    in_stock = tag(graphene.Boolean(required=True), "Products")
+    out_stock = shareable(graphene.Boolean(required=True))
+    is_listed = inaccessible(graphene.Boolean(required=True))
+
+
+class Query(graphene.ObjectType):
+    position = graphene.Field(Product)
+
+
+schema = build_schema(Query, enable_federation_2=True)
+
+query = '''
+    query getSDL {
+      _service {
+         sdl
+      }
+    }
+'''
+result = schema.execute(query)
+print(result.data)
+# {'_service': {'sdl': 'extend schema @link(url: "https://specs.apollo.dev/federation/v2.0", import: ["@inaccessible", "@shareable", "@tag"])\ntype Query {\n  position: Product\n}\n\ntype Product  {\n  id: ID!\n  inStock: Boolean! @tag(name: "Products")\n  outStock: Boolean! @shareable\n  isListed: Boolean! @inaccessible\n}'}}

graphene_federation/__init__.py

@@ -1,4 +1,9 @@
 from .main import build_schema
 from .entity import key
-from .extend import extend, external, requires
+from .extend import extend
+from .external import external
+from .requires import requires
+from .shareable import shareable
+from .inaccessible import inaccessible
 from .provides import provides
+from .override import override

graphene_federation/entity.py

@@ -8,7 +8,22 @@
 from graphene.types.schema import TypeMap
 
 from .types import _Any
-from .utils import field_name_to_type_attribute
+from .utils import (
+    field_name_to_type_attribute,
+    check_fields_exist_on_type,
+    is_valid_compound_key,
+)
+
+import collections.abc
+
+
+def update(d, u):
+    for k, v in u.items():
+        if isinstance(v, collections.abc.Mapping):
+            d[k] = update(d.get(k, {}), v)
+        else:
+            d[k] = v
+    return d
 
 
 def get_entities(schema: Schema) -> Dict[str, Any]:
@@ -24,6 +39,14 @@ def get_entities(schema: Schema) -> Dict[str, Any]:
             continue
         if getattr(type_.graphene_type, "_keys", None):
             entities[type_name] = type_.graphene_type
+
+            # Validation for compound keys
+            key_str = " ".join(type_.graphene_type._keys)
+            type_name = type_.graphene_type._meta.name
+            if "{" in key_str:  # checking for subselection to identify compound key
+                assert is_valid_compound_key(
+                    type_name, key_str, schema
+                ), f'Invalid compound key definition for type "{type_name}"'
     return entities
 
 
@@ -83,27 +106,31 @@ def resolve_entities(self, info, representations):
     return EntityQuery
 
 
-def key(fields: str) -> Callable:
+def key(fields: str, resolvable: bool = True) -> Callable:
     """
     Take as input a field that should be used as key for that entity.
     See specification: https://www.apollographql.com/docs/federation/federation-spec/#key
-
-    If the input contains a space it means it's a [compound primary key](https://www.apollographql.com/docs/federation/entities/#defining-a-compound-primary-key)
-    which is not yet supported.
     """
-    if " " in fields:
-        raise NotImplementedError("Compound primary keys are not supported.")
 
-    def decorator(Type):
+    def decorator(type_):
         # Check the provided fields actually exist on the Type.
-        assert (
-            fields in Type._meta.fields
-        ), f'Field "{fields}" does not exist on type "{Type._meta.name}"'
-
-        keys = getattr(Type, "_keys", [])
+        if " " not in fields:
+            assert (
+                fields in type_._meta.fields
+            ), f'Field "{fields}" does not exist on type "{type_._meta.name}"'
+        if "{" not in fields:
+            # Skip valid fields check if the key is a compound key. The validation for compound keys
+            # is done on calling get_entities()
+            fields_set = set(fields.replace(" ", "").split(","))
+            assert check_fields_exist_on_type(
+                fields=fields_set, type_=type_
+            ), f'Field "{fields}" does not exist on type "{type_._meta.name}"'
+
+        keys = getattr(type_, "_keys", [])
         keys.append(fields)
-        setattr(Type, "_keys", keys)
+        setattr(type_, "_keys", keys)
+        setattr(type_, "_resolvable", resolvable)
 
-        return Type
+        return type_
 
     return decorator