Allow to append comment to query (#709)

Author: dkuku

lib/postgrex.ex

@@ -67,6 +67,10 @@ defmodule Postgrex do
   @max_rows 500
   @timeout 15_000
 
+  @comment_validation_error Postgrex.Error.exception(
+                              message: "`:comment` option cannot contain sequence \"*/\""
+                            )
+
   ### PUBLIC API ###
 
   @doc """
@@ -169,6 +173,11 @@ defmodule Postgrex do
       This is useful when using Postgrex against systems that do not support composite types
       (default: `false`).
 
+    * `:comment` - When a binary string is provided, appends the given text as a comment to the
+      query.  This can be useful for tracing purposes, such as when using SQLCommenter or similar
+      tools to track query performance and behavior. Note that including a comment disables query
+      caching since each query with a different comment is treated as unique (default: `nil`).
+
   `Postgrex` uses the `DBConnection` library and supports all `DBConnection`
   options like `:idle`, `:after_connect` etc. See `DBConnection.start_link/2`
   for more information.
@@ -289,6 +298,8 @@ defmodule Postgrex do
   @spec query(conn, iodata, list, [execute_option]) ::
           {:ok, Postgrex.Result.t()} | {:error, Exception.t()}
   def query(conn, statement, params, opts \\ []) do
+    :ok = validate_comment(opts)
+
     if name = Keyword.get(opts, :cache_statement) do
       query = %Query{name: name, cache: :statement, statement: IO.iodata_to_binary(statement)}
 
@@ -319,6 +330,20 @@ defmodule Postgrex do
     end
   end
 
+  defp validate_comment(opts) do
+    case Keyword.get(opts, :comment) do
+      nil ->
+        :ok
+
+      comment when is_binary(comment) ->
+        if String.contains?(comment, "*/") do
+          raise @comment_validation_error
+        else
+          :ok
+        end
+    end
+  end
+
   @doc """
   Runs an (extended) query and returns the result or raises `Postgrex.Error` if
   there was an error. See `query/3`.
@@ -363,7 +388,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) == nil
+    opts = Keyword.put(opts, :postgrex_prepare, prepare?)
     DBConnection.prepare(conn, query, opts)
   end
 
@@ -373,7 +399,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) == nil
+    opts = Keyword.put(opts, :postgrex_prepare, prepare?)
     DBConnection.prepare!(conn, %Query{name: name, statement: statement}, opts)
   end
 

lib/postgrex/protocol.ex

@@ -347,8 +347,12 @@ defmodule Postgrex.Protocol do
     status = new_status(opts, prepare: prepare)
 
     case prepare do
-      true -> parse_describe_close(s, status, query)
-      false -> parse_describe_flush(s, status, query)
+      true ->
+        parse_describe_close(s, status, query)
+
+      false ->
+        comment = Keyword.get(opts, :comment)
+        parse_describe_flush(s, status, query, comment)
     end
   end
 
@@ -1418,9 +1422,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 +1446,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 +1471,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 +1598,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 +2094,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 +2120,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 +2411,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)
 

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)