improvement: support :no_rollback and return_query/2 callback

lib/ash/actions/create/bulk.ex

@@ -903,7 +903,7 @@ defmodule Ash.Actions.Create.Bulk do
                         must_return_records_for_changes?,
                     tenant: opts[:tenant]
                   })
-                  |> Ash.Actions.Helpers.rollback_if_in_transaction(nil)
+                  |> Ash.Actions.Helpers.rollback_if_in_transaction(resource, nil)
                   |> Enum.flat_map(fn
                     {:ok, result} ->
                       [result]
@@ -927,7 +927,7 @@ defmodule Ash.Actions.Create.Bulk do
                       tracer: opts[:tracer],
                       api: api
                     })
-                    |> Ash.Actions.Helpers.rollback_if_in_transaction(nil)
+                    |> Ash.Actions.Helpers.rollback_if_in_transaction(resource, nil)
 
                   case result do
                     {:ok, result} ->
@@ -965,7 +965,7 @@ defmodule Ash.Actions.Create.Bulk do
                       tenant: opts[:tenant]
                     }
                   )
-                  |> Ash.Actions.Helpers.rollback_if_in_transaction(nil)
+                  |> Ash.Actions.Helpers.rollback_if_in_transaction(resource, nil)
                 else
                   [changeset] = batch
                   upsert? = opts[:upsert?] || action.upsert? || false

lib/ash/actions/create/create.ex

@@ -337,7 +337,10 @@ defmodule Ash.Actions.Create do
                         opts[:upsert?] ->
                           changeset.resource
                           |> Ash.DataLayer.upsert(changeset, upsert_keys)
