fix: changeset + set_state issues

lib/ash/changeset/changeset.ex

@@ -518,8 +518,19 @@ defmodule Ash.Changeset do
   defp do_primary_key(relationship, record) when is_map(record) do
     primary_key = Ash.primary_key(relationship.destination)
 
-    if Enum.all?(primary_key, &Map.has_key?(record, &1)) do
-      pkey = Map.take(record, Ash.primary_key(relationship.destination))
+    is_pkey_map? =
+      Enum.all?(primary_key, fn key ->
+        Map.has_key?(record, key) || Map.has_key?(record, to_string(key))
+      end)
+
+    if is_pkey_map? do
+      pkey =
+        Enum.reduce(primary_key, %{}, fn key, acc ->
+          case Map.fetch(record, key) do
+            {:ok, value} -> Map.put(acc, key, value)
+            :error -> Map.put(acc, key, Map.get(record, to_string(key)))
+          end
+        end)
 
       {:ok, pkey}
     else

lib/ash/dsl/extension.ex

@@ -217,6 +217,7 @@ defmodule Ash.Dsl.Extension do
         if runtime? do
           @ash_dsl_config
         else
+          config =
             {__MODULE__, :ash_sections}
             |> Process.get([])
             |> Enum.map(fn {extension, section_path} ->
@@ -227,6 +228,10 @@ defmodule Ash.Dsl.Extension do
                )}
             end)
             |> Enum.into(%{})
+
+          Module.put_attribute(__MODULE__, :ash_dsl_config, config)
+
+          config
         end
 
       :persistent_term.put({__MODULE__, :extensions}, @extensions)
@@ -245,10 +250,6 @@ defmodule Ash.Dsl.Extension do
           transformers_to_run,
           ash_dsl_config
         )
-
-      unless runtime? do
-        Module.put_attribute(__MODULE__, :ash_dsl_config, new_dsl_config)
-      end
     end
   end
 

lib/ash/dsl/transformer.ex

@@ -86,6 +86,14 @@ defmodule Ash.Dsl.Transformer do
     |> Map.get(:entities, [])
   end
 
+  def get_option(dsl_state, path, option, extension) do
+    dsl_state
+    |> Map.get({path, extension}, %{opts: []})
+    |> Map.get(:opts)
+    |> Kernel.||([])
+    |> Keyword.get(option)
+  end
+
   @doc """
   Store a value in a special persistent term key, that will be copied to the runtime
   """