comments in query

Signed-off-by: Daniel Kukula <[email protected]>

Author: dkuku

lib/ecto/adapters/postgres/connection.ex

@@ -187,8 +187,10 @@ if Code.ensure_loaded?(Postgrex) do
       limit = limit(query, sources)
       offset = offset(query, sources)
       lock = lock(query, sources)
+      comment = comment(query)
 
       [
+        comment,
         cte,
         select,
         from,
@@ -200,7 +202,8 @@ if Code.ensure_loaded?(Postgrex) do
         combinations,
         order_by,
         limit,
-        offset | lock
+        offset
+        | lock
       ]
     end
 
@@ -214,8 +217,17 @@ if Code.ensure_loaded?(Postgrex) do
       fields = update_fields(query, sources)
       {join, wheres} = using_join(query, :update_all, "FROM", sources)
       where = where(%{query | wheres: wheres ++ query.wheres}, sources)
+      comment = comment(query)
 
-      [cte, prefix, fields, join, where | returning(query, sources)]
+      [
+        comment,
+        cte,
+        prefix,
+        fields,
+        join,
+        where
+        | returning(query, sources)
+      ]
     end
 
     @impl true
@@ -226,8 +238,19 @@ if Code.ensure_loaded?(Postgrex) do
 
       {join, wheres} = using_join(query, :delete_all, "USING", sources)
       where = where(%{query | wheres: wheres ++ query.wheres}, sources)
+      comment = comment(query)
 
-      [cte, "DELETE FROM ", from, " AS ", name, join, where | returning(query, sources)]
+      [
+        comment,
+        cte,
+        "DELETE FROM ",
+        from,
+        " AS ",
+        name,
+        join,
+        where
+        | returning(query, sources)
+      ]
     end
 
     @impl true
@@ -880,6 +903,17 @@ if Code.ensure_loaded?(Postgrex) do
     defp lock(%{lock: binary}, _sources) when is_binary(binary), do: [?\s | binary]
     defp lock(%{lock: expr} = query, sources), do: [?\s | expr(expr, sources, query)]
 
+    defp comment(%{comments: []}), do: []
+
+    defp comment(%{comments: comments}) do
+      comment =
+        Enum.map_intersperse(comments, ",", fn
+          comment when is_binary(comment) -> comment
+        end)
+
+      ["/*", comment, "*/\n"]
+    end
+
     defp boolean(_name, [], _sources, _query), do: []
 
     defp boolean(name, [%{expr: expr, op: op} | query_exprs], sources, query) do

test/ecto/adapters/postgres_test.exs

@@ -2985,6 +2985,57 @@ defmodule Ecto.Adapters.PostgresTest do
     assert SQL.ddl_logs(result) == [{:error, ~s(table "foo" exists, skipping), []}]
   end
 
+  describe "comments" do
+    test "comments before query" do
+      query =
+        from(m in Schema, update: [set: [x: 0]])
+        |> comment("comment before query")
+        |> plan(:update_all)
+
+      assert update_all(query) == ~s{/*comment before query*/\nUPDATE "schema" AS s0 SET "x" = 0}
+    end
+
+    test "with multiple comments" do
+      query =
+        Schema
+        |> select([r], r.x)
+        |> comment("comptime")
+        |> comment(:atom)
+        |> plan()
+
+      assert all(query) =~ "/*comptime,atom*/"
+    end
+
+    test "with comments in subquery" do
+      subquery =
+        Schema
+        |> select([r], r.x)
+        |> comment("subquery")
+
+      query =
+        subquery(subquery)
+        |> select([r], r.x)
+        |> comment("query")
+        |> plan()
+
+      assert all(query) ==
+               ~s'/*query*/\nSELECT s0."x" FROM ' <>
+                 ~s'(/*subquery*/\nSELECT ss0."x" AS "x" FROM "schema" AS ss0 AS s0;'
+    end
+
+    test "comments in delete_all" do
+      query = Schema |> select([r], r.x) |> comment("comment") |> plan()
+
+      assert delete_all(query) ==
+               ~s'/*comment*/\nDELETE FROM "schema" AS s0 RETURNING s0."x"'
+    end
+
+    test "comments in update_all" do
+      query = from(m in Schema, update: [set: [x: 0]]) |> comment("comment") |> plan(:update_all)
+      assert update_all(query) == ~s{/*comment*/\nUPDATE "schema" AS s0 SET "x" = 0}
+    end
+  end
+
   defp make_result(level) do
     %Postgrex.Result{
       messages: [