fix: various problems with resolving unions in relationship changes

lib/ash_graphql.ex

@@ -250,9 +250,18 @@ defmodule AshGraphql do
   end
 
   @doc false
-  def all_attributes_and_arguments(resource, already_checked \\ [], nested? \\ true) do
+  def all_attributes_and_arguments(
+        resource,
+        already_checked \\ [],
+        nested? \\ true,
+        return_new_checked? \\ false
+      ) do
     if resource in already_checked do
-      []
+      if return_new_checked? do
+        {[], already_checked}
+      else
+        []
+      end
     else
       already_checked = [resource | already_checked]
 
@@ -260,22 +269,64 @@ defmodule AshGraphql do
       |> Ash.Resource.Info.public_attributes()
       |> Enum.concat(all_arguments(resource))
       |> Enum.concat(Ash.Resource.Info.calculations(resource))
-      |> Enum.flat_map(fn %{type: type} = attr ->
+      |> Enum.reduce({[], already_checked}, fn %{type: type} = attr, {acc, already_checked} ->
-        if Ash.Type.embedded_type?(type) && nested? do
+        if nested? do
-          type = Ash.Type.NewType.subtype_of(type)
+          constraints = Map.get(attr, :constraints, [])
-
+          {nested, already_checked} = nested_attrs(type, constraints, already_checked)
-          [
+          {[attr | nested] ++ acc, already_checked}
-            attr
-            | type
-              |> unwrap_type()
-              |> all_attributes_and_arguments(already_checked, nested?)
-          ]
         else
-          [attr]
+          {[attr | acc], already_checked}
+        end
+      end)
+      |> then(fn {attrs, checked} ->
+        attrs = Enum.filter(attrs, &AshGraphql.Resource.Info.show_field?(resource, &1.name))
+
+        if return_new_checked? do
+          {attrs, checked}
+        else
+          attrs
         end
       end)
     end
-    |> Enum.filter(&AshGraphql.Resource.Info.show_field?(resource, &1.name))
+  end
+
+  defp nested_attrs(Ash.Type.Union, constraints, already_checked) do
+    Enum.reduce(
+      constraints[:types] || [],
+      {[], already_checked},
+      fn {_, config}, {attrs, already_checked} ->
+        case config[:type] do
+          {:array, type} ->
+            {new, already_checked} =
+              nested_attrs(type, config[:constraints][:items] || [], already_checked)
+
+            {attrs ++ new, already_checked}
+
+          type ->
+            {new, already_checked} =
+              nested_attrs(type, config[:constraints] || [], already_checked)
+
+            {attrs ++ new, already_checked}
+        end
+      end
+    )
+  end
+
+  defp nested_attrs(type, constraints, already_checked) do
+    cond do
+      Ash.Type.embedded_type?(type) ->
+        type
+        |> unwrap_type()
+        |> all_attributes_and_arguments(already_checked, true, true)
+
+      Ash.Type.NewType.new_type?(type) ->
+        constraints = Ash.Type.NewType.constraints(type, constraints)
+        type = Ash.Type.NewType.subtype_of(type)
+        nested_attrs(type, constraints, already_checked)
+
+      true ->
+        {[], already_checked}
+    end
   end
 
   def get_embed(type) do

lib/graphql/resolver.ex

@@ -326,7 +326,7 @@ defmodule AshGraphql.Graphql.Resolver do
       if argument do
         %{type: type, name: name, constraints: constraints} = argument
 
-        case handle_argument(type, constraints, value, name) do
+        case handle_argument(resource, action, type, constraints, value, name) do
           {:ok, value} ->
             {:cont, {:ok, Map.put(arguments, name, value)}}
 
@@ -339,10 +339,11 @@ defmodule AshGraphql.Graphql.Resolver do
     end)
   end
 
