improvement: calculation.select/2 + select calculation option

.formatter.exs

@@ -90,6 +90,7 @@ locals_without_parens = [
   required?: 1,
   resource: 1,
   resource: 2,
+  select: 1,
   sensitive?: 1,
   soft?: 1,
   sort: 1,

lib/ash/actions/load.ex

@@ -82,9 +82,14 @@ defmodule Ash.Actions.Load do
 
   defp maybe_select(query, field) do
     if query.select do
-      Ash.Query.select(query, [field])
+      Ash.Query.select(query, List.wrap(field))
     else
-      query
+      to_select =
+        query.resource
+        |> Ash.Resource.Info.attributes()
+        |> Enum.map(& &1.name)
+
+      Ash.Query.select(query, to_select)
     end
   end
 

lib/ash/calculation.ex

@@ -18,7 +18,9 @@ defmodule Ash.Calculation do
 
       def describe(opts), do: inspect({__MODULE__, opts})
 
-      defoverridable init: 1, type: 0, describe: 1
+      def select(_query, _opts), do: []
+
+      defoverridable init: 1, type: 0, describe: 1, select: 2
     end
   end
 
@@ -27,4 +29,5 @@ defmodule Ash.Calculation do
   @callback describe(Keyword.t()) :: String.t()
   @callback calculate([Ash.Resource.record()], Keyword.t(), map) ::
               {:ok, [term]} | [term] | {:error, term}
+  @callback select(Ash.Query.t(), Keyword.t()) :: list(atom)
 end

lib/ash/query/query.ex

@@ -635,19 +635,27 @@ defmodule Ash.Query do
 
             load_relationship(query, [{field, nested_query}])
 
-          calculation = Ash.Resource.Info.calculation(query.resource, field) ->
-            {module, opts} = module_and_opts(calculation.calculation)
+          resource_calculation = Ash.Resource.Info.calculation(query.resource, field) ->
+            {module, opts} = module_and_opts(resource_calculation.calculation)
 
-            with {:ok, args} <- validate_arguments(calculation, rest),
+            with {:ok, args} <- validate_arguments(resource_calculation, rest),
                  {:ok, calculation} <-
                    Calculation.new(
-                     calculation.name,
+                     resource_calculation.name,
                      module,
                      opts,
                      args
                    ) do
               calculation = %{calculation | load: field}
-              %{query | calculations: Map.put(query.calculations, field, calculation)}
+
+              fields_to_select =
+                resource_calculation.select
+                |> Kernel.||([])
+                |> Enum.concat(module.select(query, opts) || [])
+
+              query
+              |> Map.update!(:calculations, &Map.put(&1, field, calculation))
+              |> maybe_select(fields_to_select)
             end
 
           true ->
@@ -706,18 +714,26 @@ defmodule Ash.Query do
             )
         end
 
-      calculation = Ash.Resource.Info.calculation(query.resource, field) ->
+      resource_calculation = Ash.Resource.Info.calculation(query.resource, field) ->
         {module, opts} =
-          case calculation.calculation do
+          case resource_calculation.calculation do
             {module, opts} -> {module, opts}
             module -> {module, []}
           end
 
-        with {:ok, args} <- validate_arguments(calculation, %{}),
+        with {:ok, args} <- validate_arguments(resource_calculation, %{}),
              {:ok, calculation} <-
-               Calculation.new(calculation.name, module, opts, args) do
+               Calculation.new(resource_calculation.name, module, opts, args) do
           calculation = %{calculation | load: field}
-          %{query | calculations: Map.put(query.calculations, field, calculation)}
+
+          fields_to_select =
+            resource_calculation.select
+            |> Kernel.||([])
+            |> Enum.concat(module.select(query, opts) || [])
+
+          query
+          |> Map.update!(:calculations, &Map.put(&1, field, calculation))
+          |> maybe_select(fields_to_select)
         else
           {:error, error} ->
             add_error(query, :load, error)
@@ -728,6 +744,19 @@ defmodule Ash.Query do
     end
   end
 
