feat: commit_comment option for prepending a SQL comment to a COMMIT statement

lib/postgrex.ex

@@ -578,11 +578,14 @@ defmodule Postgrex do
     * `:timeout` - Transaction timeout (default: `#{@timeout}`);
     * `:mode` - Set to `:savepoint` to use savepoints instead of an SQL
     transaction, otherwise set to `:transaction` (default: `:transaction`);
+    * `:commit_comment` - When a binary string is provided, prepends the text as
+    a comment attached to the `COMMIT` statement issued to close the transaction (default: `nil`);
 
   The `:timeout` is for the duration of the transaction and all nested
   transactions and requests. This timeout overrides timeouts set by internal
   transactions and requests. The `:mode` will be used for all requests inside
-  the transaction function.
+  the transaction function. The `:commit_comment` can be helpful in distinguishing
+  between transactions in query performance monitoring tools.
 
   ## Example
 

lib/postgrex/protocol.ex

@@ -12,6 +12,9 @@ defmodule Postgrex.Protocol do
   @nonposix_errors [:closed, :timeout]
   @max_rows 500
   @text_type_oid 25
+  @commit_comment_validation_error Postgrex.QueryError.exception(
+                                     "`:comment_comment` option cannot contain sequence \"*/\""
+                                   )
 
   defstruct sock: nil,
             connection_id: nil,
@@ -537,16 +540,19 @@ defmodule Postgrex.Protocol do
           {:ok, Postgrex.Result.t(), state}
           | {DBConnection.status(), state}
           | {:disconnect, %RuntimeError{}, state}
-          | {:disconnect, %DBConnection.ConnectionError{} | Postgrex.Error.t(), state}
+          | {:disconnect,
+             %DBConnection.ConnectionError{} | Postgrex.Error.t() | Postgrex.QueryError.t(),
+             state}
   def handle_commit(_, %{postgres: {_, _}} = s) do
     lock_error(s, :commit)
   end
 
   def handle_commit(opts, %{postgres: postgres} = s) do
     case Keyword.get(opts, :mode, :transaction) do
       :transaction when postgres == :transaction ->
-        statement = "COMMIT"
-        handle_transaction(statement, opts, s)
+        with {:ok, statement} <- build_commit_statement(opts, s) do
+          handle_transaction(statement, opts, s)
+        end
 
       :savepoint when postgres == :transaction ->
         statement = "RELEASE SAVEPOINT postgrex_savepoint"
@@ -2887,6 +2893,22 @@ defmodule Postgrex.Protocol do
 
   ## transaction
 
+  defp build_commit_statement(opts, %{buffer: buffer} = s) do
+    case Keyword.get(opts, :commit_comment) do
+      comment when is_binary(comment) ->
+        if String.contains?(comment, "*/") do
+          disconnect(s, @commit_comment_validation_error, buffer)
+        else
+          statement = "/* #{comment} */\nCOMMIT"
+          {:ok, statement}
+        end
+
+      _ ->
+        statement = "COMMIT"
+        {:ok, statement}
+    end
+  end
+
   defp handle_transaction(statement, opts, %{buffer: buffer} = s) do
     status = new_status(opts, mode: :transaction)
     msgs = [msg_query(statement: statement)]
@@ -3368,6 +3390,10 @@ defmodule Postgrex.Protocol do
     {:disconnect, %{err | connection_id: connection_id}, %{s | buffer: buffer}}
   end
 
+  defp disconnect(s, %Postgrex.QueryError{} = err, buffer) do
+    {:disconnect, err, %{s | buffer: buffer}}
+  end
+
   defp disconnect(s, %RuntimeError{} = err, buffer) do
     {:disconnect, err, %{s | buffer: buffer}}
   end

test/transaction_test.exs

@@ -86,6 +86,22 @@ defmodule TransactionTest do
            end) == {:error, :rollback}
   end
 
+  @tag mode: :transaction
+  test "commit comment with possible SQL injection returns error and disconnects", context do
+    assert_raise(
+      Postgrex.QueryError,
+      "`:comment_comment` option cannot contain sequence \"*/\"",
+      fn ->
+        transaction(
+          fn conn ->
+            assert {:ok, %Postgrex.Result{rows: [[42]]}} = P.query(conn, "SELECT 42", [])
+          end,
+          commit_comment: "invalid */ comment"
+        )
+      end
+    )
+  end
+
   @tag mode: :transaction
   @tag prepare: :unnamed
   test "transaction commits with unnamed queries", context do