]> git.ipfire.org Git - thirdparty/psycopg.git/commitdiff
Added prepared statements and description
authorDaniele Varrazzo <daniele.varrazzo@gmail.com>
Mon, 16 Mar 2020 06:48:05 +0000 (19:48 +1300)
committerDaniele Varrazzo <daniele.varrazzo@gmail.com>
Mon, 16 Mar 2020 09:06:05 +0000 (22:06 +1300)
psycopg3/pq/_pq_ctypes.py
psycopg3/pq/pq_ctypes.py
tests/pq/test_exec.py
tests/pq/test_pgresult.py

index 64ed4d400bab41422fd5f08c87765d371087b8ad..ebd32059010e1a14a3b6869776cf7acf9e1ab303 100644 (file)
@@ -203,6 +203,30 @@ PQexecParams.argtypes = [
 ]
 PQexecParams.restype = PGresult_ptr
 
+PQprepare = pq.PQprepare
+PQprepare.argtypes = [PGconn_ptr, c_char_p, c_char_p, c_int, POINTER(Oid)]
+PQprepare.restype = PGresult_ptr
+
+PQexecPrepared = pq.PQexecPrepared
+PQexecPrepared.argtypes = [
+    PGconn_ptr,
+    c_char_p,
+    c_int,
+    POINTER(c_char_p),
+    POINTER(c_int),
+    POINTER(c_int),
+    c_int,
+]
+PQexecPrepared.restype = PGresult_ptr
+
+PQdescribePrepared = pq.PQdescribePrepared
+PQdescribePrepared.argtypes = [PGconn_ptr, c_char_p]
+PQdescribePrepared.restype = PGresult_ptr
+
+PQdescribePortal = pq.PQdescribePortal
+PQdescribePortal.argtypes = [PGconn_ptr, c_char_p]
+PQdescribePortal.restype = PGresult_ptr
+
 PQresultStatus = pq.PQresultStatus
 PQresultStatus.argtypes = [PGresult_ptr]
 PQresultStatus.restype = c_int
index d01edb31b5701846cc4265d1625a74c44e9597b8..64032c09fd9cd962f6aa045c503e1840823f13e8 100644 (file)
@@ -239,6 +239,82 @@ class PGconn:
             raise MemoryError("couldn't allocate PGresult")
         return PGresult(rv)
 
+    def prepare(
+        self, name, command, param_types=None,
+    ):
+        if not isinstance(name, bytes):
+            raise TypeError(
+                "'name' must be bytes, got %s instead" % type(name).__name__
+            )
+
+        if not isinstance(command, bytes):
+            raise TypeError(
+                "'command' must be bytes, got %s instead"
+                % type(command).__name__
+            )
+
+        if param_types is None:
+            nparams = 0
+            atypes = None
+        else:
+            nparams = len(param_types)
+            atypes = (impl.Oid * nparams)(*param_types)
+
+        rv = impl.PQprepare(self.pgconn_ptr, name, command, nparams, atypes,)
+        if rv is None:
+            raise MemoryError("couldn't allocate PGresult")
+        return PGresult(rv)
+
+    def exec_prepared(
+        self, name, param_values, param_formats=None, result_format=0,
+    ):
+        if not isinstance(name, bytes):
+            raise TypeError(
+                "'name' must be bytes, got %s instead" % type(name).__name__
+            )
+
+        nparams = len(param_values)
+        if nparams:
+            aparams = (c_char_p * nparams)(*param_values)
+            alenghts = (c_int * nparams)(
+                *(len(p) if p is not None else 0 for p in param_values)
+            )
+        else:
+            aparams = alenghts = None
+
+        if param_formats is None:
+            aformats = None
+        else:
+            if len(param_formats) != nparams:
+                raise ValueError(
+                    "got %d param_values but %d param_types"
+                    % (nparams, len(param_formats))
+                )
+            aformats = (c_int * nparams)(*param_formats)
+
+        rv = impl.PQexecPrepared(
+            self.pgconn_ptr,
+            name,
+            nparams,
+            aparams,
+            alenghts,
+            aformats,
+            result_format,
+        )
+        if rv is None:
+            raise MemoryError("couldn't allocate PGresult")
+        return PGresult(rv)
+
+    def describe_prepared(self, name):
+        if not isinstance(name, bytes):
+            raise TypeError(
+                "'name' must be bytes, got %s instead" % type(name).__name__
+            )
+        rv = impl.PQdescribePrepared(self.pgconn_ptr, name)
+        if rv is None:
+            raise MemoryError("couldn't allocate PGresult")
+        return PGresult(rv)
+
 
 class PGresult:
     __slots__ = ("pgresult_ptr",)
