Repository: elixir-ecto/ecto
Pull request #4221 (2023-06-29 13:51:44)

Data modifying CTE cleanup
Some small cleanups from https://github.com/elixir-ecto/ecto/pull/4220

1. We aren't able to support `insert_all` right now.
2. The cache key should take the operation into account because it will affect the underlying query

greg-rychlewski: ah have to fix the tests =/

Patch:
diff --git a/lib/ecto/query.ex b/lib/ecto/query.ex
index ad72c9ce..57207106 100644
--- a/lib/ecto/query.ex
+++ b/lib/ecto/query.ex
@@ -1364,8 +1364,8 @@ defmodule Ecto.Query do
     * `:materialized` - a boolean indicating whether the CTE should
     be materialized. If blank, the database's default behaviour
     will be used (only supported by Postgrex, for the built-in adapters)
-    * `:operation` - one of `:all`, `:update_all`, `:delete_all`, `:insert_all` 
-    indicating the operation type of the CTE query. If blank, it defaults to `:all`, 
+    * `:operation` - one of `:all`, `:update_all`, or `:delete_all`
+    indicating the operation type of the CTE query. If blank, it defaults to `:all`,
     making the CTE query a SELECT query. (only supported by Postgres built-in adapter)
 
   ## Recursive CTEs
@@ -1433,9 +1433,9 @@ defmodule Ecto.Query do
       from(cte in {"category_tree", Category}, prefix: nil)
       |> recursive_ctes(true)
       |> with_cte("category_tree", as: ^category_tree_query)
-      
+
   For Postgres built-in adapter, it is possible to define data-modifying CTE queries:
-  
+
       update_categories_query =
         Category
         |> where([c], is_nil(c.parent_id))
@@ -1445,12 +1445,12 @@ defmodule Ecto.Query do
       {"update_categories", Category}
       |> with_cte("update_categories", as: ^update_categories_query, operation: :update_all)
       |> select([c], c)
-      
-  Note: In order to retrieve the updates rows from a CTE query, the parent query 
+
+  Note: In order to retrieve the updates rows from a CTE query, the parent query
   must select rows from the CTE table instead of the table referenced by the CTE query.
   For example, `"update_categories"` will return updates rows for `"category"` table, but
   selecting from `"category"` table directly will return unaffected rows.
-  For more details see Postgres documentation on data-modifying CTEs and how these work 
+  For more details see Postgres documentation on data-modifying CTEs and how these work
   with snapshots.
 
   Keyword syntax is not supported for this feature.
