feat: loop on a recv() with timeout, rather than blocking

Author: mpetazzoni

wherobots/db/connection.py

@@ -12,6 +12,7 @@
 import websockets.sync.client
 
 from wherobots.db.constants import (
+    DEFAULT_READ_TIMEOUT_SECONDS,
     RequestKind,
     EventKind,
     ExecutionState,
@@ -80,6 +81,9 @@ def __main_loop(self):
         while self.__ws.protocol.state < websockets.protocol.State.CLOSING:
             try:
                 self.__listen()
+            except TimeoutError:
+                # Expected, retry next time
+                continue
             except websockets.exceptions.ConnectionClosedOK:
                 logging.info("Connection closed; stopping main loop.")
                 return
@@ -164,7 +168,7 @@ def __send(self, message: dict[str, Any]) -> None:
         self.__ws.send(json.dumps(message))
 
     def __recv(self) -> dict[str, Any]:
-        frame = self.__ws.recv()
+        frame = self.__ws.recv(timeout=DEFAULT_READ_TIMEOUT_SECONDS)
         if isinstance(frame, str):
             message = json.loads(frame)
         elif isinstance(frame, bytes):

wherobots/db/constants.py

@@ -8,6 +8,7 @@
 STAGING_ENDPOINT = "api.staging.wherobots.services"  # "api.staging.wherobots.com"
 DEFAULT_RUNTIME = Runtime.SEDONA
 DEFAULT_REGION = Region.AWS_US_WEST_2
+DEFAULT_READ_TIMEOUT_SECONDS = 0.5
 DEFAULT_SESSION_WAIT_TIMEOUT_SECONDS = 300
 MAX_MESSAGE_SIZE = 100 * 2**20  # 100MiB
 

wherobots/db/driver.py

@@ -126,4 +126,4 @@ def connect_direct(uri: str, headers: dict[str, str] = None) -> Connection:
         session = Connection(ws)
         return session
     except Exception as e:
-        raise InterfaceError("Failed to connect to SQL session!", e)
+        raise InterfaceError("Failed to connect to SQL session!") from e