prop/vendor/oban_pro/lib/oban/pro/migration.ex

defmodule Oban.Pro.Migration do
  @moduledoc """
  Migrations create and modify the database tables and indexes Oban Pro needs to function.

  If you use prefixes, or have multiple instances with different prefixes, you can specify the
  prefix and create multiple tables in one migration:

  ## Usage

  To use migrations in your application you'll need to generate an `Ecto.Migration` that wraps
  calls to `Oban.Pro.Migration`:

  ```bash
  mix ecto.gen.migration add_oban_pro
  ```

  Open the generated migration and delegate the `up/0` and `down/0` functions to
  `Oban.Pro.Migration`:

  ```elixir
  defmodule MyApp.Repo.Migrations.AddObanPro do
    use Ecto.Migration

    def up, do: Oban.Pro.Migration.up()

    def down, do: Oban.Pro.Migration.down()
  end
  ```

  This will run all of the necessary migrations for your database.

  Now, run the migration to create the table:

  ```bash
  mix ecto.migrate
  ```

  ## Concurrent Migrations

  For systems with many retained jobs, you may want to create indexes concurrently to avoid
  blocking writes to your tables during migration. The `only` option allows splitting migrations
  into separate steps for schema changes and index creation:

  ```bash
  mix ecto.gen.migration add_oban_pro_schemas
  mix ecto.gen.migration add_oban_pro_indexes
  ```

  The `schemas` migration will add all necessary tables and columns within a DDL transaction:

  ```elixir
  defmodule MyApp.Repo.Migrations.AddObanProSchemas do
    use Ecto.Migration

    def up, do: Oban.Pro.Migration.up(only: :schemas)
    def down, do: Oban.Pro.Migration.down(only: :schemas)
  end
  ```

  Then the `indexes` migration will add indexes concurrently, outside of a transaction:

  ```elixir
  defmodule MyApp.Repo.Migrations.AddObanProIndexes do
    use Ecto.Migration

    # Not needed with use advisory locks, i.e. `migration_lock: :pg_advisory_lock`
    @disable_migration_lock true
    @disable_ddl_transaction true

    def up, do: Oban.Pro.Migration.up(only: :indexes)
    def down, do: Oban.Pro.Migration.down(only: :indexes)
  end
  ```

  #### When to Use Split Migrations

  Use the default, combined approach when:

  1. Setting up a development or test environment
  2. Working with small tables where index creation is quick
  3. Simplicity is preferred over non-blocking operations

  Choose split migrations when:

  1. Operating on a database where blocking writes could cause downtime
  2. Working with large tables where index creation would block for an extended period

  ## Upgrading

  Like `Oban`, migrations are versioned and you may need to run the migration again when new
  `oban_pro` versions are released. To do this, generate a new migration:

  ```bash
  mix ecto.gen.migration upgrade_oban_pro
  ```

  Open the generated migration in your editor and call the `up` and `down` functions on
  `Oban.Pro.Migration`, with an optional version number:

  ```elixir
  defmodule MyApp.Repo.Migrations.UpgradeObanPro do
    use Ecto.Migration

    def up, do: Oban.Pro.Migration.up(version: "1.6.0")

    def down, do: Oban.Pro.Migration.down(version: "1.6.0")
  end
  ```

  > #### Differences from Oban : .tip
  >
  > Unlike Oban, the migration versions correspond directly to the `oban_pro` package version.
  > In addition, running `down/0` will only step back one version, not all the way down.

  ## Isolation with Prefixes

  Oban supports namespacing through PostgreSQL schemas, therefore so does Oban Pro. Specify a
  prefix within your migration:

  ```elixir
  defmodule MyApp.Repo.Migrations.AddPrefixedObanPro do
    use Ecto.Migration

    def up, do: Oban.Pro.Migration.up(prefix: "private")

    def down, do: Oban.Pro.Migration.down(prefix: "private")
  end
  ```
  """

  @behaviour Oban.Pro.Handler

  use Ecto.Migration

  require Logger

  @versions ~w(1.0.0 1.4.0 1.5.0 1.6.0 1.7.0)

  @impl Oban.Pro.Handler
  def on_start do
    event = [:oban, :supervisor, :init]

    :telemetry.attach("oban.migration", event, &__MODULE__.handle_event/4, nil)
  end

  @impl Oban.Pro.Handler
  def on_stop do
    :telemetry.detach("oban.migration")
  end

  @doc false
  def handle_event(_event, _timing, %{conf: conf}, _) do
    prefix = conf.prefix

    case :persistent_term.get(__MODULE__, %{}) do
      %{^prefix => _} ->
        :ok

      versions ->
        :persistent_term.put(__MODULE__, Map.put(versions, prefix, true))

        last = List.last(@versions)

        with_dynamic_repo(conf, fn repo ->
          with :off <- Keyword.get(repo.config(), :pool, :off),
               curr <- migrated_version(%{prefix: prefix, repo: repo}) do
            cond do
              is_nil(curr) ->
                Logger.warning(
                  "Oban Pro for prefix #{inspect(prefix)} isn't migrated. Run migrations to create " <>
                    "the tables and indexes necessary for Oban Pro to function."
                )

              not String.contains?(curr, last) ->
                Logger.warning(
                  "Oban Pro for prefix #{inspect(prefix)} is migrated to #{inspect(curr)}, but the " <>
                    "latest is #{inspect(last)}. Run migrations to upgrade to the latest version."
                )

              true ->
                :ok
            end
          end
        end)
    end
  end

  defp with_dynamic_repo(%{get_dynamic_repo: nil, repo: repo}, fun) do
    fun.(repo)
  end

  defp with_dynamic_repo(%{get_dynamic_repo: get_dyn_fun, repo: repo}, fun) do
    prev_repo =
      try do
        repo.get_dynamic_repo()
      rescue
        _ -> nil
      end

    dyna_repo =
      case get_dyn_fun do
        {mod, dyn_fun, arg} -> apply(mod, dyn_fun, arg)
        dyn_fun when is_function(dyn_fun, 0) -> dyn_fun.()
      end

    try do
      repo.put_dynamic_repo(dyna_repo)

      fun.(repo)
    after
      if prev_repo, do: repo.put_dynamic_repo(prev_repo)
    end
  end

  @doc """
  Run the `up` changes for all migrations between the initial version and the current version.

  ## Example

  Run all migrations up to the current version:

      Oban.Pro.Migration.up()

  Run migrations up to a specified version:

      Oban.Pro.Migration.up(version: "1.4.0")

  Run migrations in an alternate prefix:

      Oban.Pro.Migration.up(prefix: "payments")
  """
  def up(opts \\ []) when is_list(opts) do
    opts = with_defaults(opts)
    curr = migrated_version(opts)
    indx = migrated_version_index(curr, opts)

    cond do
      is_nil(indx) -> inner_change(0..opts.next, :up, curr, opts)
      indx < opts.next -> inner_change((indx + 1)..opts.next, :up, curr, opts)
      true -> :ok
    end
  end

  @doc """
  Run the `down` changes for migrations between the current version and the previous version.

  This behaviour differs from `Oban.Migration` by only stepping down **one version** per
  invocation. You must explicitly specify version `1.0.0` to rollback all the way and drop initial
  tables.

  ## Example

  Run migrations down from the current version to the previous:

      Oban.Pro.Migration.down()

  Run migrations down to and including a specified version:

      Oban.Pro.Migration.down(version: "1.4.0")

  Run all down migrations back to the original:

      Oban.Pro.Migration.down(version: "1.0.0")

  Run migrations in an alternate prefix:

      Oban.Pro.Migration.down(prefix: "payments")
  """
  def down(opts \\ []) when is_list(opts) do
    opts = with_defaults(opts)
    curr = migrated_version(opts)
    indx = migrated_version_index(curr, opts)

    if is_integer(indx) and indx >= opts.next do
      inner_change(indx..opts.next//-1, :down, curr, opts)
    end
  end

  @doc false
  def migrated_version(opts) when is_map(opts) do
    repo = Map.get_lazy(opts, :repo, &repo/0)
    opts = with_defaults(opts)

    query = """
    SELECT pg_catalog.obj_description(pg_class.oid, 'pg_class')
    FROM pg_class
    LEFT JOIN pg_namespace ON pg_namespace.oid = pg_class.relnamespace
    WHERE pg_class.relname = 'oban_producers'
    AND pg_namespace.nspname = '#{opts.escaped}'
    """

    case repo.query(query, [], log: false) do
      {:ok, %{rows: [[version]]}} when is_binary(version) -> version
      _ -> nil
    end
  end

  @doc false
  def partitioned?(opts) do
    query = """
    SELECT EXISTS (
      SELECT 1
      FROM pg_tables
      WHERE schemaname = '#{opts.prefix}' AND tablename = 'oban_jobs_completed'
    )
    """

    {:ok, %{rows: [[partitioned?]]}} = repo().query(query, [], log: false)

    partitioned?
  end

  defp migrated_version_index(version, opts) do
    [schemas_version, indexes_version] = normalize_recorded(version)

    case opts.only do
      [:schemas] -> find_version_index(schemas_version)
      [:indexes] -> find_version_index(indexes_version)
      _ -> find_version_index(schemas_version)
    end
  end

  defp inner_change(_curr_index..next_index//_ = range, direction, current, opts) do
    for index <- range do
      version = Enum.at(@versions, index)
      migration = Module.concat(Oban.Pro.Migrations, "V#{String.replace(version, ".", "")}")

      {do_ddl, do_idx} =
        if direction == :up do
          {:up_ddl, :up_idx}
        else
          {:down_ddl, :down_idx}
        end

      if :schemas in opts.only, do: apply(migration, do_ddl, [opts])
      if :indexes in opts.only, do: apply(migration, do_idx, [opts])

      flush()
    end

    case direction do
      :up -> record_version(next_index, current, opts)
      :down -> record_version(next_index - 1, current, opts)
    end
  end

  defp record_version(-1, _recorded, _opts), do: :ok

  defp record_version(index, recorded, opts) do
    [recorded_schemas, recorded_indexes] = normalize_recorded(recorded)

    prefix = Enum.at(@versions, index)

    versions =
      case opts.only do
        [:indexes] -> [recorded_schemas, "#{prefix}-indexes"]
        [:schemas] -> ["#{prefix}-schemas", recorded_indexes]
        _ -> ["#{prefix}-schemas", "#{prefix}-indexes"]
      end

    versions = versions |> Enum.filter(& &1) |> Enum.join(",")

    execute "COMMENT ON TABLE #{opts.quoted}.oban_producers IS '#{versions}'"
  end

  defp with_defaults(opts) do
    prefix = Access.get(opts, :prefix, "public")
    unlogged = Access.get(opts, :unlogged, true)
    version = Access.get(opts, :version, List.last(@versions))
    generated_columns = Access.get(opts, :generated_columns, false)

    only =
      opts
      |> Access.get(:only, [:schemas, :indexes])
      |> List.wrap()

    %{
      concurrently: only == [:indexes],
      escaped: String.replace(prefix, "'", "\\'"),
      generated_columns: generated_columns,
      next: find_version_index(version),
      only: only,
      prefix: prefix,
      quoted: inspect(prefix),
      unlogged: unlogged,
      version: version
    }
  end

  defp find_version_index(nil), do: nil

  defp find_version_index(version) do
    version = String.replace(version, ["-schemas", "-indexes"], "")

    Enum.find_index(@versions, &(&1 == version))
  end

  defp normalize_recorded(nil), do: [nil, nil]

  defp normalize_recorded(version) do
    case String.split(version, ",", parts: 2) do
      [schemas] when is_binary(schemas) ->
        if String.ends_with?(schemas, "-schemas") do
          [schemas, nil]
        else
          ["#{schemas}-schemas", "#{schemas}-indexes"]
        end

      recorded ->
        recorded
    end
  end
end