Wrapping PQexec and PGresult

diff --git a/psycopg3/_pq_ctypes.py b/psycopg3/_pq_ctypes.py
index 21358b72639bee57499d319382e3a74439717157..a4eec66cb9272f4094f1421a01fb6d0112d1463b 100644 (file)
--- a/psycopg3/_pq_ctypes.py
+++ b/psycopg3/_pq_ctypes.py
@@ -19,6 +19,10 @@ class PGconn_struct(Structure):
     _fields_ = []
 
 
+class PGresult_struct(Structure):
+    _fields_ = []
+
+
 class PQconninfoOption_struct(Structure):
     _fields_ = [
         ("keyword", c_char_p),
@@ -32,6 +36,7 @@ class PQconninfoOption_struct(Structure):
 
 
 PGconn_ptr = POINTER(PGconn_struct)
+PGresult_ptr = POINTER(PGresult_struct)
 PQconninfoOption_ptr = POINTER(PQconninfoOption_struct)
 
 
@@ -173,6 +178,23 @@ PQsslInUse = pq.PQsslInUse
 PQsslInUse.argtypes = [PGconn_ptr]
 PQsslInUse.restype = c_int
 
+# TODO: PQsslAttribute, PQsslAttributeNames, PQsslStruct, PQgetssl
+
+
+# 33.3. Command Execution Functions
+
+PQexec = pq.PQexec
+PQexec.argtypes = [PGconn_ptr, c_char_p]
+PQexec.restype = PGresult_ptr
+
+PQresultStatus = pq.PQresultStatus
+PQresultStatus.argtypes = [PGresult_ptr]
+PQresultStatus.restype = c_int
+
+PQclear = pq.PQclear
+PQclear.argtypes = [PGresult_ptr]
+PQclear.restype = None
+
 
 # 33.11. Miscellaneous Functions
 

diff --git a/psycopg3/pq.py b/psycopg3/pq.py
index 87daa745b7f2634a96bbdd3ebc2fc4d4e57d2f0b..4832b10542605a100589d727805143bfd66a3481 100644 (file)
--- a/psycopg3/pq.py
+++ b/psycopg3/pq.py
@@ -12,6 +12,7 @@ implementation-dependant but all the implementations share the same interface.
 from .pq_enums import (
     ConnStatus,
     PollingStatus,
+    ExecStatus,
     TransactionStatus,
     Ping,
 )
@@ -26,6 +27,7 @@ __all__ = (
     "ConnStatus",
     "PollingStatus",
     "TransactionStatus",
+    "ExecStatus",
     "Ping",
     "PGconn",
     "Conninfo",

diff --git a/psycopg3/pq_ctypes.py b/psycopg3/pq_ctypes.py
index 5469a3383850975cb576436ed63b41df0a60286e..24f3d300804728ab78c76406ffc7f50234349ba0 100644 (file)
--- a/psycopg3/pq_ctypes.py
+++ b/psycopg3/pq_ctypes.py
@@ -15,6 +15,7 @@ from ctypes import c_char_p, pointer
 from .pq_enums import (
     ConnStatus,
     PollingStatus,
+    ExecStatus,
     TransactionStatus,
     Ping,
 )
@@ -173,6 +174,12 @@ class PGconn:
     def ssl_in_use(self):
         return bool(impl.PQsslInUse(self.pgconn_ptr))
 
+    def exec_(self, command):
+        rv = impl.PQexec(self.pgconn_ptr, self._encode(command))
+        if rv is None:
+            raise MemoryError("couldn't allocate PGresult")
+        return PGresult(rv)
+
     def _encode(self, s):
         if isinstance(s, bytes):
             return s
@@ -189,6 +196,26 @@ class PGconn:
         return b.decode("utf8", "replace")
 
 
+class PGresult:
+    __slots__ = ("pgresult_ptr",)
+
+    def __init__(self, pgresult_ptr):
+        self.pgresult_ptr = pgresult_ptr
+
+    def __del__(self):
+        self.clear()
+
+    def clear(self):
+        self.pgresult_ptr, p = None, self.pgresult_ptr
+        if p is not None:
+            impl.PQclear(p)
+
+    @property
+    def status(self):
+        rv = impl.PQresultStatus(self.pgresult_ptr)
+        return ExecStatus(rv)
+
+
 ConninfoOption = namedtuple(
     "ConninfoOption", "keyword envvar compiled val label dispatcher dispsize"
 )

diff --git a/psycopg3/pq_enums.py b/psycopg3/pq_enums.py
index cafac184b189b81700fce787626d07724bc9b597..06b4ce0313293e97f3de7e859b84db0551006cef 100644 (file)
--- a/psycopg3/pq_enums.py
+++ b/psycopg3/pq_enums.py
@@ -32,6 +32,19 @@ class PollingStatus(IntEnum):
     PGRES_POLLING_ACTIVE = auto()
 
 
+class ExecStatus(IntEnum):
+    PGRES_EMPTY_QUERY = 0
+    PGRES_COMMAND_OK = auto()
+    PGRES_TUPLES_OK = auto()
+    PGRES_COPY_OUT = auto()
+    PGRES_COPY_IN = auto()
+    PGRES_BAD_RESPONSE = auto()
+    PGRES_NONFATAL_ERROR = auto()
+    PGRES_FATAL_ERROR = auto()
+    PGRES_COPY_BOTH = auto()
+    PGRES_SINGLE_TUPLE = auto()
+
+
 class TransactionStatus(IntEnum):
     PQTRANS_IDLE = 0
     PQTRANS_ACTIVE = auto()

diff --git a/tests/test_pq_exec.py b/tests/test_pq_exec.py
new file mode 100644 (file)
index 0000000..cc62ba3
--- /dev/null
+++ b/tests/test_pq_exec.py
@@ -0,0 +1,11 @@
+#!/usr/bin/env python3
+
+
+def test_exec_empty(pq, pgconn):
+    res = pgconn.exec_(b"")
+    assert res.status == pq.ExecStatus.PGRES_EMPTY_QUERY
+
+
+def test_exec_command(pq, pgconn):
+    res = pgconn.exec_("set timezone to utc")
+    assert res.status == pq.ExecStatus.PGRES_COMMAND_OK