Substra · oleobal · Aug 17, 2023 · Aug 9, 2023 · Aug 9, 2023 · Aug 10, 2023
@@ -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
 

@@ -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

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

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

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

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

@@ -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):

@@ -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):