Polymorphic Embeds

polymorphic_embed.ex

defmodule PolymorphicEmbed do
  @moduledoc """
  Allows you to embed with an unknown schema at compile time.

  The schema is decided by a `type_field` which could be an elixir module,
  but also an enum which can be converted into a module with an `embedded_schema`.
  """

  use Ecto.Type

  def type, do: :map

  def cast(%{__struct__: _} = struct) do
    {:ok, struct}
  end
  def cast(_) do
    # we don't support cast from user-supplied values,
    # as it's done in a special way in the changeset
    :error
  end

  def load(%{"__struct__" => module} = data) do
    module = String.to_existing_atom(module)

    fields =
      data
      |> Map.drop(["__struct__"])
      |> Enum.flat_map(&load_field(&1, module))

    {:ok, struct(module, fields)}
  end

  def dump(%Ecto.Changeset{} = ch) do
    dump(Ecto.Changeset.apply_changes(ch))
  end
  def dump(%struct{} = data) do
    data =
      data
      |> Map.drop([:__struct__])
      |> Enum.flat_map(&dump_field(&1, struct))
      |> Enum.into(%{})
      |> Map.put("__struct__", to_string(struct))

    {:ok, data}
  end

  def load_field({field, value}, schema) do
    field = String.to_existing_atom(field)
    case schema.__schema__(:type, field) do
      nil ->
        []

      type ->
        {:ok, value} = Ecto.Type.embedded_load(type, value, :json)
        [{field, value}]
    end
  end

  def dump_field({:__meta__, _}, _), do: []
  def dump_field({field, value}, schema) do
    type = schema.__schema__(:type, field)
    {:ok, value} = Ecto.Type.embedded_dump(type, value, :json)
    [{to_string(field), value}]
  end

  @doc """
  Similar to `Ecto.Changeset.cast_embed/3`, but it will use the `type_field` to
  decide which schema to use to cast the embed in `embed_field`.

  In the default behaviour, `cast_polymorphic_embed/4` expects the `type_field`
  to contain a schema module with an embedded_schema definition as well as a
  changeset function. However, when a function is passed into the `with_type` option,
  the function will receive the value of `type_field` and can convert it to the
  corresponding schema module.
  """
  def cast_polymorphic_embed(%{valid?: false} = changeset, _, _, _) do
    # skip casting the embed if the changeset is invalid
    changeset
  end
  def cast_polymorphic_embed(changeset, type_field, embed_field, opts \\ []) do
    import Ecto.Changeset
    {embed_type, original_embed_type} = get_embed_type(changeset, type_field, opts)

    embed_field_str = Atom.to_string(embed_field)
    embed_data = case {embed_type, original_embed_type} do
      {embed_type, embed_type} -> Map.get(changeset.data, embed_field, struct(embed_type))
      {embed_type, _nil_or_other} -> struct(embed_type)
    end
    embed_params = Map.get(changeset.params, embed_field, Map.get(changeset.params, embed_field_str, %{}))

    embed_changeset = embed_type.changeset(embed_data, embed_params)

    changeset
    |> put_change(embed_field, if(embed_changeset.valid?, do: apply_changes(embed_changeset), else: embed_changeset))
    |> Map.put(:valid?, embed_changeset.valid?)
  end

  defp get_embed_type(changeset, type_field, opts) do
    import Ecto.Changeset
    with_type_fun = Keyword.get(opts, :with_type, &(&1))
    embed_type = get_field(changeset, type_field)
    embed_type = cast_embed_type(embed_type, with_type_fun)


    original_embed_type = cast_embed_type(Map.get(changeset.data, type_field), with_type_fun)

    {embed_type, original_embed_type}
  end

  defp cast_embed_type(embed_type, fun) do
    case embed_type do
      nil -> nil
      embed_type -> fun.(embed_type)
    end
  end
end