fix: ensure calculations are selected on aggregate subqueries

Author: zachdaniel

lib/aggregate.ex

@@ -2504,7 +2504,10 @@ defmodule AshSql.Aggregate do
   def wrap_in_subquery_for_aggregates(query) do
     resource = query.__ash_bindings__.resource
     selected_by_default = Ash.Resource.Info.selected_by_default_attribute_names(resource)
-    selected_fields = extract_selected_fields(query, resource, selected_by_default)
+
+    selected_fields =
+      query.__ash_bindings__[:select] ||
+        extract_selected_fields(query, resource, selected_by_default)
 
     all_attr_names =
       resource
@@ -2541,7 +2544,14 @@ defmodule AshSql.Aggregate do
     subquery_query =
       from(row in subquery(query_with_all_attrs),
         as: ^query.__ash_bindings__.root_binding,
-        select: struct(row, ^selected_fields)
+        select:
+          struct(
+            row,
+            ^Enum.concat(
+              selected_fields,
+              query.__ash_bindings__[:select_calculations] || []
+            )
+          )
       )
 
     bindings_without_aggregates =
@@ -2560,7 +2570,6 @@ defmodule AshSql.Aggregate do
     Map.put(subquery_query, :__ash_bindings__, new_bindings)
   end
 
-  # Extract the fields that are actually selected (respects `take` clause)
   defp extract_selected_fields(
          %{select: %Ecto.Query.SelectExpr{expr: expr, take: take}},
          resource,

lib/calculation.ex

@@ -211,6 +211,7 @@ defmodule AshSql.Calculation do
         )
       end
 
-    Ecto.Query.select_merge(query, ^calcs)
+    query = Ecto.Query.select_merge(query, ^calcs)
+    put_in(query.__ash_bindings__[:select_calculations], Map.keys(calcs))
   end
 end

lib/distinct.ex

@@ -86,6 +86,14 @@ defmodule AshSql.Distinct do
                    fn ash_bindings ->
                      ash_bindings
                      |> Map.put(:__order__?, query.__ash_bindings__[:__order__?] || false)
+                     |> Map.put(
+                       :select_calculations,
+                       query.__ash_bindings__[:select_calculations]
+                     )
+                     |> Map.put(
+                       :select,
+                       query.__ash_bindings__[:select]
+                     )
                      |> Map.update(
                        :calculations_require_rewrite,
                        calculations_require_rewrite,

lib/query.ex

@@ -324,6 +324,8 @@ defmodule AshSql.Query do
                   %{
                     calculations_require_rewrite: calculations_require_rewrite,
                     aggregates_require_rewrite: aggregates_require_rewrite,
+                    select: query.__ash_bindings__[:select],
+                    select_calculations: query.__ash_bindings__[:select_calculations],
                     load_aggregates: query.__ash_bindings__[:load_aggregates]
                   },
                   fn _, v1, v2 -> Map.merge(v1, v2) end