From: Daniele Varrazzo <daniele.varrazzo@gmail.com>
Date: Sun, 15 Mar 2020 03:04:36 +0000 (+1300)
Subject: Wrapping PQexec and PGresult
X-Git-Tag: 3.0.dev0~714
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=68fc8c3c4015ad256c76df1546556dab8c30d5bb;p=thirdparty%2Fpsycopg.git

Wrapping PQexec and PGresult
---

diff --git a/psycopg3/_pq_ctypes.py b/psycopg3/_pq_ctypes.py
index 21358b726..a4eec66cb 100644
--- 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 87daa745b..4832b1054 100644
--- 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 5469a3383..24f3d3008 100644
--- 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 cafac184b..06b4ce031 100644
--- 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
index 000000000..cc62ba3b8
--- /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