Ash.Query.set_argument don't work inside an action

I have a read action that I want to do some checks first.

To do that, I want to create some preparations, the first one will fetch some data, and the others will use that data to do some checks.

Here is the preparations that will fetch data and set as an argument inside the query:

defmodule Core.Feedbacks.Template.Actions.AllFromClass.Preparations.GetClass do
  @moduledoc false

  alias Core.Feedbacks.Class

  alias Ash.Error.Changes.InvalidAttribute

  use Ash.Resource.Preparation

  def prepare(query, _opts, %{actor: actor}) do
    class_id = Ash.Query.get_argument(query, :class_id)

    case Class.get!(%{id: class_id}, actor: actor) do
      nil ->
        # TODO give the error a code
        error = InvalidAttribute.exception(message: "Can't find class")

        Ash.Query.add_error(query, error)

      class ->
        Ash.Query.set_argument(query, :class, class)
    end
  end
end

As you can see, I get a class_id argument, fetch the class with it, and then I use Ash.Query.set_argument to add the class argument.

The problem is that this call doesn´t do anything, if I dbg the query after the set_argument call, the arguments field only contains the class_id argument.