fix(protocol): handle SQL PREPARE/EXECUTE over extended protocol

integration/elixir/test/prepared_test.exs

@@ -1,6 +1,27 @@
 defmodule Pgdog.PreparedTest do
   use ExUnit.Case, async: false
 
+  test "SQL PREPARE and EXECUTE preserve the extended protocol" do
+    for options <- [[], [prepare: :unnamed]] do
+      conn = Pgdog.connect(options)
+      name = "elixir_sql_#{System.unique_integer([:positive])}"
+
+      try do
+        prepare = Postgrex.prepare!(conn, "prepare_command", "PREPARE #{name} AS SELECT $1::bigint * 2")
+        assert %Postgrex.Result{command: :prepare} = Postgrex.execute!(conn, prepare, [])
+        execute = Postgrex.prepare!(conn, "execute_command", "EXECUTE #{name}(21)")
+
+        for _ <- 1..3 do
+          assert Pgdog.one(Postgrex.execute!(conn, execute, [])) == 42
+        end
+
+        assert Pgdog.one(Postgrex.query!(conn, "SELECT 1", [])) == 1
+      after
+        GenServer.stop(conn)
+      end
+    end
+  end
+
   test "a named statement is reusable across many executions" do
     conn = Pgdog.connect()
     query = Postgrex.prepare!(conn, "elixir_echo", "SELECT $1::bigint")

integration/python/test_extended_sql_prepare.py

