From e0cbe591ee12aeafcdb5c907a5c5c68410302a95 Mon Sep 17 00:00:00 2001
From: Zach Daniel <zach@zachdaniel.dev>
Date: Sun, 4 Dec 2022 13:35:52 -0500
Subject: [PATCH] improvement: set context once, early

---
 lib/ash/actions/read.ex |  7 +------
 lib/ash/query/query.ex  | 12 ++++++------
 2 files changed, 7 insertions(+), 12 deletions(-)

diff --git a/lib/ash/actions/read.ex b/lib/ash/actions/read.ex
index d8ad90e0..cd44fa39 100644
--- a/lib/ash/actions/read.ex
+++ b/lib/ash/actions/read.ex
@@ -862,6 +862,7 @@ defmodule Ash.Actions.Read do
                 :distinct,
                 :select
               ])
+              |> Ash.Query.set_context(%{action: ash_query.action})
               |> Ash.Query.data_layer_query(only_validate_filter?: true)
 
             ash_query =
@@ -901,12 +902,6 @@ defmodule Ash.Actions.Read do
                      {path, ash_query.tenant, data}
                    ),
                  filter <- update_aggregate_filters(filter, data, path),
-                 {:ok, query} <-
-                   Ash.DataLayer.set_context(
-                     ash_query.resource,
-                     query,
-                     Map.put(ash_query.context, :action, ash_query.action)
-                   ),
                  {:ok, query} <-
                    Ash.DataLayer.select(
                      query,
diff --git a/lib/ash/query/query.ex b/lib/ash/query/query.ex
index 44c1aaaf..45da0196 100644
--- a/lib/ash/query/query.ex
+++ b/lib/ash/query/query.ex
@@ -1946,6 +1946,12 @@ defmodule Ash.Query do
     aggregates = Enum.uniq_by(filter_aggregates ++ sort_aggregates, & &1.name)
 
     with {:ok, query} <-
+           Ash.DataLayer.set_context(
+             resource,
+             query,
+             Map.put(ash_query.context, :action, ash_query.action)
+           ),
+         {:ok, query} <-
            add_tenant(query, ash_query),
          {:ok, query} <-
            add_aggregates(query, ash_query, aggregates),
@@ -1957,12 +1963,6 @@ defmodule Ash.Query do
            Ash.DataLayer.limit(query, ash_query.limit, resource),
          {:ok, query} <-
            Ash.DataLayer.offset(query, ash_query.offset, resource),
-         {:ok, query} <-
-           Ash.DataLayer.set_context(
-             resource,
-             query,
-             Map.put(ash_query.context, :action, ash_query.action)
-           ),
          {:ok, query} <- Ash.DataLayer.select(query, ash_query.select, ash_query.resource) do
       if opts[:no_modify?] || !ash_query.action || !ash_query.action.modify_query do
         {:ok, query}