Fix wait_for() in blocking client to raise internal TimeoutError

edgedb · May 29, 2024 · b3407c1 · b3407c1
1 parent ee3c5ce
commit b3407c1
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 8 deletions.
diff --git a/edgedb/blocking_client.py b/edgedb/blocking_client.py
@@ -135,6 +135,9 @@ async def close(self, timeout=None):
                     await self._protocol.wait_for(
                         self._protocol.wait_for_disconnect(), timeout
                     )
+            except TimeoutError:
+                self.terminate()
+                raise errors.QueryTimeoutError()
             except Exception:
                 self.terminate()
                 raise

diff --git a/edgedb/protocol/blocking_proto.pyx b/edgedb/protocol/blocking_proto.pyx
@@ -63,18 +63,14 @@ cdef class BlockingIOProtocol(protocol.SansIOProtocolBackwardsCompatible):
     async def wait_for_message(self):
         cdef float timeout
         if self.deadline > 0:
-            timeout = self.deadline - time.monotonic()
-            if timeout <= 0:
-                self.abort()
-                raise errors.QueryTimeoutError()
             while not self.buffer.take_message():
+                timeout = self.deadline - time.monotonic()
+                if timeout <= 0:
+                    self.abort()
+                    raise TimeoutError
                 try:
                     self.sock.settimeout(timeout)
                     data = self.sock.recv(RECV_BUF)
-                    timeout = self.deadline - time.monotonic()
-                    if timeout <= 0:
-                        self.abort()
-                        raise TimeoutError
                 except OSError as e:
                     self._disconnect()
                     raise con_utils.wrap_error(e) from e