Skip to content

Commit

Permalink
feat: commit_comment option for prepending a SQL comment to a `COMM…
Browse files Browse the repository at this point in the history
…IT` statement (#669)
  • Loading branch information
bceskavich authored Mar 30, 2024
1 parent fcad612 commit ff417bd
Show file tree
Hide file tree
Showing 3 changed files with 42 additions and 3 deletions.
5 changes: 4 additions & 1 deletion lib/postgrex.ex
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
24 changes: 22 additions & 2 deletions lib/postgrex/protocol.ex
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,9 @@ defmodule Postgrex.Protocol do
@nonposix_errors [:closed, :timeout]
@max_rows 500
@text_type_oid 25
@commit_comment_validation_error Postgrex.Error.exception(
message: "`:commit_comment` option cannot contain sequence \"*/\""
)

defstruct sock: nil,
connection_id: nil,
Expand Down Expand Up @@ -545,8 +548,9 @@ defmodule Postgrex.Protocol do
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"
Expand Down Expand Up @@ -2887,6 +2891,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)]
Expand Down
16 changes: 16 additions & 0 deletions test/transaction_test.exs
Original file line number Diff line number Diff line change
Expand Up @@ -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.Error,
"`:commit_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
Expand Down

0 comments on commit ff417bd

Please sign in to comment.