Allow using type/2 with virtual fields

Author: v0idpwn

lib/ecto/query/planner.ex

@@ -1181,7 +1181,7 @@ defmodule Ecto.Query.Planner do
 
   defp prewalk({:type, _, [arg, type]}, kind, query, expr, acc, adapter) do
     {arg, acc} = prewalk(arg, kind, query, expr, acc, adapter)
-    type = field_type!(kind, query, expr, type)
+    type = field_type!(kind, query, expr, type, true)
     {%Ecto.Query.Tagged{value: arg, tag: type, type: Ecto.Type.type(type)}, acc}
   end
 
@@ -1700,25 +1700,31 @@ defmodule Ecto.Query.Planner do
     end
   end
 
-  defp field_type!(kind, query, expr, {composite, {ix, field}}) when is_integer(ix) do
-    {composite, type!(kind, query, expr, ix, field)}
+  defp field_type!(kind, query, expr, type, allow_virtuals? \\ false)
+
+  defp field_type!(kind, query, expr, {composite, {ix, field}}, allow_virtuals?) when is_integer(ix) do
+    {composite, type!(kind, query, expr, ix, field, allow_virtuals?)}
   end
-  defp field_type!(kind, query, expr, {ix, field}) when is_integer(ix) do
-    type!(kind, query, expr, ix, field)
+
+  defp field_type!(kind, query, expr, {ix, field}, allow_virtuals?) when is_integer(ix) do
+    type!(kind, query, expr, ix, field, allow_virtuals?)
   end
-  defp field_type!(_kind, _query, _expr, type) do
+
+  defp field_type!(_kind, _query, _expr, type, _) do
     type
   end
 
-  defp type!(_kind, _query, _expr, nil, _field), do: :any
+  defp type!(kind, query, expr, schema, field, allow_virtuals? \\ false)
 
-  defp type!(kind, query, expr, ix, field) when is_integer(ix) do
+  defp type!(_kind, _query, _expr, nil, _field, _allow_virtuals?), do: :any
+
+  defp type!(kind, query, expr, ix, field, allow_virtuals?) when is_integer(ix) do
     case get_source!(kind, query, ix) do
       {:fragment, _, _} ->
         :any
 
       {_, schema, _} ->
-        type!(kind, query, expr, schema, field)
+        type!(kind, query, expr, schema, field, allow_virtuals?)
 
       %Ecto.SubQuery{select: select} ->
         case subquery_type_for(select, field) do
@@ -1728,11 +1734,14 @@ defmodule Ecto.Query.Planner do
     end
   end
 
-  defp type!(kind, query, expr, schema, field) when is_atom(schema) do
+  defp type!(kind, query, expr, schema, field, allow_virtuals?) when is_atom(schema) do
     cond do
       type = schema.__schema__(:type, field) ->
         type
 
+      type = allow_virtuals? && schema.__schema__(:virtual_type, field) ->
+        type
+
       Map.has_key?(schema.__struct__(), field) ->
         error! query, expr, "field `#{field}` in `#{kind}` is a virtual field in schema #{inspect schema}"
 

lib/ecto/schema.ex

@@ -473,6 +473,7 @@ defmodule Ecto.Schema do
 
       Module.register_attribute(__MODULE__, :ecto_primary_keys, accumulate: true)
       Module.register_attribute(__MODULE__, :ecto_fields, accumulate: true)
+      Module.register_attribute(__MODULE__, :ecto_virtual_fields, accumulate: true)
       Module.register_attribute(__MODULE__, :ecto_query_fields, accumulate: true)
       Module.register_attribute(__MODULE__, :ecto_field_sources, accumulate: true)
       Module.register_attribute(__MODULE__, :ecto_assocs, accumulate: true)
@@ -600,6 +601,7 @@ defmodule Ecto.Schema do
         autoupdate = @ecto_autoupdate |> Enum.reverse
         fields = @ecto_fields |> Enum.reverse
         query_fields = @ecto_query_fields |> Enum.reverse
+        virtual_fields = @ecto_virtual_fields |> Enum.reverse
         field_sources = @ecto_field_sources |> Enum.reverse
         assocs = @ecto_assocs |> Enum.reverse
         embeds = @ecto_embeds |> Enum.reverse
@@ -629,6 +631,7 @@ defmodule Ecto.Schema do
         def __schema__(:autoupdate), do: unquote(Macro.escape(autoupdate))
         def __schema__(:loaded), do: unquote(Macro.escape(loaded))
         def __schema__(:redact_fields), do: unquote(redacted_fields)
+        def __schema__(:virtual_fields), do: unquote(Enum.map(virtual_fields, &elem(&1, 0)))
 
         def __schema__(:query) do
           %Ecto.Query{
@@ -639,7 +642,7 @@ defmodule Ecto.Schema do
           }
         end
 
-        for clauses <- Ecto.Schema.__schema__(fields, field_sources, assocs, embeds),
+        for clauses <- Ecto.Schema.__schema__(fields, field_sources, assocs, embeds, virtual_fields),
             {args, body} <- clauses do
           def __schema__(unquote_splicing(args)), do: unquote(body)
         end
@@ -1919,7 +1922,9 @@ defmodule Ecto.Schema do
       Module.put_attribute(mod, :ecto_redact_fields, name)
     end
 
-    unless virtual? do
+    if virtual? do
+      Module.put_attribute(mod, :ecto_virtual_fields, {name, type})
+    else
       source = opts[:source] || Module.get_attribute(mod, :field_source_mapper).(name)
 
       if name != source do
@@ -2082,7 +2087,7 @@ defmodule Ecto.Schema do
   end
 
   @doc false
-  def __schema__(fields, field_sources, assocs, embeds) do
+  def __schema__(fields, field_sources, assocs, embeds, virtual_fields) do
     load =
       for {name, type} <- fields do
         if alias = field_sources[name] do
@@ -2107,6 +2112,11 @@ defmodule Ecto.Schema do
         {[:type, name], Macro.escape(type)}
       end
 
+    virtual_types_quoted =
+      for {name, type} <- virtual_fields do
+        {[:virtual_type, name], Macro.escape(type)}
+      end
+
     assoc_quoted =
       for {name, refl} <- assocs do
         {[:association, name], Macro.escape(refl)}
@@ -2131,6 +2141,7 @@ defmodule Ecto.Schema do
     catch_all = [
       {[:field_source, quote(do: _)], nil},
       {[:type, quote(do: _)], nil},
+      {[:virtual_type, quote(do: _)], nil},
       {[:association, quote(do: _)], nil},
       {[:embed, quote(do: _)], nil}
     ]
@@ -2139,6 +2150,7 @@ defmodule Ecto.Schema do
       single_arg,
       field_sources_quoted,
       types_quoted,
+      virtual_types_quoted,
       assoc_quoted,
       embed_quoted,
       catch_all

test/ecto/query/planner_test.exs

@@ -998,6 +998,11 @@ defmodule Ecto.Query.PlannerTest do
     end
   end
 
+  test "normalize: allow virtual fields in type/2" do
+    query = from(Comment, []) |> select([c], type(fragment("1"), c.temp))
+    normalize(query)
+  end
+
   test "normalize: validate fields in left side of in expressions" do
     query = from(Post, []) |> where([p], p.id in [1, 2, 3])
     normalize(query)