allow to append comment to query

lib/postgrex.ex

@@ -289,6 +289,9 @@ defmodule Postgrex do
   @spec query(conn, iodata, list, [execute_option]) ::
           {:ok, Postgrex.Result.t()} | {:error, Exception.t()}
   def query(conn, statement, params, opts \\ []) do
+    prepare? = !Keyword.get(opts, :comment)
+    opts = Keyword.put(opts, :postgrex_prepare, prepare?)
+
     if name = Keyword.get(opts, :cache_statement) do
       query = %Query{name: name, cache: :statement, statement: IO.iodata_to_binary(statement)}
 
@@ -363,7 +366,8 @@ defmodule Postgrex do
           {:ok, Postgrex.Query.t()} | {:error, Exception.t()}
   def prepare(conn, name, statement, opts \\ []) do
     query = %Query{name: name, statement: statement}
-    opts = Keyword.put(opts, :postgrex_prepare, true)
+    prepare? = !Keyword.get(opts, :comment)
+    opts = Keyword.put(opts, :postgrex_prepare, prepare?)
     DBConnection.prepare(conn, query, opts)
   end
 
@@ -373,7 +377,8 @@ defmodule Postgrex do
   """
   @spec prepare!(conn, iodata, iodata, [option]) :: Postgrex.Query.t()
   def prepare!(conn, name, statement, opts \\ []) do
-    opts = Keyword.put(opts, :postgrex_prepare, true)
+    prepare? = !Keyword.get(opts, :comment)
+    opts = Keyword.put(opts, :postgrex_prepare, prepare?)
     DBConnection.prepare!(conn, %Query{name: name, statement: statement}, opts)
   end
 

lib/postgrex/protocol.ex

@@ -16,6 +16,9 @@ defmodule Postgrex.Protocol do
                                      message:
                                        "`:commit_comment` option cannot contain sequence \"*/\""
                                    )
+  @comment_validation_error Postgrex.Error.exception(
+                              message: "`:comment` option cannot contain sequence \"*/\""
+                            )
 
   defstruct sock: nil,
             connection_id: nil,
@@ -343,12 +346,19 @@ defmodule Postgrex.Protocol do
   end
 
   def handle_prepare(%Query{name: ""} = query, opts, s) do
-    prepare = Keyword.get(opts, :postgrex_prepare, false)
-    status = new_status(opts, prepare: prepare)
+    prepare? = Keyword.get(opts, :postgrex_prepare, false)
+    status = new_status(opts, prepare: prepare?)
 
-    case prepare do
-      true -> parse_describe_close(s, status, query)
-      false -> parse_describe_flush(s, status, query)
+    if prepare? do
+      parse_describe_close(s, status, query)
+    else
+      comment = Keyword.get(opts, :comment)
+
+      if is_binary(comment) && String.contains?(comment, "*/") do
+        raise @comment_validation_error
+      else
+        parse_describe_flush(s, status, query, comment)
+      end
     end
   end
 
@@ -361,11 +371,11 @@ defmodule Postgrex.Protocol do
     if new_query = cached_query(s, query) do
       {:ok, new_query, s}
     else
-      prepare = Keyword.get(opts, :postgrex_prepare, false)
-      status = new_status(opts, prepare: prepare)
+      prepare? = Keyword.get(opts, :postgrex_prepare, false)
+      status = new_status(opts, prepare: prepare?)
 
       result =
-        case prepare do
+        case prepare? do
           true -> close_parse_describe(s, status, query)
           false -> close_parse_describe_flush(s, status, query)
         end
@@ -1418,9 +1428,9 @@ defmodule Postgrex.Protocol do
     parse_describe(s, status, query)
   end
 
-  defp parse_describe_flush(s, %{mode: :transaction} = status, query) do
+  defp parse_describe_flush(s, %{mode: :transaction} = status, query, comment) do
     %{buffer: buffer} = s
-    msgs = parse_describe_msgs(query, [msg_flush()])
+    msgs = parse_describe_comment_msgs(query, comment, [msg_flush()])
 
     with :ok <- msg_send(%{s | buffer: nil}, msgs, buffer),
          {:ok, %Query{ref: ref} = query, %{postgres: postgres} = s, buffer} <-
@@ -1442,11 +1452,12 @@ defmodule Postgrex.Protocol do
   defp parse_describe_flush(
          %{postgres: :transaction, buffer: buffer} = s,
          %{mode: :savepoint} = status,
-         query
+         query,
+         comment
        ) do
     msgs =
       [msg_query(statement: "SAVEPOINT postgrex_query")] ++
-        parse_describe_msgs(query, [msg_flush()])
+        parse_describe_comment_msgs(query, comment, [msg_flush()])
 
     with :ok <- msg_send(%{s | buffer: nil}, msgs, buffer),
          {:ok, _, %{buffer: buffer} = s} <- recv_transaction(s, status, buffer),
@@ -1466,7 +1477,7 @@ defmodule Postgrex.Protocol do
     end
   end
 
-  defp parse_describe_flush(%{postgres: postgres} = s, %{mode: :savepoint}, _)
+  defp parse_describe_flush(%{postgres: postgres} = s, %{mode: :savepoint}, _, _)
        when postgres in [:idle, :error] do
     transaction_error(s, postgres)
   end
@@ -1593,6 +1604,16 @@ defmodule Postgrex.Protocol do
     transaction_error(s, postgres)
   end
 
+  defp parse_describe_comment_msgs(query, comment, tail) when is_binary(comment) do
+    statement = query.statement <> "/* #{comment} */"
+    query = %{query | statement: statement}
+    parse_describe_msgs(query, tail)
+  end
+
+  defp parse_describe_comment_msgs(query, _comment, tail) do
+    parse_describe_msgs(query, tail)
+  end
+
   defp parse_describe_msgs(query, tail) do
     %Query{name: name, statement: statement, param_oids: param_oids} = query
     type_oids = param_oids || []
@@ -2079,7 +2100,7 @@ defmodule Postgrex.Protocol do
 
       _ ->
         # flush awaiting execute or declare
-        parse_describe_flush(s, status, query)
+        parse_describe_flush(s, status, query, nil)
     end
   end
 
@@ -2105,7 +2126,7 @@ defmodule Postgrex.Protocol do
   defp handle_prepare_execute(%Query{name: ""} = query, params, opts, s) do
     status = new_status(opts)
 
-    case parse_describe_flush(s, status, query) do
+    case parse_describe_flush(s, status, query, nil) do
       {:ok, query, s} ->
         bind_execute_close(s, status, query, params)
 
@@ -2396,7 +2417,7 @@ defmodule Postgrex.Protocol do
   defp handle_prepare_bind(%Query{name: ""} = query, params, res, opts, s) do
     status = new_status(opts)
 
-    case parse_describe_flush(s, status, query) do
+    case parse_describe_flush(s, status, query, nil) do
       {:ok, query, s} ->
         bind(s, status, query, params, res)
 
@@ -3371,6 +3392,7 @@ defmodule Postgrex.Protocol do
 
   defp msg_send(s, msgs, buffer) when is_list(msgs) do
     binaries = Enum.reduce(msgs, [], &[&2 | maybe_encode_msg(&1)])
+
     do_send(s, binaries, buffer)
   end
 

test/query_test.exs

@@ -1851,6 +1851,15 @@ defmodule QueryTest do
     assert [["1", "2"], ["3", "4"]] = query("COPY (VALUES (1, 2), (3, 4)) TO STDOUT", [], opts)
   end
 
+  test "comment", context do
+    assert [[123]] = query("select 123", [], comment: "query comment goes here")
+
+    assert_raise Postgrex.Error, fn ->
+      query("select 123", [], comment: "*/ DROP TABLE 123 --")
+    end
+  end
+
+  @tag :big_binary
   test "receive packet with remainder greater than 64MB", context do
     # to ensure remainder is more than 64MB use 64MBx2+1
     big_binary = :binary.copy(<<1>>, 128 * 1024 * 1024 + 1)