Allow take with tuple fragment sources (#4666)

Author: greg-rychlewski

integration_test/cases/repo.exs

@@ -782,6 +782,19 @@ defmodule Ecto.Integration.RepoTest do
     assert_raise Ecto.NoResultsError, fn -> query |> last |> TestRepo.one! end
   end
 
+  test "fragment source mapped to schema" do
+    query = from f in {fragment("select 1 as num"), Barebone}
+    assert %Barebone{num: 1} = TestRepo.one(query)
+  end
+
+  test "fragment source mapped to schema with take" do
+    query = from f in {fragment("select 1 as visits"), Post}, select: struct(f, [:visits])
+    assert %Post{visits: 1} = TestRepo.one(query)
+
+    query = from f in {fragment("select 1 as visits"), Post}, select: map(f, [:visits])
+    assert TestRepo.one(query) == %{visits: 1}
+  end
+
   test "exists?" do
     TestRepo.insert!(%Post{title: "1", visits: 2})
     TestRepo.insert!(%Post{title: "2", visits: 1})

lib/ecto/query/planner.ex

@@ -2220,11 +2220,15 @@ defmodule Ecto.Query.Planner do
         error!(query, "struct/2 in select expects a source with a schema")
 
       {{:ok, {kind, fields}}, {source, schema, prefix}} when is_binary(source) ->
-        dumper = if schema, do: schema.__schema__(:dump), else: %{}
+        {types, fields} = select_dump_for_schema(schema, List.wrap(fields), ix, drop)
         schema = if kind == :map, do: nil, else: schema
-        {types, fields} = select_dump(List.wrap(fields), dumper, ix, drop)
         {{:source, {source, schema}, prefix || query.prefix, types}, fields}
 
+      {{:ok, {kind, fields}}, {{:fragment, _, _} = source, schema, prefix}}  ->
+        {types, fields} = select_dump_for_schema(schema, List.wrap(fields), ix, drop)
+        schema = if kind == :map, do: nil, else: schema
+        {{:source, {source, schema}, prefix, types}, fields}
+
       {{:ok, {_, fields}}, _} ->
         {{:map, Enum.map(fields, &{&1, {:value, :any}})},
          Enum.map(fields, &select_field(&1, ix, :always))}
@@ -2258,6 +2262,11 @@ defmodule Ecto.Query.Planner do
     end
   end
 
+  defp select_dump_for_schema(schema, fields, ix, drop) do
+    dumper = if schema, do: schema.__schema__(:dump), else: %{}
+    select_dump(List.wrap(fields), dumper, ix, drop)
+  end
+
   defp select_dump(fields, dumper, ix, drop) do
     fields
     |> Enum.reverse()

test/ecto/query/planner_test.exs

@@ -958,6 +958,23 @@ defmodule Ecto.Query.PlannerTest do
     assert [:all, {:from, {{:fragment, _, _}, Barebone, _, _}, []}] = cache_key
   end
 
+  test "plan: tuple source with fragment and take" do
+    {query, cast_params, dump_params, cache_key} =
+      plan(from f in {fragment("? as text", ^"hi"), Post}, select: struct(f, [:text]))
+
+    assert query.select.take == %{0 => {:struct, [:text]}}
+    assert {{{:fragment, [], _}, Post, nil}} = query.sources
+    assert cast_params == ["hi"]
+    assert dump_params == ["hi"]
+
+    assert [
+             :all,
+             {:take, %{0 => {:struct, [:text]}}},
+             {:from, {{:fragment, _, _}, Post, _, _}, []},
+             {:select, {:&, [], [0]}}
+           ] = cache_key
+  end
+
   describe "plan: CTEs" do
     test "with uncacheable queries are uncacheable" do
       {_, _, _, cache} =
@@ -2601,6 +2618,18 @@ defmodule Ecto.Query.PlannerTest do
     assert query.select.fields == [{{:., [writable: :always], [{:&, [], [0]}, :num]}, [], []}]
   end
 
+  test "normalize: tuple source with fragment and take" do
+    {query, _, _, select} =
+      normalize_with_params(
+        from f in {fragment("? as text", ^"hi"), Post}, select: struct(f, [:text])
+      )
+
+    %{from: {_, {:source, {{:fragment, _, _}, Post}, nil, types}}} = select
+    assert types == [text: :string]
+    assert {{:fragment, _, _}, Post} = query.from.source
+    assert query.select.fields == [{{:., [writable: :always], [{:&, [], [0]}, :text]}, [], []}]
+  end
+
   describe "normalize: subqueries in boolean expressions" do
     test "replaces {:subquery, index} with an Ecto.SubQuery struct" do
       subquery = from(p in Post, select: p.visits)