@@ -310,6 +386,13 @@ class PGresult:
             else:
                 return b""
 
+    @property
+    def nparams(self):
+        return impl.PQnparams(self.pgresult_ptr)
+
+    def param_type(self, param_number):
+        return impl.PQparamtype(self.pgresult_ptr, param_number)
+
 
 ConninfoOption = namedtuple(
     "ConninfoOption", "keyword envvar compiled val label dispatcher dispsize"
index cfbe6de2dd0d55c0e4aaccaece3e794a2fbff0ba..73d26db0e4ff4fa1ff3bf86e4b9ab91becff52bf 100644 (file)
@@ -70,3 +70,47 @@ def test_exec_params_binary_out(pq, pgconn, fmt, out):
     )
     assert res.status == pq.ExecStatus.PGRES_TUPLES_OK
     assert res.get_value(0, 0) == out
+
+
+def test_prepare(pq, pgconn):
+    res = pgconn.prepare(b"prep", b"select $1::int + $2::int")
+    assert res.status == pq.ExecStatus.PGRES_COMMAND_OK, res.error_message
+
+    res = pgconn.exec_prepared(b"prep", [b"3", b"5"])
+    assert res.get_value(0, 0) == b"8"
+
+
+def test_prepare_types(pq, pgconn):
+    res = pgconn.prepare(b"prep", b"select $1 + $2", [23, 23])
+    assert res.status == pq.ExecStatus.PGRES_COMMAND_OK, res.error_message
+
+    res = pgconn.exec_prepared(b"prep", [b"3", b"5"])
+    assert res.get_value(0, 0) == b"8"
+
+
+def test_exec_prepared_binary_in(pq, pgconn):
+    val = b"foo\00bar"
+    res = pgconn.prepare(b"", b"select length($1::bytea), length($2::bytea)")
+
+    res = pgconn.exec_prepared(b"", [val, val], param_formats=[0, 1])
+    assert res.status == pq.ExecStatus.PGRES_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_exec_prepared_binary_out(pq, pgconn, fmt, out):
+    val = b"foo\00bar"
+    res = pgconn.prepare(b"", b"select $1::bytea")
+    assert res.status == pq.ExecStatus.PGRES_COMMAND_OK, res.error_message
+
+    res = pgconn.exec_prepared(
+        b"", [val], param_formats=[1], result_format=fmt
+    )
+    assert res.status == pq.ExecStatus.PGRES_TUPLES_OK
+    assert res.get_value(0, 0) == out
index 530a9748ba347fd6babd8ef367181759e49e4bdc..e151d87992bddc2b423a50bb04ddaf4035979719 100644 (file)
@@ -111,3 +111,15 @@ def test_get_value(pq, pgconn):
     assert res.get_value(0, 0) == b"a"
     assert res.get_value(0, 1) == b""
     assert res.get_value(0, 2) is None
+
+
+def test_nparams_types(pq, pgconn):
+    res = pgconn.prepare(b"", b"select $1::int, $2::text")
+    assert res.status == pq.ExecStatus.PGRES_COMMAND_OK, res.error_message
+
+    res = pgconn.describe_prepared(b"")
+    assert res.status == pq.ExecStatus.PGRES_COMMAND_OK, res.error_message
+
+    assert res.nparams == 2
+    assert res.param_type(0) == 23
+    assert res.param_type(1) == 25