Added PQconn.reset sync and async

diff --git a/psycopg3/_pq_ctypes.py b/psycopg3/_pq_ctypes.py
index d310b629ade8e5bd13ca08e3617f9e54a114e8af..0a09c52d3aba5c388a8d03f34186c0aeb873e025 100644 (file)
--- a/psycopg3/_pq_ctypes.py
+++ b/psycopg3/_pq_ctypes.py
@@ -74,6 +74,18 @@ PQfinish = pq.PQfinish
 PQfinish.argtypes = [PGconn_ptr]
 PQfinish.restype = None
 
+PQreset = pq.PQreset
+PQreset.argtypes = [PGconn_ptr]
+PQreset.restype = None
+
+PQresetStart = pq.PQresetStart
+PQresetStart.argtypes = [PGconn_ptr]
+PQresetStart.restype = c_int
+
+PQresetPoll = pq.PQresetPoll
+PQresetPoll.argtypes = [PGconn_ptr]
+PQresetPoll.restype = c_int
+
 
 # 33.2. Connection Status Functions
 

diff --git a/psycopg3/pq_ctypes.py b/psycopg3/pq_ctypes.py
index 55dd598728d99770d040c6fcdcc4d959a3a0a16e..1e1d02bc57c704003454bd0430a5bca7590213df 100644 (file)
--- a/psycopg3/pq_ctypes.py
+++ b/psycopg3/pq_ctypes.py
@@ -14,6 +14,10 @@ from .pq_enums import ConnStatus, PostgresPollingStatus
 from . import _pq_ctypes as impl
 
 
+class PQerror(Exception):
+    pass
+
+
 class PGconn:
     __slots__ = ("pgconn_ptr",)
 
@@ -71,6 +75,18 @@ class PGconn:
         finally:
             impl.PQconninfoFree(opts)
 
+    def reset(self):
+        impl.PQreset(self.pgconn_ptr)
+
+    def reset_start(self):
+        rv = impl.PQresetStart(self.pgconn_ptr)
+        if rv == 0:
+            raise PQerror("couldn't reset connection")
+
+    def reset_poll(self):
+        rv = impl.PQresetPoll(self.pgconn_ptr)
+        return PostgresPollingStatus(rv)
+
     @property
     def status(self):
         rv = impl.PQstatus(self.pgconn_ptr)

diff --git a/tests/test_pq.py b/tests/test_pq.py
index 464d0f138ab0150b3da733279dc338b2f81570a1..3fd94ba3af9b999ad02b4f4de68973d3a88c120a 100644 (file)
--- a/tests/test_pq.py
+++ b/tests/test_pq.py
@@ -92,3 +92,29 @@ def test_info(pq, dsn):
     assert dbname.label == "Database-Name"
     assert dbname.dispatcher == ""
     assert dbname.dispsize == 20
+
+
+def test_reset(pq, dsn):
+    conn = pq.PGconn.connect(dsn)
+    assert conn.status == ConnStatus.CONNECTION_OK
+    # TODO: break it
+    conn.reset()
+    assert conn.status == ConnStatus.CONNECTION_OK
+
+
+def test_reset_async(pq, dsn):
+    conn = pq.PGconn.connect(dsn)
+    assert conn.status == ConnStatus.CONNECTION_OK
+    # TODO: break it
+    conn.reset_start()
+    while 1:
+        rv = conn.connect_poll()
+        if rv == PostgresPollingStatus.PGRES_POLLING_READING:
+            select([conn.socket], [], [])
+        elif rv == PostgresPollingStatus.PGRES_POLLING_WRITING:
+            select([], [conn.socket], [])
+        else:
+            break
+
+    assert rv == PostgresPollingStatus.PGRES_POLLING_OK
+    assert conn.status == ConnStatus.CONNECTION_OK