From 89929dcda51c642b0541d78772e01a7b11e30b4f Mon Sep 17 00:00:00 2001
From: Simon Fishel <simon@wherobots.com>
Date: Mon, 20 Apr 2026 15:06:02 -0700
Subject: [PATCH] fix: add retry logic for WebSocket connection to SQL session

Adds tenacity-based retry with exponential backoff to the initial
WebSocket connection in connect_direct(). This fixes intermittent
failures in the SQL session canary tests caused by transient
connection errors during the WebSocket handshake.

Bump version to 0.26.1.
---
 pyproject.toml         |  2 +-
 wherobots/db/driver.py | 25 +++++++++++++++++++++----
 2 files changed, 22 insertions(+), 5 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index e954d45..161e431 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [project]
 name = "wherobots-python-dbapi"
-version = "0.26.0"
+version = "0.26.1"
 description = "Python DB-API driver for Wherobots DB"
 authors = [{ name = "Maxime Petazzoni", email = "max@wherobots.com" }]
 requires-python = ">=3.10, <4"
diff --git a/wherobots/db/driver.py b/wherobots/db/driver.py
index 0da2eb6..b18ff84 100644
--- a/wherobots/db/driver.py
+++ b/wherobots/db/driver.py
@@ -13,6 +13,7 @@
 import tenacity
 from typing import Final, Union, Dict
 import urllib.parse
+import websockets.exceptions
 import websockets.sync.client
 import certifi
 
@@ -200,17 +201,33 @@ def connect_direct(
     geometry_representation: Union[GeometryRepresentation, None] = None,
 ) -> Connection:
     uri_with_protocol = f"{uri}/{protocol}"
+    ssl_context = ssl.create_default_context()
+    ssl_context.load_verify_locations(certifi.where())
 
-    try:
+    @tenacity.retry(
+        stop=tenacity.stop_after_attempt(5),
+        wait=tenacity.wait_exponential(multiplier=1, min=1, max=5),
+        retry=tenacity.retry_if_exception_type(
+            (
+                ConnectionRefusedError,
+                ConnectionResetError,
+                TimeoutError,
+                websockets.exceptions.InvalidHandshake,
+            )
+        ),
+        reraise=True,
+    )
+    def ws_connect() -> websockets.sync.client.ClientConnection:
         logging.info("Connecting to SQL session at %s ...", uri_with_protocol)
-        ssl_context = ssl.create_default_context()
-        ssl_context.load_verify_locations(certifi.where())
-        ws = websockets.sync.client.connect(
+        return websockets.sync.client.connect(
             uri=uri_with_protocol,
             additional_headers=headers,
             max_size=MAX_MESSAGE_SIZE,
             ssl=ssl_context,
         )
+
+    try:
+        ws = ws_connect()
     except Exception as e:
         raise InterfaceError("Failed to connect to SQL session!") from e