From af1878e2d1d71471f3e8a6ad345e5ea1566ba5a8 Mon Sep 17 00:00:00 2001 From: Jia Yu Date: Sun, 1 Nov 2026 23:37:37 -0800 Subject: [PATCH 1/2] feat(connection): stop the SQL session on request Closing a connection leaves its SQL session running until its inactivity timeout passes, which a batch job pays for after it has finished. Connection.stop_session() closes the connection, then asks the API to stop the session (DELETE /sql/session/{id}), which stops within seconds. Connection.session_key reports the key the API created the session under, read from the session's stored request while connect() waits for READY, so a caller can check that a session is its own before stopping it. A connection from connect_direct() has no API to ask: stop_session() raises NotSupportedError and leaves it open. --- README.md | 33 ++++++++++++++ tests/test_disconnect.py | 54 ++++++++++++++++++++++- tests/test_driver.py | 89 ++++++++++++++++++++++++++++++++++++-- wherobots/db/connection.py | 41 ++++++++++++++++++ wherobots/db/driver.py | 45 ++++++++++++++++--- 5 files changed, 251 insertions(+), 11 deletions(-) diff --git a/README.md b/README.md index 813c704..aad3c04 100644 --- a/README.md +++ b/README.md @@ -243,6 +243,39 @@ including whether a write committed, is unknown. The error's (When `execute()` itself raises `ConnectionLostError`, the connection was already lost and the statement was never sent.) +### Stopping the SQL session + +Closing a connection leaves its SQL session running until +`shutdown_after_inactive_seconds` pass with no connection on it, so that +a later connection can reuse it. To stop it right away instead, for +example at the end of a batch job, call `stop_session()`. It closes the +connection, then asks the API to stop the session, which stops within +seconds: + +```python +conn = connect(api_key=..., session_key="nightly-build-1234") +try: + ... # run the job's queries +finally: + if conn.session_key == "nightly-build-1234": + conn.stop_session() + else: + conn.close() +``` + +The session stops for every connection on it. Other clients may share a +session, so check `conn.session_key` before stopping it: it's the key the +API created the session under, which is your `session_key` only when the +session is yours alone. 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). + +`stop_session()` raises `InterfaceError` if the API refuses or can't be +reached; the connection is closed anyway, and the session stops after its +inactivity timeout instead. Stopping a session that already stopped +succeeds. A connection from `connect_direct()` has no API to ask, so its +`stop_session()` raises `NotSupportedError` and leaves it open. + ### Runtime and region selection You can chose the Wherobots runtime you want to use using the `runtime` diff --git a/tests/test_disconnect.py b/tests/test_disconnect.py index f1021ca..f6da310 100644 --- a/tests/test_disconnect.py +++ b/tests/test_disconnect.py @@ -22,7 +22,12 @@ from wherobots.db._transport import abort_connection from wherobots.db.connection import Connection, Query from wherobots.db.driver import connect_direct -from wherobots.db.errors import OperationalError, ProgrammingError +from wherobots.db.errors import ( + InterfaceError, + NotSupportedError, + OperationalError, + ProgrammingError, +) from wherobots.db.types import ExecutionState @@ -899,3 +904,50 @@ def test_connection_reports_its_session_and_fallback(): assert plain.session_fallback is None finally: plain.close() + + +def test_connection_reports_the_session_key_it_was_given(): + conn = Connection(Transport(), session_key="dbt-run-1") + try: + assert conn.session_key == "dbt-run-1" + finally: + conn.close() + plain = Connection(Transport()) + try: + assert plain.session_key is None + finally: + plain.close() + + +def test_stop_session_closes_the_connection_then_asks_the_api(): + ws = Transport() + closed_when_asked = [] + conn = Connection( + ws, + session_id="session-1", + session_stopper=lambda: closed_when_asked.append(conn._Connection__closed), + ) + conn.stop_session() + assert closed_when_asked == [True] + assert ws.closed.is_set() + assert not conn._Connection__thread.is_alive() + + +def test_a_refused_stop_raises_and_leaves_the_connection_closed(): + def refuse(): + raise InterfaceError("Failed to stop SQL session session-1: 500") + + conn = Connection(Transport(), session_id="session-1", session_stopper=refuse) + with pytest.raises(InterfaceError, match="session-1"): + conn.stop_session() + assert conn._Connection__closed + + +def test_a_direct_connection_cannot_stop_its_session_and_stays_open(): + conn = Connection(Transport()) + try: + with pytest.raises(NotSupportedError): + conn.stop_session() + assert not conn._Connection__closed + finally: + conn.close() diff --git a/tests/test_driver.py b/tests/test_driver.py index 3b51d51..8c68b5d 100644 --- a/tests/test_driver.py +++ b/tests/test_driver.py @@ -38,13 +38,18 @@ def _redirect(status_code, headers=None): def _mock_ready_session( - mock_post, mock_get, session_url=None, redirect_headers=None, history=None + mock_post, + mock_get, + session_url=None, + redirect_headers=None, + history=None, + api_payload=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. + that whole history. ``api_payload`` is the session's stored request. """ post_resp = MagicMock() post_resp.status_code = 200 @@ -63,6 +68,8 @@ def _mock_ready_session( "status": "READY", "appMeta": {"url": "https://compute.example.com/sql/org/session-id"}, } + if api_payload is not None: + get_resp.json.return_value["apiPayload"] = api_payload mock_get.return_value = get_resp @@ -72,6 +79,7 @@ def _run_connect_full( session_url=None, redirect_headers=None, history=None, + api_payload=None, **connect_kwargs, ): """Drive a successful connect(). @@ -80,7 +88,9 @@ 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, history) + _mock_ready_session( + mock_post, mock_get, session_url, redirect_headers, history, api_payload + ) with patch("wherobots.db.driver.connect_direct") as mock_cd: mock_cd.return_value = MagicMock() @@ -161,6 +171,79 @@ def test_the_connection_gets_the_session_and_fallback( assert conn_kwargs["session_fallback"] == "quota" +class TestStopSession: + """connect() hands the connection the session's key and a way to stop it.""" + + @pytest.mark.parametrize( + ("api_payload", "session_key"), + [ + ({"sessionKey": "dbt-run-1", "sessionType": "multi"}, "dbt-run-1"), + ({"sessionKey": None}, None), + ({}, None), + (None, None), + ], + ids=["keyed", "keyless", "no key field", "no stored request"], + ) + @patch("wherobots.db.driver.requests.get") + @patch("wherobots.db.driver.requests.post") + def test_the_key_the_api_stored_reaches_the_connection( + self, mock_post, mock_get, api_payload, session_key + ): + _, cd_kwargs = _run_connect_full( + mock_post, mock_get, api_payload=api_payload, session_key="dbt-run-1" + ) + assert cd_kwargs["session_key"] == session_key + + @patch("wherobots.db.driver.requests.delete") + @patch("wherobots.db.driver.requests.get") + @patch("wherobots.db.driver.requests.post") + def test_the_stopper_deletes_the_session_with_the_connect_headers( + self, mock_post, mock_get, mock_delete + ): + _, cd_kwargs = _run_connect_full(mock_post, mock_get) + mock_delete.assert_not_called() + cd_kwargs["session_stopper"]() + mock_delete.assert_called_once() + args, kwargs = mock_delete.call_args + assert args == ("https://api.example.com/sql/session/test-id",) + assert kwargs["headers"]["X-API-Key"] == "test-key" + assert kwargs["timeout"] == DEFAULT_HTTP_TIMEOUT + + @pytest.mark.parametrize( + "failure", + [requests.HTTPError("404 Client Error"), requests.ConnectionError("down")], + ids=["refused", "unreachable"], + ) + @patch("wherobots.db.driver.requests.delete") + @patch("wherobots.db.driver.requests.get") + @patch("wherobots.db.driver.requests.post") + def test_a_failed_stop_is_an_interface_error_naming_the_session( + self, mock_post, mock_get, mock_delete, failure + ): + _, cd_kwargs = _run_connect_full(mock_post, mock_get) + if isinstance(failure, requests.HTTPError): + mock_delete.return_value.raise_for_status.side_effect = failure + else: + mock_delete.side_effect = failure + with pytest.raises(InterfaceError, match="test-id"): + cd_kwargs["session_stopper"]() + + @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_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_ws.return_value = MagicMock() + with patch("wherobots.db.driver.Connection") as mock_conn: + connect(api_key="test-key", session_key="k") + _, conn_kwargs = mock_conn.call_args + assert conn_kwargs["session_key"] == "k" + assert callable(conn_kwargs["session_stopper"]) + + class TestConnectRegionRuntime: """region/runtime accept enum|str and are omitted when not provided.""" diff --git a/wherobots/db/connection.py b/wherobots/db/connection.py index 5eff23a..6c05552 100644 --- a/wherobots/db/connection.py +++ b/wherobots/db/connection.py @@ -286,6 +286,8 @@ def __init__( Callable[[float], websockets.sync.client.ClientConnection] | None ) = None, reconnect_timeout_seconds: float | None = DEFAULT_RECONNECT_TIMEOUT_SECONDS, + session_key: str | None = None, + session_stopper: Callable[[], None] | None = None, ): """ :param transport_factory: Opens a new WebSocket to the same SQL session, given @@ -297,6 +299,10 @@ def __init__( :param reconnect_timeout_seconds: How long to keep trying to reconnect after losing the transport. ``None``, or an invalid value (non-numeric, non-finite, or not positive), disables automatic reconnection. + :param session_key: See :attr:`session_key`. + :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``. """ self.__ws = ws self.__transport_factory = transport_factory @@ -343,6 +349,8 @@ def __init__( self.__session_id = session_id self.__session_fallback = session_fallback + self.__session_key = session_key + self.__session_stopper = session_stopper self.__lock = threading.Lock() self.__send_lock = threading.Lock() self.__shutdown_done = threading.Event() @@ -370,6 +378,39 @@ def session_fallback(self) -> str | None: another of the user's sessions.""" return self.__session_fallback + @property + def session_key(self) -> str | None: + """The session key the API created this connection's SQL session + under, as the API reports it, or None for a session without a key. + It differs from the ``session_key`` passed to ``connect()`` when the + API shared another key's session (see :attr:`session_fallback`), or + ignored keys for the organization. None for a connection from + ``connect_direct()``, which never asks the API.""" + return self.__session_key + + 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. + + :raises NotSupportedError: The connection came from + ``connect_direct()``, which has no API to ask. The connection + stays open. + :raises InterfaceError: The API refused or couldn't be reached. The + connection is closed, and the session stops after its inactivity + timeout instead. + """ + if self.__session_stopper is None: + raise NotSupportedError( + "Only a connection from connect() can stop its SQL session" + ) + self.close() + self.__session_stopper() + def __enter__(self): return self diff --git a/wherobots/db/driver.py b/wherobots/db/driver.py index b8786fc..b8a1367 100644 --- a/wherobots/db/driver.py +++ b/wherobots/db/driver.py @@ -12,7 +12,7 @@ import requests import tenacity import threading -from typing import Any, Final, Union, Dict +from typing import Any, Callable, Final, Union, Dict import urllib.parse import websockets.exceptions import websockets.sync.client @@ -282,7 +282,7 @@ def connect( before_sleep=lambda _: _check_cancelled(cancel_event), reraise=True, ) - def get_session_uri() -> str: + def get_ready_session() -> Dict[str, Any]: _check_cancelled(cancel_event) r = requests.get(session_id_url, headers=headers, timeout=DEFAULT_HTTP_TIMEOUT) r.raise_for_status() @@ -292,20 +292,43 @@ def get_session_uri() -> str: if status.is_starting(): raise tenacity.TryAgain("SQL Session is not ready yet") elif status == AppStatus.READY: - return payload["appMeta"]["url"] + return payload else: logging.error("SQL session creation failed: %s; should not retry.", status) raise OperationalError(f"Failed to create SQL session: {status}") try: logging.info("Getting SQL session status from %s ...", session_id_url) - session_uri = get_session_uri() + session = get_ready_session() + session_uri = session["appMeta"]["url"] logging.debug("SQL session URI from app status: %s", session_uri) except InterfaceError: raise except Exception as e: raise InterfaceError("Could not acquire SQL session!", e) + session_id = ( + urllib.parse.urlparse(session_id_url).path.rstrip("/").rsplit("/", 1)[-1] + ) + # The key the API created the session under, which differs from the one + # asked for after a fallback, or when the organization ignores keys. + api_payload = session.get("apiPayload") + session_key_on_api = ( + api_payload.get("sessionKey") if isinstance(api_payload, dict) else None + ) + # Snapshot the headers, as for reconnecting: a stop sends what this + # connection sent even if the caller later mutates them. + stop_headers = dict(headers) + + def stop_session() -> None: + try: + resp = requests.delete( + session_id_url, headers=stop_headers, timeout=DEFAULT_HTTP_TIMEOUT + ) + resp.raise_for_status() + except requests.RequestException as e: + raise InterfaceError(f"Failed to stop SQL session {session_id}: {e}") from e + return connect_direct( uri=http_to_ws(session_uri), headers=headers, @@ -314,12 +337,12 @@ def get_session_uri() -> str: data_compression=data_compression, geometry_representation=geometry_representation, cancel_event=cancel_event, - session_id=urllib.parse.urlparse(session_id_url) - .path.rstrip("/") - .rsplit("/", 1)[-1], + session_id=session_id, stale_query_probe_seconds=stale_query_probe_seconds, session_fallback=session_fallback, reconnect_timeout_seconds=reconnect_timeout_seconds, + session_key=session_key_on_api, + session_stopper=stop_session, ) @@ -351,6 +374,8 @@ def connect_direct( stale_query_probe_seconds: Union[float, None] = DEFAULT_STALE_QUERY_PROBE_SECONDS, session_fallback: str | None = None, reconnect_timeout_seconds: Union[float, None] = DEFAULT_RECONNECT_TIMEOUT_SECONDS, + session_key: str | None = None, + session_stopper: Union[Callable[[], None], None] = None, ) -> Connection: """Connect to a SQL session's WebSocket endpoint directly. @@ -361,6 +386,10 @@ def connect_direct( :param stale_query_probe_seconds: See :func:`connect`. :param session_fallback: See :attr:`Connection.session_fallback`. :param reconnect_timeout_seconds: See :func:`connect`. + :param session_key: See :attr:`Connection.session_key`. + :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. """ uri_with_protocol = f"{uri}/{protocol}" ssl_context = ssl.create_default_context() @@ -435,4 +464,6 @@ def open_transport(timeout: float) -> websockets.sync.client.ClientConnection: session_fallback=session_fallback, transport_factory=open_transport, reconnect_timeout_seconds=reconnect_timeout_seconds, + session_key=session_key, + session_stopper=session_stopper, ) From 3a791d88a763351889292a7d7118f4de5fc02944 Mon Sep 17 00:00:00 2001 From: Jia Yu Date: Sun, 1 Nov 2026 23:52:38 -0800 Subject: [PATCH 2/2] docs: require a key unique to the run before stopping a session Two jobs that send the same session_key can share one session, so a matching key proves the session is the run's own only when the key is unique to the run. The example now builds the key from a run ID. --- README.md | 19 +++++++++++-------- 1 file changed, 11 insertions(+), 8 deletions(-) diff --git a/README.md b/README.md index aad3c04..7e8820e 100644 --- a/README.md +++ b/README.md @@ -253,22 +253,25 @@ connection, then asks the API to stop the session, which stops within seconds: ```python -conn = connect(api_key=..., session_key="nightly-build-1234") +run_key = f"nightly-build-{run_id}" # unique to this run of the job +conn = connect(api_key=..., session_key=run_key) try: ... # run the job's queries finally: - if conn.session_key == "nightly-build-1234": + if conn.session_key == run_key: conn.stop_session() else: conn.close() ``` -The session stops for every connection on it. Other clients may share a -session, so check `conn.session_key` before stopping it: it's the key the -API created the session under, which is your `session_key` only when the -session is yours alone. 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 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). `stop_session()` raises `InterfaceError` if the API refuses or can't be reached; the connection is closed anyway, and the session stops after its