@@ -0,0 +1,149 @@
+import uuid
+
+import asyncpg
+import pytest
+from globals import admin, no_out_of_sync, normal_async, sharded_async
+
+
+@pytest.fixture
+def full_prepared_statements():
+    with admin() as connection:
+        connection.execute("SET prepared_statements TO 'full'")
+        try:
+            yield
+        finally:
+            connection.execute("RELOAD")
+
+
+@pytest.fixture
+def rewritten_prepared_statements(full_prepared_statements):
+    with admin() as connection:
+        connection.execute("SET rewrite_enabled TO true")
+    yield
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("connect", [normal_async, sharded_async])
+@pytest.mark.parametrize("extended_prepare", [False, True])
+async def test_sql_prepare_execute_extended(connect, extended_prepare, full_prepared_statements):
+    connection = await connect()
+    name = "extended_sql_" + uuid.uuid4().hex
+    try:
+        sql = f"PREPARE {name} AS SELECT $1::integer * 2"
+        if extended_prepare:
+            outer = await connection.prepare(sql, timeout=5)
+            assert outer.get_parameters() == ()
+            assert outer.get_attributes() == ()
+            for _ in range(3):
+                assert await outer.fetch(timeout=5) == []
+                assert outer.get_statusmsg() == "PREPARE"
+        else:
+            assert await connection.execute(sql, timeout=5) == "PREPARE"
+
+        execute = await connection.prepare(f"EXECUTE {name}(21)", timeout=5)
+        for _ in range(3):
+            assert [tuple(row) for row in await execute.fetch(timeout=5)] == [(42,)]
+        assert await connection.fetchval("SELECT 1", timeout=5) == 1
+        no_out_of_sync()
+    finally:
+        await connection.close(timeout=5)
+
+
+@pytest.mark.asyncio
+async def test_extended_sql_prepare_error_recovers(full_prepared_statements):
+    connection = await normal_async()
+    name = "extended_sql_error_" + uuid.uuid4().hex
+    try:
+        outer = await connection.prepare(
+            f"PREPARE {name} AS SELECT nonexistent_issue1403_column", timeout=5
+        )
+        with pytest.raises(asyncpg.UndefinedColumnError):
+            await outer.fetch(timeout=5)
+        assert await connection.fetchval("SELECT 1", timeout=5) == 1
+        no_out_of_sync()
+    finally:
+        await connection.close(timeout=5)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("extended_prepare", [False, True])
+async def test_extended_execute_limit_offset(extended_prepare, rewritten_prepared_statements):
+    connection = await sharded_async()
+    schema = "extended_limit_" + uuid.uuid4().hex
+    try:
+        await connection.execute(f'CREATE SCHEMA "{schema}"')
+        await connection.execute(f'CREATE TABLE "{schema}".sharded (id BIGINT PRIMARY KEY)')
+        for value in range(1, 56):
+            await connection.execute(f'INSERT INTO "{schema}".sharded VALUES ($1)', value)
+        cases = [
+            ("", "LIMIT 5 OFFSET $1", "(10)", list(range(45, 40, -1))),
+            ("", "LIMIT $2 OFFSET $1", "(5, 10)", list(range(50, 40, -1))),
+            ("", "LIMIT $1 OFFSET $2", "(5, 10)", list(range(45, 40, -1))),
+            ("", "LIMIT 10 OFFSET 5", "", list(range(50, 40, -1))),
+            ("WHERE id < $2", "LIMIT $3 OFFSET $1", "(5, 25, 10)", list(range(19, 9, -1))),
+            ("WHERE id = 35", "LIMIT 1 OFFSET 0", "", [35]),
+        ]
+        for index, (predicate, clause, arguments, expected) in enumerate(cases):
+            name = f"{schema}_{index}"
+            sql = f'PREPARE {name} AS SELECT id FROM "{schema}".sharded {predicate} ORDER BY id DESC {clause}'
+            if extended_prepare:
+                prepare = await connection.prepare(sql, timeout=5)
+                await prepare.fetch(timeout=5)
+            else:
+                await connection.execute(sql, timeout=5)
+            execute = await connection.prepare(f"EXECUTE {name}{arguments}", timeout=5)
+            for _ in range(3):
+                assert [row[0] for row in await execute.fetch(timeout=5)] == expected
+        no_out_of_sync()
+    finally:
+        await connection.execute(f'DROP SCHEMA IF EXISTS "{schema}" CASCADE')
+        await connection.close(timeout=5)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("extended_prepare", [False, True])
+async def test_extended_execute_generated_values(extended_prepare, rewritten_prepared_statements):
+    connection = await sharded_async()
+    name = "extended_id_" + uuid.uuid4().hex
+    try:
+        sql = f"PREPARE {name} AS SELECT pgdog.unique_id()"
+        if extended_prepare:
+            prepare = await connection.prepare(sql, timeout=5)
+            await prepare.fetch(timeout=5)
+        else:
+            await connection.execute(sql, timeout=5)
+        execute = await connection.prepare(f"EXECUTE {name}", timeout=5)
+        values = [await execute.fetchval(timeout=5) for _ in range(10)]
+        assert len(set(values)) == len(values), "each execution must generate a fresh ID"
+        assert all(isinstance(value, int) for value in values)
+        no_out_of_sync()
+    finally:
+        await connection.close(timeout=5)
+
+
+@pytest.mark.asyncio
+async def test_sql_prepare_registers_each_clients_name(full_prepared_statements):
+    connections = [await normal_async(), await normal_async()]
+    name = "shared_sql_" + uuid.uuid4().hex
+    try:
+        for value, connection in enumerate(connections):
+            prepare = await connection.prepare(f"PREPARE {name} AS SELECT $1::integer", timeout=5)
+            await prepare.fetch(timeout=5)
+            execute = await connection.prepare(f"EXECUTE {name}({value})", timeout=5)
+            assert await execute.fetchval(timeout=5) == value
+    finally:
+        for connection in connections:
+            await connection.close(timeout=5)
+
+
+@pytest.mark.asyncio
+async def test_repeated_unnamed_sql_execute(full_prepared_statements):
+    connection = await normal_async()
+    name = "unnamed_sql_" + uuid.uuid4().hex
+    try:
+        await connection.execute(f"PREPARE {name} AS SELECT 42::integer")
+        execute = await connection.prepare(f"EXECUTE {name}", name="", timeout=5)
+        for _ in range(3):
+            assert await execute.fetchval(timeout=5) == 42
+    finally:
+        await connection.close(timeout=5)

integration/ruby/prepared_disabled/prepared_spec.rb

