Allow data-modifying CTEs in Ecto (#532)

lib/ecto/adapters/myxql/connection.ex

@@ -312,16 +312,26 @@ if Code.ensure_loaded?(MyXQL) do
       error!(query, "MySQL adapter does not support materialized CTEs")
     end
 
-    defp cte_expr({name, _opts, cte}, sources, query) do
-      [quote_name(name), " AS ", cte_query(cte, sources, query)]
+    defp cte_expr({name, opts, cte}, sources, query) do
+      operation_opt = Map.get(opts, :operation)
+
+      [quote_name(name), " AS ", cte_query(cte, sources, query, operation_opt)]
     end
 
-    defp cte_query(%Ecto.Query{} = query, sources, parent_query) do
+    defp cte_query(query, sources, parent_query, nil) do
+      cte_query(query, sources, parent_query, :all)
+    end
+
+    defp cte_query(%Ecto.Query{} = query, sources, parent_query, :all) do
       query = put_in(query.aliases[@parent_as], {parent_query, sources})
       ["(", all(query, subquery_as_prefix(sources)), ")"]
     end
 
-    defp cte_query(%QueryExpr{expr: expr}, sources, query) do
+    defp cte_query(%Ecto.Query{} = query, _sources, _parent_query, operation) do
+      error!(query, "MySQL adapter does not support data-modifying CTEs (operation: #{operation})")
+    end
+
+    defp cte_query(%QueryExpr{expr: expr}, sources, query, _operation) do
       expr(expr, sources, query)
     end
 

lib/ecto/adapters/postgres/connection.ex

@@ -421,16 +421,36 @@ if Code.ensure_loaded?(Postgrex) do
           true -> "MATERIALIZED"
           false -> "NOT MATERIALIZED"
         end
+
+      operation_opt = Map.get(opts, :operation)
 
-      [quote_name(name), " AS ", materialized_opt, cte_query(cte, sources, query)]
+      [quote_name(name), " AS ", materialized_opt, cte_query(cte, sources, query, operation_opt)]
     end
 
-    defp cte_query(%Ecto.Query{} = query, sources, parent_query) do
+    defp cte_query(query, sources, parent_query, nil) do
+      cte_query(query, sources, parent_query, :all)
+    end
+
+    defp cte_query(%Ecto.Query{} = query, sources, parent_query, :update_all) do
+      query = put_in(query.aliases[@parent_as], {parent_query, sources})
+      ["(", update_all(query), ")"]
+    end
+
+    defp cte_query(%Ecto.Query{} = query, sources, parent_query, :delete_all) do
+      query = put_in(query.aliases[@parent_as], {parent_query, sources})
+      ["(", delete_all(query), ")"]
+    end
+
+    defp cte_query(%Ecto.Query{} = query, _sources, _parent_query, :insert_all) do
+      error!(query, "Postgres adapter does not support CTE operation :insert_all")
+    end
+
+    defp cte_query(%Ecto.Query{} = query, sources, parent_query, :all) do
       query = put_in(query.aliases[@parent_as], {parent_query, sources})
       ["(", all(query, subquery_as_prefix(sources)), ")"]
     end
 
-    defp cte_query(%QueryExpr{expr: expr}, sources, query) do
+    defp cte_query(%QueryExpr{expr: expr}, sources, query, _operation) do
       expr(expr, sources, query)
     end
 

lib/ecto/adapters/tds/connection.ex

@@ -438,8 +438,10 @@ if Code.ensure_loaded?(Tds) do
       error!(query, "Tds adapter does not support materialized CTEs")
     end
 
-    defp cte_expr({name, _opts, cte}, sources, query) do
-      [quote_name(name), cte_header(cte, query), " AS ", cte_query(cte, sources, query)]
+    defp cte_expr({name, opts, cte}, sources, query) do
+      operation_opt = Map.get(opts, :operation)
+
+      [quote_name(name), cte_header(cte, query), " AS ", cte_query(cte, sources, query, operation_opt)]
     end
 
     defp cte_header(%QueryExpr{}, query) do
@@ -467,10 +469,18 @@ if Code.ensure_loaded?(Tds) do
       ]
     end
 