diff --git a/lib/ecto/query/builder/cte.ex b/lib/ecto/query/builder/cte.ex
index 9ef39ba7..a253fd29 100644
--- a/lib/ecto/query/builder/cte.ex
+++ b/lib/ecto/query/builder/cte.ex
@@ -35,7 +35,7 @@ defmodule Ecto.Query.Builder.CTE do
   If possible, it does all calculations at compile time to avoid
   runtime work.
   """
-  @spec build(Macro.t, Macro.t, Macro.t, nil | boolean(), nil | :all | :update_all | :delete_all | :insert_all, Macro.Env.t) :: Macro.t
+  @spec build(Macro.t, Macro.t, Macro.t, nil | boolean(), nil | :all | :update_all | :delete_all , Macro.Env.t) :: Macro.t
   def build(query, name, cte, materialized, operation, env) do
     Builder.apply_query(query, __MODULE__, [escape(name, env), build_cte(name, cte, env), materialized, operation], env)
   end
@@ -73,17 +73,17 @@ defmodule Ecto.Query.Builder.CTE do
   @doc """
   The callback applied by `build/4` to build the query.
   """
-  @spec apply(Ecto.Queryable.t, bitstring, Ecto.Queryable.t, nil | boolean(), nil | :all | :update_all | :delete_all | :insert_all) :: Ecto.Query.t
+  @spec apply(Ecto.Queryable.t, bitstring, Ecto.Queryable.t, nil | boolean(), nil | :all | :update_all | :delete_all) :: Ecto.Query.t
   # Runtime
   def apply(query, name, with_query, materialized, nil) do
     apply(query, name, with_query, materialized, :all)
   end
-  
-  def apply(_query, _name, _with_query, _materialized, operation) 
-    when operation not in [:all, :update_all, :delete_all, :insert_all] do
-    Builder.error!("`operation` option must be one of :all, :update_all, :delete_all, :insert_all")
+
+  def apply(_query, _name, _with_query, _materialized, operation)
+    when operation not in [:all, :update_all, :delete_all] do
+    Builder.error!("`operation` option must be one of :all, :update_all, or :delete_all")
   end
-  
+
   def apply(%Ecto.Query{with_ctes: with_expr} = query, name, %_{} = with_query, materialized, operation) do
     %{query | with_ctes: apply_cte(with_expr, name, with_query, materialized, operation)}
   end
diff --git a/lib/ecto/query/planner.ex b/lib/ecto/query/planner.ex
index 5de6d846..bee99244 100644
--- a/lib/ecto/query/planner.ex
+++ b/lib/ecto/query/planner.ex
@@ -757,11 +757,11 @@ defmodule Ecto.Query.Planner do
     Enum.reduce queries, cache_and_params, fn
       {name, opts, %Ecto.Query{} = query}, {cache, params} ->
         {_, params, inner_cache} = traverse_cache(query, :all, {[], params}, adapter)
-        {merge_cache({key, name, opts[:materialized], inner_cache}, cache, inner_cache != :nocache), params}
+        {merge_cache({key, name, opts[:materialized], opts[:operation], inner_cache}, cache, inner_cache != :nocache), params}
 
       {name, opts, %Ecto.Query.QueryExpr{} = query_expr}, {cache, params} ->
         {params, cacheable?} = cast_and_merge_params(:with_cte, query, query_expr, params, adapter)
-        {merge_cache({key, name, opts[:materialized], expr_to_cache(query_expr)}, cache, cacheable?), params}
+        {merge_cache({key, name, opts[:materialized], opts[:operation], expr_to_cache(query_expr)}, cache, cacheable?), params}
     end
   end
 


Test patch:
diff --git a/test/ecto/query/planner_test.exs b/test/ecto/query/planner_test.exs
index 74cc6da0..0f804d45 100644
--- a/test/ecto/query/planner_test.exs
+++ b/test/ecto/query/planner_test.exs
@@ -738,8 +738,7 @@ defmodule Ecto.Query.PlannerTest do
       assert [
         :all,
         {:from, {"comments", Comment, _, nil}, []},
-        {:non_recursive_cte, "cte",
-        nil,
+        {:non_recursive_cte, "cte", nil, :all,
          [:all, {:prefix, "another"}, {:from, {"comments", Comment, _, nil}, []}, {:select, {:&, _, [0]}}]}
       ] = cache
 
@@ -760,7 +759,7 @@ defmodule Ecto.Query.PlannerTest do
       %{queries: [{"cte", %{}, query_expr}]} = with_expr
       expr = {:fragment, [], [raw: "SELECT * FROM comments WHERE id = ", expr: {:^, [], [0]}, raw: ""]}
       assert expr == query_expr.expr
-      assert [:all, {:from, {"comments", Comment, _, nil}, []}, {:recursive_cte, "cte", nil, ^expr}] = cache
+      assert [:all, {:from, {"comments", Comment, _, nil}, []}, {:recursive_cte, "cte", nil, :all, ^expr}] = cache
     end
 
     test "on update_all" do
@@ -784,7 +783,7 @@ defmodule Ecto.Query.PlannerTest do
       assert {{"comments", Comment, "another"}} = cte.sources
       assert %{expr: {:^, [], [0]}, params: [{500, :integer}]} = cte.limit
 
-      assert [:update_all, _, _, _, _, {:non_recursive_cte, "recent_comments", nil, cte_cache}] = cache
+      assert [:update_all, _, _, _, _, {:non_recursive_cte, "recent_comments", nil, :all, cte_cache}] = cache
       assert [
                :all,
                {:prefix, "another"},
@@ -816,7 +815,7 @@ defmodule Ecto.Query.PlannerTest do
       assert {{"comments", Comment, "another"}} = cte.sources
       assert %{expr: {:^, [], [0]}, params: [{500, :integer}]} = cte.limit
 
-      assert [:delete_all, _, _, _, {:non_recursive_cte, "recent_comments", nil, cte_cache}] = cache
+      assert [:delete_all, _, _, _, {:non_recursive_cte, "recent_comments", nil, :all, cte_cache}] = cache
       assert [
                :all,
                {:prefix, "another"},