improvement: calculation values from requests

lib/ash/actions/helpers.ex

@@ -2,6 +2,19 @@ defmodule Ash.Actions.Helpers do
   @moduledoc false
   require Logger
 
+  def validate_calculation_load!(%Ash.Query{}, module) do
+    raise """
+    `#{inspect(module)}.load/3` returned a query.
+
+    Returning a query from the `load/3` callback of a calculation is now deprecated.
+    Instead, return the load statement itself, i.e instead of `Ash.Query.load(query, [...])`,
+    just return `[...]`. This is so that Ash can examine the requirements of just this single
+    calculation to ensure that all required values are present
+    """
+  end
+
+  def validate_calculation_load!(other, _), do: other
+
   def warn_missed!(action, result) do
     case Map.get(result, :resource_notifications, []) do
       empty when empty in [nil, []] ->

lib/ash/actions/load.ex

@@ -908,10 +908,14 @@ defmodule Ash.Actions.Load do
 
   defp load_for_calcs(query) do
     Enum.reduce(query.calculations || %{}, query, fn {_, calc}, query ->
-      calc.module.load(
+      Ash.Query.load(
         query,
-        calc.opts,
+        calc.module.load(
-        Map.put(calc.context, :context, query.context)
+          query,
+          calc.opts,
+          Map.put(calc.context, :context, query.context)
+        )
+        |> Ash.Actions.Helpers.validate_calculation_load!(calc.module)
       )
     end)
   end

lib/ash/actions/read.ex

@@ -319,8 +319,14 @@ defmodule Ash.Actions.Read do
                     Map.get(fetched_data, :aggregates_in_query) || []
                   )
                   |> add_calculation_values(
+                    query.resource,
+                    api,
+                    action,
+                    error_path,
+                    path,
                     Map.get(fetched_data, :ultimate_query) || query,
-                    Map.get(fetched_data, :calculations_at_runtime) || []
+                    Map.get(fetched_data, :calculations_at_runtime) || [],
+                    get_in(context, path ++ [:calculation_results]) || :error
                   )
                   |> case do
                     {:ok, values} ->
@@ -335,6 +341,9 @@ defmodule Ash.Actions.Read do
                       )
                       |> add_query(Map.get(fetched_data, :ultimate_query), request_opts)
                       |> unwrap_for_get(get?, query.resource)
+
+                    {:requests, requests} ->
+                      {:requests, Enum.map(requests, &{&1, :data})}
                   end
 
                 deps ->
@@ -588,14 +597,16 @@ defmodule Ash.Actions.Read do
         {calculations_in_query, calculations_at_runtime} =
           if Ash.DataLayer.data_layer_can?(ash_query.resource, :expression_calculation) &&
                !request_opts[:initial_data] do
-            Enum.split_with(ash_query.calculations, fn {_name, calculation} ->
+            ash_query.calculations
+            |> Map.values()
+            |> Enum.split_with(fn calculation ->
               Enum.find(used_calculations, &(&1.name == calculation.name)) ||
                 calculation.name in Enum.map(ash_query.sort || [], &elem(&1, 0)) ||
                 (:erlang.function_exported(calculation.module, :expression, 2) &&
                    !calculation.allow_async?)
             end)
           else
-            {[], ash_query.calculations}
+            {[], Map.values(ash_query.calculations)}
           end
 
         {aggregate_auth_requests, aggregate_value_requests, aggregates_in_query} =
@@ -603,7 +614,7 @@ defmodule Ash.Actions.Read do
             ash_query,
             can_be_in_query?,
             authorizing?,
-            Enum.map(calculations_in_query, &elem(&1, 1)),
+            calculations_in_query,
             path
           )
 
@@ -640,7 +651,7 @@ defmodule Ash.Actions.Read do
                 |> Ash.Query.data_layer_query(only_validate_filter?: true)
 
               ash_query =
-                if ash_query.select || calculations_at_runtime == %{} do
+                if ash_query.select || calculations_at_runtime == [] do
                   ash_query
                 else
                   to_select =
@@ -652,7 +663,7 @@ defmodule Ash.Actions.Read do
                 end
 
               ash_query =
