Added async query prepare functions in pq wrapper

diff --git a/psycopg3/pq/_pq_ctypes.py b/psycopg3/pq/_pq_ctypes.py
index 0d232466db1cce011cba14a65a47770c6b960c26..b292189771be5f34763975ea1d12cfc13153bc96 100644 (file)
--- a/psycopg3/pq/_pq_ctypes.py
+++ b/psycopg3/pq/_pq_ctypes.py
@@ -407,8 +407,23 @@ PQsendQueryParams.argtypes = [
 ]
 PQsendQueryParams.restype = c_int
 
-# TODO: PQsendPrepare PQsendQueryPrepared
-#       PQsendDescribePrepared PQsendDescribePortal
+PQsendPrepare = pq.PQsendPrepare
+PQsendPrepare.argtypes = [PGconn_ptr, c_char_p, c_char_p, c_int, POINTER(Oid)]
+PQsendPrepare.restype = c_int
+
+PQsendQueryPrepared = pq.PQsendQueryPrepared
+PQsendQueryPrepared.argtypes = [
+    PGconn_ptr,
+    c_char_p,
+    c_int,
+    POINTER(c_char_p),
+    POINTER(c_int),
+    POINTER(c_int),
+    c_int,
+]
+PQsendQueryPrepared.restype = c_int
+
+# TODO: PQsendDescribePrepared PQsendDescribePortal
 
 PQgetResult = pq.PQgetResult
 PQgetResult.argtypes = [PGconn_ptr]

diff --git a/psycopg3/pq/_pq_ctypes.pyi b/psycopg3/pq/_pq_ctypes.pyi
index a47e4530e2e012a77ebf42e4a92210d5ad6fd0bc..1a976fabd105101894e9f04476bdce9dd8c4569e 100644 (file)
--- a/psycopg3/pq/_pq_ctypes.pyi
+++ b/psycopg3/pq/_pq_ctypes.pyi
@@ -45,6 +45,22 @@ def PQgetvalue(
     arg1: Optional[PGresult_struct], arg2: int, arg3: int
 ) -> pointer[c_char]: ...
 def PQcmdTuples(arg1: Optional[PGresult_struct]) -> bytes: ...
+def PQsendPrepare(
+    arg1: Optional[PGconn_struct],
+    arg2: bytes,
+    arg3: bytes,
+    arg4: int,
+    arg5: Optional[Array[c_uint]],
+) -> int: ...
+def PQsendQueryPrepared(
+    arg1: Optional[PGconn_struct],
+    arg2: bytes,
+    arg3: int,
+    arg4: Optional[Array[c_char_p]],
+    arg5: Optional[Array[c_int]],
+    arg6: Optional[Array[c_int]],
+    arg7: int,
+) -> int: ...
 
 # fmt: off
 # autogenerated: start

diff --git a/psycopg3/pq/pq_ctypes.py b/psycopg3/pq/pq_ctypes.py
index af0df71b908f528685c369b8dc4858d18dc809a3..f36cf9f34b7c4d880532793a5803be4e24a2011a 100644 (file)
--- a/psycopg3/pq/pq_ctypes.py
+++ b/psycopg3/pq/pq_ctypes.py
@@ -232,6 +232,48 @@ class PGconn:
                 f"sending query and params failed: {error_message(self)}"
             )
 
