feat(connection): stop the SQL session on request

README.md

@@ -243,6 +243,42 @@ 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
+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 == run_key:
+        conn.stop_session()
+    else:
+        conn.close()
+```
+
+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
+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`

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

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

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