feat(connection): keep a session to its key, and report whether the API did

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,21 @@ 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's
+  `None` for a session without a key, and it differs from your key when
+  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
+  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 +369,14 @@ 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
+    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
     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.

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"

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 = []

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

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()
@@ -372,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
@@ -388,14 +393,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

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(
@@ -264,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(
@@ -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,
     )