@@ -1,12 +1,15 @@
 # frozen_string_literal: true
 
 require_relative '../rspec_helper'
+require_relative '../sql_prepare_examples'
 
 # With prepared_statements = "disabled" pgdog forwards protocol messages as-is
 # without rewriting or caching.
 describe 'prepared_statements = disabled' do
   after { ensure_done }
 
+  it_behaves_like 'SQL PREPARE over extended protocol', 'pgdog_session'
+
   # Anonymous statements (empty name) are a single Parse+Bind+Execute+Sync
   # cycle on one backend — no state needs to survive across cycles.
   it 'executes anonymous parameterized queries' do

integration/ruby/prepared_extended/prepared_spec.rb

@@ -1,6 +1,7 @@
 # frozen_string_literal: true
 
 require_relative '../rspec_helper'
+require_relative '../sql_prepare_examples'
 
 # Uses the main integration pgdog.toml which sets prepared_statements = "extended".
 # "extended" rewrites and replays named extended-protocol statements (Parse/Bind)
@@ -10,6 +11,8 @@
 describe 'prepared_statements = extended' do
   after { ensure_done }
 
+  it_behaves_like 'SQL PREPARE over extended protocol', 'pgdog_session'
+
   # Anonymous statements (empty name) are a single Parse+Bind+Execute+Sync
   # cycle on one backend — no state needs to survive across cycles.
   it 'executes anonymous parameterized queries' do

integration/ruby/prepared_full/prepared_spec.rb

@@ -1,10 +1,13 @@
 # frozen_string_literal: true
 
 require_relative '../rspec_helper'
+require_relative '../sql_prepare_examples'
 
 describe 'prepared_statements = full' do
   after { ensure_done }
 
+  it_behaves_like 'SQL PREPARE over extended protocol', 'pgdog'
+
   # Mirror of disabled suite: anonymous statements carry no per-backend state,
   # so they work identically regardless of the prepared_statements setting.
   it 'executes anonymous parameterized queries' do

integration/ruby/sql_prepare_examples.rb

@@ -0,0 +1,29 @@
+# frozen_string_literal: true
+
+shared_examples 'SQL PREPARE over extended protocol' do |user|
+  it 'executes SQL PREPARE and EXECUTE through named statements' do
+    conn = connect('pgdog', user)
+    name = "sql_extended_#{SecureRandom.hex(6)}"
+    conn.prepare('prepare_command', "PREPARE #{name} AS SELECT $1::bigint * 2 AS val")
+    expect(conn.exec_prepared('prepare_command', []).cmd_status).to eq('PREPARE')
+    conn.prepare('execute_command', "EXECUTE #{name}(21)")
+    3.times do
+      expect(conn.exec_prepared('execute_command', [])[0]['val'].to_i).to eq(42)
+    end
+    expect(conn.exec('SELECT 1')[0].values).to eq(['1'])
+  ensure
+    conn.close if conn && !conn.finished?
+  end
+
+  it 'executes SQL PREPARE and EXECUTE through unnamed statements' do
+    conn = connect('pgdog', user)
+    name = "sql_unnamed_#{SecureRandom.hex(6)}"
+    expect(conn.exec_params("PREPARE #{name} AS SELECT $1::bigint * 2 AS val", []).cmd_status).to eq('PREPARE')
+    3.times do
+      expect(conn.exec_params("EXECUTE #{name}(21)", [])[0]['val'].to_i).to eq(42)
+    end
+    expect(conn.exec('SELECT 1')[0].values).to eq(['1'])
+  ensure
+    conn.close if conn && !conn.finished?
+  end
+end

integration/rust/tests/integration/simple_prepared.rs

@@ -34,7 +34,6 @@ async fn test_simple_prepared_ttl() {
 }
 
 /// <https://github.com/pgdogdev/pgdog/issues/1383>
-/// TODO: will need to support extended-protocol `Bind`s later for the re-write
 #[tokio::test]
 async fn test_simple_prepared_limit() {
     let mut conn =