-                          |> Ash.Actions.Helpers.rollback_if_in_transaction(changeset)
+                          |> Ash.Actions.Helpers.rollback_if_in_transaction(
+                            changeset.resource,
+                            changeset
+                          )
                           |> add_tenant(changeset)
                           |> manage_relationships(api, changeset,
                             actor: opts[:actor],
@@ -348,7 +351,10 @@ defmodule Ash.Actions.Create do
                         true ->
                           changeset.resource
                           |> Ash.DataLayer.create(changeset)
-                          |> Ash.Actions.Helpers.rollback_if_in_transaction(changeset)
+                          |> Ash.Actions.Helpers.rollback_if_in_transaction(
+                            changeset.resource,
+                            changeset
+                          )
                           |> add_tenant(changeset)
                           |> manage_relationships(api, changeset,
                             actor: opts[:actor],

lib/ash/actions/destroy/destroy.ex

@@ -182,7 +182,7 @@ defmodule Ash.Actions.Destroy do
                 else
                   changeset.resource
                   |> Ash.DataLayer.destroy(changeset)
-                  |> Ash.Actions.Helpers.rollback_if_in_transaction(changeset)
+                  |> Ash.Actions.Helpers.rollback_if_in_transaction(changeset.resource, changeset)
                   |> case do
                     :ok ->
                       {:ok, data} = Ash.Changeset.apply_attributes(changeset, force?: true)

lib/ash/actions/helpers.ex

@@ -3,23 +3,28 @@ defmodule Ash.Actions.Helpers do
   require Logger
   require Ash.Flags
 
-  def rollback_if_in_transaction({:error, error}, changeset) do
+  def rollback_if_in_transaction({:error, error}, resource, changeset) do
-    if Ash.DataLayer.in_transaction?(changeset.resource) do
+    if Ash.DataLayer.in_transaction?(resource) do
-      if changeset do
+      case changeset do
-        Ash.DataLayer.rollback(changeset.resource, Ash.Changeset.add_error(changeset, error))
+        %Ash.Changeset{} = changeset ->
-      else
+          Ash.DataLayer.rollback(resource, Ash.Changeset.add_error(changeset, error))
-        Ash.DataLayer.rollback(changeset.resource, Ash.Error.to_error_class(error))
+
+        %Ash.Query{} = query ->
+          Ash.DataLayer.rollback(resource, Ash.Query.add_error(query, error))
+
+        _ ->
+          Ash.DataLayer.rollback(resource, Ash.Error.to_error_class(error))
       end
     else
       {:error, error}
     end
   end
 
-  def rollback_if_in_transaction({:error, :no_rollback, error}, _changeset) do
+  def rollback_if_in_transaction({:error, :no_rollback, error}, _, _changeset) do
     {:error, error}
   end
 
-  def rollback_if_in_transaction(success, _), do: success
+  def rollback_if_in_transaction(success, _, _), do: success
 
   def validate_calculation_load!(%{__struct__: Ash.Query}, module) do
     raise """

lib/ash/actions/read/read.ex

@@ -1405,7 +1405,10 @@ defmodule Ash.Actions.Read do
                    ])
                    |> Map.put(:context, ash_query.context)
                    |> Ash.Query.set_context(%{action: ash_query.action})
-                   |> Ash.Query.data_layer_query(only_validate_filter?: true),
+                   |> Ash.Query.data_layer_query(
+                     only_validate_filter?: true,
+                     run_return_query?: false
+                   ),
                  {:ok, filter} <-
                    filter_with_related(
                      Enum.map(filter_requests, & &1.path),
@@ -1450,7 +1453,11 @@ defmodule Ash.Actions.Read do
                      ash_query.resource
                    ),
                  {:ok, query} <-
-                   Ash.DataLayer.sort(query, ash_query.sort, ash_query.resource),
+                   Ash.DataLayer.sort(
+                     query,
+                     ash_query.sort,
+                     ash_query.resource
+                   ),
                  {:ok, query} <-
                    Ash.DataLayer.distinct_sort(query, ash_query.distinct_sort, ash_query.resource),
                  {:ok, query} <-
@@ -2947,6 +2954,7 @@ defmodule Ash.Actions.Read do
     else
       query
       |> Ash.DataLayer.run_query(resource)
+      |> Helpers.rollback_if_in_transaction(ash_query.resource, ash_query)
       |> Helpers.select(ash_query)
       |> Helpers.load_runtime_types(ash_query, load_attributes?)
     end

lib/ash/actions/update/update.ex

@@ -272,7 +272,10 @@ defmodule Ash.Actions.Update do
 
                         changeset.resource
                         |> Ash.DataLayer.update(changeset)
-                        |> Ash.Actions.Helpers.rollback_if_in_transaction(changeset)
+                        |> Ash.Actions.Helpers.rollback_if_in_transaction(
+                          changeset.resource,
+                          changeset
+                        )
                         |> add_tenant(changeset)
                         |> manage_relationships(api, changeset,
                           actor: opts[:actor],

lib/ash/api/api.ex

@@ -1030,7 +1030,10 @@ defmodule Ash.Api do
                       |> Ash.Query.data_layer_query()
                       |> case do
                         {:ok, data_layer_query} ->
-                          case Ash.DataLayer.run_query(data_layer_query, query.resource) do
+                          data_layer_query
+                          |> Ash.DataLayer.run_query(query.resource)
+                          |> Ash.Actions.Helpers.rollback_if_in_transaction(query.resource, query)
+                          |> case do
                             {:ok, results} ->
                               if Enum.count(results) == Enum.count(data) do
                                 {:ok, true}
@@ -1065,7 +1068,10 @@ defmodule Ash.Api do
                     |> Ash.Query.data_layer_query()
                     |> case do
                       {:ok, data_layer_query} ->
-                        case Ash.DataLayer.run_query(data_layer_query, resource) do
+                        data_layer_query
+                        |> Ash.DataLayer.run_query(resource)
+                        |> Ash.Actions.Helpers.rollback_if_in_transaction(query.resource, query)
+                        |> case do
                           {:ok, [_]} ->
                             {:ok, true}
 

lib/ash/data_layer/data_layer.ex

@@ -378,13 +378,13 @@ defmodule Ash.DataLayer do
   end
 
   @spec update(Ash.Resource.t(), Ash.Changeset.t()) ::
-          {:ok, Ash.Resource.record()} | {:error, term}
+          {:ok, Ash.Resource.record()} | {:error, term} | {:error, :no_rollback, term}
   def update(resource, changeset) do
     Ash.DataLayer.data_layer(resource).update(resource, changeset)
   end
 
   @spec create(Ash.Resource.t(), Ash.Changeset.t()) ::
-          {:ok, Ash.Resource.record()} | {:error, term}
+          {:ok, Ash.Resource.record()} | {:error, term} | {:error, :no_rollback, term}
   def create(resource, changeset) do
     Ash.DataLayer.data_layer(resource).create(resource, changeset)
   end
@@ -409,11 +409,13 @@ defmodule Ash.DataLayer do
           :ok
           | {:ok, Enumerable.t(Ash.Resource.record())}
           | {:error, Ash.Error.t()}
+          | {:error, :no_rollback, Ash.Error.t()}
   def bulk_create(resource, changesets, options) do
     Ash.DataLayer.data_layer(resource).bulk_create(resource, changesets, options)
   end
 
-  @spec destroy(Ash.Resource.t(), Ash.Changeset.t()) :: :ok | {:error, term}
+  @spec destroy(Ash.Resource.t(), Ash.Changeset.t()) ::
+          :ok | {:error, term} | {:error, :no_rollback, term}
   def destroy(resource, changeset) do
     Ash.DataLayer.data_layer(resource).destroy(resource, changeset)
   end
@@ -615,7 +617,7 @@ defmodule Ash.DataLayer do
   end
 
   @spec run_query(data_layer_query(), central_resource :: Ash.Resource.t()) ::
-          {:ok, list(Ash.Resource.record())} | {:error, term}
+          {:ok, list(Ash.Resource.record())} | {:error, term} | {:error, :no_rollback, term}
   def run_query(query, central_resource) do
     Ash.DataLayer.data_layer(central_resource).run_query(query, central_resource)
   end

lib/ash/engine/request.ex

@@ -919,6 +919,7 @@ defmodule Ash.Engine.Request do
       {:ok, data_layer_query} ->
         data_layer_query
         |> Ash.DataLayer.run_query(request.resource)
+        |> Ash.Actions.Helpers.rollback_if_in_transaction(request.resource, new_query)
         |> case do
           {:ok, results} ->
             pkey = Ash.Resource.Info.primary_key(request.resource)
@@ -971,6 +972,10 @@ defmodule Ash.Engine.Request do
       {:ok, data_layer_query} ->
         data_layer_query
         |> Ash.DataLayer.run_query(request.resource)
+        |> Ash.Actions.Helpers.rollback_if_in_transaction(
+          request.resource,
+          query_with_pkey_filter
+        )
         |> case do
           {:ok, []} ->
             {:error,

lib/ash/query/aggregate.ex

@@ -588,7 +588,8 @@ defmodule Ash.Query.Aggregate do
                      Ash.DataLayer.run_query(
                        data_layer_query,
                        query.resource
-                     ) do
+                     )
+                     |> Ash.Actions.Helpers.rollback_if_in_transaction(query.resource, query) do
                 loaded_aggregates =
                   aggregates
                   |> Enum.map(& &1.load)

lib/ash/query/query.ex

@@ -2627,7 +2627,7 @@ defmodule Ash.Query do
          {:ok, query} <-
            Ash.DataLayer.offset(query, ash_query.offset, resource),
          {:ok, query} <- Ash.DataLayer.lock(query, ash_query.lock, resource),
-         {:ok, query} <- Ash.DataLayer.return_query(query, resource) do
+         {:ok, query} <- maybe_return_query(query, resource, opts) do
       if opts[:no_modify?] || !ash_query.action || !ash_query.action.modify_query do
         {:ok, query}
       else
@@ -2644,6 +2644,14 @@ defmodule Ash.Query do
     end
   end
 
+  defp maybe_return_query(query, resource, opts) do
+    if Keyword.get(opts, :run_return_query?, true) do
+      Ash.DataLayer.return_query(query, resource)
+    else
+      {:ok, query}
+    end
+  end
+
   defp add_tenant(query, ash_query) do
     with :context <- Ash.Resource.Info.multitenancy_strategy(ash_query.resource),
          tenant when not is_nil(tenant) <- ash_query.tenant,