From eb0e6e5abc43bc8456d70e9fd015beaa9bd14250 Mon Sep 17 00:00:00 2001 From: Jia Yu Date: Fri, 6 Nov 2026 12:18:23 -0800 Subject: [PATCH 1/3] feat(connection): keep a session to its key, and report whether the API did Over a quota, the API can give a keyed session to a connection under another key, and the session's own connection can't tell: its session_key is still its own and it got no fallback. Stopping the session then cuts the other connection off. connect(session_key_exclusive=True) sends sessionKeyExclusive with the session key, which asks the API to keep the session to that key: it never shares it with another key, even over a quota. The flag requires a session_key; without one connect() raises ValueError before any request. Connection.session_key_exclusive reports whether the API did, read from the session's stored request like session_key. It is True only when the API says so, so it is False from an API that predates key-exclusive sessions, which ignores the flag. A caller stops a session only when both session_key is its own and session_key_exclusive is True; the README's example now does. --- README.md | 29 +++++++++++++---- tests/test_disconnect.py | 13 ++++++++ tests/test_driver.py | 67 ++++++++++++++++++++++++++++++++++++-- wherobots/db/connection.py | 20 ++++++++++-- wherobots/db/driver.py | 19 +++++++++++ 5 files changed, 136 insertions(+), 12 deletions(-) diff --git a/README.md b/README.md index 7e8820e..03e1995 100644 --- a/README.md +++ b/README.md @@ -254,11 +254,11 @@ seconds: ```python run_key = f"nightly-build-{run_id}" # unique to this run of the job -conn = connect(api_key=..., session_key=run_key) +conn = connect(api_key=..., session_key=run_key, session_key_exclusive=True) try: ... # run the job's queries finally: - if conn.session_key == run_key: + if conn.session_key == run_key and conn.session_key_exclusive: conn.stop_session() else: conn.close() @@ -267,11 +267,19 @@ finally: The session stops for every connection on it, so stop only a session no other client uses. Every connection that asks with the same `session_key` can share a session, so give each run of the job a key of its own, for -example one that includes a run ID. Then check `conn.session_key` before -stopping: it's the key the API created the session under. It equals your -key when the API kept the run on its own sessions, it's `None` for a -session without a key, and it differs from your key when the API shared -another session (see `session_fallback` below). +example one that includes a run ID. Over a quota, the API can also give +your session to a connection under another key, and your connection +can't tell; `session_key_exclusive=True` asks the API to keep the session +to your key instead. Then check both before stopping: + +* `conn.session_key` is the key the API created the session under. It + equals your key when the API kept the run on its own sessions, it's + `None` for a session without a key, and it differs from your key when + the API shared another session (see `session_fallback` below). +* `conn.session_key_exclusive` is `True` only when the API reports that + it keeps the session to that key. It's `False` from an API that + predates key-exclusive sessions, and for a shared session that isn't + kept to its key. `stop_session()` raises `InterfaceError` if the API refuses or can't be reached; the connection is closed anyway, and the session stops after its @@ -359,6 +367,13 @@ users may find useful: 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. +* `session_key_exclusive`: with `session_key`, keeps this connection's + session to that key even over a quota: the API never shares it with a + connection under another key, so you can stop it without cutting off + another client. Over a quota, the connection itself can still be given + another key's session, and `session_fallback` says so. + `connection.session_key_exclusive` says whether the API kept the + session to its key. * `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 cade315..6132dd0 100644 --- a/tests/test_disconnect.py +++ b/tests/test_disconnect.py @@ -930,6 +930,19 @@ def test_connection_reports_the_session_key_it_was_given(): plain.close() +def test_connection_reports_whether_its_session_is_kept_to_its_key(): + conn = Connection(Transport(), session_key="dbt-run-1", session_key_exclusive=True) + try: + assert conn.session_key_exclusive is True + finally: + conn.close() + plain = Connection(Transport(), session_key="dbt-run-1") + try: + assert plain.session_key_exclusive is False + finally: + plain.close() + + def test_stop_session_closes_the_connection_then_asks_the_api(): ws = Transport() closed_when_asked = [] diff --git a/tests/test_driver.py b/tests/test_driver.py index 8c68b5d..b01cd03 100644 --- a/tests/test_driver.py +++ b/tests/test_driver.py @@ -116,6 +116,31 @@ 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"] + @patch("wherobots.db.driver.requests.get") + @patch("wherobots.db.driver.requests.post") + def test_session_key_exclusive_is_sent_in_the_body(self, mock_post, mock_get): + kwargs = _run_connect( + mock_post, mock_get, session_key="dbt-run-1", session_key_exclusive=True + ) + assert kwargs["json"]["sessionKey"] == "dbt-run-1" + assert kwargs["json"]["sessionKeyExclusive"] is True + + @pytest.mark.parametrize("session_key", ["dbt-run-1", None]) + @patch("wherobots.db.driver.requests.get") + @patch("wherobots.db.driver.requests.post") + def test_session_key_exclusive_is_not_sent_unless_asked_for( + self, mock_post, mock_get, session_key + ): + """A request without it is the same as before it existed.""" + kwargs = _run_connect(mock_post, mock_get, session_key=session_key) + assert "sessionKeyExclusive" not in kwargs["json"] + + @patch("wherobots.db.driver.requests.post") + def test_session_key_exclusive_requires_a_session_key(self, mock_post): + with pytest.raises(ValueError, match="session_key"): + connect(api_key="test-key", session_key_exclusive=True) + mock_post.assert_not_called() + @pytest.mark.parametrize( ("redirect_headers", "fallback"), [ @@ -194,6 +219,39 @@ def test_the_key_the_api_stored_reaches_the_connection( ) assert cd_kwargs["session_key"] == session_key + @pytest.mark.parametrize( + ("api_payload", "exclusive"), + [ + ({"sessionKey": "dbt-run-1", "sessionKeyExclusive": True}, True), + ({"sessionKey": "dbt-run-1", "sessionKeyExclusive": False}, False), + ({"sessionKey": "dbt-run-1", "sessionKeyExclusive": "true"}, False), + ({"sessionKey": "dbt-run-1"}, False), + (None, False), + ], + ids=[ + "kept to its key", + "not kept", + "not a boolean", + "older API", + "no stored request", + ], + ) + @patch("wherobots.db.driver.requests.get") + @patch("wherobots.db.driver.requests.post") + def test_whether_the_api_keeps_the_session_to_its_key_reaches_the_connection( + self, mock_post, mock_get, api_payload, exclusive + ): + """Only the API's word counts, never what the request asked for: an + API without key-exclusive sessions ignores the request's flag.""" + _, cd_kwargs = _run_connect_full( + mock_post, + mock_get, + api_payload=api_payload, + session_key="dbt-run-1", + session_key_exclusive=True, + ) + assert cd_kwargs["session_key_exclusive"] is exclusive + @patch("wherobots.db.driver.requests.delete") @patch("wherobots.db.driver.requests.get") @patch("wherobots.db.driver.requests.post") @@ -235,12 +293,17 @@ def test_the_connection_gets_the_key_and_the_stopper( self, mock_post, mock_get, mock_ws ): """Through the real connect_direct, not around it.""" - _mock_ready_session(mock_post, mock_get, api_payload={"sessionKey": "k"}) + _mock_ready_session( + mock_post, + mock_get, + api_payload={"sessionKey": "k", "sessionKeyExclusive": True}, + ) mock_ws.return_value = MagicMock() with patch("wherobots.db.driver.Connection") as mock_conn: - connect(api_key="test-key", session_key="k") + connect(api_key="test-key", session_key="k", session_key_exclusive=True) _, conn_kwargs = mock_conn.call_args assert conn_kwargs["session_key"] == "k" + assert conn_kwargs["session_key_exclusive"] is True assert callable(conn_kwargs["session_stopper"]) diff --git a/wherobots/db/connection.py b/wherobots/db/connection.py index 6c05552..c311f51 100644 --- a/wherobots/db/connection.py +++ b/wherobots/db/connection.py @@ -288,6 +288,7 @@ def __init__( reconnect_timeout_seconds: float | None = DEFAULT_RECONNECT_TIMEOUT_SECONDS, session_key: str | None = None, session_stopper: Callable[[], None] | None = None, + session_key_exclusive: bool = False, ): """ :param transport_factory: Opens a new WebSocket to the same SQL session, given @@ -303,6 +304,7 @@ def __init__( :param session_stopper: Asks the API to stop this connection's SQL session, raising ``InterfaceError`` if it can't. Without it, :meth:`stop_session` raises ``NotSupportedError``. + :param session_key_exclusive: See :attr:`session_key_exclusive`. """ self.__ws = ws self.__transport_factory = transport_factory @@ -350,6 +352,7 @@ def __init__( self.__session_id = session_id self.__session_fallback = session_fallback self.__session_key = session_key + self.__session_key_exclusive = session_key_exclusive self.__session_stopper = session_stopper self.__lock = threading.Lock() self.__send_lock = threading.Lock() @@ -388,14 +391,25 @@ def session_key(self) -> str | None: ``connect_direct()``, which never asks the API.""" return self.__session_key + @property + def session_key_exclusive(self) -> bool: + """Whether the API keeps this connection's SQL session to its + :attr:`session_key`: it never shares the session with another key, + even over a quota, as ``connect(session_key_exclusive=True)`` asks. + True only when the API reports it, so False also for an API that + predates key-exclusive sessions, and for a connection from + ``connect_direct()``.""" + return self.__session_key_exclusive + def stop_session(self) -> None: """Close this connection, then ask the API to stop its SQL session. The session stops within seconds instead of after its inactivity timeout, for every connection on it: another connection that shares - the session loses it too. Check :attr:`session_key` first if other - clients may share it. Stopping a session that already stopped - succeeds. + the session loses it too. Unless :attr:`session_key_exclusive` is + True, another client may have been given the session over a quota, + even when :attr:`session_key` is your own key. Stopping a session + that already stopped succeeds. :raises NotSupportedError: The connection came from ``connect_direct()``, which has no API to ask. The connection diff --git a/wherobots/db/driver.py b/wherobots/db/driver.py index b8a1367..88b85c2 100644 --- a/wherobots/db/driver.py +++ b/wherobots/db/driver.py @@ -131,6 +131,7 @@ def connect( stale_query_probe_seconds: Union[float, None] = DEFAULT_STALE_QUERY_PROBE_SECONDS, session_key: Union[str, None] = None, reconnect_timeout_seconds: Union[float, None] = DEFAULT_RECONNECT_TIMEOUT_SECONDS, + session_key_exclusive: bool = False, ) -> Connection: """Create a connection to a Wherobots SQL session. @@ -169,11 +170,18 @@ def connect( it can't reconnect in time do they fail, with ``ConnectionLostError``. ``None`` disables reconnecting, as does an invalid value (non-numeric, non-finite, or not positive): pending queries then fail at once. + :param session_key_exclusive: Keeps this connection's session to its + ``session_key`` even over a quota: the API never shares it with + another key, so the connection can stop it without cutting off + another client. Requires ``session_key``. The connection's + ``session_key_exclusive`` says whether the API did. """ if not token and not api_key: raise ValueError("At least one of `token` or `api_key` is required") if token and api_key: raise ValueError("`token` and `api_key` can't be both provided") + if session_key_exclusive and session_key is None: + raise ValueError("`session_key_exclusive` requires a `session_key`") headers = gen_user_agent_header() # Merge caller-supplied headers first so the driver's own auth headers, @@ -221,6 +229,8 @@ def connect( } if session_key is not None: body["sessionKey"] = session_key + if session_key_exclusive: + body["sessionKeyExclusive"] = True try: resp = requests.post( @@ -316,6 +326,11 @@ def get_ready_session() -> Dict[str, Any]: session_key_on_api = ( api_payload.get("sessionKey") if isinstance(api_payload, dict) else None ) + # Only the API's word counts: an API without key-exclusive sessions + # ignores the request's flag. + session_key_exclusive_on_api = ( + isinstance(api_payload, dict) and api_payload.get("sessionKeyExclusive") is True + ) # Snapshot the headers, as for reconnecting: a stop sends what this # connection sent even if the caller later mutates them. stop_headers = dict(headers) @@ -343,6 +358,7 @@ def stop_session() -> None: reconnect_timeout_seconds=reconnect_timeout_seconds, session_key=session_key_on_api, session_stopper=stop_session, + session_key_exclusive=session_key_exclusive_on_api, ) @@ -376,6 +392,7 @@ def connect_direct( reconnect_timeout_seconds: Union[float, None] = DEFAULT_RECONNECT_TIMEOUT_SECONDS, session_key: str | None = None, session_stopper: Union[Callable[[], None], None] = None, + session_key_exclusive: bool = False, ) -> Connection: """Connect to a SQL session's WebSocket endpoint directly. @@ -390,6 +407,7 @@ def connect_direct( :param session_stopper: Asks the API to stop the SQL session; see :meth:`Connection.stop_session`. ``connect()`` supplies it. Without it, the connection can't stop its session. + :param session_key_exclusive: See :attr:`Connection.session_key_exclusive`. """ uri_with_protocol = f"{uri}/{protocol}" ssl_context = ssl.create_default_context() @@ -466,4 +484,5 @@ def open_transport(timeout: float) -> websockets.sync.client.ClientConnection: reconnect_timeout_seconds=reconnect_timeout_seconds, session_key=session_key, session_stopper=session_stopper, + session_key_exclusive=session_key_exclusive, ) From 2c74261b9052fe9df486717b4c94f40f45f1c581 Mon Sep 17 00:00:00 2001 From: Jia Yu Date: Fri, 6 Nov 2026 12:32:26 -0800 Subject: [PATCH 2/3] chore: bump version to 0.32.0 --- pyproject.toml | 2 +- uv.lock | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index feb8837..92d54fa 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "wherobots-python-dbapi" -version = "0.31.0" +version = "0.32.0" description = "Python DB-API driver for Wherobots DB" authors = [{ name = "Maxime Petazzoni", email = "max@wherobots.com" }] requires-python = ">=3.10, <4" diff --git a/uv.lock b/uv.lock index 5340ec6..6affa01 100644 --- a/uv.lock +++ b/uv.lock @@ -1858,7 +1858,7 @@ wheels = [ [[package]] name = "wherobots-python-dbapi" -version = "0.31.0" +version = "0.32.0" source = { editable = "." } dependencies = [ { name = "cbor2" }, From dad6d07dd7882f72b53ad1ce818093ac2e79c229 Mon Sep 17 00:00:00 2001 From: Jia Yu Date: Fri, 6 Nov 2026 21:03:21 -0800 Subject: [PATCH 3/3] docs(connection): a quota fallback can share a session under the caller's own key Over a quota, the API can give a session_key_exclusive request a session opened under the same key without the flag. session_fallback is then "quota" while session_key matches, so the README no longer says a matching key means the run's own session, and session_fallback and the fallback warning no longer say the session is always another key's. --- README.md | 11 +++++++---- wherobots/db/connection.py | 10 ++++++---- wherobots/db/driver.py | 2 +- 3 files changed, 14 insertions(+), 9 deletions(-) diff --git a/README.md b/README.md index 03e1995..8e89858 100644 --- a/README.md +++ b/README.md @@ -272,10 +272,12 @@ your session to a connection under another key, and your connection can't tell; `session_key_exclusive=True` asks the API to keep the session to your key instead. Then check both before stopping: -* `conn.session_key` is the key the API created the session under. It - equals your key when the API kept the run on its own sessions, it's +* `conn.session_key` is the key the API created the session under. It's `None` for a session without a key, and it differs from your key when - the API shared another session (see `session_fallback` below). + the API shared another key's session (see `session_fallback` below). + It can equal your key on a shared session too: over a quota, the API + can give you a session opened under your key without + `session_key_exclusive`. * `conn.session_key_exclusive` is `True` only when the API reports that it keeps the session to that key. It's `False` from an API that predates key-exclusive sessions, and for a shared session that isn't @@ -371,7 +373,8 @@ users may find useful: session to that key even over a quota: the API never shares it with a connection under another key, so you can stop it without cutting off another client. Over a quota, the connection itself can still be given - another key's session, and `session_fallback` says so. + a shared session, even one under its own key; `session_fallback` then + says so, and `session_key_exclusive` is `False`. `connection.session_key_exclusive` says whether the API kept the session to its key. * `stale_query_probe_seconds`: how long a running query may go without diff --git a/wherobots/db/connection.py b/wherobots/db/connection.py index c311f51..d82c86b 100644 --- a/wherobots/db/connection.py +++ b/wherobots/db/connection.py @@ -375,10 +375,12 @@ def session_id(self) -> str | None: @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.""" + """Why the API gave this connection a shared session instead of one of + its key's own, 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: one under another key, or, for + ``connect(session_key_exclusive=True)``, one opened under the same key + without it.""" return self.__session_fallback @property diff --git a/wherobots/db/driver.py b/wherobots/db/driver.py index 88b85c2..9b32bf4 100644 --- a/wherobots/db/driver.py +++ b/wherobots/db/driver.py @@ -274,7 +274,7 @@ def connect( ) if session_fallback is not None: logging.warning( - "The API shared a SQL session under another key (%s).", session_fallback + "The API gave this connection a shared SQL session (%s).", session_fallback ) @tenacity.retry(