-    defp cte_query(%Ecto.Query{} = query, sources, parent_query) do
+    defp cte_query(query, sources, parent_query, nil) do
+      cte_query(query, sources, parent_query, :all)
+    end
+
+    defp cte_query(%Ecto.Query{} = query, sources, parent_query, :all) do
       query = put_in(query.aliases[@parent_as], {parent_query, sources})
       [?(, all(query, subquery_as_prefix(sources)), ?)]
     end
+
+    defp cte_query(%Ecto.Query{} = query, _sources, _parent_query, operation) do
+      error!(query, "Tds adapter does not support data-modifying CTEs (operation: #{operation})")
+    end
 
     defp update_fields(%Query{updates: updates} = query, sources) do
       for(

mix.exs

@@ -77,7 +77,7 @@ defmodule EctoSQL.MixProject do
     if path = System.get_env("ECTO_PATH") do
       {:ecto, path: path}
     else
-      {:ecto, "~> 3.10.2"}
+      {:ecto, git: "https://github.com/elixir-ecto/ecto.git"}
     end
   end
 

mix.lock

@@ -5,7 +5,7 @@
   "decimal": {:hex, :decimal, "2.1.1", "5611dca5d4b2c3dd497dec8f68751f1f1a54755e8ed2a966c2633cf885973ad6", [:mix], [], "hexpm", "53cfe5f497ed0e7771ae1a475575603d77425099ba5faef9394932b35020ffcc"},
   "deep_merge": {:hex, :deep_merge, "0.2.0", "c1050fa2edf4848b9f556fba1b75afc66608a4219659e3311d9c9427b5b680b3", [:mix], [], "hexpm", "e3bf435a54ed27b0ba3a01eb117ae017988804e136edcbe8a6a14c310daa966e"},
   "earmark_parser": {:hex, :earmark_parser, "1.4.31", "a93921cdc6b9b869f519213d5bc79d9e218ba768d7270d46fdcf1c01bacff9e2", [:mix], [], "hexpm", "317d367ee0335ef037a87e46c91a2269fef6306413f731e8ec11fc45a7efd059"},
-  "ecto": {:hex, :ecto, "3.10.2", "6b887160281a61aa16843e47735b8a266caa437f80588c3ab80a8a960e6abe37", [:mix], [{:decimal, "~> 1.6 or ~> 2.0", [hex: :decimal, repo: "hexpm", optional: false]}, {:jason, "~> 1.0", [hex: :jason, repo: "hexpm", optional: true]}, {:telemetry, "~> 0.4 or ~> 1.0", [hex: :telemetry, repo: "hexpm", optional: false]}], "hexpm", "6a895778f0d7648a4b34b486af59a1c8009041fbdf2b17f1ac215eb829c60235"},
+  "ecto": {:git, "https://github.com/elixir-ecto/ecto.git", "cf379688df5c786b4f6e5b5cbf283489972b26b3", []},
   "ex_doc": {:hex, :ex_doc, "0.29.4", "6257ecbb20c7396b1fe5accd55b7b0d23f44b6aa18017b415cb4c2b91d997729", [:mix], [{:earmark_parser, "~> 1.4.31", [hex: :earmark_parser, repo: "hexpm", optional: false]}, {:makeup_elixir, "~> 0.14", [hex: :makeup_elixir, repo: "hexpm", optional: false]}, {:makeup_erlang, "~> 0.1", [hex: :makeup_erlang, repo: "hexpm", optional: false]}], "hexpm", "2c6699a737ae46cb61e4ed012af931b57b699643b24dabe2400a8168414bc4f5"},
   "jason": {:hex, :jason, "1.4.0", "e855647bc964a44e2f67df589ccf49105ae039d4179db7f6271dfd3843dc27e6", [:mix], [{:decimal, "~> 1.0 or ~> 2.0", [hex: :decimal, repo: "hexpm", optional: true]}], "hexpm", "79a3791085b2a0f743ca04cec0f7be26443738779d09302e01318f97bdb82121"},
   "makeup": {:hex, :makeup, "1.1.0", "6b67c8bc2882a6b6a445859952a602afc1a41c2e08379ca057c0f525366fc3ca", [:mix], [{:nimble_parsec, "~> 1.2.2 or ~> 1.3", [hex: :nimble_parsec, repo: "hexpm", optional: false]}], "hexpm", "0a45ed501f4a8897f580eabf99a2e5234ea3e75a4373c8a52824f6e873be57a6"},

test/ecto/adapters/myxql_test.exs

@@ -281,6 +281,24 @@ defmodule Ecto.Adapters.MyXQLTest do
       ~s{SELECT GROUP_CONCAT(st0.`id` SEPARATOR ' / ') AS `breadcrumbs` FROM `tree` AS st0) AS s1 ON TRUE}
   end
 
