From fe3398dfcf7c8a2d61833f4dfdcc7fbf6771cfd2 Mon Sep 17 00:00:00 2001
From: Daniele Varrazzo <daniele.varrazzo@gmail.com>
Date: Sat, 14 Mar 2020 18:30:03 +1300
Subject: [PATCH] Added PQconn.reset sync and async

---
 psycopg3/_pq_ctypes.py | 12 ++++++++++++
 psycopg3/pq_ctypes.py  | 16 ++++++++++++++++
 tests/test_pq.py       | 26 ++++++++++++++++++++++++++
 3 files changed, 54 insertions(+)

diff --git a/psycopg3/_pq_ctypes.py b/psycopg3/_pq_ctypes.py
index d310b629a..0a09c52d3 100644
--- 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 55dd59872..1e1d02bc5 100644
--- 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 464d0f138..3fd94ba3a 100644
--- 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
-- 
2.47.3