Refactor primary key clauses handling

CHANGELOG.md

@@ -21,6 +21,7 @@
 * [`PowPersistentSession.Plug.Cookie`] Changed default cookie name to `persistent_session`
 * [`PowPersistentSession.Plug.Cookie`] Removed renewal of cookie as the token will always expire
 * [`PowPersistentSession.Plug.Cookie`] No longer expires invalid cookies
+* [`Pow.Operations`] Added `Pow.Operations.fetch_primary_key_values/2`
 
 ## Removed
 

lib/extensions/email_confirmation/plug.ex

@@ -3,7 +3,7 @@ defmodule PowEmailConfirmation.Plug do
   Plug helper methods.
   """
   alias Plug.Conn
-  alias Pow.Plug
+  alias Pow.{Operations, Plug}
   alias PowEmailConfirmation.Ecto.Context
 
   @doc """
@@ -56,10 +56,18 @@ defmodule PowEmailConfirmation.Plug do
     end
   end
 
-  defp maybe_renew_conn(conn, %{id: user_id} = user, config) do
-    case Plug.current_user(conn, config) do
-      %{id: ^user_id} -> Plug.get_plug(config).do_create(conn, user, config)
-      _any            -> conn
+  defp maybe_renew_conn(conn, user, config) do
+    case equal_user?(user, Plug.current_user(conn, config), config) do
+      true  -> Plug.get_plug(config).do_create(conn, user, config)
+      false -> conn
     end
   end
+
+  defp equal_user?(_user, nil, _config), do: false
+  defp equal_user?(user, current_user, config) do
+    {:ok, clauses1} = Operations.fetch_primary_key_values(user, config)
+    {:ok, clauses2} = Operations.fetch_primary_key_values(current_user, config)
+
+    Keyword.equal?(clauses1, clauses2)
+  end
 end

lib/extensions/persistent_session/plug/cookie.ex

@@ -92,7 +92,7 @@ defmodule PowPersistentSession.Plug.Cookie do
     {store, store_config} = store(config)
     cookie_key            = cookie_key(config)
     key                   = cookie_id(config)
-    value                 = persistent_session_value(conn, user)
+    value                 = persistent_session_value(conn, user, config)
     opts                  = cookie_opts(config)
 
     store.put(store_config, key, value)
@@ -102,8 +102,8 @@ defmodule PowPersistentSession.Plug.Cookie do
     |> Conn.put_resp_cookie(cookie_key, key, opts)
   end
 
-  defp persistent_session_value(conn, user) do
-    clauses  = user_to_get_by_clauses(user)
+  defp persistent_session_value(conn, user, config) do
+    clauses  = user_to_get_by_clauses!(user, config)
     metadata =
       conn.private
       |> Map.get(:pow_persistent_session_metadata, [])
@@ -112,7 +112,12 @@ defmodule PowPersistentSession.Plug.Cookie do
     {clauses, metadata}
   end
 
-  defp user_to_get_by_clauses(%{id: id}), do: [id: id]
+  defp user_to_get_by_clauses!(user, config) do
+    case Operations.fetch_primary_key_values(user, config) do
+      {:ok, clauses}  -> clauses
+      {:error, error} -> raise error
+    end
+  end
 
   defp maybe_put_fingerprint_in_session_metadata(metadata, conn) do
     conn.private

lib/pow/ecto/context.ex

@@ -222,12 +222,20 @@ defmodule Pow.Ecto.Context do
   defp reload_after_write({:ok, struct}, config) do
     # When ecto updates/inserts, has_many :through associations are set to nil.
     # So we'll just reload when writes happen.
-    opts = repo_opts(config, [:prefix])
-    struct = Config.repo!(config).get!(struct.__struct__, struct.id, opts)
+    opts    = repo_opts(config, [:prefix])
+    clauses = fetch_primary_key_values!(struct, config)
+    struct  = Config.repo!(config).get_by!(struct.__struct__, clauses, opts)
 
     {:ok, struct}
   end
 
+  defp fetch_primary_key_values!(struct, config) do
+    case Operations.fetch_primary_key_values(struct, config) do
+      {:error, error} -> raise error
+      {:ok, clauses}  -> clauses
+    end
+  end
+
   # TODO: Remove by 1.1.0
   @deprecated "Use `Pow.Config.repo!/1` instead"
   defdelegate repo(config), to: Config, as: :repo!

lib/pow/operations.ex