+  test "CTE with update statement" do
+    cte_query =
+      "categories"
+      |> where([c], is_nil(c.parent_id))
+      |> update([c], set: [desc: "Root category"])
+      |> select([c], %{id: c.id, desc: c.desc})
+
+    query =
+      "update_categories"
+      |> with_cte("update_categories", as: ^cte_query, operation: :update_all)
+      |> select([c], %{id: c.id, desc: c.desc})
+      |> plan()
+
+      assert_raise Ecto.QueryError, ~r/MySQL adapter does not support data-modifying CTEs/, fn ->
+        all(query)
+      end
+  end
+
   test "select" do
     query = Schema |> select([r], {r.x, r.y}) |> plan()
     assert all(query) == ~s{SELECT s0.`x`, s0.`y` FROM `schema` AS s0}

test/ecto/adapters/postgres_test.exs

@@ -383,6 +383,42 @@ defmodule Ecto.Adapters.PostgresTest do
       ~s{SELECT STRING_AGG(st0."id", ' / ') AS "breadcrumbs" FROM "tree" AS st0) AS s1 ON TRUE}
   end
 
+  test "CTE with update statement" do
+      cte_query =
+        "categories"
+        |> where([c], is_nil(c.parent_id))
+        |> update([c], set: [desc: "Root category"])
+        |> select([c], %{id: c.id, desc: c.desc})
+
+      query =
+        "update_categories"
+        |> with_cte("search_categories", as: ^cte_query)
+        |> with_cte("delete_categories", as: ^cte_query, operation: :delete_all)
+        |> with_cte("update_categories", as: ^cte_query, operation: :update_all)
+        |> select([c], %{id: c.id, desc: c.desc})
+        |> plan()
+
+      assert all(query) ==
+        ~s{WITH "search_categories" AS } <>
+        ~s{(} <>
+        ~s{SELECT sc0."id" AS "id", sc0."desc" AS "desc" } <>
+        ~s{FROM "categories" AS sc0 } <>
+        ~s{WHERE (sc0."parent_id" IS NULL)} <>
+        ~s{), } <>
+        ~s{"delete_categories" AS (} <>
+        ~s{DELETE FROM "categories" AS c0 } <>
+        ~s{WHERE (c0."parent_id" IS NULL) } <>
+        ~s{RETURNING c0."id" AS "id", c0."desc" AS "desc"} <>
+        ~s{), } <>
+        ~s{"update_categories" AS (} <>
+        ~s{UPDATE "categories" AS c0 SET "desc" = 'Root category' } <>
+        ~s{WHERE (c0."parent_id" IS NULL) } <>
+        ~s{RETURNING c0."id" AS "id", c0."desc" AS "desc"} <>
+        ~s{) } <>
+        ~s{SELECT u0."id", u0."desc" } <>
+        ~s{FROM "update_categories" AS u0}
+  end
+
   test "select" do
     query = Schema |> select([r], {r.x, r.y}) |> plan()
     assert all(query) == ~s{SELECT s0."x", s0."y" FROM "schema" AS s0}

test/ecto/adapters/tds_test.exs

@@ -364,6 +364,24 @@ defmodule Ecto.Adapters.TdsTest do
       ~s{SELECT STRING_AGG(st0.[id], ' / ') AS [breadcrumbs] FROM [tree] AS st0) AS s1}
   end
 
+  test "CTE with update statement" do
+    cte_query =
+      "categories"
+      |> where([c], is_nil(c.parent_id))
+      |> update([c], set: [desc: "Root category"])
+      |> select([c], %{id: c.id, desc: c.desc})
+
+    query =
+      "update_categories"
+      |> with_cte("update_categories", as: ^cte_query, operation: :update_all)
+      |> select([c], %{id: c.id, desc: c.desc})
+      |> plan()
+
+      assert_raise Ecto.QueryError, ~r/Tds adapter does not support data-modifying CTEs/, fn ->
+        all(query)
+      end
+  end
+
   test "select" do
     query = Schema |> select([r], {r.x, r.y}) |> plan()
     assert all(query) == ~s{SELECT s0.[x], s0.[y] FROM [schema] AS s0}