16470769502

Committed 23 Jul 2025 12:34PM UTC coverage: 57.067% (+1.7%) from 55.355%

Build # 16470769502

Build Type

Pull #694

github

Committed by

web-flow

Commit Message

Merge 78a9c0b2c into deaa48192

Pull Request Pull Request #694: feat: improved named prepared statements support

Run Details

175 of 217 new or added lines in 11 files covered. (80.65%)

16 existing lines in 4 files now uncovered.

1292 of 2264 relevant lines covered (57.07%)

1126.08 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

93.44

/lib/supavisor/protocol/prepared_statements.ex

defmodule Supavisor.Protocol.PreparedStatements do
  @moduledoc """
  Handles prepared statement binary packet transformations.
  """

  alias Supavisor.Protocol.PreparedStatements.PreparedStatement
  alias Supavisor.Protocol.PreparedStatements.Storage

  @type statement_map() :: Storage.t()

  @type statement_name() :: String.t()

  @type pkt() :: binary()

  @type handled_pkt() ::
          {:parse_pkt, statement_name(), pkt()}
          | {:bind_pkt, statement_name(), bind_pkt :: pkt(), parse_pkt :: pkt()}
          | {:close_pkt, statement_name(), pkt()}
          | pkt()

  @client_limit 100
  @backend_limit 200
  @client_memory_limit_bytes 1_000_000

  @doc """
  Upper limit of prepared statements from the client
  """
  @spec client_limit() :: pos_integer()
  def client_limit, do: @client_limit

  @doc """
  Upper limit of prepared statements backend-side.

  Should rotate prepared statements to avoid surpassing it.
  """
  @spec backend_limit() :: pos_integer()
  def backend_limit, do: @backend_limit

  @doc """
  Upper limit of prepared statements memory from the client in bytes
  """
  @spec client_memory_limit_bytes() :: pos_integer()
  def client_memory_limit_bytes, do: @client_memory_limit_bytes

  @doc """
  Receives a statement name and returns a close packet for it
  """
  @spec build_close_pkt(statement_name) :: pkt()
  def build_close_pkt(statement_name) do
    len = byte_size(statement_name)
    <<?C, len + 6::32, ?S, statement_name::binary, 0>>
  end

  @doc """
  Handles prepared statement packets and returns appropriate tuples for packets
  that need special treatment according to the protocol.
  """
  @spec handle_pkt(statement_map(), pkt()) ::
          {:ok, statement_map(), handled_pkt()}
          | {:error, :max_prepared_statements}
          | {:error, :max_prepared_statements_memory}
          | {:error, :prepared_statement_on_simple_query}
          | {:error, :duplicate_prepared_statement, statement_name()}
          | {:error, :prepared_statement_not_found, statement_name()}
  def handle_pkt(client_statements, binary) do
    case binary do
      # Parse message (P)
      <<?P, len::32, rest::binary>> ->
        handle_parse_message(client_statements, binary, len, rest)

      # Bind message (B)
      <<?B, len::32, rest::binary>> ->
        handle_bind_message(client_statements, binary, len, rest)

      # Close message (C)
      <<?C, len::32, ?S, rest::binary>> ->
        handle_close_message(client_statements, len, rest)

      # Describe message (D)
      <<?D, len::32, ?S, rest::binary>> ->
        handle_describe_message(client_statements, len, rest)

      # Query message (Q)
      <<?Q, len::32, rest::binary>> ->
        handle_simple_query_message(client_statements, binary, len, rest)

      # All other messages pass through unchanged
      _ ->
        {:ok, client_statements, binary}
    end
  end

  defp handle_parse_message(client_statements, original_bin, len, rest) do
    case extract_null_terminated_string(rest) do
      # Unnamed prepared statements are passed through unchanged
      {"", _} ->
        {:ok, client_statements, original_bin}

      {client_side_name, remaining} ->
        cond do
          Storage.statement_count(client_statements) >= @client_limit ->
            {:error, :max_prepared_statements}

          Storage.statement_memory(client_statements) > @client_memory_limit_bytes ->
            {:error, :max_prepared_statements_memory}

          Storage.get(client_statements, client_side_name) ->
            {:error, :duplicate_prepared_statement, client_side_name}

          true ->
            server_side_name = gen_server_side_name(rest)

            new_len = len + (byte_size(server_side_name) - byte_size(client_side_name))
            new_bin = <<?P, new_len::32, server_side_name::binary, 0, remaining::binary>>

            prepared_statement = %PreparedStatement{
              name: server_side_name,
              parse_pkt: new_bin
            }

            new_client_statements =
              Storage.put(client_statements, client_side_name, prepared_statement)

            {:ok, new_client_statements, {:parse_pkt, server_side_name, new_bin}}
        end
    end
  end

  defp handle_bind_message(client_statements, bin, len, rest) do
    {_portal_name, after_portal} = extract_null_terminated_string(rest)

    case extract_null_terminated_string(after_portal) do
      {"", _} ->
        {:ok, client_statements, bin}

      {client_side_name, packet_after_client_name} ->
        case Storage.get(client_statements, client_side_name) do
          %PreparedStatement{name: server_side_name, parse_pkt: parse_pkt} ->
            new_len = len + (byte_size(server_side_name) - byte_size(client_side_name))

            new_bin =
              <<?B, new_len::32, 0, server_side_name::binary, 0,
                packet_after_client_name::binary>>

            {:ok, client_statements, {:bind_pkt, server_side_name, new_bin, parse_pkt}}

          nil ->
            {:error, :prepared_statement_not_found}
        end
    end
  end

  defp handle_close_message(client_statements, len, rest) do
    {client_side_name, _} = extract_null_terminated_string(rest)

    {prepared_statement, new_client_statements} = Storage.pop(client_statements, client_side_name)

    server_side_name =
      case prepared_statement do
        %PreparedStatement{name: name} -> name
        nil -> "supavisor_none"
      end

    new_len = len + (byte_size(server_side_name) - byte_size(client_side_name))
    new_bin = <<?C, new_len::32, ?S, server_side_name::binary, 0>>

    {:ok, new_client_statements, {:close_pkt, server_side_name, new_bin}}
  end

  defp handle_describe_message(client_statements, len, rest) do
    {client_side_name, _} = extract_null_terminated_string(rest)

    server_side_name =
      case Storage.get(client_statements, client_side_name) do
        %PreparedStatement{name: name} -> name
        nil -> ""
      end

    new_len = len + (byte_size(server_side_name) - byte_size(client_side_name))
    new_bin = <<?D, new_len::32, ?S, server_side_name::binary, 0>>

    {:ok, client_statements, {:describe_pkt, server_side_name, new_bin}}
  end

  defp handle_simple_query_message(client_statements, binary, _len, rest) do
    case rest do
      "PREPARE" <> _ ->
        {:error, :prepared_statement_on_simple_query}

      _ ->
        {:ok, client_statements, binary}
    end
  end

  defp extract_null_terminated_string(binary) do
    case :binary.split(binary, <<0>>) do
      [string, rest] -> {string, rest}
      [string] -> {string, <<>>}
    end
  end

  defp gen_server_side_name(binary) do
    {_name, rest} = extract_null_terminated_string(binary)
    {query, _rest} = extract_null_terminated_string(rest)

    fingerprint =
      case Supavisor.PgParser.fingerprint(query) do
        {:ok, fingerprint} -> fingerprint
        {:error, _} -> System.unique_integer()
      end

    "supavisor_#{fingerprint}"
  end
