bulk_upsert

defmodule BulkUpsert do
  @moduledoc "Bulk upsert a list of Ecto structs and their nested associations in one call."

  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`)

  - `:chunk_size` - The number of parent attrs items to insert into the database in a single
  query. Can be increased or decreased as needed to avoid hitting Postgres max item limit for a
  single query. (Default: `1000`)

  - `: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 or source 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}]}`
    - A `many_to_many` join table is keyed by its source, e.g. `%{"persons_topics" => [...]}`.

  - `:placeholders` - Set fields from shared values that are sent to the database once instead of
  once per row, using the `:placeholders` feature of Ecto's `insert_all/3`. This option is a map
  whose key is the schema or source being upserted, and the value is a map of `field => value`.
  The fields are set after changeset validation, so they do not need to appear in the attrs.
  (Default: `%{}`)
    - Example: `%{YourProject.Persons.Person => %{inserted_at: DateTime.utc_now()}}`
    - Placeholder fields bypass the changeset, so they are not cast or validated.
    - Do not include a placeholder field in the changeset's `validate_required/2`. The value is
    absent during validation, so the changeset would be invalid and the row would be skipped.

  - `: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.Person => [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

  ## Associations

  Nested associations are upserted in the same call as the parent:

  - `has_many` and `has_one`: the associated records are upserted into their own table. Each child
  must include its foreign key in its attrs, since it is upserted directly via `insert_all/3`.

  - `many_to_many`: the associated records are upserted into their own table, and the join table
  rows linking each parent to its associations are upserted as well. Duplicate records and links
  are removed automatically.

  - `embeds_one` and `embeds_many`: embedded data has no table of its own, so it is stored inline
  on the parent row as part of the parent upsert.

  ## Known limitations

  - Nested `belongs_to` associations are not upserted. To associate with a `belongs_to` parent,
  include its foreign key field in the attrs (e.g. `category_id`).
  """
  def bulk_upsert(repo_module, schema_module, attrs_list, opts \\ []) do
    changeset_function_atom = Keyword.get(opts, :changeset_function_atom, :changeset)
    chunk_size = Keyword.get(opts, :chunk_size, 1000)
    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(chunk_size)
    # 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

  # Upsert a list of entries into a schema or source, applying any configured placeholders and
  # chunking large payloads to stay within Postgres bulk limits.
  defp insert_all_entries(entries, schema_or_source, insert_all_opts, context) do
    placeholder_values = context.placeholders[schema_or_source] || %{}

    # Placeholder fields are set here, after changeset validation, because a `{:placeholder, key}`
    # tuple cannot pass through a changeset. Each placeholder value is sent to Postgres once.
    {entries, insert_all_opts} =
      if placeholder_values == %{} do
        {entries, insert_all_opts}
      else
        placeholder_attrs =
          Map.new(placeholder_values, fn {field, _value} -> {field, {:placeholder, field}} end)

        {
          Enum.map(entries, &Map.merge(&1, placeholder_attrs)),
          Keyword.put(insert_all_opts, :placeholders, placeholder_values)
        }
      end

    entries
    |> Enum.chunk_every(context.chunk_size)
    |> Enum.each(fn entries_chunk ->
      apply(context.insert_all_function_module, context.insert_all_function_atom, [
        schema_or_source,
        entries_chunk,
        insert_all_opts
      ])
    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, [])
    chunk_size = Keyword.get(opts, :chunk_size, 1000)
    timeout = Keyword.get(opts, :timeout, @default_timeout)
    placeholders = Keyword.get(opts, :placeholders, %{})

    # Bundle the values that every `insert_all` call site shares
    insert_all_context = %{
      insert_all_function_module: insert_all_function_module,
      insert_all_function_atom: insert_all_function_atom,
      chunk_size: chunk_size,
      placeholders: placeholders
    }

    # 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
          # Drop all assoc data from the changeset (assocs are handled separately in a later step)
          |> Enum.map(fn %Ecto.Changeset{} = changeset ->
            changeset
            |> Map.put(
              :changes,
              Map.drop(changeset.changes, schema_module.__schema__(:associations))
            )
          end)
          |> Enum.map(&attrs_from_changeset/1)

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

        insert_all_entries(attrs_list, schema_module, parent_insert_all_opts, insert_all_context)

        # Perform bulk upsert for all `has_many` and `has_one` associations
        for association <- get_schema_associations(schema_module, :has) do
          association_attrs_list =
            changesets
            |> Enum.map(& &1.changes)
            |> Enum.map(&Map.get(&1, association))
            # `has_many` changes are a list of changesets; `has_one` changes are a single
            # changeset; an absent association is `nil`. `List.wrap/1` normalizes each into a
            # (possibly empty) list, which `flat_map` concatenates.
            |> Enum.flat_map(&List.wrap/1)
            |> 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] || []
            )

          insert_all_entries(
            association_attrs_list,
            association_schema_module,
            association_insert_all_opts,
            insert_all_context
          )
        end

        # Perform bulk upsert for all `many_to_many` associations
        for association <- get_schema_associations(schema_module, :many_to_many) do
          %Ecto.Association.ManyToMany{
            related: related_schema_module,
            join_through: join_through,
            join_keys: [{owner_join_key, owner_key}, {related_join_key, related_key}]
          } = schema_module.__changeset__()[association] |> elem(1)

          # Pair each parent's primary key with each of its related changesets, so the related
          # records and the join rows can both be derived from the same data.
          parent_related_pairs =
            Enum.flat_map(changesets, fn parent_changeset ->
              parent_changeset.changes
              |> Map.get(association)
              |> List.wrap()
              |> Enum.map(fn related_changeset -> {parent_changeset, related_changeset} end)
            end)

          # Upsert the related records into their own table. The same record may be referenced by
          # multiple parents, so duplicates are removed to avoid conflicting twice in one query.
          related_attrs_list =
            parent_related_pairs
            |> Enum.map(fn {_parent_changeset, related_changeset} ->
              attrs_from_changeset(related_changeset)
            end)
            |> Enum.uniq_by(&Map.take(&1, related_schema_module.__schema__(:primary_key)))

          related_insert_all_opts =
            Keyword.merge(
              [
                on_conflict:
                  {:replace_all_except,
                   related_schema_module.__schema__(:primary_key) ++ replace_all_except},
                conflict_target: related_schema_module.__schema__(:primary_key),
                timeout: timeout
              ],
              insert_all_opts[related_schema_module] || []
            )

          insert_all_entries(
            related_attrs_list,
            related_schema_module,
            related_insert_all_opts,
            insert_all_context
          )

          # Upsert the join table rows that link each parent to its related records. The same link
          # may be listed more than once, so duplicate rows are removed for the same reason.
          join_attrs_list =
            parent_related_pairs
            |> Enum.map(fn {parent_changeset, related_changeset} ->
              %{
                owner_join_key => Ecto.Changeset.fetch_field!(parent_changeset, owner_key),
                related_join_key => Ecto.Changeset.fetch_field!(related_changeset, related_key)
              }
            end)
            |> Enum.uniq()

          join_insert_all_opts =
            Keyword.merge(
              [
                on_conflict: :nothing,
                conflict_target: [owner_join_key, related_join_key],
                timeout: timeout
              ],
              insert_all_opts[join_through] || []
            )

          insert_all_entries(
            join_attrs_list,
            join_through,
            join_insert_all_opts,
            insert_all_context
          )
        end
      end,
      timeout: timeout
    )
  end

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

  # Get all `many_to_many` associations for a given schema.
  defp get_schema_associations(schema_module, :many_to_many) do
    schema_module.__changeset__()
    |> Enum.filter(fn {_k, v} ->
      match?({:assoc, %Ecto.Association.ManyToMany{}}, 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_association_attrs =
      schema_module
      |> get_schema_associations(:has)
      # 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]
          # `has_one` changes are a single changeset; `has_many` changes are a list of changesets.
          |> List.wrap()
          |> 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_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