Handle ErrorResponse in ping (#393)

* Handle IdleSessionTimeoutError in blocking client reconnect
edgedb · Nov 23, 2022 · 8b28947 · 8b28947
1 parent 4b8bec6
commit 8b28947
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 2 deletions.
diff --git a/edgedb/blocking_client.py b/edgedb/blocking_client.py
@@ -152,8 +152,8 @@ async def raw_query(self, query_context: abstract.QueryContext):
                 time.monotonic() - self._protocol.last_active_timestamp
                 > self._ping_wait_time
             ):
-                await self._protocol._sync()
-        except errors.ClientConnectionError:
+                await self._protocol.ping()
+        except (errors.IdleSessionTimeoutError, errors.ClientConnectionError):
             await self.connect()
 
         return await super().raw_query(query_context)

diff --git a/edgedb/protocol/protocol.pyx b/edgedb/protocol/protocol.pyx
@@ -711,6 +711,27 @@ cdef class SansIOProtocol:
             else:
                 self.fallthrough()
 
+    async def ping(self):
+        cdef char mtype
+        self.write(WriteBuffer.new_message(SYNC_MSG).end_message())
+        exc = None
+        while True:
+            if not self.buffer.take_message():
+                await self.wait_for_message()
+            mtype = self.buffer.get_message_type()
+
+            if mtype == READY_FOR_COMMAND_MSG:
+                self.parse_sync_message()
+                break
+            elif mtype == ERROR_RESPONSE_MSG:
+                exc = self.parse_error_message()
+                self.buffer.finish_message()
+                break
+            else:
+                self.fallthrough()
+        if exc is not None:
+            raise exc
+
     async def restore(self, bytes header, data_gen):
         cdef:
             WriteBuffer buf