@@ -104,4 +104,31 @@ defmodule Pow.Operations do
   defp context_module(config) do
     Config.get(config, :users_context, Context)
   end
+
+  @doc """
+  Retrieve a keyword list of primary key value(s) from the provided struct.
+
+  The keys will be fetched from the `__schema__/1` method in the struct module.
+  If no `__schema__/1` method exists, then it's expected that the struct has
+  `:id` as its only primary key.
+  """
+  @spec fetch_primary_key_values(struct(), Config.t()) :: {:ok, keyword()} | {:error, term()}
+  def fetch_primary_key_values(%mod{} = struct, _config) do
+    mod
+    |> function_exported?(:__schema__, 1)
+    |> case do
+      true  -> mod.__schema__(:primary_key)
+      false -> [:id]
+    end
+    |> map_primary_key_values(struct, [])
+  end
+
+  defp map_primary_key_values([], %mod{}, []), do: {:error, "No primary keys found for #{inspect mod}"}
+  defp map_primary_key_values([key | rest], %mod{} = struct, acc) do
+    case Map.get(struct, key) do
+      nil   -> {:error, "Primary key value for key `#{inspect key}` in #{inspect mod} can't be `nil`"}
+      value -> map_primary_key_values(rest, struct, acc ++ [{key, value}])
+    end
+  end
+  defp map_primary_key_values([], _struct, acc), do: {:ok, acc}
 end

lib/pow/store/credentials_cache.ex

@@ -12,7 +12,7 @@ defmodule Pow.Store.CredentialsCache do
     * `users/2` - to list all current users
     * `sessions/2` - to list all current sessions
   """
-  alias Pow.{Config, Store.Base}
+  alias Pow.{Config, Operations, Store.Base}
 
   use Base,
     ttl: :timer.minutes(30),
@@ -40,7 +40,7 @@ defmodule Pow.Store.CredentialsCache do
 
   # TODO: Refactor by 1.1.0
   defp fetch_sessions(config, backend_config, user) do
-    {struct, id} = user_to_struct_id(user)
+    {struct, id} = user_to_struct_id!(user, [])
 
     config
     |> Base.all(backend_config, [struct, :user, id, :session, :_])
@@ -66,7 +66,7 @@ defmodule Pow.Store.CredentialsCache do
   """
   @impl true
   def put(config, session_id, {user, metadata}) do
