From dc70980f2ef03d68603fb6660f2209bcfffdd085 Mon Sep 17 00:00:00 2001
From: Kaden Wilkinson <wilkinson.kaden@gmail.com>
Date: Thu, 15 Jul 2021 17:08:18 -0600
Subject: [PATCH] Don't add _Entity/_entities when no @key types found

---
 .../federation/schema/entities_field.ex       | 62 +++++++++--------
 .../federation/schema/entity_union.ex         | 59 +++++-----------
 lib/absinthe/federation/schema/phase.ex       | 43 +++++++++---
 lib/absinthe/federation/schema/utils.ex       | 40 +++++++++++
 .../federation/schema/entities_field_test.exs | 68 +++++++++++++++----
 .../federation/schema/entity_union_test.exs   | 41 +++++++++--
 6 files changed, 214 insertions(+), 99 deletions(-)
 create mode 100644 lib/absinthe/federation/schema/utils.ex

diff --git a/lib/absinthe/federation/schema/entities_field.ex b/lib/absinthe/federation/schema/entities_field.ex
index 1f25026..e18cfba 100644
--- a/lib/absinthe/federation/schema/entities_field.ex
+++ b/lib/absinthe/federation/schema/entities_field.ex
@@ -9,6 +9,8 @@ defmodule Absinthe.Federation.Schema.EntitiesField do
   alias Absinthe.Blueprint.TypeReference.NonNull
   alias Absinthe.Schema.Notation
 
+  alias Absinthe.Federation.Schema.Utils
+
   # TODO: Fix __reference__ typespec upstream in absinthe
   @type input_value_definition :: %InputValueDefinition{
           name: String.t(),
@@ -51,34 +53,40 @@ defmodule Absinthe.Federation.Schema.EntitiesField do
           __private__: []
         }
 
