fix: Federation V2 Scalar Support (#14)

Co-authored-by: Patrick Arminio <patrick.arminio@gmail.com>
Co-authored-by: Arun Suresh Kumar <arun@strollby.com>
Co-authored-by: Adarsh Divakaran <adarshdevamritham@gmail.com>
Co-authored-by: Arun Suresh Kumar <89654966+arun-sureshkumar@users.noreply.github.com>

Author: 4 people

examples/entities.py

@@ -6,10 +6,18 @@ def get_file_by_id(id):
     return File(**{'id': id, 'name': 'test_name'})
 
 
+class Author(graphene.ObjectType):
+    id = graphene.ID(required=True)
+    name = graphene.String(required=True)
+
+
 @key(fields='id')
+@key(fields='id author { name }')
+@key(fields='id author { id name }')
 class File(graphene.ObjectType):
     id = graphene.Int(required=True)
     name = graphene.String()
+    author = graphene.Field(Author, required=True)
 
     def resolve_id(self, info, **kwargs):
         return 1
@@ -25,7 +33,7 @@ class Query(graphene.ObjectType):
     file = graphene.Field(File)
 
     def resolve_file(self, **kwargs):
-        return None   # no direct access
+        return None  # no direct access
 
 
 schema = build_schema(Query)
@@ -41,7 +49,7 @@ def resolve_file(self, **kwargs):
 print(result.data)
 # {'_service': {'sdl': 'type Query {\n  file: File\n}\n\ntype File @key(fields: "id") {\n  id: Int!\n  name: String\n}'}}
 
-query ='''
+query = '''
     query entities($_representations: [_Any!]!) {
       _entities(representations: $_representations) {
         ... on File {
@@ -55,10 +63,10 @@ def resolve_file(self, **kwargs):
 
 result = schema.execute(query, variables={
     "_representations": [
-      {
-        "__typename": "File",
-        "id": 1
-      }
+        {
+            "__typename": "File",
+            "id": 1
+        }
     ]
 })
 print(result.data)

graphene_federation/entity.py

@@ -121,7 +121,7 @@ def decorator(type_):
         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(","))
+            fields_set = set(fields.split(" "))
             assert check_fields_exist_on_type(
                 fields=fields_set, type_=type_
             ), f'Field "{fields}" does not exist on type "{type_._meta.name}"'

graphene_federation/extend.py

@@ -43,7 +43,7 @@ def decorator(type_):
         if "{" not in fields:  # Check for compound keys
             # Skip valid fields check if the key is a compound key. The validation for compound keys
             # is done on calling get_extended_types()
-            fields_set = set(fields.replace(" ", "").split(","))
+            fields_set = set(fields.split(" "))
             assert check_fields_exist_on_type(
                 fields=fields_set, type_=type_
             ), f'Field "{fields}" does not exist on type "{type_._meta.name}"'

graphene_federation/tests/test_scalar.py

@@ -0,0 +1,57 @@
+from typing import Any
+
+import graphene
+from graphene import Scalar, String, ObjectType
+from graphql import graphql_sync
+
+from graphene_federation import build_schema, shareable, inaccessible
+
+
+def test_custom_scalar():
+    class AddressScalar(Scalar):
+        base = String
+
+        @staticmethod
+        def coerce_address(value: Any):
+            ...
+
+        serialize = coerce_address
+        parse_value = coerce_address
+
+        @staticmethod
+        def parse_literal(ast):
+            ...
+
+    @shareable
+    class TestScalar(graphene.ObjectType):
+        test_shareable_scalar = shareable(String(x=AddressScalar()))
+        test_inaccessible_scalar = inaccessible(String(x=AddressScalar()))
+
+    class Query(ObjectType):
+        test = String(x=AddressScalar())
+        test2 = graphene.List(AddressScalar, required=True)
+
+    schema = build_schema(query=Query, enable_federation_2=True, types=(TestScalar,))
+    query = """
+        query {
+            _service {
+                sdl
+            }
+        }
+        """
+    result = graphql_sync(schema.graphql_schema, query)
+    assert (
+        result.data["_service"]["sdl"].strip()
+        == """extend schema @link(url: "https://specs.apollo.dev/federation/v2.0", import: ["@inaccessible", "@shareable"])
+type TestScalar  @shareable {
+  testShareableScalar(x: AddressScalar): String @shareable
+  testInaccessibleScalar(x: AddressScalar): String @inaccessible
+}
+
+scalar AddressScalar
+
+type Query {
+  test(x: AddressScalar): String
+  test2: [AddressScalar]!
+}"""
+    )

graphene_federation/utils.py

@@ -3,9 +3,10 @@
 import graphene
 from graphene import Schema, ObjectType
 from graphene.types.definitions import GrapheneObjectType
+from graphene.types.scalars import ScalarOptions
 from graphene.types.union import UnionOptions
 from graphene.utils.str_converters import to_camel_case
-from graphql import parse, GraphQLScalarType
+from graphql import parse, GraphQLScalarType, GraphQLNonNull
 
 
 def field_name_to_type_attribute(schema: Schema, model: Any) -> Callable[[str], str]:
@@ -47,9 +48,11 @@ def is_valid_compound_key(type_name: str, key: str, schema: Schema):
 
     while key_nodes:
         selection_node, parent_object_type = key_nodes[0]
-
-        for field in selection_node.selection_set.selections:
+        if isinstance(parent_object_type, GraphQLNonNull):
+            parent_type_fields = parent_object_type.of_type.fields
+        else:
             parent_type_fields = parent_object_type.fields
+        for field in selection_node.selection_set.selections:
             if schema.auto_camelcase:
                 field_name = to_camel_case(field.name.value)
             else:
@@ -62,14 +65,20 @@ def is_valid_compound_key(type_name: str, key: str, schema: Schema):
             if field.selection_set:
                 # If the field has sub-selections, add it to node mappings to check for valid subfields
 
-                if isinstance(field_type, GraphQLScalarType):
+                if isinstance(field_type, GraphQLScalarType) or (
+                    isinstance(field_type, GraphQLNonNull)
+                    and isinstance(field_type.of_type, GraphQLScalarType)
+                ):
                     # sub-selections are added to a scalar type, key is not valid
                     return False
 
                 key_nodes.append((field, field_type))
             else:
                 # If there are no sub-selections for a field, it should be a scalar
-                if not isinstance(field_type, GraphQLScalarType):
+                if not isinstance(field_type, GraphQLScalarType) and not (
+                    isinstance(field_type, GraphQLNonNull)
+                    and isinstance(field_type.of_type, GraphQLScalarType)
+                ):
                     return False
 
         key_nodes.pop(0)  # Remove the current node as it is fully processed
@@ -80,8 +89,10 @@ def is_valid_compound_key(type_name: str, key: str, schema: Schema):
 def get_attributed_fields(attribute: str, schema: Schema):
     fields = {}
     for type_name, type_ in schema.graphql_schema.type_map.items():
-        if not hasattr(type_, "graphene_type") or isinstance(
-            type_.graphene_type._meta, UnionOptions
+        if (
+            not hasattr(type_, "graphene_type")
+            or isinstance(type_.graphene_type._meta, UnionOptions)
+            or isinstance(type_.graphene_type._meta, ScalarOptions)
         ):
             continue
         for field in list(type_.graphene_type._meta.fields):