-    {struct, id} = user_to_struct_id(user)
+    {struct, id} = user_to_struct_id!(user, [])
     user_key     = [struct, :user, id]
     session_key  = [struct, :user, id, :session, session_id]
     records      = [
@@ -129,33 +129,27 @@ defmodule Pow.Store.CredentialsCache do
     end
   end
 
-  defp user_to_struct_id(%struct{} = user) do
-    key_value = case function_exported?(struct, :__schema__, 1) do
-      true  -> key_value_from_primary_keys(user)
-      false -> primary_keys_to_keyword_list!([:id], user)
-    end
-
-    {struct, key_value}
-  end
-  defp user_to_struct_id(_user), do: raise "Only structs can be stored as credentials"
-
-  defp key_value_from_primary_keys(%struct{} = user) do
-    :primary_key
-    |> struct.__schema__()
-    |> Enum.sort()
-    |> primary_keys_to_keyword_list!(user)
-  end
+  @spec user_to_struct_id!(any(), Config.t()) :: {atom(), keyword()} | no_return()
+  defp user_to_struct_id!(%mod{} = user, config) do
+    key_values =
+      user
+      |> fetch_primary_key_values!(config)
+      |> Enum.sort(&elem(&1, 0) < elem(&2, 0))
+      |> case do
+        [id: id] -> id
+        clauses  -> clauses
+      end
 
-  defp primary_keys_to_keyword_list!([], %struct{}), do: raise "No primary keys found for #{inspect struct}"
-  defp primary_keys_to_keyword_list!([key], user), do: get_primary_key_value!(user, key)
-  defp primary_keys_to_keyword_list!(keys, user) do
-    Enum.map(keys, &{&1, get_primary_key_value!(user, &1)})
+    {mod, key_values}
   end
-
-  defp get_primary_key_value!(%struct{} = user, key) do
-    case Map.get(user, key) do
-      nil -> raise "Primary key value for key `#{inspect key}` in #{inspect struct} can't be `nil`"
-      val -> val
+  defp user_to_struct_id!(_user, _config), do: raise "Only structs can be stored as credentials"
+
+  defp fetch_primary_key_values!(user, config) do
+    user
+    |> Operations.fetch_primary_key_values(config)
+    |> case do
+      {:error, error} -> raise error
+      {:ok, clauses}  -> clauses
     end
   end
 

test/extensions/email_confirmation/phoenix/controllers/confirmation_controller_test.exs

@@ -1,6 +1,8 @@
 defmodule PowEmailConfirmation.Phoenix.ConfirmationControllerTest do
   use PowEmailConfirmation.TestWeb.Phoenix.ConnCase
 
+  alias PowEmailConfirmation.Test.Users.User
+
   @session_key "auth"
 
   describe "show/2" do
@@ -51,7 +53,7 @@ defmodule PowEmailConfirmation.Phoenix.ConfirmationControllerTest do
       session_id = conn.private[:plug_session][@session_key]
       conn       =
         conn
-        |> Pow.Plug.assign_current_user(%{id: 1}, [])
+        |> Pow.Plug.assign_current_user(%User{id: 1}, [])
         |> get(Routes.pow_email_confirmation_confirmation_path(conn, :show, "valid"))
 
       assert redirected_to(conn) == Routes.pow_registration_path(conn, :edit)
@@ -63,7 +65,7 @@ defmodule PowEmailConfirmation.Phoenix.ConfirmationControllerTest do
       session_id = conn.private[:plug_session][@session_key]
       conn       =
         conn
-        |> Pow.Plug.assign_current_user(%{id: 2}, [])
+        |> Pow.Plug.assign_current_user(%User{id: 2}, [])
         |> get(Routes.pow_email_confirmation_confirmation_path(conn, :show, "valid"))
 
       assert redirected_to(conn) == Routes.pow_registration_path(conn, :edit)

test/extensions/persistent_session/plug/cookie_test.exs

@@ -220,6 +220,12 @@ defmodule PowPersistentSession.Plug.CookieTest do
     assert %{max_age: 1, path: "/"} = conn.resp_cookies[@cookie_key]
   end
 
+  test "create/3 handles clause error", %{conn: conn, config: config} do
+    assert_raise RuntimeError, "Primary key value for key `:id` in #{inspect User} can't be `nil`", fn ->
+      Cookie.create(conn, %User{id: nil}, config)
+    end
+  end
+
   test "create/3 with custom cookie options", %{conn: conn, config: config} do
     config = Keyword.put(config, :persistent_session_cookie_opts, @custom_cookie_opts)
     conn   = Cookie.create(conn, %User{id: 1}, config)

test/pow/operations_test.exs

@@ -1,4 +1,59 @@
 defmodule Pow.OperationsTest do
   use ExUnit.Case
   doctest Pow.Operations
+
+  alias Pow.Operations
+
+  defmodule PrimaryFieldUser do
+    use Ecto.Schema
+
+    schema "users" do
+      timestamps()
+    end
+  end
+
+  defmodule NoPrimaryFieldUser do
+    use Ecto.Schema
+
+    @primary_key false
+    schema "users" do
+      timestamps()
+    end
+  end
+
+  defmodule CompositePrimaryFieldsUser do
+    use Ecto.Schema
+
+    @primary_key false
+    schema "users" do
+      field :some_id, :integer, primary_key: true
+      field :another_id, :integer, primary_key: true
+
+      timestamps()
+    end
+  end
+
+  defmodule NonEctoUser do
+    defstruct [:id]
+  end
+
+  @config []
+
+  describe "fetch_primary_key_values/2" do
+    test "handles nil primary key value" do
+      assert Operations.fetch_primary_key_values(%PrimaryFieldUser{id: nil}, @config) == {:error, "Primary key value for key `:id` in #{inspect PrimaryFieldUser} can't be `nil`"}
+      assert Operations.fetch_primary_key_values(%CompositePrimaryFieldsUser{}, @config) == {:error, "Primary key value for key `:some_id` in #{inspect CompositePrimaryFieldsUser} can't be `nil`"}
+      assert Operations.fetch_primary_key_values(%NonEctoUser{}, @config) == {:error, "Primary key value for key `:id` in #{inspect NonEctoUser} can't be `nil`"}
+    end
+
+    test "requires primary key" do
+      assert Operations.fetch_primary_key_values(%NoPrimaryFieldUser{}, @config) == {:error, "No primary keys found for #{inspect NoPrimaryFieldUser}"}
+    end
+
+    test "returns keyword list" do
+      assert Operations.fetch_primary_key_values(%PrimaryFieldUser{id: 1}, @config) == {:ok, id: 1}
+      assert Operations.fetch_primary_key_values(%CompositePrimaryFieldsUser{some_id: 1, another_id: 2}, @config) == {:ok, some_id: 1, another_id: 2}
+      assert Operations.fetch_primary_key_values(%NonEctoUser{id: 1}, @config) == {:ok, id: 1}
+    end
+  end
 end

test/pow/store/credentials_cache_test.exs

@@ -75,23 +75,6 @@ defmodule Pow.Store.CredentialsCacheTest do
     assert CredentialsCache.get(@config, "key_3") == {user, fingerprint: 1}
   end
 
-  test "raises for nil primary key value" do
-    user_1 = %User{id: nil}
-
-    assert_raise RuntimeError, "Primary key value for key `:id` in Pow.Test.Ecto.Users.User can't be `nil`", fn ->
-      CredentialsCache.put(@config, "key_1", {user_1, a: 1})
-    end
-  end
-
-  defmodule NoPrimaryFieldUser do
-    use Ecto.Schema
-
-    @primary_key false
-    schema "users" do
-      timestamps()
-    end
-  end
-
   defmodule CompositePrimaryFieldsUser do
     use Ecto.Schema
 
@@ -104,36 +87,13 @@ defmodule Pow.Store.CredentialsCacheTest do
     end
   end
 
-  test "handles custom primary fields" do
-    assert_raise RuntimeError, "No primary keys found for Pow.Store.CredentialsCacheTest.NoPrimaryFieldUser", fn ->
-      CredentialsCache.put(@config, "key_1", {%NoPrimaryFieldUser{}, a: 1})
-    end
-
-    assert_raise RuntimeError, "Primary key value for key `:another_id` in Pow.Store.CredentialsCacheTest.CompositePrimaryFieldsUser can't be `nil`", fn ->
-      CredentialsCache.put(@config, "key_1", {%CompositePrimaryFieldsUser{}, a: 1})
-    end
-
+  test "sorts composite primary keys" do
     user = %CompositePrimaryFieldsUser{some_id: 1, another_id: 2}
 
     CredentialsCache.put(@config, "key_1", {user, a: 1})
 
     assert CredentialsCache.users(@config, CompositePrimaryFieldsUser) == [user]
-  end
-
-  defmodule NonEctoUser do
-    defstruct [:id]
-  end
-
-  test "handles non-ecto user struct" do
-    assert_raise RuntimeError, "Primary key value for key `:id` in Pow.Store.CredentialsCacheTest.NonEctoUser can't be `nil`", fn ->
-      CredentialsCache.put(@config, "key_1", {%NonEctoUser{}, a: 1})
-    end
-
-    user = %NonEctoUser{id: 1}
-
-    assert CredentialsCache.put(@config, "key_1", {user, a: 1})
-
-    assert CredentialsCache.users(@config, NonEctoUser) == [user]
+    assert EtsCacheMock.get(@backend_config, [CompositePrimaryFieldsUser, :user, [another_id: 2, some_id: 1]]) == user
   end
 
   # TODO: Remove by 1.1.0

test/support/extensions/email_confirmation/repo_mock.ex

@@ -78,7 +78,7 @@ defmodule PowEmailConfirmation.Test.RepoMock do
     {:ok, user}
   end
 
-  def get!(User, 1, _opts), do: Process.get({:user, 1})
+  def get_by!(User, [id: 1], _opts), do: Process.get({:user, 1})
 
   defmodule Invitation do
     @moduledoc false
@@ -89,6 +89,6 @@ defmodule PowEmailConfirmation.Test.RepoMock do
 
     def update(changeset, opts), do: RepoMock.update(changeset, opts)
 
-    def get!(User, 1, _opts), do: Process.get({:user, 1})
+    def get_by!(User, [id: 1], _opts), do: Process.get({:user, 1})
   end
 end

test/support/extensions/invitation/repo_mock.ex

@@ -27,7 +27,7 @@ defmodule PowInvitation.Test.RepoMock do
   end
   def insert(%{valid?: false} = changeset, _opts), do: {:error, %{changeset | action: :insert}}
 
-  def get!(User, 1, _opts), do: Process.get({:user, 1})
+  def get_by!(User, [id: 1], _opts), do: Process.get({:user, 1})
 
   def get_by(User, [invitation_token: "valid"], _opts), do: %{@user | invitation_token: "valid"}
   def get_by(User, [invitation_token: "valid_but_accepted"], _opts), do: %{@user | invitation_accepted_at: :now}