-  @spec build() :: field_definition()
-  def build() do
-    %FieldDefinition{
-      __reference__: Notation.build_reference(__ENV__),
-      description: """
-      Returns a non-nullable list of _Entity types
-      and have a single argument with an argument name of representations
-      and type [_Any!]! (non-nullable list of non-nullable _Any scalars).
-      The _entities field on the query root must allow a list of _Any scalars
-      which are "representations" of entities from external services.
-      These representations should be validated with the following rules:
-
-      - Any representation without a __typename: String field is invalid.
-      - Representations must contain at least the fields defined in the fieldset of a @key directive on the base type.
-      """,
-      identifier: :_entities,
-      module: __MODULE__,
-      name: "_entities",
-      type: %NonNull{
-        of_type: %ListType{
-          of_type: %Name{
-            name: "_Entity"
-          }
+  @spec build(Blueprint.t()) :: field_definition() | nil
+  def build(blueprint) do
+    case Utils.key_field_types(blueprint) do
+      [] ->
+        nil
+
+      _found_types ->
+        %FieldDefinition{
+          __reference__: Notation.build_reference(__ENV__),
+          description: """
+          Returns a non-nullable list of _Entity types
+          and have a single argument with an argument name of representations
+          and type [_Any!]! (non-nullable list of non-nullable _Any scalars).
+          The _entities field on the query root must allow a list of _Any scalars
+          which are "representations" of entities from external services.
+          These representations should be validated with the following rules:
+
+          - Any representation without a __typename: String field is invalid.
+          - Representations must contain at least the fields defined in the fieldset of a @key directive on the base type.
+          """,
+          identifier: :_entities,
+          module: __MODULE__,
+          name: "_entities",
+          type: %NonNull{
+            of_type: %ListType{
+              of_type: %Name{
+                name: "_Entity"
+              }
+            }
+          },
+          middleware: [{Absinthe.Resolution, &__MODULE__.resolver/3}],
+          arguments: build_arguments()
         }
-      },
-      middleware: [{Absinthe.Resolution, &__MODULE__.resolver/3}],
-      arguments: build_arguments()
-    }
+    end
   end
 
   def resolver(parent, %{representations: representations}, resolution) do
diff --git a/lib/absinthe/federation/schema/entity_union.ex b/lib/absinthe/federation/schema/entity_union.ex
index 423ad2d..9590408 100644
--- a/lib/absinthe/federation/schema/entity_union.ex
+++ b/lib/absinthe/federation/schema/entity_union.ex
@@ -3,20 +3,27 @@ defmodule Absinthe.Federation.Schema.EntityUnion do
 
   alias Absinthe.Blueprint
   alias Absinthe.Blueprint.Schema.UnionTypeDefinition
-  alias Absinthe.Blueprint.TypeReference.Name
   alias Absinthe.Schema.Notation
-  alias Absinthe.Type
 
+  alias Absinthe.Federation.Schema.Utils
+
+  @spec build(Blueprint.t()) :: UnionTypeDefinition.t() | nil
   def build(blueprint) do
-    %UnionTypeDefinition{
-      __reference__: Notation.build_reference(__ENV__),
-      description: "a union of all types that use the @key directive",
-      identifier: :_entity,
-      module: __MODULE__,
-      name: "_Entity",
-      types: types(blueprint),
-      resolve_type: &Absinthe.Federation.Schema.EntityUnion.resolve_type/2
-    }
+    case Utils.key_field_types(blueprint) do
+      [] ->
+        nil
+
+      found_types ->
+        %UnionTypeDefinition{
+          __reference__: Notation.build_reference(__ENV__),
+          description: "a union of all types that use the @key directive",
+          identifier: :_entity,
+          module: __MODULE__,
+          name: "_Entity",
+          types: found_types,
+          resolve_type: &Absinthe.Federation.Schema.EntityUnion.resolve_type/2
+        }
+    end
   end
 
   # TODO: This is a very naive approach to resolve the union type and should be replaced by something better
@@ -40,34 +47,4 @@ defmodule Absinthe.Federation.Schema.EntityUnion do
     |> Macro.underscore()
     |> String.to_existing_atom()
   end
-
-  defp types(node) do
-    {_node, types} = Blueprint.postwalk(node, [], &collect_types/2)
-
-    types
-  end
-
-  defp collect_types(
-         %{name: name, __private__: _private} = node,
-         types
-       ) do
-    if has_key_directive?(node) do
-      {node, [%Name{name: name} | types]}
-    else
-      {node, types}
-    end
-  end
-
-  defp collect_types(node, acc), do: {node, acc}
-
-  defp has_key_directive?(node) do
-    meta = Type.meta(node)
-    has_meta_key = Map.has_key?(meta, :key_fields)
-    node_directives = Map.get(node, :directives, [])
-    has_key_directive = Enum.any?(node_directives, &is_key_directive?/1)
-    has_meta_key or has_key_directive
-  end
-
-  defp is_key_directive?(%{name: "key"} = _directive), do: true
-  defp is_key_directive?(_directive), do: false
 end
diff --git a/lib/absinthe/federation/schema/phase.ex b/lib/absinthe/federation/schema/phase.ex
index 2656d2a..7952c8a 100644
--- a/lib/absinthe/federation/schema/phase.ex
+++ b/lib/absinthe/federation/schema/phase.ex
@@ -14,29 +14,50 @@ defmodule Absinthe.Federation.Schema.Phase do
   @dialyzer {:nowarn_function, add_directive: 2}
 
   def run(%Blueprint{} = blueprint, _) do
-    blueprint = Blueprint.postwalk(blueprint, &collect_types/1)
+    blueprint = Blueprint.postwalk(blueprint, &collect_types(&1, blueprint))
     {:ok, blueprint}
   end
 
-  @spec collect_types(Blueprint.node_t()) :: Blueprint.node_t()
-  defp collect_types(%Schema.SchemaDefinition{type_definitions: type_definitions} = node) do
-    entity_union = EntityUnion.build(node)
+  @spec collect_types(Blueprint.node_t(), Blueprint.t()) :: Blueprint.node_t()
+  defp collect_types(%Schema.SchemaDefinition{type_definitions: type_definitions} = node, blueprint) do
+    case EntityUnion.build(blueprint) do
+      nil ->
+        node
 
-    %{node | type_definitions: [entity_union | type_definitions]}
+      entity_union ->
+        %{node | type_definitions: [entity_union | type_definitions]}
+    end
   end
 
-  defp collect_types(%Schema.ObjectTypeDefinition{identifier: :query, fields: fields} = node) do
-    service_field = ServiceField.build()
-    entities_field = EntitiesField.build()
-    %{node | fields: [service_field, entities_field] ++ fields}
+  defp collect_types(%Schema.ObjectTypeDefinition{identifier: :query, fields: fields} = node, blueprint) do
+    new_fields =
+      fields
+      |> add_service_field()
+      |> maybe_add_entities_field(blueprint)
+
+    %{node | fields: new_fields}
   end
 
-  defp collect_types(%{__private__: _private} = node) do
+  defp collect_types(%{__private__: _private} = node, _) do
     meta = Type.meta(node)
     maybe_add_directives(node, meta)
   end
 
-  defp collect_types(node), do: node
+  defp collect_types(node, _), do: node
+
+  defp add_service_field(fields) do
+    [ServiceField.build() | fields]
+  end
+
+  defp maybe_add_entities_field(fields, node) do
+    case EntitiesField.build(node) do
+      nil ->
+        fields
+
+      entities_field ->
+        [entities_field | fields]
+    end
+  end
 
   @spec maybe_add_directives(term(), any()) :: term()
   defp maybe_add_directives(node, meta) do
diff --git a/lib/absinthe/federation/schema/utils.ex b/lib/absinthe/federation/schema/utils.ex
new file mode 100644
index 0000000..91c21fc
--- /dev/null
+++ b/lib/absinthe/federation/schema/utils.ex
@@ -0,0 +1,40 @@
+defmodule Absinthe.Federation.Schema.Utils do
+  @moduledoc false
+
+  alias Absinthe.Blueprint
+  alias Absinthe.Blueprint.TypeReference.Name
+  alias Absinthe.Type
+
+  @spec key_field_types(Blueprint.t()) :: list(Blueprint.node_t())
+  def key_field_types(blueprint) do
+    {_blueprint, types} = Blueprint.postwalk(blueprint, [], &collect_key_field_types/2)
+
+    types
+  end
+
+  defp collect_key_field_types(
+         %{name: name, __private__: _private} = node,
+         types
+       ) do
+    if has_key_directive?(node) do
+      {node, [%Name{name: name} | types]}
+    else
+      {node, types}
+    end
+  end
+
+  defp collect_key_field_types(node, acc), do: {node, acc}
+
+  @spec has_key_directive?(struct()) :: boolean()
+  def has_key_directive?(node) do
+    meta = Type.meta(node)
+    has_meta_key = Map.has_key?(meta, :key_fields)
+    node_directives = Map.get(node, :directives, [])
+    has_key_directive = Enum.any?(node_directives, &is_key_directive?/1)
+    has_meta_key or has_key_directive
+  end
+
+  @spec is_key_directive?(Blueprint.Directive.t()) :: boolean()
+  def is_key_directive?(%{name: "key"} = _directive), do: true
+  def is_key_directive?(_directive), do: false
+end
diff --git a/test/absinthe/federation/schema/entities_field_test.exs b/test/absinthe/federation/schema/entities_field_test.exs
index 25e109e..8395630 100644
--- a/test/absinthe/federation/schema/entities_field_test.exs
+++ b/test/absinthe/federation/schema/entities_field_test.exs
@@ -8,23 +8,40 @@ defmodule Absinthe.Federation.Schema.EntitiesFieldTest do
 
   alias Absinthe.Federation.Schema.EntitiesField
 
+  defmodule EntitiesSchema do
+    use Absinthe.Schema
+    use Absinthe.Federation.Schema
+
+    query do
+    end
+
+    object :foo do
+      key_fields("id")
+      field :id, :id
+    end
+  end
+
+  setup do
+    {:ok, blueprint: EntitiesSchema.__absinthe_blueprint__()}
+  end
+
   describe "build" do
-    test "builds field definition" do
-      assert %FieldDefinition{} = EntitiesField.build()
+    test "builds field definition", %{blueprint: blueprint} do
+      assert %FieldDefinition{} = EntitiesField.build(blueprint)
     end
 
-    test "builds field definition with name" do
-      field_definition = EntitiesField.build()
+    test "builds field definition with name", %{blueprint: blueprint} do
+      field_definition = EntitiesField.build(blueprint)
       assert field_definition.name == "_entities"
     end
 
-    test "builds field definition with identifier" do
-      field_definition = EntitiesField.build()
+    test "builds field definition with identifier", %{blueprint: blueprint} do
+      field_definition = EntitiesField.build(blueprint)
       assert field_definition.identifier == :_entities
     end
 
-    test "builds field definition with type" do
-      field_definition = EntitiesField.build()
+    test "builds field definition with type", %{blueprint: blueprint} do
+      field_definition = EntitiesField.build(blueprint)
 
       assert %NonNull{
                of_type: %List{
@@ -35,8 +52,8 @@ defmodule Absinthe.Federation.Schema.EntitiesFieldTest do
              } = field_definition.type
     end
 
-    test "builds field definition with middleware" do
-      field_definition = EntitiesField.build()
+    test "builds field definition with middleware", %{blueprint: blueprint} do
+      field_definition = EntitiesField.build(blueprint)
       assert Enum.count(field_definition.middleware) == 1
     end
   end
@@ -73,18 +90,41 @@ defmodule Absinthe.Federation.Schema.EntitiesFieldTest do
   end
 
   describe "sdl" do
-    defmodule SDLSchema do
+    defmodule SDLWithKeyFieldsSchema do
       use Absinthe.Schema
       use Absinthe.Federation.Schema
 
       query do
-        field :test, :string
+        field :user, :user
+      end
+
+      object :user do
+        key_fields("id")
+        field :id, non_null(:id)
       end
     end
 
-    test "renders correctly in sdl" do
-      sdl = Absinthe.Schema.to_sdl(SDLSchema)
+    test "renders correctly in sdl with @key" do
+      sdl = Absinthe.Schema.to_sdl(SDLWithKeyFieldsSchema)
       assert sdl =~ "_entities(representations: [_Any!]!): [_Entity]!"
     end
+
+    defmodule SDLWithoutKeyFieldsSchema do
+      use Absinthe.Schema
+      use Absinthe.Federation.Schema
+
+      query do
+        field :test, :string
+      end
+
+      object :user do
+        field :id, non_null(:id)
+      end
+    end
+
+    test "does not render in sdl without @key" do
+      sdl = Absinthe.Schema.to_sdl(SDLWithoutKeyFieldsSchema)
+      refute sdl =~ "_entities(representations: [_Any!]!): [_Entity]!"
+    end
   end
 end
diff --git a/test/absinthe/federation/schema/entity_union_test.exs b/test/absinthe/federation/schema/entity_union_test.exs
index 45524d0..0c6ee11 100644
--- a/test/absinthe/federation/schema/entity_union_test.exs
+++ b/test/absinthe/federation/schema/entity_union_test.exs
@@ -8,8 +8,15 @@ defmodule Absinthe.Federation.Schema.EntityUnionTest do
   describe "build" do
     defmodule EntityUnionSchema do
       use Absinthe.Schema
+      use Absinthe.Federation.Schema
 
       query do
+        field :user, :user
+      end
+
+      object :user do
+        key_fields("id")
+        field :id, non_null(:id)
       end
     end
 
@@ -35,7 +42,9 @@ defmodule Absinthe.Federation.Schema.EntityUnionTest do
     test "builds field definition with type", %{blueprint: blueprint} do
       union_type_definition = EntityUnion.build(blueprint)
 
-      assert union_type_definition.types == []
+      assert union_type_definition.types == [
+               %Absinthe.Blueprint.TypeReference.Name{errors: [], name: "User", schema_node: nil}
+             ]
     end
   end
 
@@ -87,7 +96,7 @@ defmodule Absinthe.Federation.Schema.EntityUnionTest do
   end
 
   describe "sdl" do
-    defmodule SDLSchema do
+    defmodule SDLWithKeyFieldsSchema do
       use Absinthe.Schema
       use Absinthe.Federation.Schema
 
@@ -103,11 +112,13 @@ defmodule Absinthe.Federation.Schema.EntityUnionTest do
     end
 
     test "renders union correctly in sdl based schema" do
-      sdl = Absinthe.Schema.to_sdl(SDLSchema)
+      sdl = Absinthe.Schema.to_sdl(SDLWithKeyFieldsSchema)
       assert sdl =~ "union _Entity = User"
     end
+  end
 
-    defmodule MacroSchema do
+  describe "macro" do
+    defmodule MacroWithKeyFieldsSchema do
       use Absinthe.Schema
       use Absinthe.Federation.Schema
 
@@ -121,9 +132,27 @@ defmodule Absinthe.Federation.Schema.EntityUnionTest do
       end
     end
 
-    test "renders union correctly in macro based schema" do
-      sdl = Absinthe.Schema.to_sdl(MacroSchema)
+    test "renders union correctly in macro based schema with @key types" do
+      sdl = Absinthe.Schema.to_sdl(MacroWithKeyFieldsSchema)
       assert sdl =~ "union _Entity = User"
     end
+
+    defmodule MacroWithoutKeyFields do
+      use Absinthe.Schema
+      use Absinthe.Federation.Schema
+
+      query do
+        field :me, :user
+      end
+
+      object :user do
+        field :id, non_null(:id)
+      end
+    end
+
+    test "renders no union in schema without @key types" do
+      sdl = Absinthe.Schema.to_sdl(MacroWithoutKeyFields)
+      refute sdl =~ "union _Entity = User"
+    end
   end
 end