From e2916b9e0bfd04aa7b2a60ed08b8dd0ca53147aa Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Olivier=20L=C3=A9obal?= Date: Thu, 17 Aug 2023 12:18:40 +0200 Subject: [PATCH] feat: add Client.logout and context manager (#381) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add a Client.logout function, which deletes a token obtained by Client.login (if any) Add support for using Client within a context manager Signed-off-by: Olivier Léobal Co-authored-by: Thibault Fouqueray --- CHANGELOG.md | 13 +++++++++ substra/sdk/backends/base.py | 4 +++ substra/sdk/backends/local/backend.py | 3 +++ substra/sdk/backends/local/dal.py | 4 +++ substra/sdk/backends/remote/backend.py | 3 +++ substra/sdk/backends/remote/rest_client.py | 19 +++++++++++++ substra/sdk/client.py | 26 ++++++++++++++++++ tests/sdk/test_client.py | 31 +++++++++++++++++----- 8 files changed, 96 insertions(+), 7 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 01fc5c4c..a2425157 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,10 +12,23 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - `wait_completion` param on `get_performances`, `list_task_output_assets` and `get_task_output_asset` to block execution until execution is over ([#368](https://github.com/Substra/substra/pull/368)) - `list_task_output_assets` and `get_task_output_asset` wait that the compute task is over before getting assets ([#369](https://github.com/Substra/substra/pull/369)) - warning and help message when logging in with username/password rather than token ([#378](https://github.com/Substra/substra/pull/378)) +- new `Client.logout` function, mirroring `Client.login` ([#381](https://github.com/Substra/substra/pull/381)) +- `Client` can now be used within a context manager ([#381](https://github.com/Substra/substra/pull/381)) + ```python + with Client( + client_name="org-1", + backend_type="remote", + url="http://substra-backend.org-1.com:8000", + username="org-1", + password="p@sswr0d44", + ) as client: + pass + ``` ### Changed - change how API responses are parsed to match server changes ([#379](https://github.com/Substra/substra/pull/379)) +- `Client` will now terminate the sessions it starts when given username & password ([#381](https://github.com/Substra/substra/pull/381)) ### Fixed diff --git a/substra/sdk/backends/base.py b/substra/sdk/backends/base.py index b6002a49..573c8404 100644 --- a/substra/sdk/backends/base.py +++ b/substra/sdk/backends/base.py @@ -14,6 +14,10 @@ def backend_mode(self) -> BackendType: def login(self, username, password): pass + @abc.abstractmethod + def logout(self): + pass + @abc.abstractmethod def get(self, asset_type, key): raise NotImplementedError diff --git a/substra/sdk/backends/local/backend.py b/substra/sdk/backends/local/backend.py index 529aa3fc..e32e53f8 100644 --- a/substra/sdk/backends/local/backend.py +++ b/substra/sdk/backends/local/backend.py @@ -74,6 +74,9 @@ def backend_mode(self) -> schemas.BackendType: def login(self, username, password): self._db.login(username, password) + def logout(self): + self._db.logout() + def get(self, asset_type, key): return self._db.get(asset_type, key) diff --git a/substra/sdk/backends/local/dal.py b/substra/sdk/backends/local/dal.py index bfc264e2..135a19fa 100644 --- a/substra/sdk/backends/local/dal.py +++ b/substra/sdk/backends/local/dal.py @@ -53,6 +53,10 @@ def login(self, username, password): if self._remote: self._remote.login(username, password) + def logout(self): + if self._remote: + self._remote.logout() + def add(self, asset): return self._db.add(asset) diff --git a/substra/sdk/backends/remote/backend.py b/substra/sdk/backends/remote/backend.py index c05ac8d9..f1fca6ec 100644 --- a/substra/sdk/backends/remote/backend.py +++ b/substra/sdk/backends/remote/backend.py @@ -42,6 +42,9 @@ def backend_mode(self) -> schemas.BackendType: def login(self, username, password): return self._client.login(username, password) + def logout(self): + return self._client.logout() + def get(self, asset_type, key): """Get an asset by key.""" asset = self._client.get(asset_type.to_server(), key) diff --git a/substra/sdk/backends/remote/rest_client.py b/substra/sdk/backends/remote/rest_client.py index fd2c4401..2b3a5e3c 100644 --- a/substra/sdk/backends/remote/rest_client.py +++ b/substra/sdk/backends/remote/rest_client.py @@ -52,6 +52,7 @@ def __init__(self, url, insecure, token): if not url: raise exceptions.SDKException("url required to connect to the Substra server") self._base_url = url[:-1] if url.endswith("/") else url + self._token_id = None # filled in by self.login def login(self, username, password): # we do not use self._headers in order to avoid existing tokens to be sent alongside the @@ -90,6 +91,7 @@ def login(self, username, password): try: rj = r.json() token = rj["token"] + self._token_id = rj["id"] except json.decoder.JSONDecodeError: # sometimes requests seem to be fine, but the json is not being found # this might be if the url seems to be correct (in the syntax) @@ -108,6 +110,23 @@ def login(self, username, password): return token + def logout(self) -> None: + if self._token_id is None: + return + try: + r = requests.delete( + f"{self._base_url}/active-api-tokens/", params={"id": self._token_id}, headers=self._headers + ) + r.raise_for_status() + self._token_id = None + logger.info("Successfully logged out") + except requests.exceptions.ConnectionError as e: + raise exceptions.ConnectionError.from_request_exception(e) + except requests.exceptions.Timeout as e: + raise exceptions.Timeout.from_request_exception(e) + except requests.exceptions.HTTPError as e: + raise exceptions.HTTPError.from_request_exception(e) + # TODO: '__request' is too complex, consider refactoring def __request(self, request_name, url, **request_kwargs): # noqa: C901 """Base request helper.""" diff --git a/substra/sdk/client.py b/substra/sdk/client.py index baf605bb..57f850bb 100644 --- a/substra/sdk/client.py +++ b/substra/sdk/client.py @@ -199,6 +199,12 @@ class Client: Defaults to None. username (str, optional): Username to authenticate to the Substra platform. Used in conjunction with a password to generate a token if not given, using the `login` function. + + If using username/password, you should use a context manager to ensure the session terminates as intended: + ``` + with Client(username, password) as client: + ... + ``` Not stored. Defaults to None. password (str, optional): Password to authenticate to the Substra platform. @@ -311,6 +317,15 @@ def __init__( ) self._token = self.login(config_dict["username"].value, config_dict["password"].value) + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_value, traceback): + self.logout() + + def __del__(self): + self.logout() + def _get_backend(self, backend_type: schemas.BackendType): # Three possibilities: # - remote: get a remote backend @@ -369,6 +384,17 @@ def login(self, username, password): self._token = self._backend.login(username, password) return self._token + @logit + def logout(self) -> None: + """ + Log out from a remote server, if Client.login was used + (otherwise, nothing happens) + """ + if not self._backend: + raise exceptions.SDKException("No backend found") + self._backend.logout() + self._token = None + @staticmethod def _get_spec(asset_type, data): if isinstance(data, asset_type): diff --git a/tests/sdk/test_client.py b/tests/sdk/test_client.py index 7ca63a2b..5ac0f9cb 100644 --- a/tests/sdk/test_client.py +++ b/tests/sdk/test_client.py @@ -91,14 +91,31 @@ def test_client_should_raise_when_missing_name(): def test_client_with_password(mocker): mocker.patch("substra.sdk.Client.login", side_effect=stub_login) - client = Client( - backend_type="remote", - url="example.com", - token=None, - username="org-1", - password="password1", - ) + rest_client_logout = mocker.patch("substra.sdk.backends.remote.rest_client.Client.logout") + client_args = { + "backend_type": "remote", + "url": "example.com", + "token": None, + "username": "org-1", + "password": "password1", + } + + client = Client(**client_args) assert client._token == "token1" + client.logout() + assert client._token is None + rest_client_logout.assert_called_once() + del client + + rest_client_logout.reset_mock() + client = Client(**client_args) + del client + rest_client_logout.assert_called_once() + + rest_client_logout.reset_mock() + with Client(**client_args) as client: + assert client._token == "token1" + rest_client_logout.assert_called_once() def test_client_token_supercedes_password(mocker):