improvement: prune calculations made unnecessary by field policies

closes #1356

lib/ash/actions/read/calculations.ex

@@ -1731,7 +1731,17 @@ defmodule Ash.Actions.Read.Calculations do
   end
 
   defp find_equivalent_calculation(query, calculation, authorize?) do
-    if !authorize? || match?({:__calc_dep__, _}, calculation.name) do
+    reusable? =
+      if authorize? && calculation.module.has_expression?() do
+        calculation.module.expression(calculation.opts, calculation.context)
+        |> Ash.Filter.list_refs(false, false, true, true)
+        |> Enum.any?(&(&1.relationship_path != []))
+        |> Kernel.not()
+      else
+        true
+      end
+
+    if reusable? do
       Enum.find_value(query.calculations, :error, fn {_, other_calc} ->
         if other_calc.module == calculation.module and other_calc.opts == calculation.opts and
              other_calc.context.arguments == calculation.context.arguments do
@@ -1891,10 +1901,11 @@ defmodule Ash.Actions.Read.Calculations do
   # Deselect fields that we know statically cannot be seen
   # The field may be reselected later as a calculation dependency
   # this is an optimization not a guarantee
-  def deselect_known_forbidden_fields(ash_query, calculations_at_runtime) do
+  def deselect_known_forbidden_fields(ash_query, calculations_at_runtime, calculations_in_query) do
     depended_on_fields = ash_query.context[:private][:depended_on_fields] || []
 
     calculations_at_runtime
+    |> Enum.concat(calculations_in_query)
     |> Enum.reduce([], fn
       %{
         name: {:__ash_fields_are_visible__, fields},
@@ -1927,10 +1938,12 @@ defmodule Ash.Actions.Read.Calculations do
     end)
     |> Enum.uniq()
     |> Kernel.--(depended_on_fields)
-    |> then(&unload_forbidden_fields(ash_query, &1))
+    |> then(
+      &unload_forbidden_fields(ash_query, &1, calculations_at_runtime, calculations_in_query)
+    )
   end
 
-  defp unload_forbidden_fields(ash_query, fields) do
+  defp unload_forbidden_fields(ash_query, fields, calculations_at_runtime, calculations_in_query) do
     fields
     |> Enum.group_by(fn field ->
       cond do
@@ -1944,17 +1957,17 @@ defmodule Ash.Actions.Read.Calculations do
           :calculation
       end
     end)
-    |> Enum.reduce(ash_query, fn
-      {:attribute, fields}, ash_query ->
-        ash_query
-        |> Ash.Query.deselect(fields)
-        |> unload_attribute_calculations(fields)
+    |> Enum.reduce({ash_query, calculations_at_runtime, calculations_in_query}, fn
+      {:attribute, fields}, {ash_query, calculations_at_runtime, calculations_in_query} ->
+        {ash_query
+         |> Ash.Query.deselect(fields)
+         |> unload_attribute_calculations(fields), calculations_at_runtime, calculations_in_query}
 
-      {:aggregate, fields}, ash_query ->
-        unload_aggregates(ash_query, fields)
+      {:aggregate, fields}, {ash_query, calculations_at_runtime, calculations_in_query} ->
+        {unload_aggregates(ash_query, fields), calculations_at_runtime, calculations_in_query}
 
-      {:calculation, fields}, ash_query ->
-        unload_calculations(ash_query, fields)
+      {:calculation, fields}, {ash_query, calculations_at_runtime, calculations_in_query} ->
+        unload_calculations(ash_query, fields, calculations_at_runtime, calculations_in_query)
     end)
   end
 
@@ -1990,7 +2003,7 @@ defmodule Ash.Actions.Read.Calculations do
     %{ash_query | calculations: Map.drop(ash_query.calculations, drop)}
   end
 
-  defp unload_calculations(ash_query, fields) do
+  defp unload_calculations(ash_query, fields, calculations_at_runtime, calculations_in_query) do
     drop =
       ash_query.calculations
       |> Enum.flat_map(fn
@@ -2002,6 +2015,36 @@ defmodule Ash.Actions.Read.Calculations do
           end
       end)
 
-    %{ash_query | calculations: Map.drop(ash_query.calculations, drop)}
+    Enum.reduce(
+      drop,
+      {ash_query, calculations_at_runtime, calculations_in_query},
+      fn drop, {ash_query, calculations_at_runtime, calculations_in_query} ->
+        if Enum.any?(ash_query.context[:calculation_dependencies], fn {_source, dest} ->
+             drop in dest
+           end) do
+          {%{ash_query | calculations: Map.delete(ash_query.calculations, drop)},
+           calculations_at_runtime, calculations_in_query}
+        else
+          {%{ash_query | calculations: Map.delete(ash_query.calculations, drop)},
+           Enum.reject(calculations_at_runtime, &(&1.name == drop)),
+           Enum.reject(calculations_in_query, &(&1 == drop))}
+        end
+      end
+    )
+    |> remove_unreferenced_calculations()
+  end
+
+  defp remove_unreferenced_calculations(
+         {ash_query, calculations_at_runtime, calculations_in_query}
+       ) do
+    {ash_query, Enum.filter(calculations_at_runtime, &used?(ash_query, &1.name)),
+     Enum.filter(calculations_in_query, &used?(ash_query, &1.name))}
+  end
+
+  defp used?(ash_query, name) do
+    Map.has_key?(ash_query, name) ||
+      Enum.any?(ash_query.context[:calculation_dependencies], fn {_source, dest} ->
+        name in dest
+      end)
   end
 end

lib/ash/actions/read/read.ex

@@ -251,14 +251,19 @@ defmodule Ash.Actions.Read do
 
       {data_result, query_ran} =
         case data_result do
-          {:ok, result, query} -> {{:ok, result, nil}, query}
-          {:ok, result, count, query} -> {{:ok, result, count}, query}
-          {{:error, _, _} = error, query} -> {error, query}
-          {{:error, _} = error, query} -> {error, query}
-          other -> {other, query}
+          {:ok, _result, _count, _calculations_at_runtime, _calculations_in_query, query} =
+              data_result ->
+            {data_result, query}
+
+          {{:error, _} = data_result, query} ->
+            {data_result, query}
+
+          data_result ->
+            {data_result, query}
         end
 
-      with {:ok, data, count} <- data_result,
+      with {:ok, data, count, calculations_at_runtime, calculations_in_query, _query} <-
+             data_result,
            data = add_tenant(data, query),
            {:ok, data} <-
              load_through_attributes(
@@ -358,12 +363,13 @@ defmodule Ash.Actions.Read do
          },
          pre_authorization_query <- query,
          {:ok, query} <- authorize_query(query, opts) do
-      maybe_in_transaction(query, opts, fn ->
+      maybe_in_transaction(query, opts, fn notify_callback ->
         with query_before_pagination <- query,
-             query <-
+             {query, calculations_at_runtime, calculations_in_query} <-
                Ash.Actions.Read.Calculations.deselect_known_forbidden_fields(
                  query,
-                 calculations_at_runtime ++ calculations_in_query
+                 calculations_at_runtime,
+                 calculations_in_query
                ),
              {:ok, data_layer_calculations} <- hydrate_calculations(query, calculations_in_query),
              {:ok, query} <-
@@ -480,10 +486,12 @@ defmodule Ash.Actions.Read do
              {:ok, results} <- run_authorize_results(query, results),
              {:ok, results, after_notifications} <- run_after_action(query, results),
              {:ok, count} <- maybe_await(count) do
-          {:ok, results, count, query, before_notifications ++ after_notifications}
+          notify_callback.(query, before_notifications ++ after_notifications)
+          {:ok, results, count, calculations_at_runtime, calculations_in_query, query}
         else
           {%{valid?: false} = query, before_notifications} ->
-            {{:error, query, before_notifications}, query}
+            notify_callback.(before_notifications)
+            {{:error, query}, query}
 
           {{:error, %Ash.Query{} = query}, query} ->
             {:error, query}
@@ -790,7 +798,9 @@ defmodule Ash.Actions.Read do
         query.action.transaction? ->
           Ash.DataLayer.transaction(
             [query.resource | query.action.touches_resources],
-            func,
+            fn ->
+              func.(&notify_or_store(&1, &2, notify?))
+            end,
             query.timeout,
             %{
               type: :read,
@@ -802,36 +812,25 @@ defmodule Ash.Actions.Read do
               data_layer_context: query.context[:data_layer]
             }
           )
+          |> case do
+            {:error, {:error, error}} -> {:error, error}
+            {:error, error} -> {:error, error}
+            {:ok, result} -> {:ok, result}
+          end
 
         query.timeout ->
-          {:ok,
-           Ash.ProcessHelpers.task_with_timeout(
-             func,
-             query.resource,
-             query.timeout,
-             "#{inspect(query.resource)}.#{query.action.name}",
-             opts[:tracer]
-           )}
+          Ash.ProcessHelpers.task_with_timeout(
+            fn ->
+              func.(&notify_or_store(&1, &2, notify?))
+            end,
+            query.resource,
+            query.timeout,
+            "#{inspect(query.resource)}.#{query.action.name}",
+            opts[:tracer]
+          )
 
         true ->
-          {:ok, func.()}
-      end
-      |> case do
-        {:ok, {:ok, result, count, query_ran, notifications}} ->
-          notify_or_store(query, notifications, notify?)
-
-          {:ok, result, count, query_ran}
-
-        {:ok, {:error, error, notifications}} ->
-          notify_or_store(query, notifications, notify?)
-
-          {:error, error}
-
-        {:ok, value} ->
-          value
-
-        other ->
-          other
+          func.(&notify_or_store(&1, &2, notify?))
       end
     after
       if notify? do
@@ -879,22 +878,24 @@ defmodule Ash.Actions.Read do
        ) do
     must_be_reselected = List.wrap(query.select) -- Ash.Resource.Info.primary_key(query.resource)
 
-    query =
+    {query, calculations_at_runtime, calculations_in_query} =
       Ash.Actions.Read.Calculations.deselect_known_forbidden_fields(
         query,
-        calculations_at_runtime ++ calculations_in_query
+        calculations_at_runtime,
+        calculations_in_query
       )
 
     if missing_pkeys? ||
          (Enum.empty?(must_be_reselected) && Enum.empty?(query.aggregates) &&
             Enum.empty?(calculations_in_query)) do
-      {:ok, initial_data, query}
+      {:ok, initial_data, 0, calculations_at_runtime, calculations_in_query, query}
     else
       reselect_and_load(
         initial_data,
         query,
         must_be_reselected,
         calculations_in_query,
+        calculations_at_runtime,
         opts
       )
     end
@@ -905,6 +906,7 @@ defmodule Ash.Actions.Read do
          query,
          must_be_reselected,
          calculations_in_query,
+         calculations_at_runtime,
          opts
        ) do
     primary_key = Ash.Resource.Info.primary_key(query.resource)
@@ -1004,7 +1006,7 @@ defmodule Ash.Actions.Read do
       |> compute_expression_at_runtime_for_missing_records(query, data_layer_calculations)
       |> case do
         {:ok, result} ->
-          {:ok, result, 0, query}
+          {:ok, result, 0, calculations_at_runtime, calculations_in_query, query}
 
         {:error, error} ->
           {:error, error}

test/policy/field_policy_pruning_test.exs

@@ -1,12 +1,17 @@
-defmodule Ash.Test.Policy.FieldPolicyTest do
+defmodule Ash.Test.Policy.FieldPolicyPruningTest do
   @doc false
   use ExUnit.Case
 
-  defmodule App.Core.TestResource do
+  defmodule TestResource do
     use Ash.Resource,
-      domain: App.Core,
+      domain: Ash.Test.Domain,
+      data_layer: Ash.DataLayer.Ets,
       authorizers: [Ash.Policy.Authorizer]
 
+    ets do
+      private? true
+    end
+
     policies do
       policy always() do
         authorize_if always()
@@ -18,7 +23,11 @@ defmodule Ash.Test.Policy.FieldPolicyTest do
         authorize_if always()
       end
 
-      field_policy :graphs do
+      field_policy :calc do
+        forbid_if always()
+      end
+
+      field_policy :calc2 do
         forbid_if always()
       end
     end
@@ -28,13 +37,21 @@ defmodule Ash.Test.Policy.FieldPolicyTest do
     end
 
     actions do
-      read :read do
-        primary? true
-      end
+      defaults [:create, :read, :update, :destroy]
     end
 
     calculations do
-      calculate :graphs, :map do
+      calculate :calc2, :map do
+        public? true
+
+        calculation fn records, _ ->
+          Enum.map(records, & &1.id)
+        end
+      end
+
+      calculate :calc, :map do
+        load :calc2
+
         calculation fn records, _ ->
           raise "shouldn't get here!"
         end
@@ -45,5 +62,10 @@ defmodule Ash.Test.Policy.FieldPolicyTest do
   end
 
   test "field policies prune unnecessary calculations" do
+    Ash.create!(TestResource, %{}, authorize?: false)
+
+    TestResource
+    |> Ash.Query.load([:calc, :calc2])
+    |> Ash.read!()
   end
 end