+  defp maybe_select(query, field) do
+    if query.select do
+      Ash.Query.select(query, List.wrap(field))
+    else
+      to_select =
+        query.resource
+        |> Ash.Resource.Info.attributes()
+        |> Enum.map(& &1.name)
+
+      Ash.Query.select(query, to_select)
+    end
+  end
+
   defp validate_arguments(calculation, args) do
     Enum.reduce_while(calculation.arguments, {:ok, %{}}, fn argument, {:ok, arg_values} ->
       value = default(Map.get(args, argument.name), argument.default)

lib/ash/resource/calculation/calculation.ex

@@ -1,6 +1,15 @@
 defmodule Ash.Resource.Calculation do
   @moduledoc "Represents a named calculation on a resource"
-  defstruct [:name, :type, :calculation, :arguments, :description, :private?, :allow_nil?]
+  defstruct [
+    :name,
+    :type,
+    :calculation,
+    :arguments,
+    :description,
+    :private?,
+    :allow_nil?,
+    :select
+  ]
 
   @schema [
     name: [
@@ -27,6 +36,11 @@ defmodule Ash.Resource.Calculation do
       doc:
         "Whether or not the calculation will appear in any interfaces created off of this resource, e.g AshJsonApi and AshGraphql"
     ],
+    select: [
+      type: {:list, :atom},
+      default: [],
+      doc: "A list of fields to ensure selected in the case that the calculation is run."
+    ],
     allow_nil?: [
       type: :boolean,
       default: true,

lib/ash/resource/calculation/concat.ex

@@ -10,6 +10,10 @@ defmodule Ash.Resource.Calculation.Concat do
     end
   end
 
+  def select(_query, opts) do
+    opts[:keys]
+  end
+
   def calculate(records, opts, _) do
     Enum.map(records, fn record ->
       Enum.map_join(opts[:keys], opts[:separator] || "", fn key ->

lib/ash/resource/dsl.ex

@@ -845,14 +845,20 @@ defmodule Ash.Resource.Dsl do
 
     Takes a module that must adopt the `Ash.Calculation` behaviour. See that module
     for more information.
+
+    To ensure that the necessary fields are selected:
+
+    1.) Specifying the `select` option on a calculation in the resource.
+    2.) Define a `select/2` callback in the calculation module
+    3.) Set `always_select?` on the attribute in question
     """,
     examples: [
-      "calculate :full_name, :string, MyApp.MyResource.FullName",
+      "calculate :full_name, :string, MyApp.MyResource.FullName, select: [:first_name, :last_name]",
       {
         "`Ash.Calculation` implementation example:",
-        "calculate :full_name, :string, {MyApp.FullName, keys: [:first_name, :last_name]}"
+        "calculate :full_name, :string, {MyApp.FullName, keys: [:first_name, :last_name]}, select: [:first_name, :last_name]"
       },
-      "calculate :full_name, :string, full_name([:first_name, :last_name])"
+      "calculate :full_name, :string, full_name([:first_name, :last_name]), select: [:first_name, :last_name]"
     ],
     target: Ash.Resource.Calculation,
     args: [:name, :type, :calculation],

test/actions/create_test.exs

@@ -614,7 +614,7 @@ defmodule Ash.Test.Actions.CreateTest do
         |> replace_relationship(:author, author)
         |> Api.create!()
 
-      assert post.author == author
+      assert post.author.id == author.id
     end
   end
 

test/actions/update_test.exs

@@ -634,8 +634,8 @@ defmodule Ash.Test.Actions.UpdateTest do
       |> replace_relationship(:author, author2)
       |> Api.update!()
 
-      assert Api.get!(Author, author2.id, load: [:posts]).posts == [
-               Api.get!(Post, post.id)
+      assert Enum.map(Api.get!(Author, author2.id, load: [:posts]).posts, & &1.id) == [
+               Api.get!(Post, post.id).id
              ]
     end
 
@@ -658,8 +658,8 @@ defmodule Ash.Test.Actions.UpdateTest do
 
       updated_post = post |> new() |> replace_relationship(:author, author2) |> Api.update!()
 
-      assert updated_post.author ==
-               Api.get!(Author, author2.id)
+      assert updated_post.author.id ==
+               Api.get!(Author, author2.id).id
     end
   end
 

test/calculation_test.exs

@@ -45,6 +45,7 @@ defmodule Ash.Test.CalculationTest do
 
     calculations do
       calculate :full_name, :string, {Concat, keys: [:first_name, :last_name]} do
+        select [:first_name, :last_name]
         # We currently need to use the [allow_empty?: true, trim?: false] constraints here.
         # As it's an empty string, the separator would otherwise be trimmed and set to `nil`.
         argument :separator, :string,
@@ -97,6 +98,18 @@ defmodule Ash.Test.CalculationTest do
     assert full_names == ["brian - cranston", "zach - daniel"]
   end
 
+  test "fields are selected if necessary for the calculation" do
+    full_names =
+      User
+      |> Ash.Query.select(:first_name)
+      |> Ash.Query.load(full_name: %{separator: " - "})
+      |> Api.read!()
+      |> Enum.map(& &1.full_name)
+      |> Enum.sort()
+
+    assert full_names == ["brian - cranston", "zach - daniel"]
+  end
+
   test "custom calculations can be added to a query" do
     full_names =
       User