From d568ac18a887a8ebd261c24d3e510817883063f7 Mon Sep 17 00:00:00 2001
From: Zach Daniel <zach@zachdaniel.dev>
Date: Fri, 4 Aug 2023 18:53:09 -0400
Subject: [PATCH] improvement: support new `from_many?` option improvement:
 subquery after distinct to handle distinct

---
 lib/data_layer.ex | 10 +++++++++-
 lib/join.ex       |  3 ++-
 2 files changed, 11 insertions(+), 2 deletions(-)

diff --git a/lib/data_layer.ex b/lib/data_layer.ex
index 5a4b486..5f433eb 100644
--- a/lib/data_layer.ex
+++ b/lib/data_layer.ex
@@ -1856,7 +1856,15 @@ defmodule AshPostgres.DataLayer do
 
   @impl true
   def lock(query, :for_update, _) do
-    {:ok, Ecto.Query.lock(query, [{^0, a}], fragment("FOR UPDATE OF ?", a))}
+    if query.distinct do
+      new_query =
+        Ecto.Query.lock(%{query | distinct: nil}, [{^0, a}], fragment("FOR UPDATE OF ?", a))
+
+      q = from(row in subquery(new_query), [])
+      {:ok, %{q | distinct: query.distinct}}
+    else
+      {:ok, Ecto.Query.lock(query, [{^0, a}], fragment("FOR UPDATE OF ?", a))}
+    end
   end
 
   @locks [
diff --git a/lib/join.ex b/lib/join.ex
index 87bd6ca..b67c74d 100644
--- a/lib/join.ex
+++ b/lib/join.ex
@@ -430,7 +430,8 @@ defmodule AshPostgres.Join do
   def get_binding(_, _, _, _), do: nil
 
   defp add_distinct(relationship, _join_type, joined_query) do
-    if !joined_query.__ash_bindings__.in_group? && relationship.cardinality == :many &&
+    if !joined_query.__ash_bindings__.in_group? &&
+         (relationship.cardinality == :many || Map.get(relationship, :from_many?)) &&
          !joined_query.distinct do
       from(row in joined_query,
         distinct: ^Ash.Resource.Info.primary_key(joined_query.__ash_bindings__.resource)