diff --git a/lib/postgrex.ex b/lib/postgrex.ex index f76ab3c8..0bd1d280 100644 --- a/lib/postgrex.ex +++ b/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 diff --git a/lib/postgrex/protocol.ex b/lib/postgrex/protocol.ex index 81556365..17d8aa05 100644 --- a/lib/postgrex/protocol.ex +++ b/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) diff --git a/test/query_test.exs b/test/query_test.exs index 72240411..376962d8 100644 --- a/test/query_test.exs +++ b/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)