From 30cd113ac26b275b42e13516970f1308a5520617 Mon Sep 17 00:00:00 2001 From: Carl Robben Date: Fri, 29 May 2026 17:17:36 +1200 Subject: [PATCH 1/2] feat: pool connections via a per-client requests.Session - BaseClient and IfmClient each build one requests.Session at construction and route every send through it - Session-wide verify default set once; per-call Authorization headers still merged at request time - _build_session helper centralises the verify default so new endpoints can't forget the flag - verify_ssl tests repointed at client._session.request to reflect the new dispatch site --- datamasque/client/base.py | 34 ++++++++++++++++++++++++++++------ datamasque/client/ifm.py | 14 +++++++------- tests/test_base.py | 27 +++++++++------------------ 3 files changed, 44 insertions(+), 31 deletions(-) diff --git a/datamasque/client/base.py b/datamasque/client/base.py index ca2ca15..e828746 100644 --- a/datamasque/client/base.py +++ b/datamasque/client/base.py @@ -24,6 +24,21 @@ FileOrContent = Union[str, bytes, TextIOBase, BufferedIOBase, Path] _T = TypeVar("_T", bound=BaseModel) + +def _build_session(verify_ssl: bool) -> requests.Session: + """ + Build a configured `requests.Session` for one client's lifetime. + + Centralises the `verify` default so every call site inherits it + automatically — keeping the per-call code free of boilerplate and removing + the risk of forgetting the flag on a new endpoint. + """ + + session = requests.Session() + session.verify = verify_ssl + return session + + # Substrings (case-insensitive) that mark a key whose value should be redacted # before logging on an error path, so that passwords, API tokens, and similar secrets don't # end up in user-visible logs when a request fails. @@ -71,6 +86,15 @@ class BaseClient: Holds the connection config, cached auth token, and the core `make_request` dispatcher used by all per-feature mixins that compose `DataMasqueClient`. + + Uses a single `requests.Session` for the lifetime of the client so that + per-host TCP / TLS connections are pooled across calls (paginated list + endpoints and tight polling loops benefit most). Session-wide defaults + (`verify`) are set once on construction; per-call headers like + `Authorization` are merged at request time. + + `requests.Session` is not thread-safe; do not share a client between + threads. Construct one per worker. """ token: str = "" @@ -86,6 +110,7 @@ def __init__(self, connection_config: DataMasqueInstanceConfig) -> None: self.password = connection_config.password self.verify_ssl = connection_config.verify_ssl self.token_source = connection_config.token_source + self._session = _build_session(self.verify_ssl) @contextmanager def _maybe_suppress_insecure_warning(self) -> Iterator[None]: @@ -186,23 +211,20 @@ def make_request( url = urljoin(self.base_url, path) def send() -> Response: - headers: Optional[dict] = {"Authorization": self.token} if requires_authorization else None + headers = {"Authorization": self.token} if requires_authorization else None try: with self._maybe_suppress_insecure_warning(): if files: files_payload = {f.field_name: (f.filename, f.content, f.content_type or "") for f in files} - return requests.request( + return self._session.request( method, url, data=data, params=params, headers=headers, files=files_payload, - verify=self.verify_ssl, ) - return requests.request( - method, url, json=data, params=params, headers=headers, verify=self.verify_ssl - ) + return self._session.request(method, url, json=data, params=params, headers=headers) except requests.RequestException as e: raise DataMasqueTransportError(f"Failed to reach DataMasque server at {url}: {e}") from e diff --git a/datamasque/client/ifm.py b/datamasque/client/ifm.py index 374d46e..11a041d 100644 --- a/datamasque/client/ifm.py +++ b/datamasque/client/ifm.py @@ -17,7 +17,7 @@ from pydantic import BaseModel from requests import Response -from datamasque.client.base import suppress_insecure_warning_if_needed +from datamasque.client.base import _build_session, suppress_insecure_warning_if_needed from datamasque.client.exceptions import ( DataMasqueApiError, DataMasqueNotReadyError, @@ -82,6 +82,9 @@ def __init__(self, connection_config: DataMasqueIfmInstanceConfig) -> None: self.password = connection_config.password self.verify_ssl = connection_config.verify_ssl self.token_source = connection_config.token_source + # One session for both admin-server (JWT login/refresh) and IFM (data plane) + # traffic -- different hosts, but a single session handles per-host pooling. + self._session = _build_session(self.verify_ssl) def authenticate(self) -> None: """Obtain an access (and refresh) token from the admin server, or via `token_source`.""" @@ -95,10 +98,9 @@ def authenticate(self) -> None: login_url = urljoin(self.admin_server_base_url, "/api/auth/jwt/login/") try: with self._maybe_suppress_insecure_warning(): - response = requests.post( + response = self._session.post( login_url, json={"username": self.username, "password": self.password}, - verify=self.verify_ssl, ) except requests.RequestException as e: raise DataMasqueTransportError(f"Failed to reach admin server at {login_url}: {e}") from e @@ -122,10 +124,9 @@ def _refresh_or_reauth(self) -> None: refresh_url = urljoin(self.admin_server_base_url, "/api/auth/jwt/refresh/") try: with self._maybe_suppress_insecure_warning(): - response = requests.post( + response = self._session.post( refresh_url, json={"refresh": self.refresh_token}, - verify=self.verify_ssl, ) except requests.RequestException as e: raise DataMasqueTransportError(f"Failed to reach admin server at {refresh_url}: {e}") from e @@ -187,13 +188,12 @@ def _make_request( def send() -> Response: try: with self._maybe_suppress_insecure_warning(): - return requests.request( + return self._session.request( method, url, json=json_body, params=params, headers={"Authorization": f"Bearer {self.access_token}"}, - verify=self.verify_ssl, ) except requests.RequestException as e: raise DataMasqueTransportError(f"Failed to reach IFM server at {url}: {e}") from e diff --git a/tests/test_base.py b/tests/test_base.py index 00bf32f..e72e8e2 100644 --- a/tests/test_base.py +++ b/tests/test_base.py @@ -74,8 +74,12 @@ def test_healthcheck_transport_failure(client): @pytest.mark.parametrize("verify_ssl", [True, False]) -def test_make_request_verify_ssl_true_by_default(config, verify_ssl): - """Verifies SSL setting is passed through to the `requests` call.""" +def test_session_verify_reflects_config(config, verify_ssl): + """ + `verify_ssl` is applied to the client's `requests.Session` once at construction. + + Every outgoing call then inherits it without per-call boilerplate. + """ config_with_ssl = DataMasqueInstanceConfig( base_url=config.base_url, username=config.username, @@ -84,24 +88,14 @@ def test_make_request_verify_ssl_true_by_default(config, verify_ssl): ) client = DataMasqueClient(config_with_ssl) - with patch( - "datamasque.client.base.requests.request", - return_value=make_ok_response(), - ) as mock_request: - client.make_request("GET", "/api/test/") - - _, kwargs = mock_request.call_args - assert kwargs["verify"] is verify_ssl + assert client._session.verify is verify_ssl def test_make_request_verify_ssl_true_does_not_touch_global_warning_filter(client): """With `verify_ssl=True`, the client should not modify `warnings.filters`.""" filters_before = list(warnings.filters) - with patch( - "datamasque.client.base.requests.request", - return_value=make_ok_response(), - ): + with patch.object(client._session, "request", return_value=make_ok_response()): client.make_request("GET", "/api/test/") assert warnings.filters == filters_before @@ -125,10 +119,7 @@ def raise_insecure_warning_then_respond(*_args, **_kwargs): with warnings.catch_warnings(record=True) as captured: warnings.simplefilter("always") # ensure we'd otherwise see the warning - with patch( - "datamasque.client.base.requests.request", - side_effect=raise_insecure_warning_then_respond, - ): + with patch.object(client._session, "request", side_effect=raise_insecure_warning_then_respond): client.make_request("GET", "/api/test/") # The warning raised inside the request call was suppressed by the client. From 79f462f489aacac02e3c8da13d303d8fa7fe6a36 Mon Sep 17 00:00:00 2001 From: Carl Robben Date: Fri, 29 May 2026 17:18:14 +1200 Subject: [PATCH 2/2] fix: gate 401 re-auth retry on requires_authorization - A 401 on a requires_authorization=False call now surfaces as-is instead of triggering a re-auth and replay - Anonymous 401s (e.g. admin-install on an already-configured instance) no longer misreport as bad credentials - Avoids a wasted login round-trip on calls the caller explicitly marked as not needing auth - Test asserts a single request and no /api/auth/token/login/ replay on an anonymous 401 Closes #9 --- datamasque/client/base.py | 9 ++++++++- tests/test_base.py | 31 +++++++++++++++++++++++++++++++ 2 files changed, 39 insertions(+), 1 deletion(-) diff --git a/datamasque/client/base.py b/datamasque/client/base.py index e828746..442e867 100644 --- a/datamasque/client/base.py +++ b/datamasque/client/base.py @@ -229,7 +229,14 @@ def send() -> Response: raise DataMasqueTransportError(f"Failed to reach DataMasque server at {url}: {e}") from e response = send() - if response.status_code == 401: + if response.status_code == 401 and requires_authorization: + # Token-expiry recovery: re-auth and replay. Only meaningful when the + # caller actually sent a token; on `requires_authorization=False` + # calls a 401 means the server itself is rejecting anonymous access + # (e.g. admin-install on an already-configured instance), and + # re-authing with whatever creds the client happens to hold would + # both misdiagnose the failure and emit a misleading + # "credentials are incorrect" error to the user. logger.debug("Re-authenticating") self.authenticate() # Reset file pointers so the retry doesn't send empty files diff --git a/tests/test_base.py b/tests/test_base.py index e72e8e2..6d1d4b0 100644 --- a/tests/test_base.py +++ b/tests/test_base.py @@ -280,6 +280,37 @@ def test_token_source_called_again_on_401_retry(): assert client.token == "Token t2" +def test_401_does_not_retry_when_requires_authorization_is_false(client): + """ + A 401 on an anonymous request must surface as-is, not trigger a re-auth retry. + + `/api/users/admin-install/` returns 401 once any user exists -- the endpoint + is gated on "no user has been created yet" and DRF treats it as a normal + auth-required endpoint thereafter. Re-authing on that 401 would both + misdiagnose the failure ("login credentials are correct") and waste a + round-trip on a call the caller said doesn't need auth. + """ + with requests_mock.Mocker() as m: + m.post( + "http://test-server/api/users/admin-install/", + status_code=401, + json={"detail": "Authentication credentials were not provided."}, + ) + + with pytest.raises(DataMasqueApiError) as excinfo: + client.make_request( + "POST", + "/api/users/admin-install/", + data={"email": "x@y", "username": "x", "password": "p", "re_password": "p", "allowed_hosts": []}, + requires_authorization=False, + ) + + assert excinfo.value.response.status_code == 401 + # Exactly one request: no re-auth roundtrip to /api/auth/token/login/ and no replay. + assert m.call_count == 1 + assert m.request_history[0].path == "/api/users/admin-install/" + + def test_token_source_callable_exception_propagates(): """Errors from `token_source` are surfaced to the caller, not swallowed."""