-                Enum.reduce(calculations_at_runtime, ash_query, fn {_, calculation}, ash_query ->
+                Enum.reduce(calculations_at_runtime, ash_query, fn calculation, ash_query ->
                   if calculation.select do
                     Ash.Query.select(ash_query, calculation.select || [])
                   else
@@ -690,7 +701,7 @@ defmodule Ash.Actions.Read do
                      add_calculations(
                        query,
                        ash_query,
-                       Enum.map(calculations_in_query, &elem(&1, 1))
+                       calculations_in_query
                      ),
                    {:ok, query} <-
                      Ash.DataLayer.filter(
@@ -1072,58 +1083,178 @@ defmodule Ash.Actions.Read do
     end
   end
 
-  defp add_calculation_values(results, query, calculations) do
+  defp calculation_request(
-    {can_be_runtime, require_query} =
+         all_calcs,
-      calculations
+         resource,
-      |> Enum.map(&elem(&1, 1))
+         api,
-      |> Enum.split_with(fn calculation ->
+         action,
-        :erlang.function_exported(calculation.module, :calculate, 3)
+         error_path,
+         path,
+         results,
+         calculation,
+         query
+       ) do
+    dependencies =
+      query
+      |> calculation.module.load(calculation.opts, calculation.context)
+      |> Ash.Actions.Helpers.validate_calculation_load!(calculation.module)
+      |> Enum.map(fn
+        {key, _} ->
+          key
+
+        key ->
+          key
+      end)
+      |> Enum.filter(fn key ->
+        key in all_calcs
+      end)
+      |> Enum.map(fn key ->
+        path ++ [:calculation_results, key]
       end)
 
-    can_be_runtime
+    if function_exported?(calculation.module, :calculate, 3) do
-    |> Enum.reduce_while({:ok, %{}, require_query}, fn calculation,
+      Request.new(
-                                                       {:ok, calculation_results, require_query} ->
+        resource: resource,
-      case calculation.module.calculate(results, calculation.opts, calculation.context) do
+        api: api,
-        results when is_list(results) ->
+        action: action,
-          {:cont, {:ok, Map.put(calculation_results, calculation, results), require_query}}
+        error_path: error_path,
+        query: query,
+        authorize?: false,
+        async?: true,
+        data:
+          Request.resolve(dependencies, fn data ->
+            temp_results =
+              Enum.reduce(get_in(data, path ++ [:calculation_results]) || %{}, results, fn {key,
+                                                                                            config},
+                                                                                           results ->
+                add_calc_to_results(results, key, config[:data])
+              end)
 
-        :unknown ->
+            case calculation.module.calculate(temp_results, calculation.opts, calculation.context) do
-          {:cont, {:ok, calculation_results, [calculation | require_query]}}
+              {:ok, values} ->
+                {:ok,
+                 %{
+                   values: values,
+                   load?: not is_nil(calculation.load)
+                 }}
 
-        {:ok, results} ->
+              {:error, error} ->
-          {:cont, {:ok, Map.put(calculation_results, calculation, results), require_query}}
+                {:error, error}
 
-        {:error, error} ->
+              values ->
-          {:halt, {:error, error}}
+                {:ok,
-      end
+                 %{
-    end)
+                   values: values,
-    |> case do
+                   load?: not is_nil(calculation.load)
-      {:ok, calculation_results, require_query} ->
+                 }}
-        Enum.reduce(calculation_results, results, fn {calculation, values}, records ->
+            end
-          if calculation.load do
+          end),
-            :lists.zipwith(
+        path: path ++ [:calculation_results, calculation.name],
-              fn record, value -> Map.put(record, calculation.name, value) end,
+        name: "#{inspect(path)} calculation #{calculation.name}"
-              records,
+      )
-              values
+    else
-            )
+      Request.new(
-          else
+        resource: resource,
-            :lists.zipwith(
+        api: api,
-              fn record, value ->
+        action: action,
-                %{record | calculations: Map.put(record.calculations, calculation.name, value)}
+        error_path: error_path,
-              end,
+        query: query,
-              records,
+        authorize?: false,
-              values
+        async?: true,
-            )
+        data:
-          end
+          Request.resolve([], fn _ ->
-        end)
+            case run_calculation_query(results, [calculation], query) do
-        |> run_calculation_query(require_query, query)
+              {:ok, results_with_calc} ->
+                {:ok,
+                 %{
+                   values:
+                     Enum.map(results_with_calc, fn record ->
+                       if calculation.load do
+                         Map.get(record, calculation.name)
+                       else
+                         Map.get(record.calculations, calculation.name)
+                       end
+                     end),
+                   load?: not is_nil(calculation.load)
+                 }}
 
-      {:error, error} ->
+              other ->
-        {:error, error}
+                other
+            end
+          end),
+        path: path ++ [:calculation_results, calculation.name],
+        name: "#{inspect(path)} calculation #{calculation.name}"
+      )
     end
   end
 
-  defp run_calculation_query(results, [], _), do: {:ok, results}
+  defp add_calculation_values(
+         results,
+         resource,
+         api,
+         action,
+         error_path,
+         path,
+         query,
+         calculations,
+         :error
+       )
+       when calculations != [] do
+    {:requests,
+     Enum.map(
+       calculations,
+       &calculation_request(
+         calculations,
+         resource,
+         api,
+         action,
+         error_path,
+         path,
+         results,
+         &1,
+         query
+       )
+     )}
+  end
+
+  defp add_calculation_values(
+         results,
+         _resource,
+         _api,
+         _action,
+         _error_path,
+         _path,
+         _query,
+         _calculations,
+         calculation_values
+       ) do
+    if calculation_values == :error do
+      {:ok, results}
+    else
+      {:ok,
+       Enum.reduce(calculation_values, results, fn {name, config}, results ->
+         add_calc_to_results(results, name, config[:data])
+       end)}
+    end
+  end
+
+  defp add_calc_to_results(results, name, config) do
+    if config[:load?] do
+      :lists.zipwith(
+        fn record, value -> Map.put(record, name, value) end,
+        results,
+        config[:values]
+      )
+    else
+      :lists.zipwith(
+        fn record, value ->
+          %{record | calculations: Map.put(record.calculations, name, value)}
+        end,
+        results,
+        config[:values]
+      )
+    end
+  end
 
   defp run_calculation_query(results, calculations, query) do
     pkey = Ash.Resource.Info.primary_key(query.resource)
@@ -1141,38 +1272,11 @@ defmodule Ash.Actions.Read do
          query <- Ash.Query.filter(query, ^[or: pkey_filter]),
          {:ok, data_layer_query} <- Ash.Query.data_layer_query(query),
          {:ok, data_layer_query} <-
-           add_calculations(data_layer_query, query, calculations),
+           add_calculations(data_layer_query, query, calculations) do
-         {:ok, calculation_results} <-
+      Ash.DataLayer.run_query(
-           Ash.DataLayer.run_query(
+        data_layer_query,
-             data_layer_query,
+        query.resource
-             query.resource
+      )
-           ) do
-      results_with_calculations =
-        results
-        |> Enum.map(fn result ->
-          result_pkey = Map.take(result, pkey)
-          match = Enum.find(calculation_results, &(Map.take(&1, pkey) == result_pkey))
-
-          if match do
-            Enum.reduce(calculations, result, fn calculation, result ->
-              if calculation.load do
-                Map.put(result, calculation.load, Map.get(match, calculation.load))
-              else
-                Map.update!(result, :calculations, fn calculations ->
-                  Map.put(
-                    calculations,
-                    calculation.name,
-                    Map.get(match, :calculations)[calculation.name]
-                  )
-                end)
-              end
-            end)
-          else
-            result
-          end
-        end)
-
-      {:ok, results_with_calculations}
     end
   end
 

lib/ash/actions/sort.ex

@@ -208,7 +208,9 @@ defmodule Ash.Actions.Sort do
              module,
              opts,
              type,
-             Map.put(input, :context, context)
+             Map.put(input, :context, context),
+             calc.filterable?,
+             calc.load
            ) do
       {sorts ++ [{calc, order}], errors}
     else

lib/ash/calculation.ex

@@ -15,7 +15,7 @@ defmodule Ash.Calculation do
 
       def describe(opts), do: "##{inspect(__MODULE__)}<#{inspect(opts)}>"
 
-      def load(query, _opts, _context), do: query
+      def load(_query, _opts, _context), do: []
 
       def select(_query, _opts, _context), do: []
 
@@ -28,7 +28,7 @@ defmodule Ash.Calculation do
   @callback calculate([Ash.Resource.record()], Keyword.t(), map) ::
               {:ok, [term]} | [term] | {:error, term} | :unknown
   @callback expression(Keyword.t(), map) :: any
-  @callback load(Ash.Query.t(), Keyword.t(), map) :: Ash.Query.t()
+  @callback load(Ash.Query.t(), Keyword.t(), map) :: Keyword.t()
   @callback select(Ash.Query.t(), Keyword.t(), map) :: list(atom)
 
   @optional_callbacks expression: 2, calculate: 3

lib/ash/filter/filter.ex

@@ -1979,7 +1979,8 @@ defmodule Ash.Filter do
                  opts,
                  resource_calculation.type,
                  Map.put(args, :context, context.query_context),
-                 resource_calculation.filterable?
+                 resource_calculation.filterable?,
+                 resource_calculation.load
                ) do
           case parse_predicates(nested_statement, calculation, context) do
             {:ok, nested_statement} ->
@@ -2149,7 +2150,8 @@ defmodule Ash.Filter do
                  opts,
                  resource_calculation.type,
                  Map.put(args, :context, Map.get(context, :query_context, %{})),
-                 resource_calculation.filterable?
+                 resource_calculation.filterable?,
+                 resource_calculation.load
                ) do
           {:ok,
            %Ref{
@@ -2252,7 +2254,8 @@ defmodule Ash.Filter do
                      opts,
                      resource_calculation.type,
                      Map.put(args, :context, context.query_context),
-                     resource_calculation.filterable?
+                     resource_calculation.filterable?,
+                     resource_calculation.load
                    ) do
               {:ok, %{ref | attribute: calculation, resource: related}}
             else

lib/ash/query/calculation.ex

@@ -8,6 +8,7 @@ defmodule Ash.Query.Calculation do
     :load,
     :type,
     context: %{},
+    required_loads: [],
     select: [],
     sequence: 0,
     allow_async?: false,
@@ -16,7 +17,7 @@ defmodule Ash.Query.Calculation do
 
   @type t :: %__MODULE__{}
 
-  def new(name, module, opts, type, context \\ %{}, filterable? \\ true) do
+  def new(name, module, opts, type, context \\ %{}, filterable? \\ true, required_loads \\ []) do
     case module.init(opts) do
       {:ok, opts} ->
         {:ok,
@@ -26,6 +27,7 @@ defmodule Ash.Query.Calculation do
            type: type,
            opts: opts,
            context: context,
+           required_loads: required_loads,
            filterable?: filterable?
          }}
 

lib/ash/query/query.ex

@@ -939,7 +939,9 @@ defmodule Ash.Query do
                      module,
                      opts,
                      resource_calculation.type,
-                     Map.put(args, :context, query.context)
+                     Map.put(args, :context, query.context),
+                     resource_calculation.filterable?,
+                     resource_calculation.load
                    ) do
               fields_to_select =
                 resource_calculation.select
