feat(DB-773): take a SQL session key and report the API's session fallback

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.

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()

tests/test_driver.py

@@ -30,17 +30,30 @@ 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 _redirect(status_code, headers=None):
+    response = MagicMock()
+    response.status_code = status_code
+    response.headers = headers or {}
+    return response
 
-    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.
+
+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; ``history`` replaces
+    that whole history.
     """
     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.headers = {}
+    post_resp.history = (
+        history if history is not None else [_redirect(303, redirect_headers)]
+    )
     mock_post.return_value = post_resp
 
     get_resp = MagicMock()
@@ -52,6 +65,23 @@ 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,
+    history=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, history)
+
     with patch("wherobots.db.driver.connect_direct") as mock_cd:
         mock_cd.return_value = MagicMock()
         connect(api_key="test-key", **connect_kwargs)
@@ -61,6 +91,76 @@ 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.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")
+    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 +258,8 @@ 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 = []
+        post_resp.headers = {}
         mock_post.return_value = post_resp
 
         # GET returns INITIALIZING (triggers TryAgain)
@@ -205,6 +307,8 @@ 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 = []
+        post_resp.headers = {}
         mock_post.return_value = post_resp
 
         get_resp = MagicMock()
@@ -365,6 +469,8 @@ 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 = []
+        post_resp.headers = {}
         mock_post.return_value = post_resp
 
         get_resp = MagicMock()

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
 

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")

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,21 @@ 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 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
+        )
 
     @tenacity.retry(
         stop=tenacity.stop_after_delay(wait_timeout),
@@ -280,6 +306,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 +336,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 +388,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,
     )