From be74cce511ad217bf9d55952a318c54e9dfb862a Mon Sep 17 00:00:00 2001 From: Daniel Kukula Date: Sat, 12 Oct 2024 11:32:58 +0200 Subject: [PATCH 1/8] implement comment in postgrex --- lib/postgrex.ex | 6 +++-- lib/postgrex/protocol.ex | 48 +++++++++++++++++++++++++++++----------- test/query_test.exs | 9 ++++++++ 3 files changed, 48 insertions(+), 15 deletions(-) diff --git a/lib/postgrex.ex b/lib/postgrex.ex index f76ab3c8..0e73447d 100644 --- a/lib/postgrex.ex +++ b/lib/postgrex.ex @@ -363,7 +363,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 +374,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 diff --git a/lib/postgrex/protocol.ex b/lib/postgrex/protocol.ex index 81556365..d2b462f0 100644 --- a/lib/postgrex/protocol.ex +++ b/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 @@ -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 = "/* #{comment} */\n" <> query.statement + 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 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) From 409d07d3f93ed32d9e151d7f77a22da0fb85bd61 Mon Sep 17 00:00:00 2001 From: Daniel Kukula Date: Sat, 12 Oct 2024 11:42:16 +0200 Subject: [PATCH 2/8] move at the end of query --- lib/postgrex/protocol.ex | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/lib/postgrex/protocol.ex b/lib/postgrex/protocol.ex index d2b462f0..b98d3e81 100644 --- a/lib/postgrex/protocol.ex +++ b/lib/postgrex/protocol.ex @@ -1605,7 +1605,7 @@ defmodule Postgrex.Protocol do end defp parse_describe_comment_msgs(query, comment, tail) when is_binary(comment) do - statement = "/* #{comment} */\n" <> query.statement + statement = query.statement <> "/* #{comment} */" query = %{query | statement: statement} parse_describe_msgs(query, tail) end From a5121fb84d404531302dbe91e4a267d4660908e0 Mon Sep 17 00:00:00 2001 From: Daniel Kukula Date: Sun, 13 Oct 2024 13:24:06 +0200 Subject: [PATCH 3/8] set postgrex_prepare in query --- lib/postgrex.ex | 3 +++ lib/postgrex/protocol.ex | 6 +++--- 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/lib/postgrex.ex b/lib/postgrex.ex index 0e73447d..845f947d 100644 --- a/lib/postgrex.ex +++ b/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)} diff --git a/lib/postgrex/protocol.ex b/lib/postgrex/protocol.ex index b98d3e81..6aa18795 100644 --- a/lib/postgrex/protocol.ex +++ b/lib/postgrex/protocol.ex @@ -371,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 From 78a806934311dff588f208eb9fe27628ecdc57ae Mon Sep 17 00:00:00 2001 From: Daniel Kukula Date: Sun, 13 Oct 2024 20:17:29 +0200 Subject: [PATCH 4/8] cr suggestions --- lib/postgrex.ex | 7 ++----- lib/postgrex/protocol.ex | 33 +++++++++++++++++---------------- 2 files changed, 19 insertions(+), 21 deletions(-) diff --git a/lib/postgrex.ex b/lib/postgrex.ex index 845f947d..8f53b5ba 100644 --- a/lib/postgrex.ex +++ b/lib/postgrex.ex @@ -289,9 +289,6 @@ 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)} @@ -366,7 +363,7 @@ defmodule Postgrex do {:ok, Postgrex.Query.t()} | {:error, Exception.t()} def prepare(conn, name, statement, opts \\ []) do query = %Query{name: name, statement: statement} - prepare? = !Keyword.get(opts, :comment) + prepare? = Keyword.get(opts, :comment) == nil opts = Keyword.put(opts, :postgrex_prepare, prepare?) DBConnection.prepare(conn, query, opts) end @@ -377,7 +374,7 @@ defmodule Postgrex do """ @spec prepare!(conn, iodata, iodata, [option]) :: Postgrex.Query.t() def prepare!(conn, name, statement, opts \\ []) do - prepare? = !Keyword.get(opts, :comment) + 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 6aa18795..e2628ea5 100644 --- a/lib/postgrex/protocol.ex +++ b/lib/postgrex/protocol.ex @@ -346,19 +346,21 @@ 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) - if prepare? do - parse_describe_close(s, status, query) - else - comment = Keyword.get(opts, :comment) + case prepare do + true -> + parse_describe_close(s, status, query) - if is_binary(comment) && String.contains?(comment, "*/") do - raise @comment_validation_error - else - parse_describe_flush(s, status, query, comment) - end + false -> + 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 @@ -371,11 +373,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 @@ -1605,7 +1607,7 @@ defmodule Postgrex.Protocol do end defp parse_describe_comment_msgs(query, comment, tail) when is_binary(comment) do - statement = query.statement <> "/* #{comment} */" + statement = query.statement <> "/*#{comment}*/" query = %{query | statement: statement} parse_describe_msgs(query, tail) end @@ -3392,7 +3394,6 @@ 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 From 7f2f5db4005d0522f5954c73995d20bb3ad4ad12 Mon Sep 17 00:00:00 2001 From: Daniel Kukula Date: Sun, 13 Oct 2024 21:14:48 +0200 Subject: [PATCH 5/8] add docs --- lib/postgrex.ex | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/lib/postgrex.ex b/lib/postgrex.ex index 8f53b5ba..f4953988 100644 --- a/lib/postgrex.ex +++ b/lib/postgrex.ex @@ -169,6 +169,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. From 518ee8fff1a2ff5874205f98f9fd3c6607c80761 Mon Sep 17 00:00:00 2001 From: Daniel Kukula Date: Tue, 15 Oct 2024 20:20:09 +0200 Subject: [PATCH 6/8] cr suggestions --- lib/postgrex/protocol.ex | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/lib/postgrex/protocol.ex b/lib/postgrex/protocol.ex index e2628ea5..dfa3eb69 100644 --- a/lib/postgrex/protocol.ex +++ b/lib/postgrex/protocol.ex @@ -1607,7 +1607,7 @@ defmodule Postgrex.Protocol do end defp parse_describe_comment_msgs(query, comment, tail) when is_binary(comment) do - statement = query.statement <> "/*#{comment}*/" + statement = [query.statement, "/*", comment, "*/"] query = %{query | statement: statement} parse_describe_msgs(query, tail) end From 70fd3e3af9e9cc9eb397937f95368a29aa60b619 Mon Sep 17 00:00:00 2001 From: Daniel Kukula Date: Wed, 23 Oct 2024 08:06:46 +0200 Subject: [PATCH 7/8] remove coment validation --- lib/postgrex/protocol.ex | 10 +--------- test/query_test.exs | 4 ---- 2 files changed, 1 insertion(+), 13 deletions(-) diff --git a/lib/postgrex/protocol.ex b/lib/postgrex/protocol.ex index dfa3eb69..17d8aa05 100644 --- a/lib/postgrex/protocol.ex +++ b/lib/postgrex/protocol.ex @@ -16,9 +16,6 @@ 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, @@ -355,12 +352,7 @@ defmodule Postgrex.Protocol do false -> 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 + parse_describe_flush(s, status, query, comment) end end diff --git a/test/query_test.exs b/test/query_test.exs index 376962d8..bc3d895c 100644 --- a/test/query_test.exs +++ b/test/query_test.exs @@ -1853,10 +1853,6 @@ defmodule QueryTest do 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 From f783aa5d8e9a2cbc5f5791d522f866e67b853345 Mon Sep 17 00:00:00 2001 From: Daniel Kukula Date: Wed, 23 Oct 2024 19:19:47 +0200 Subject: [PATCH 8/8] validate comment in Postgrex module --- lib/postgrex.ex | 20 ++++++++++++++++++++ test/query_test.exs | 4 ++++ 2 files changed, 24 insertions(+) diff --git a/lib/postgrex.ex b/lib/postgrex.ex index f4953988..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 """ @@ -294,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)} @@ -324,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`. diff --git a/test/query_test.exs b/test/query_test.exs index bc3d895c..376962d8 100644 --- a/test/query_test.exs +++ b/test/query_test.exs @@ -1853,6 +1853,10 @@ defmodule QueryTest do 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