+    def send_prepare(
+        self,
+        name: bytes,
+        command: bytes,
+        param_types: Optional[Sequence[int]] = None,
+    ) -> None:
+        atypes: Optional[Array[impl.Oid]]
+        if param_types is None:
+            nparams = 0
+            atypes = None
+        else:
+            nparams = len(param_types)
+            atypes = (impl.Oid * nparams)(*param_types)
+
+        self._ensure_pgconn()
+        if not impl.PQsendPrepare(
+            self.pgconn_ptr, name, command, nparams, atypes
+        ):
+            raise PQerror(
+                f"sending query and params failed: {error_message(self)}"
+            )
+
+    def send_query_prepared(
+        self,
+        name: bytes,
+        param_values: Sequence[Optional[bytes]],
+        param_formats: Optional[Sequence[Format]] = None,
+        result_format: Format = Format.TEXT,
+    ) -> None:
+        # repurpose this function with a cheeky replacement of query with name,
+        # drop the param_types from the result
+        args = self._query_params_args(
+            name, param_values, None, param_formats, result_format
+        )
+        args = args[:3] + args[4:]
+
+        self._ensure_pgconn()
+        if not impl.PQsendQueryPrepared(*args):
+            raise PQerror(
+                f"sending prepared query failed: {error_message(self)}"
+            )
+
     def _query_params_args(
         self,
         command: bytes,

diff --git a/tests/pq/test_async.py b/tests/pq/test_async.py
index 7cedc9574ec5ccd703c81376792b163160c24034..969e406c725f0c05ae35cef9981b30f3dcfeaee5 100644 (file)
--- a/tests/pq/test_async.py
+++ b/tests/pq/test_async.py
@@ -77,7 +77,7 @@ def test_send_query_compact_test(pq, pgconn):
 
 
 def test_send_query_params(pq, pgconn):
-    res = pgconn.send_query_params(b"select $1::int + $2", [b"5", b"3"])
+    pgconn.send_query_params(b"select $1::int + $2", [b"5", b"3"])
     (res,) = psycopg3.waiting.wait(psycopg3.generators.execute(pgconn))
     assert res.status == pq.ExecStatus.TUPLES_OK
     assert res.get_value(0, 0) == b"8"
@@ -85,3 +85,62 @@ def test_send_query_params(pq, pgconn):
     pgconn.finish()
     with pytest.raises(psycopg3.OperationalError):
         pgconn.send_query_params(b"select $1", [b"1"])
+
+
+def test_send_prepare(pq, pgconn):
+    pgconn.send_prepare(b"prep", b"select $1::int + $2::int")
+    (res,) = psycopg3.waiting.wait(psycopg3.generators.execute(pgconn))
+    assert res.status == pq.ExecStatus.COMMAND_OK, res.error_message
+
+    pgconn.send_query_prepared(b"prep", [b"3", b"5"])
+    (res,) = psycopg3.waiting.wait(psycopg3.generators.execute(pgconn))
+    assert res.get_value(0, 0) == b"8"
+
+    pgconn.finish()
+    with pytest.raises(psycopg3.OperationalError):
+        pgconn.send_prepare(b"prep", b"select $1::int + $2::int")
+    with pytest.raises(psycopg3.OperationalError):
+        pgconn.send_query_prepared(b"prep", [b"3", b"5"])
+
+
+def test_send_prepare_types(pq, pgconn):
+    pgconn.send_prepare(b"prep", b"select $1 + $2", [23, 23])
+    (res,) = psycopg3.waiting.wait(psycopg3.generators.execute(pgconn))
+    assert res.status == pq.ExecStatus.COMMAND_OK, res.error_message
+
+    pgconn.send_query_prepared(b"prep", [b"3", b"5"])
+    (res,) = psycopg3.waiting.wait(psycopg3.generators.execute(pgconn))
+    assert res.get_value(0, 0) == b"8"
+
+
+def test_send_prepared_binary_in(pq, pgconn):
+    val = b"foo\00bar"
+    pgconn.send_prepare(b"", b"select length($1::bytea), length($2::bytea)")
+    (res,) = psycopg3.waiting.wait(psycopg3.generators.execute(pgconn))
+    assert res.status == pq.ExecStatus.COMMAND_OK, res.error_message
+
+    pgconn.send_query_prepared(b"", [val, val], param_formats=[0, 1])
+    (res,) = psycopg3.waiting.wait(psycopg3.generators.execute(pgconn))
+    assert res.status == pq.ExecStatus.TUPLES_OK
+    assert res.get_value(0, 0) == b"3"
+    assert res.get_value(0, 1) == b"7"
+
+    with pytest.raises(ValueError):
+        pgconn.exec_params(b"select $1::bytea", [val], param_formats=[1, 1])
+
+
+@pytest.mark.parametrize(
+    "fmt, out", [(0, b"\\x666f6f00626172"), (1, b"foo\00bar")]
+)
+def test_send_prepared_binary_out(pq, pgconn, fmt, out):
+    val = b"foo\00bar"
+    pgconn.send_prepare(b"", b"select $1::bytea")
+    (res,) = psycopg3.waiting.wait(psycopg3.generators.execute(pgconn))
+    assert res.status == pq.ExecStatus.COMMAND_OK, res.error_message
+
+    pgconn.send_query_prepared(
+        b"", [val], param_formats=[1], result_format=fmt
+    )
+    (res,) = psycopg3.waiting.wait(psycopg3.generators.execute(pgconn))
+    assert res.status == pq.ExecStatus.TUPLES_OK
+    assert res.get_value(0, 0) == out

diff --git a/tests/pq/test_exec.py b/tests/pq/test_exec.py
index d575d6b063957b2530ea8b94453c6c1a9dde3d62..f090e0afd7214da112dab288c163c0f194c9e730 100644 (file)
--- a/tests/pq/test_exec.py
+++ b/tests/pq/test_exec.py
@@ -107,6 +107,7 @@ def test_prepare_types(pq, pgconn):
 def test_exec_prepared_binary_in(pq, pgconn):
     val = b"foo\00bar"
     res = pgconn.prepare(b"", b"select length($1::bytea), length($2::bytea)")
+    assert res.status == pq.ExecStatus.COMMAND_OK, res.error_message
 
     res = pgconn.exec_prepared(b"", [val, val], param_formats=[0, 1])
     assert res.status == pq.ExecStatus.TUPLES_OK