From 27761731616b6061898f7ade8499976cd8b274a6 Mon Sep 17 00:00:00 2001 From: Kevin Date: Thu, 15 Oct 2026 14:49:07 -0700 Subject: [PATCH 1/2] feat: take a SQL session key and report the API's session fallback connect() takes session_key, sent as the session request's sessionKey: the API reuses a session only for requests with the same key, so a caller can keep sessions of its own (Waldo keeps one per chat). When a session of the key's own would exceed a quota, the API shares another of the user's sessions and names why in the X-Wherobots-Sql-Session-Fallback header of its redirect. requests follows that redirect, so the driver reads the header from the response history and exposes it as Connection.session_fallback, next to a now-public Connection.session_id. Co-Authored-By: Claude Opus 5.5 --- README.md | 8 ++++ tests/test_disconnect.py | 16 ++++++++ tests/test_driver.py | 84 +++++++++++++++++++++++++++++++++++--- wherobots/db/connection.py | 15 +++++++ wherobots/db/constants.py | 3 ++ wherobots/db/driver.py | 36 ++++++++++++---- 6 files changed, 150 insertions(+), 12 deletions(-) diff --git a/README.md b/README.md index d99cfb0..ece9305 100644 --- a/README.md +++ b/README.md @@ -277,6 +277,14 @@ users may find useful: caller needs to interrupt it (e.g. on client disconnect or timeout). The event is checked before each HTTP request, between retry attempts, and before the WebSocket handshake. +* `session_key`: keeps this connection on SQL sessions of its own. The + API reuses a session only for connections that ask with the same key, + and never one opened without a key, so a heavy query elsewhere never + holds the session your queries run on. When a new session would exceed + your organization's quota, the API shares another of your sessions + under another key instead, with or without `session_key`, and the + connection's `session_fallback` is `"quota"` (it is `None` otherwise). + `connection.session_id` names the session. * `stale_query_probe_seconds`: how long a running query may go without any message from the SQL session before the driver asks the session for its state again, in case the query's completion event was lost. diff --git a/tests/test_disconnect.py b/tests/test_disconnect.py index da70d8c..f1021ca 100644 --- a/tests/test_disconnect.py +++ b/tests/test_disconnect.py @@ -883,3 +883,19 @@ def test_direct_connection_uses_explicit_session_id(): conn.close() with pytest.raises(OperationalError, match="session=session-1"): cursor.fetchall() + + +def test_connection_reports_its_session_and_fallback(): + ws = Transport() + conn = Connection(ws, session_id="session-1", session_fallback="quota") + try: + assert conn.session_id == "session-1" + assert conn.session_fallback == "quota" + finally: + conn.close() + plain = Connection(Transport()) + try: + assert plain.session_id is None + assert plain.session_fallback is None + finally: + plain.close() diff --git a/tests/test_driver.py b/tests/test_driver.py index 3295cd0..929dd4f 100644 --- a/tests/test_driver.py +++ b/tests/test_driver.py @@ -30,17 +30,20 @@ def _run_connect(mock_post, mock_get, **connect_kwargs): return kwargs -def _run_connect_full(mock_post, mock_get, session_url=None, **connect_kwargs): - """Drive a successful connect(). +def _mock_ready_session(mock_post, mock_get, session_url=None, redirect_headers=None): + """Answer the session request with a followed 303 and a READY status. - Returns a tuple of (kwargs passed to requests.post, kwargs passed to the - patched connect_direct) so tests can assert on both the HTTP POST and the - WebSocket upgrade path. + ``redirect_headers`` are the headers of the API's 303 answer, which + `requests` keeps in the final response's history. """ + redirect = MagicMock() + redirect.status_code = 303 + redirect.headers = redirect_headers or {} post_resp = MagicMock() post_resp.status_code = 200 post_resp.url = session_url or "https://api.example.com/sql/session/test-id" post_resp.raise_for_status = MagicMock() + post_resp.history = [redirect] mock_post.return_value = post_resp get_resp = MagicMock() @@ -52,6 +55,18 @@ def _run_connect_full(mock_post, mock_get, session_url=None, **connect_kwargs): } mock_get.return_value = get_resp + +def _run_connect_full( + mock_post, mock_get, session_url=None, redirect_headers=None, **connect_kwargs +): + """Drive a successful connect(). + + Returns a tuple of (kwargs passed to requests.post, kwargs passed to the + patched connect_direct) so tests can assert on both the HTTP POST and the + WebSocket upgrade path. + """ + _mock_ready_session(mock_post, mock_get, session_url, redirect_headers) + with patch("wherobots.db.driver.connect_direct") as mock_cd: mock_cd.return_value = MagicMock() connect(api_key="test-key", **connect_kwargs) @@ -61,6 +76,62 @@ def _run_connect_full(mock_post, mock_get, session_url=None, **connect_kwargs): return post_kwargs, cd_kwargs +class TestConnectSessionKey: + """session_key reaches the API, and the API's fallback reaches the caller.""" + + @patch("wherobots.db.driver.requests.get") + @patch("wherobots.db.driver.requests.post") + def test_session_key_is_sent_in_the_body(self, mock_post, mock_get): + kwargs = _run_connect(mock_post, mock_get, session_key="waldo:chat-1") + assert kwargs["json"]["sessionKey"] == "waldo:chat-1" + + @patch("wherobots.db.driver.requests.get") + @patch("wherobots.db.driver.requests.post") + def test_omitted_session_key_is_not_sent(self, mock_post, mock_get): + kwargs = _run_connect(mock_post, mock_get) + assert "sessionKey" not in kwargs["json"] + + @pytest.mark.parametrize( + ("redirect_headers", "fallback"), + [ + ({"X-Wherobots-Sql-Session-Fallback": "quota"}, "quota"), + ({}, None), + ], + ids=["fallback", "own session"], + ) + @patch("wherobots.db.driver.requests.get") + @patch("wherobots.db.driver.requests.post") + def test_the_apis_fallback_reaches_the_connection( + self, mock_post, mock_get, redirect_headers, fallback + ): + _, cd_kwargs = _run_connect_full( + mock_post, + mock_get, + redirect_headers=redirect_headers, + session_key="waldo:chat-2", + ) + assert cd_kwargs["session_fallback"] == fallback + + @patch("wherobots.db.driver.websockets.sync.client.connect") + @patch("wherobots.db.driver.requests.get") + @patch("wherobots.db.driver.requests.post") + def test_the_connection_gets_the_session_and_fallback( + self, mock_post, mock_get, mock_ws + ): + """Through the real connect_direct, not around it.""" + _mock_ready_session( + mock_post, + mock_get, + redirect_headers={"X-Wherobots-Sql-Session-Fallback": "quota"}, + ) + mock_ws.return_value = MagicMock() + with patch("wherobots.db.driver.Connection") as mock_conn: + connect(api_key="test-key", session_key="waldo:chat-2") + _, conn_kwargs = mock_conn.call_args + assert conn_kwargs["session_id"] == "test-id" + assert conn_kwargs["session_fallback"] == "quota" + + class TestConnectRegionRuntime: """region/runtime accept enum|str and are omitted when not provided.""" @@ -158,6 +229,7 @@ def test_cancel_during_polling(self, mock_post, mock_get): post_resp.status_code = 200 post_resp.url = "https://api.example.com/sql/session/test-id" post_resp.raise_for_status = MagicMock() + post_resp.history = [] mock_post.return_value = post_resp # GET returns INITIALIZING (triggers TryAgain) @@ -205,6 +277,7 @@ def test_http_timeout_on_get(self, mock_post, mock_get): post_resp.status_code = 200 post_resp.url = "https://api.example.com/sql/session/test-id" post_resp.raise_for_status = MagicMock() + post_resp.history = [] mock_post.return_value = post_resp get_resp = MagicMock() @@ -365,6 +438,7 @@ def test_extra_headers_cannot_spoof_bearer_token(self, mock_post, mock_get): post_resp.status_code = 200 post_resp.url = "https://api.example.com/sql/session/test-id" post_resp.raise_for_status = MagicMock() + post_resp.history = [] mock_post.return_value = post_resp get_resp = MagicMock() diff --git a/wherobots/db/connection.py b/wherobots/db/connection.py index 5d53298..3d1f68a 100644 --- a/wherobots/db/connection.py +++ b/wherobots/db/connection.py @@ -214,6 +214,7 @@ def __init__( geometry_representation: GeometryRepresentation | None = None, session_id: str | None = None, stale_query_probe_seconds: float | None = DEFAULT_STALE_QUERY_PROBE_SECONDS, + session_fallback: str | None = None, ): self.__ws = ws self.__results_format = results_format @@ -238,6 +239,7 @@ def __init__( self.__recv_timeout = check self.__session_id = session_id + self.__session_fallback = session_fallback self.__lock = threading.Lock() self.__send_lock = threading.Lock() self.__shutdown_done = threading.Event() @@ -252,6 +254,19 @@ def __init__( ) self.__thread.start() + @property + def session_id(self) -> str | None: + """The id of the SQL session this connection is on, when known.""" + return self.__session_id + + @property + def session_fallback(self) -> str | None: + """Why the API gave this connection a session under another key than + the one it asked for, or None when it did not. ``"quota"``: a session + of the key's own would have exceeded a quota, so the connection shares + another of the user's sessions.""" + return self.__session_fallback + def __enter__(self): return self diff --git a/wherobots/db/constants.py b/wherobots/db/constants.py index 965bd59..5296ee9 100644 --- a/wherobots/db/constants.py +++ b/wherobots/db/constants.py @@ -19,6 +19,9 @@ DEFAULT_SESSION_WAIT_TIMEOUT_SECONDS: float = 900 # Silence, per query, before the driver re-asks the session for its state. DEFAULT_STALE_QUERY_PROBE_SECONDS: float = 30.0 +# The API sets this on its answer to a session request when it gave the request +# a session under another key than its own, naming why ("quota"). +SESSION_FALLBACK_HEADER: str = "X-Wherobots-Sql-Session-Fallback" MAX_MESSAGE_SIZE: int = 100 * 2**20 # 100MiB PROTOCOL_VERSION: Version = Version("1.0.0") diff --git a/wherobots/db/driver.py b/wherobots/db/driver.py index 399bced..baec9df 100644 --- a/wherobots/db/driver.py +++ b/wherobots/db/driver.py @@ -12,7 +12,7 @@ import requests import tenacity import threading -from typing import Final, Union, Dict +from typing import Any, Final, Union, Dict import urllib.parse import websockets.exceptions import websockets.sync.client @@ -29,6 +29,7 @@ MAX_MESSAGE_SIZE, PARAM_STYLE, PROTOCOL_VERSION, + SESSION_FALLBACK_HEADER, ) from .errors import ( InterfaceError, @@ -127,6 +128,7 @@ def connect( cancel_event: Union[threading.Event, None] = None, extra_headers: Union[Dict[str, str], None] = None, stale_query_probe_seconds: Union[float, None] = DEFAULT_STALE_QUERY_PROBE_SECONDS, + session_key: Union[str, None] = None, ) -> Connection: """Create a connection to a Wherobots SQL session. @@ -154,6 +156,11 @@ def connect( back off to at most 8 times this interval. ``None`` disables probing; an invalid value (non-numeric, non-finite, or not positive) also disables it. + :param session_key: Keeps this connection on SQL sessions of its own: the + API reuses a session only for connections that ask with the same key, + and never one opened without a key. When a session of the key's own + would exceed a quota, the API shares one under another key instead; + the connection's ``session_fallback`` then says so. """ if not token and not api_key: raise ValueError("At least one of `token` or `api_key` is required") @@ -198,18 +205,22 @@ def connect( _check_cancelled(cancel_event) + body: Dict[str, Any] = { + "runtimeId": runtime_id, + "shutdownAfterInactiveSeconds": shutdown_after_inactive_seconds, + "version": version, + "sessionType": session_type.value, + } + if session_key is not None: + body["sessionKey"] = session_key + try: resp = requests.post( url=f"{host}/sql/session", # `requests` omits query params whose value is None, so an omitted # region is simply not sent and the API applies the org default. params={"region": region_name, "force_new": force_new}, - json={ - "runtimeId": runtime_id, - "shutdownAfterInactiveSeconds": shutdown_after_inactive_seconds, - "version": version, - "sessionType": session_type.value, - }, + json=body, headers=headers, timeout=DEFAULT_HTTP_TIMEOUT, ) @@ -228,6 +239,13 @@ def connect( # At this point we've been redirected to /sql/session/{session_id}, which we'll need to keep polling until the # session is in READY state. session_id_url = resp.url + # The fallback header is on the API's redirect, which `requests` followed. + answer = resp.history[0] if resp.history else resp + session_fallback = answer.headers.get(SESSION_FALLBACK_HEADER) or None + if session_fallback is not None: + logging.warning( + "The API shared a SQL session under another key (%s).", session_fallback + ) @tenacity.retry( stop=tenacity.stop_after_delay(wait_timeout), @@ -280,6 +298,7 @@ def get_session_uri() -> str: .path.rstrip("/") .rsplit("/", 1)[-1], stale_query_probe_seconds=stale_query_probe_seconds, + session_fallback=session_fallback, ) @@ -309,10 +328,12 @@ def connect_direct( cancel_event: Union[threading.Event, None] = None, session_id: str | None = None, stale_query_probe_seconds: Union[float, None] = DEFAULT_STALE_QUERY_PROBE_SECONDS, + session_fallback: str | None = None, ) -> Connection: """Connect to a SQL session's WebSocket endpoint directly. :param stale_query_probe_seconds: See :func:`connect`. + :param session_fallback: See :attr:`Connection.session_fallback`. """ uri_with_protocol = f"{uri}/{protocol}" ssl_context = ssl.create_default_context() @@ -359,4 +380,5 @@ def ws_connect() -> websockets.sync.client.ClientConnection: geometry_representation=geometry_representation, session_id=session_id, stale_query_probe_seconds=stale_query_probe_seconds, + session_fallback=session_fallback, ) From bf42a1f3acac29c4d2b42d48bd64252ed3639115 Mon Sep 17 00:00:00 2001 From: Kevin Date: Thu, 15 Oct 2026 16:59:48 -0700 Subject: [PATCH 2/2] fix: read the session fallback header from every redirect hop Salty Hambot: when the session POST is first redirected by a 307 or 308 (to a canonical endpoint, say) before the API's 303, requests keeps both in resp.history, and reading history[0] missed the header, so the connection reported no fallback. Read every hop, latest first. Co-Authored-By: Claude Opus 5.5 --- tests/test_driver.py | 48 +++++++++++++++++++++++++++++++++++------- wherobots/db/driver.py | 14 +++++++++--- 2 files changed, 51 insertions(+), 11 deletions(-) diff --git a/tests/test_driver.py b/tests/test_driver.py index 929dd4f..f2909e7 100644 --- a/tests/test_driver.py +++ b/tests/test_driver.py @@ -30,20 +30,30 @@ def _run_connect(mock_post, mock_get, **connect_kwargs): return kwargs -def _mock_ready_session(mock_post, mock_get, session_url=None, redirect_headers=None): +def _redirect(status_code, headers=None): + response = MagicMock() + response.status_code = status_code + response.headers = headers or {} + return response + + +def _mock_ready_session( + mock_post, mock_get, session_url=None, redirect_headers=None, history=None +): """Answer the session request with a followed 303 and a READY status. ``redirect_headers`` are the headers of the API's 303 answer, which - `requests` keeps in the final response's history. + `requests` keeps in the final response's history; ``history`` replaces + that whole history. """ - redirect = MagicMock() - redirect.status_code = 303 - redirect.headers = redirect_headers or {} post_resp = MagicMock() post_resp.status_code = 200 post_resp.url = session_url or "https://api.example.com/sql/session/test-id" post_resp.raise_for_status = MagicMock() - post_resp.history = [redirect] + post_resp.headers = {} + post_resp.history = ( + history if history is not None else [_redirect(303, redirect_headers)] + ) mock_post.return_value = post_resp get_resp = MagicMock() @@ -57,7 +67,12 @@ def _mock_ready_session(mock_post, mock_get, session_url=None, redirect_headers= def _run_connect_full( - mock_post, mock_get, session_url=None, redirect_headers=None, **connect_kwargs + mock_post, + mock_get, + session_url=None, + redirect_headers=None, + history=None, + **connect_kwargs, ): """Drive a successful connect(). @@ -65,7 +80,7 @@ def _run_connect_full( patched connect_direct) so tests can assert on both the HTTP POST and the WebSocket upgrade path. """ - _mock_ready_session(mock_post, mock_get, session_url, redirect_headers) + _mock_ready_session(mock_post, mock_get, session_url, redirect_headers, history) with patch("wherobots.db.driver.connect_direct") as mock_cd: mock_cd.return_value = MagicMock() @@ -112,6 +127,20 @@ def test_the_apis_fallback_reaches_the_connection( ) assert cd_kwargs["session_fallback"] == fallback + @patch("wherobots.db.driver.requests.get") + @patch("wherobots.db.driver.requests.post") + def test_the_fallback_is_read_past_an_earlier_redirect(self, mock_post, mock_get): + """A 307 to a canonical endpoint before the API's 303 hides nothing.""" + _, cd_kwargs = _run_connect_full( + mock_post, + mock_get, + history=[ + _redirect(307), + _redirect(303, {"X-Wherobots-Sql-Session-Fallback": "quota"}), + ], + ) + assert cd_kwargs["session_fallback"] == "quota" + @patch("wherobots.db.driver.websockets.sync.client.connect") @patch("wherobots.db.driver.requests.get") @patch("wherobots.db.driver.requests.post") @@ -230,6 +259,7 @@ def test_cancel_during_polling(self, mock_post, mock_get): post_resp.url = "https://api.example.com/sql/session/test-id" post_resp.raise_for_status = MagicMock() post_resp.history = [] + post_resp.headers = {} mock_post.return_value = post_resp # GET returns INITIALIZING (triggers TryAgain) @@ -278,6 +308,7 @@ def test_http_timeout_on_get(self, mock_post, mock_get): post_resp.url = "https://api.example.com/sql/session/test-id" post_resp.raise_for_status = MagicMock() post_resp.history = [] + post_resp.headers = {} mock_post.return_value = post_resp get_resp = MagicMock() @@ -439,6 +470,7 @@ def test_extra_headers_cannot_spoof_bearer_token(self, mock_post, mock_get): post_resp.url = "https://api.example.com/sql/session/test-id" post_resp.raise_for_status = MagicMock() post_resp.history = [] + post_resp.headers = {} mock_post.return_value = post_resp get_resp = MagicMock() diff --git a/wherobots/db/driver.py b/wherobots/db/driver.py index baec9df..06fe601 100644 --- a/wherobots/db/driver.py +++ b/wherobots/db/driver.py @@ -239,9 +239,17 @@ def connect( # At this point we've been redirected to /sql/session/{session_id}, which we'll need to keep polling until the # session is in READY state. session_id_url = resp.url - # The fallback header is on the API's redirect, which `requests` followed. - answer = resp.history[0] if resp.history else resp - session_fallback = answer.headers.get(SESSION_FALLBACK_HEADER) or None + # The fallback header is on the API's 303, which `requests` followed. Other + # redirects (a 307 or 308 to a canonical endpoint) may come before it, so + # every hop is read, latest first. + session_fallback = next( + ( + hop.headers[SESSION_FALLBACK_HEADER] + for hop in reversed([*resp.history, resp]) + if hop.headers.get(SESSION_FALLBACK_HEADER) + ), + None, + ) if session_fallback is not None: logging.warning( "The API shared a SQL session under another key (%s).", session_fallback