From: Daniele Varrazzo <daniele.varrazzo@gmail.com>
Date: Sun, 28 Feb 2021 00:36:34 +0000 (+0100)
Subject: Add Connection.fileno()
X-Git-Tag: 3.0.dev0~103
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=4fba62cd14495dbd5863cbedf94733c39c4105df;p=thirdparty%2Fpsycopg.git

Add Connection.fileno()
---

diff --git a/docs/connection.rst b/docs/connection.rst
index a7c25f196..c82c7bf65 100644
--- a/docs/connection.rst
+++ b/docs/connection.rst
@@ -166,6 +166,8 @@ The `!Connection` class
         TODO
 
 
+    .. automethod:: fileno
+
     .. autoattribute:: prepare_threshold
         :annotation: Optional[int]
 
diff --git a/psycopg3/psycopg3/connection.py b/psycopg3/psycopg3/connection.py
index 9c701fdd3..d724cda51 100644
--- a/psycopg3/psycopg3/connection.py
+++ b/psycopg3/psycopg3/connection.py
@@ -209,6 +209,18 @@ class BaseConnection(AdaptContext):
         # implement the AdaptContext protocol
         return self
 
+    def fileno(self) -> int:
+        """Return the file descriptor of the connection.
+
+        This function allows to use the connection as file-like object in
+        functions waiting for readiness, such as the ones defined in the
+        `selectors` module.
+        """
+        try:
+            return self.pgconn.socket
+        except pq.PQerror as exc:
+            raise e.OperationalError(str(exc))
+
     def cancel(self) -> None:
         """Cancel the current operation on the connection."""
         c = self.pgconn.get_cancel()
diff --git a/tests/test_concurrency.py b/tests/test_concurrency.py
index 193896423..2f9347bf1 100644
--- a/tests/test_concurrency.py
+++ b/tests/test_concurrency.py
@@ -7,6 +7,7 @@ import sys
 import time
 import queue
 import pytest
+import selectors
 import threading
 import subprocess as sp
 
@@ -170,3 +171,26 @@ def test_cancel(conn):
     # still working
     conn.rollback()
     assert cur.execute("select 1").fetchone()[0] == 1
+
+
+@pytest.mark.slow
+def test_identify_closure(conn, dsn):
+    conn2 = psycopg3.connect(dsn)
+
+    def closer():
+        time.sleep(0.3)
+        conn2.execute(
+            "select pg_terminate_backend(%s)", [conn.pgconn.backend_pid]
+        )
+
+    t0 = time.time()
+    sel = selectors.DefaultSelector()
+    sel.register(conn, selectors.EVENT_READ)
+    t = threading.Thread(target=closer)
+    t.start()
+
+    assert sel.select(timeout=1.0)
+    with pytest.raises(psycopg3.OperationalError):
+        conn.execute("select 1")
+    t1 = time.time()
+    assert 0.3 < t1 - t0 < 0.5
diff --git a/tests/test_concurrency_async.py b/tests/test_concurrency_async.py
index 6b21db1d3..446db1fb9 100644
--- a/tests/test_concurrency_async.py
+++ b/tests/test_concurrency_async.py
@@ -127,3 +127,26 @@ async def test_cancel(aconn):
     cur = aconn.cursor()
     await cur.execute("select 1")
     assert await cur.fetchone() == (1,)
+
+
+@pytest.mark.slow
+async def test_identify_closure(aconn, dsn):
+    conn2 = await psycopg3.AsyncConnection.connect(dsn)
+
+    async def closer():
+        await asyncio.sleep(0.3)
+        await conn2.execute(
+            "select pg_terminate_backend(%s)", [aconn.pgconn.backend_pid]
+        )
+
+    t0 = time.time()
+    ev = asyncio.Event()
+    loop = asyncio.get_event_loop()
+    loop.add_reader(aconn.fileno(), ev.set)
+    asyncio.ensure_future(closer())
+
+    await asyncio.wait_for(ev.wait(), 1.0)
+    with pytest.raises(psycopg3.OperationalError):
+        await aconn.execute("select 1")
+    t1 = time.time()
+    assert 0.3 < t1 - t0 < 0.5
diff --git a/tests/test_connection.py b/tests/test_connection.py
index 8ab62ee77..83328f71a 100644
--- a/tests/test_connection.py
+++ b/tests/test_connection.py
@@ -511,3 +511,10 @@ def test_str(conn):
     assert "[IDLE]" in str(conn)
     conn.close()
     assert "[BAD]" in str(conn)
+
+
+def test_fileno(conn):
+    assert conn.fileno() == conn.pgconn.socket
+    conn.close()
+    with pytest.raises(psycopg3.OperationalError):
+        conn.fileno()
diff --git a/tests/test_connection_async.py b/tests/test_connection_async.py
index 9ec28560e..956cff151 100644
--- a/tests/test_connection_async.py
+++ b/tests/test_connection_async.py
@@ -529,3 +529,10 @@ async def test_str(aconn):
     assert "[IDLE]" in str(aconn)
     await aconn.close()
     assert "[BAD]" in str(aconn)
+
+
+async def test_fileno(aconn):
+    assert aconn.fileno() == aconn.pgconn.socket
+    await aconn.close()
+    with pytest.raises(psycopg3.OperationalError):
+        aconn.fileno()