@@ -954,11 +956,14 @@ defmodule Ash.Query do
               }
 
               query =
-                query
+                Ash.Query.load(
-                |> module.load(
+                  query,
-                  opts,
+                  module.load(
-                  calculation.context
+                    query,
-                  |> Map.put(:context, query.context)
+                    opts,
+                    Map.put(calculation.context, :context, query.context)
+                  )
+                  |> Ash.Actions.Helpers.validate_calculation_load!(module)
                 )
 
               query
@@ -1038,7 +1043,9 @@ defmodule Ash.Query do
                  module,
                  opts,
                  resource_calculation.type,
-                 Map.put(args, :context, query.context)
+                 Map.put(args, :context, query.context),
+                 resource_calculation.filterable?,
+                 resource_calculation.load
                ) do
           calculation = %{calculation | load: field}
 
@@ -1048,11 +1055,14 @@ defmodule Ash.Query do
             |> Enum.concat(module.select(query, opts, calculation.context) || [])
 
           query =
-            query
+            Ash.Query.load(
-            |> module.load(
+              query,
-              opts,
+              module.load(
-              calculation.context
+                query,
-              |> Map.put(:context, query.context)
+                opts,
+                Map.put(calculation.context, :context, query.context)
+              )
+              |> Ash.Actions.Helpers.validate_calculation_load!(module)
             )
             |> Ash.Query.load(resource_calculation.load || [])
 
@@ -1477,8 +1487,6 @@ defmodule Ash.Query do
 
   The `module_and_opts` argument accepts either a `module` or a `{module, opts}`. For more information
   on what that module should look like, see `Ash.Calculation`.
-
-  More features for calculations, like passing anonymous functions, will be supported in the future.
   """
   def calculate(query, name, module_and_opts, type, context \\ %{}) do
     query = to_query(query)
@@ -1494,11 +1502,14 @@ defmodule Ash.Query do
         fields_to_select = module.select(query, opts, calculation.context) || []
 
         query =
-          query
+          Ash.Query.load(
-          |> module.load(
+            query,
-            opts,
+            module.load(
-            calculation.context
+              query,
-            |> Map.put(:context, query.context)
+              opts,
+              Map.put(calculation.context, :context, query.context)
+            )
+            |> Ash.Actions.Helpers.validate_calculation_load!(module)
           )
 
         calculation = %{calculation | select: fields_to_select}
@@ -1832,11 +1843,14 @@ defmodule Ash.Query do
             calculation = %{calculation | load: name, select: fields_to_select}
 
             query =
-              query
+              Ash.Query.load(
-              |> module.load(
+                query,
-                opts,
+                module.load(
-                calculation.context
+                  query,
-                |> Map.put(:context, query.context)
+                  opts,
+                  Map.put(calculation.context, :context, query.context)
+                )
+                |> Ash.Actions.Helpers.validate_calculation_load!(module)
               )
 
             Ash.Query.load(query, resource_load)

lib/ash/resource/calculation/expression.ex

@@ -101,12 +101,10 @@ defmodule Ash.Resource.Calculation.Expression do
           end)
           |> Enum.map(& &1.name)
 
-        names = Enum.uniq(aggs_from_calcs ++ aggs_from_this_calc)
+        Enum.uniq(aggs_from_calcs ++ aggs_from_this_calc)
-
-        Ash.Query.load(query, names)
 
       {:error, _} ->
-        query
+        []
     end
   end