bulk_upsert

defmodule BulkUpsert do
  @moduledoc "Enable bulk upsert functionality when working with Ecto."

  require Logger

  @default_timeout 15_000

  @doc """
  Validate a list of attrs maps (`attrs_list`) by passing them through an Ecto changeset,
  then upsert the valid items to the database that corresponds to a given Ecto `repo_module` (e.g.
  `YourProject.Repo`).

  Using a changeset serves two purposes:
    1. The changeset can be used to validate and transform the data.
    2. Using a changeset allows this function to perform bulk upserts with nested associations.

  For validation, each list item in the `attrs_list` is converted to a changeset for a given
  `schema_module`. By default, this function expects the schema module to contain a 2-arity
  function called `:changeset`. (See the `#options` section for more info.)

  ## Basic example

      iex> BulkUpsert.bulk_upsert(
      ...>   YourProject.Repo,
      ...>   YourProject.Persons.Person,
      ...>   _attrs_list = [
      ...>     %{id: 1, name: "Alice", age: 25, phone_number: "555-1234"},
      ...>     %{id: 2, name: "Bob", age: 35, phone_number: "555-2345"},
      ...>   ]
      ...> )
      :ok

  ## Options

  - `:changeset_function_atom` - The name of the 2-arity changeset function to apply for the given
  `schema_module` (Default: `:changeset`)

  - `:insert_all_function_module` - Instead of using the `:insert_all` function in the given
  `repo_module`, you may specify the name of a custom module to use instead. (Default:
  Inherited from the value specified in the `repo_module` function argument, e.g.
  `YourProject.Repo`))
    - Example: `YourProject.OtherRepo`

  - `:insert_all_function_atom` - Instead of using your repo module's `:insert_all`
  function, you may pass a compatible equivalent that accepts the same arguments. (Default:
  `:insert_all`)
    - Example: `:insert_all_with_autogenerated_timestamps`

  - `:insert_all_opts` - Pass custom `opts` to the `insert_all/3` function. This option consists
  of a map whose key is the schema module that may have items being upserted, and the value is the
  `YourProject.Repo.insert_all/3` opts that will be applied when items for that schema are being
  upserted. By default, this function is configured to replace all values in a given struct,
  except for the primary key(s) and the insert timestamp. (Default: `%{}`)
    - Example: `%{YourProject.Persons.Person => [on_conflict: {:nothing}]}`

  - `:recover_changeset_errors` - If the given fields in a changeset have errors, then replace
  them with a custom fallback value. (Default: `%{}`)
    - Example: `%{YourProject.Persons.Person => %{phone_number: "INVALID"}}`

  - `:replace_all_except` - If a row already exists, then all fields will be replaced except the
  primary key, and any fields specified here. (Default: `[]`)
    - Example: `[:field, :other_field]`

  - `:timeout` - The maximum timeout for a transaction. (Default: `#{@default_timeout}`)
    - Example: `60_000`

  ## Examples

  Upsert a list of Person attrs using the changeset function
  `YourProject.Persons.Person.upsert_changeset/2` to validate the attrs:

      iex> attrs_list = [%{id: 1, name: "Alice", ...}]

      iex> BulkUpsert.bulk_upsert(
      ...>   YourProject.Repo,
      ...>   YourProject.Persons.Person,
      ...>   attrs_list,
      ...>   changeset_function_atom: :upsert_changeset
      ...> )
      :ok

  Upsert a list of attrs, but overwrite the `:name` field if there is a conflict.

  If using this option, you must declare each schema that will get a customized `:insert_all_opts`
  keyword list. Any schemas that are not given custom `:insert_all_opts` will overwrite all fields
  except the primary key:

      iex> insert_all_opts = %{
      ...>   YourProject.Persons => [on_conflict: {:replace, [:name]}]
      ...> }

      iex> BulkUpsert.bulk_upsert(
      ...>   YourProject.Repo,
      ...>   YourProject.Persons.Person,
      ...>   _attrs_list = [%{id: 1, name: "Alicia"}],
      ...>   insert_all_opts: insert_all_opts
      ...> )
      :ok

  ## Known limitations

  - This function will not currently work with the `:placeholders` option of
  Ecto's `insert_all/3` function. This is because the attrs are passed directly to the changesets
  for validation, so the placeholder values will not be parsed correctly. This functionality can
  be added later if needed.
  """
  def bulk_upsert(repo_module, schema_module, attrs_list, opts \\ []) do
    changeset_function_atom = Keyword.get(opts, :changeset_function_atom, :changeset)
    recover_changeset_errors = Keyword.get(opts, :recover_changeset_errors, %{})

    attrs_list
    # Convert to changesets so the data can be validated before upsertion
    |> Enum.map(fn attrs -> apply(schema_module, changeset_function_atom, [attrs]) end)
    |> then(&handle_invalid_changesets(schema_module, &1, recover_changeset_errors))
    # Work around Postgres bulk limits by chunking large payloads
    |> Enum.chunk_every(2_000)
    # Use `Enum.map/2` instead of `Task.async_stream/2`. (This slightly decreases performance, but
    # prevents issues when using the Ecto sandbox (i.e. in the `:test` configuration environment)
    # since other functions may also call `Task.async_stream/2` before calling this function.
    # These nested async calls cause issues with the sandbox. If the additional performance is
    # required, the caller may be able to pass in its PID to the `Repo.insert_all/3` opts to work
    # around this issue, at the cost of additional complexity in the codebase)
    |> Enum.map(&do_bulk_upsert(repo_module, schema_module, &1, opts))

    :ok
  end

  defp attrs_from_changeset(changeset) do
    struct = Ecto.Changeset.apply_action!(changeset, :build_for_bulk_upsert)

    struct
    |> Map.from_struct()
    |> Map.reject(fn {k, _v} -> k not in Map.keys(changeset.changes) end)
  end

  defp do_bulk_upsert(repo_module, schema_module, changesets, opts) do
    insert_all_function_module = Keyword.get(opts, :insert_all_function_module, repo_module)
    insert_all_function_atom = Keyword.get(opts, :insert_all_function_atom, :insert_all)
    insert_all_opts = Keyword.get(opts, :insert_all_opts, [])
    replace_all_except = Keyword.get(opts, :replace_all_except, [])
    timeout = Keyword.get(opts, :timeout, @default_timeout)

    # Wrap all bulk upserts in a transaction so that any failures will roll back all changes made
    # to the parent and all of its associations
    repo_module.transaction(
      fn ->
        # Perform bulk upsert for all parent attrs
        attrs_list = changesets |> Enum.map(&attrs_from_changeset/1)

        # Build `insert_all` opts for the parent schema
        insert_all_opts =
          Keyword.merge(
            _default_parent_insert_all_opts = [
              on_conflict:
                {:replace_all_except,
                 schema_module.__schema__(:primary_key) ++ replace_all_except},
              conflict_target: schema_module.__schema__(:primary_key),
              timeout: timeout
            ],
            insert_all_opts[schema_module] || []
          )

        apply(insert_all_function_module, insert_all_function_atom, [
          schema_module,
          attrs_list,
          insert_all_opts
        ])

        # Perform bulk upsert for all 'has_many' associations
        for association <- get_schema_associations(schema_module, :has_many) do
          association_attrs_list =
            changesets
            |> Enum.map(& &1.changes)
            |> Enum.map(&Map.get(&1, association))
            # Reject associations that do not have nested changesets (FIXME: Is this necessary?)
            |> Enum.reject(&is_nil/1)
            |> List.flatten()
            |> Enum.map(&attrs_from_changeset/1)

          association_schema_module =
            schema_module.__changeset__()[association] |> elem(1) |> Map.fetch!(:related)

          # Build `insert_all` opts for the association schema
          association_insert_all_opts =
            Keyword.merge(
              _default_association_insert_all_opts = [
                on_conflict:
                  {:replace_all_except,
                   association_schema_module.__schema__(:primary_key) ++ replace_all_except},
                conflict_target: association_schema_module.__schema__(:primary_key),
                timeout: timeout
              ],
              insert_all_opts[association_schema_module] || []
            )

          apply(insert_all_function_module, insert_all_function_atom, [
            association_schema_module,
            association_attrs_list,
            association_insert_all_opts
          ])
        end

        # FIXME: Add logic for other bulk upsert for other associations as needed: `has_one`,
        # `many_to_many`, `embeds_one`, `embeds_many`
      end,
      timeout: timeout
    )
  end

  # Get all `:has_many` associations for a given schema.
  defp get_schema_associations(schema_module, :has_many) do
    schema_module.__changeset__()
    |> Enum.filter(fn {_k, v} ->
      match?({:assoc, %Ecto.Association.Has{cardinality: :many}}, v)
    end)
    |> Keyword.keys()
  end

  defp handle_invalid_changesets(schema_module, changesets, recover_changeset_errors) do
    changesets
    |> recover_changesets_with_recoverable_errors(recover_changeset_errors)
    |> then(&reject_invalid_changesets(schema_module, &1))
  end

  defp log_on_bulk_upsert_changeset_error(schema_module, changeset) do
    item_id_or_ids =
      schema_module.__schema__(:primary_key)
      |> Enum.reduce(%{}, fn primary_key_field, acc ->
        acc |> Map.put(primary_key_field, changeset.changes[primary_key_field])
      end)

    invalid_parent_attrs =
      changeset.errors
      |> Enum.reduce(%{}, fn {k, _v}, acc -> Map.put(acc, k, changeset.changes[k]) end)
      # If a parent has an error in an association, the error will appear as a changeset, which
      # clutters up the logs. So, remove association errors from the invalid attrs map. The error
      # message for the field will still appear in the logs, so the information about the error
      # will still get passed along
      |> Map.new(fn {k, v} ->
        if k in schema_module.__schema__(:associations),
          do: {k, :changesets_hidden_to_keep_logs_shorter},
          else: {k, v}
      end)

    invalid_has_many_association_attrs =
      schema_module
      |> get_schema_associations(:has_many)
      # Only check associations that are present in the changeset's changes (i.e. they aren't nil)
      |> Enum.reject(&is_nil(changeset.changes[&1]))
      |> Enum.reduce(%{}, fn association, acc ->
        association_error_items =
          changeset.changes[association]
          |> Enum.reject(fn changeset -> Enum.empty?(changeset.errors) end)
          |> Enum.reduce([], fn changeset, acc ->
            changeset_error_items =
              changeset.errors
              |> Keyword.keys()
              |> Enum.reduce([], fn key, acc ->
                acc |> Keyword.put(key, changeset.changes[key])
              end)

            changeset_error_items ++ acc
          end)

        if Enum.empty?(association_error_items),
          do: acc,
          else: acc |> Map.put(association, association_error_items)
      end)

    invalid_attrs = Map.merge(invalid_parent_attrs, invalid_has_many_association_attrs)

    Logger.debug(
      """
      This changeset has one or more unrecoverable errors. The item associated with this \
      changeset will not be upserted.\
      """,
      reason: :bulk_upsert_changeset_error,
      schema_module: inspect(schema_module),
      item_id_or_ids: item_id_or_ids,
      # NOTE: If one item in an array contains an invalid value, the whole array will be logged
      fields_with_invalid_attrs: Map.keys(invalid_attrs),
      changeset_errors: changeset.errors
    )
  end

  defp recover_changesets_with_recoverable_errors(changesets, recover_changeset_errors)
       when changesets == [] or recover_changeset_errors == %{} do
    changesets
  end

  defp recover_changesets_with_recoverable_errors(changesets, recover_changeset_errors) do
    changeset_schema_module = List.first(changesets).data.__struct__

    with {_matching_schema_module, recoverable_items} <-
           Enum.find(recover_changeset_errors, fn {schema_module, _recoverable_items} ->
             schema_module == changeset_schema_module
           end) do
      do_recover_changesets_with_recoverable_errors(changesets, recoverable_items)
    else
      _ -> changesets
    end
  end

  defp do_recover_changesets_with_recoverable_errors(changesets, recoverable_items) do
    recoverable_fields = Map.keys(recoverable_items)

    changesets
    |> Enum.map(fn changeset ->
      if changeset.valid? do
        changeset
      else
        changeset_error_fields = Keyword.keys(changeset.errors)

        recoverable_changeset_error_fields =
          Enum.filter(changeset_error_fields, &(&1 in recoverable_fields)) |> Enum.uniq()

        all_errors_in_changeset_are_recoverable? =
          changeset_error_fields |> Enum.all?(&(&1 in recoverable_fields))

        if all_errors_in_changeset_are_recoverable? do
          # Recover all errors in this changeset

          recoverable_changeset_error_fields
          |> Enum.reduce(changeset, fn recoverable_changeset_error_field, acc_changeset ->
            {_field, recover_to_value} =
              recoverable_items
              |> Enum.find(fn {field, _recover_to_value} ->
                field == recoverable_changeset_error_field
              end)

            changes_with_recovered_change =
              acc_changeset.changes
              |> Map.put(recoverable_changeset_error_field, recover_to_value)

            Logger.debug(
              (
                primary_key_info =
                  acc_changeset.data.__struct__.__schema__(:primary_key)
                  |> Keyword.new(fn primary_key_field ->
                    {primary_key_field, Map.fetch!(acc_changeset.changes, primary_key_field)}
                  end)

                struct_name = Macro.to_string(acc_changeset.data.__struct__)

                """
                Recovered changeset error for struct #{struct_name} with primary key(s) \
                `#{inspect(primary_key_info)}` in the field \
                `#{recoverable_changeset_error_field}`.\
                """
              )
            )

            acc_changeset |> Map.put(:changes, changes_with_recovered_change)
          end)
          # Clear the changeset's errors and mark the changeset as valid
          |> Map.merge(%{errors: [], valid?: true})
        else
          # The changeset has errors that are not in the list of recoverable error fields. It will
          # be removed later in the pipeline
          changeset
        end
      end
    end)
  end

  defp reject_invalid_changesets(schema_module, changesets) do
    changesets
    |> Enum.reject(fn changeset ->
      if changeset.valid? do
        _reject_changeset? = false
      else
        log_on_bulk_upsert_changeset_error(schema_module, changeset)
        _reject_changeset? = true
      end
    end)
  end
end