From 129ab24a1f66f0a5b9bb4678274fbfee4e416204 Mon Sep 17 00:00:00 2001
From: Zach Daniel <zachary.s.daniel@gmail.com>
Date: Mon, 13 Jul 2020 00:37:02 -0400
Subject: [PATCH] fix: changeset + set_state issues

---
 lib/ash/changeset/changeset.ex | 15 +++++++++++++--
 lib/ash/dsl/extension.ex       | 29 +++++++++++++++--------------
 lib/ash/dsl/transformer.ex     |  8 ++++++++
 3 files changed, 36 insertions(+), 16 deletions(-)

diff --git a/lib/ash/changeset/changeset.ex b/lib/ash/changeset/changeset.ex
index 2218d4f5..04e9f45c 100644
--- a/lib/ash/changeset/changeset.ex
+++ b/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
diff --git a/lib/ash/dsl/extension.ex b/lib/ash/dsl/extension.ex
index ea32b975..a25006f3 100644
--- a/lib/ash/dsl/extension.ex
+++ b/lib/ash/dsl/extension.ex
@@ -217,16 +217,21 @@ defmodule Ash.Dsl.Extension do
         if runtime? do
           @ash_dsl_config
         else
-          {__MODULE__, :ash_sections}
-          |> Process.get([])
-          |> Enum.map(fn {extension, section_path} ->
-            {{section_path, extension},
-             Process.get(
-               {__MODULE__, :ash, section_path},
-               []
-             )}
-          end)
-          |> Enum.into(%{})
+          config =
+            {__MODULE__, :ash_sections}
+            |> Process.get([])
+            |> Enum.map(fn {extension, section_path} ->
+              {{section_path, extension},
+               Process.get(
+                 {__MODULE__, :ash, section_path},
+                 []
+               )}
+            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
 
diff --git a/lib/ash/dsl/transformer.ex b/lib/ash/dsl/transformer.ex
index 3576d219..4a9d88dd 100644
--- a/lib/ash/dsl/transformer.ex
+++ b/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
   """