-  defp handle_argument({:array, type}, constraints, value, name) when is_list(value) do
+  defp handle_argument(resource, action, {:array, type}, constraints, value, name)
+       when is_list(value) do
     value
     |> Enum.reduce_while({:ok, []}, fn value, {:ok, acc} ->
-      case handle_argument(type, constraints[:items], value, name) do
+      case handle_argument(resource, action, type, constraints[:items], value, name) do
         {:ok, value} ->
           {:cont, {:ok, [value | acc]}}
 
@@ -356,14 +357,99 @@ defmodule AshGraphql.Graphql.Resolver do
     end
   end
 
-  defp handle_argument(Ash.Type.Union, constraints, value, name) do
+  defp handle_argument(_resource, _action, Ash.Type.Union, constraints, value, name) do
     handle_union_type(value, constraints, name)
   end
 
-  defp handle_argument(type, constraints, value, name) do
+  defp handle_argument(resource, action, type, constraints, value, name) do
     cond do
+      AshGraphql.Resource.Info.managed_relationship(resource, action, %{name: name}) &&
+          is_map(value) ->
+        managed_relationship =
+          AshGraphql.Resource.Info.managed_relationship(resource, action, %{name: name})
+
+        opts = AshGraphql.Resource.find_manage_change(%{name: name}, action, resource)
+
+        relationship =
+          Ash.Resource.Info.relationship(resource, opts[:relationship]) ||
+            raise """
+            No relationship found when building managed relationship input: #{opts[:relationship]}
+            """
+
+        manage_opts_schema =
+          if opts[:opts][:type] do
+            defaults = Ash.Changeset.manage_relationship_opts(opts[:opts][:type])
+
+            Enum.reduce(defaults, Ash.Changeset.manage_relationship_schema(), fn {key, value},
+                                                                                 manage_opts ->
+              Spark.OptionsHelpers.set_default!(manage_opts, key, value)
+            end)
+          else
+            Ash.Changeset.manage_relationship_schema()
+          end
+
+        manage_opts = Spark.OptionsHelpers.validate!(opts[:opts], manage_opts_schema)
+
+        fields =
+          manage_opts
+          |> AshGraphql.Resource.manage_fields(
+            managed_relationship,
+            relationship,
+            __MODULE__
+          )
+          |> Enum.reject(fn
+            {_, :__primary_key, _} ->
+              true
+
+            {_, {:identity, _}, _} ->
+              true
+
+            _ ->
+              false
+          end)
+          |> Map.new(fn {_, _, %{identifier: identifier}} = field ->
+            {identifier, field}
+          end)
+
+        Enum.reduce_while(value, {:ok, %{}}, fn {key, value}, {:ok, acc} ->
+          field_name =
+            resource
+            |> AshGraphql.Resource.Info.field_names()
+            |> Enum.map(fn {l, r} -> {r, l} end)
+            |> Keyword.get(key, key)
+
+          case Map.get(fields, field_name) do
+            nil ->
+              {:cont, {:ok, Map.put(acc, key, value)}}
+
+            {resource, action, _} ->
+              action = Ash.Resource.Info.action(resource, action)
+              attributes = Ash.Resource.Info.public_attributes(resource)
+
+              argument =
+                Enum.find(action.arguments, &(&1.name == field_name)) ||
+                  Enum.find(attributes, &(&1.name == field_name))
+
+              if argument do
+                %{type: type, name: name, constraints: constraints} = argument
+
+                case handle_argument(resource, action, type, constraints, value, name) do
+                  {:ok, value} ->
+                    {:cont, {:ok, Map.put(acc, key, value)}}
+
+                  {:error, error} ->
+                    {:halt, {:error, error}}
+                end
+              else
+                {:cont, {:ok, Map.put(acc, key, value)}}
+              end
+          end
+        end)
+
       Ash.Type.NewType.new_type?(type) ->
         handle_argument(
+          resource,
+          action,
           Ash.Type.NewType.subtype_of(type),
           Ash.Type.NewType.constraints(type, constraints),
           value,
@@ -412,7 +498,14 @@ defmodule AshGraphql.Graphql.Resolver do
             end)
 
           if field do
-            case handle_argument(field.type, field.constraints, value, "#{name}.#{key}") do
+            case handle_argument(
+                   resource,
+                   action,
+                   field.type,
+                   field.constraints,
+                   value,
+                   "#{name}.#{key}"
+                 ) do
               {:ok, value} ->
                 {:cont, {:ok, Map.put(acc, key, value)}}
 

lib/resource/resource.ex

@@ -1007,7 +1007,8 @@ defmodule AshGraphql.Resource do
       end
   end
 
-  defp find_manage_change(argument, action, resource) do
+  @doc false
+  def find_manage_change(argument, action, resource) do
     if AshGraphql.Resource.Info.managed_relationship(resource, action, argument) do
       Enum.find_value(action.changes, fn
         %{change: {Ash.Resource.Change.ManageRelationship, opts}} ->
@@ -1447,11 +1448,7 @@ defmodule AshGraphql.Resource do
 
     manage_opts = Spark.OptionsHelpers.validate!(opts[:opts], manage_opts_schema)
 
-    fields =
+    fields = manage_fields(manage_opts, managed_relationship, relationship, schema)
-      on_match_fields(manage_opts, relationship, schema) ++
-        on_no_match_fields(manage_opts, relationship, schema) ++
-        on_lookup_fields(manage_opts, relationship, schema) ++
-        manage_pkey_fields(manage_opts, managed_relationship, relationship, schema)
 
     type = managed_relationship.type_name || default_managed_type_name(resource, action, argument)
 
@@ -1477,6 +1474,14 @@ defmodule AshGraphql.Resource do
     }
   end
 
+  @doc false
+  def manage_fields(manage_opts, managed_relationship, relationship, schema) do
+    on_match_fields(manage_opts, relationship, schema) ++
+      on_no_match_fields(manage_opts, relationship, schema) ++
+      on_lookup_fields(manage_opts, relationship, schema) ++
+      manage_pkey_fields(manage_opts, managed_relationship, relationship, schema)
+  end
+
   defp check_for_conflicts!(fields, managed_relationship, resource) do
     {ok, errors} =
       fields
@@ -3460,8 +3465,10 @@ defmodule AshGraphql.Resource do
          function_exported?(type, :graphql_unnested_unions, 1) do
       unnested_types = type.graphql_unnested_unions(constraints)
 
-      [{AshGraphql.Graphql.Resolver, :resolve_union},
+      [
-       {name, type, field, resource, unnested_types}]
+        {AshGraphql.Graphql.Resolver, :resolve_union},
+        {name, type, field, resource, unnested_types}
+      ]
     else
       []
     end