end

1	defmodule Supavisor.Protocol.PreparedStatements do
2	@moduledoc """
3	Handles prepared statement binary packet transformations.
4	"""
5
6	alias Supavisor.Protocol.PreparedStatements.PreparedStatement
7	alias Supavisor.Protocol.PreparedStatements.Storage
8
9	@type statement_map() :: Storage.t()
10
11	@type statement_name() :: String.t()
12
13	@type pkt() :: binary()
14
15	@type handled_pkt() ::
16	{:parse_pkt, statement_name(), pkt()}
17	\| {:bind_pkt, statement_name(), bind_pkt :: pkt(), parse_pkt :: pkt()}
18	\| {:close_pkt, statement_name(), pkt()}
19	\| pkt()
20
21	@client_limit 100
22	@backend_limit 200
23	@client_memory_limit_bytes 1_000_000
24
25	@doc """
26	Upper limit of prepared statements from the client
27	"""
28	@spec client_limit() :: pos_integer()
29	def client_limit, do: @client_limit	4✔
30
31	@doc """
32	Upper limit of prepared statements backend-side.
33
34	Should rotate prepared statements to avoid surpassing it.
35	"""
36	@spec backend_limit() :: pos_integer()
37	def backend_limit, do: @backend_limit	3,593✔
38
39	@doc """
40	Upper limit of prepared statements memory from the client in bytes
41	"""
42	@spec client_memory_limit_bytes() :: pos_integer()
43	def client_memory_limit_bytes, do: @client_memory_limit_bytes	1✔
44
45	@doc """
46	Receives a statement name and returns a close packet for it
47	"""
48	@spec build_close_pkt(statement_name) :: pkt()
49	def build_close_pkt(statement_name) do
NEW 50	len = byte_size(statement_name)	×
NEW 51	<<?C, len + 6::32, ?S, statement_name::binary, 0>>	×
52	end
53
54	@doc """
55	Handles prepared statement packets and returns appropriate tuples for packets
56	that need special treatment according to the protocol.
57	"""
58	@spec handle_pkt(statement_map(), pkt()) ::
59	{:ok, statement_map(), handled_pkt()}
60	\| {:error, :max_prepared_statements}
61	\| {:error, :max_prepared_statements_memory}
62	\| {:error, :prepared_statement_on_simple_query}
63	\| {:error, :duplicate_prepared_statement, statement_name()}
64	\| {:error, :prepared_statement_not_found, statement_name()}
65	def handle_pkt(client_statements, binary) do
66	case binary do	13,761✔
67	# Parse message (P)
68	<<?P, len::32, rest::binary>> ->
69	handle_parse_message(client_statements, binary, len, rest)	2,037✔
70
71	# Bind message (B)
72	<<?B, len::32, rest::binary>> ->
73	handle_bind_message(client_statements, binary, len, rest)	2,217✔
74
75	# Close message (C)
76	<<?C, len::32, ?S, rest::binary>> ->
77	handle_close_message(client_statements, len, rest)	1,334✔
78
79	# Describe message (D)
80	<<?D, len::32, ?S, rest::binary>> ->
81	handle_describe_message(client_statements, len, rest)	2,032✔
82
83	# Query message (Q)
84	<<?Q, len::32, rest::binary>> ->
85	handle_simple_query_message(client_statements, binary, len, rest)	194✔
86
87	# All other messages pass through unchanged
88	_ ->
89	{:ok, client_statements, binary}	5,947✔
90	end
91	end
92
93	defp handle_parse_message(client_statements, original_bin, len, rest) do
94	case extract_null_terminated_string(rest) do	2,037✔
95	# Unnamed prepared statements are passed through unchanged
96	{"", _} ->
97	{:ok, client_statements, original_bin}	706✔
98
99	{client_side_name, remaining} ->
100	cond do	1,331✔
101	Storage.statement_count(client_statements) >= @client_limit ->	2✔
102	{:error, :max_prepared_statements}
103
104	Storage.statement_memory(client_statements) > @client_memory_limit_bytes ->	1,329✔
105	{:error, :max_prepared_statements_memory}
106
107	Storage.get(client_statements, client_side_name) ->	1,327✔
108	{:error, :duplicate_prepared_statement, client_side_name}	1✔
109
110	true ->	1,326✔
111	server_side_name = gen_server_side_name(rest)	1,326✔
112
113	new_len = len + (byte_size(server_side_name) - byte_size(client_side_name))	1,326✔
114	new_bin = <<?P, new_len::32, server_side_name::binary, 0, remaining::binary>>	1,326✔
115
116	prepared_statement = %PreparedStatement{	1,326✔
117	name: server_side_name,
118	parse_pkt: new_bin
119	}
120
121	new_client_statements =	1,326✔
122	Storage.put(client_statements, client_side_name, prepared_statement)
123
124	{:ok, new_client_statements, {:parse_pkt, server_side_name, new_bin}}	1,326✔
125	end
126	end
127	end
128
129	defp handle_bind_message(client_statements, bin, len, rest) do
130	{_portal_name, after_portal} = extract_null_terminated_string(rest)	2,217✔
131
132	case extract_null_terminated_string(after_portal) do	2,217✔
133	{"", _} ->
134	{:ok, client_statements, bin}	687✔
135
136	{client_side_name, packet_after_client_name} ->
137	case Storage.get(client_statements, client_side_name) do	1,530✔
138	%PreparedStatement{name: server_side_name, parse_pkt: parse_pkt} ->
139	new_len = len + (byte_size(server_side_name) - byte_size(client_side_name))	1,529✔
140
141	new_bin =	1,529✔
142	<<?B, new_len::32, 0, server_side_name::binary, 0,
143	packet_after_client_name::binary>>
144
145	{:ok, client_statements, {:bind_pkt, server_side_name, new_bin, parse_pkt}}	1,529✔
146
147	nil ->	1✔
148	{:error, :prepared_statement_not_found}
149	end
150	end
151	end
152
153	defp handle_close_message(client_statements, len, rest) do
154	{client_side_name, _} = extract_null_terminated_string(rest)	1,334✔
155
156	{prepared_statement, new_client_statements} = Storage.pop(client_statements, client_side_name)	1,334✔
157
158	server_side_name =	1,334✔
159	case prepared_statement do
160	%PreparedStatement{name: name} -> name	1✔
161	nil -> "supavisor_none"	1,333✔
162	end
163
164	new_len = len + (byte_size(server_side_name) - byte_size(client_side_name))	1,334✔
165	new_bin = <<?C, new_len::32, ?S, server_side_name::binary, 0>>	1,334✔
166
167	{:ok, new_client_statements, {:close_pkt, server_side_name, new_bin}}	1,334✔
168	end
169
170	defp handle_describe_message(client_statements, len, rest) do
171	{client_side_name, _} = extract_null_terminated_string(rest)	2,032✔
172
173	server_side_name =	2,032✔
174	case Storage.get(client_statements, client_side_name) do
175	%PreparedStatement{name: name} -> name	1,326✔
176	nil -> ""	706✔
177	end
178
179	new_len = len + (byte_size(server_side_name) - byte_size(client_side_name))	2,032✔
180	new_bin = <<?D, new_len::32, ?S, server_side_name::binary, 0>>	2,032✔
181
182	{:ok, client_statements, {:describe_pkt, server_side_name, new_bin}}	2,032✔
183	end
184
185	defp handle_simple_query_message(client_statements, binary, _len, rest) do
186	case rest do	194✔
187	"PREPARE" <> _ ->	1✔
188	{:error, :prepared_statement_on_simple_query}
189
190	_ ->
191	{:ok, client_statements, binary}	193✔
192	end
193	end
194
195	defp extract_null_terminated_string(binary) do
196	case :binary.split(binary, <<0>>) do	12,489✔
197	[string, rest] -> {string, rest}	12,489✔
NEW 198	[string] -> {string, <<>>}	×
199	end
200	end
201
202	defp gen_server_side_name(binary) do
203	{_name, rest} = extract_null_terminated_string(binary)	1,326✔
204	{query, _rest} = extract_null_terminated_string(rest)	1,326✔
205
206	fingerprint =	1,326✔
207	case Supavisor.PgParser.fingerprint(query) do
208	{:ok, fingerprint} -> fingerprint	1,326✔
NEW 209	{:error, _} -> System.unique_integer()	×
210	end
211
212	"supavisor_#{fingerprint}"	1,326✔
213	end
214	end

supabase / supavisor / 16470769502

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous