[wk-libs] increase network timeout to 30 mins, allow custom timeout (#577)

* increase default timeout to 2min, allow custom timeout
* add changelog entry
* Merge branch 'master' into custom-timeouts
* bump timout to 30min, fallback to previous wk_url, allow ds download without org
* formatting
* docs
* make timeout default a constant, fix test by specifying url explicitly

Author: jstriebel

webknossos/Changelog.md

@@ -12,8 +12,10 @@ For upgrade instructions, please check the respective *Breaking Changes* section
 ### Breaking Changes
 
 ### Added
+- A custom network request timeout can be set using `webknossos_context(…, timeout=300)` or `export WK_TIMEOUT="300"`. [#577](https://github.com/scalableminds/webknossos-libs/pull/577)
 
 ### Changed
+- The default network request timeout changed from ½min to 30 min. [#577](https://github.com/scalableminds/webknossos-libs/pull/577)
 
 ### Fixed
 

webknossos/tests/test_generated_client.py

@@ -1,5 +1,6 @@
 import pytest
 
+from webknossos.client._defaults import DEFAULT_WEBKNOSSOS_URL
 from webknossos.client._generated.api.default import (
     annotation_info,
     build_info,
@@ -99,7 +100,7 @@ def test_user_list(auth_client: Client) -> None:
 
 @pytest.mark.vcr()
 def test_dataset_info() -> None:
-    with webknossos_context():
+    with webknossos_context(url=DEFAULT_WEBKNOSSOS_URL):
         client = _get_generated_client()
     response = dataset_info.sync(
         organization_name="scalable_minds",

webknossos/webknossos/client/_defaults.py

@@ -1 +1,2 @@
 DEFAULT_WEBKNOSSOS_URL = "https://webknossos.org"
+DEFAULT_HTTP_TIMEOUT = 1800

webknossos/webknossos/client/_download_dataset.py

@@ -23,14 +23,19 @@
 
 def download_dataset(
     dataset_name: str,
-    organization_name: str,
+    organization_name: Optional[str] = None,
     bbox: Optional[BoundingBox] = None,
     layers: Optional[List[str]] = None,
     mags: Optional[List[Mag]] = None,
     path: Optional[Union[PathLike, str]] = None,
     exist_ok: bool = False,
 ) -> Dataset:
     client = _get_generated_client()
+    context = _get_context()
+
+    if organization_name is None:
+        organization_name = context.organization
+
     dataset_info_response = dataset_info.sync_detailed(
         organization_name=organization_name,
         data_set_name=dataset_name,
@@ -40,10 +45,8 @@ def download_dataset(
     parsed = dataset_info_response.parsed
     assert parsed is not None
 
-    datastore_client = _get_context().get_generated_datastore_client(
-        parsed.data_store.url
-    )
-    optional_datastore_token = _get_context().datastore_token
+    datastore_client = context.get_generated_datastore_client(parsed.data_store.url)
+    optional_datastore_token = context.datastore_token
 
     actual_path = Path(dataset_name) if path is None else Path(path)
     if actual_path.exists():

webknossos/webknossos/client/context.py

@@ -8,7 +8,8 @@
 [https://webknossos.org/auth/token](https://webknossos.org/auth/token)
 
 Using the same methods, you can also specify the webknossos-server if you
-are not using the default [webknossos.org](https://webknossos.org) instance.
+are not using the default [webknossos.org](https://webknossos.org) instance,
+as well as a timeout for network requests (default is 30 minutes).
 
 There are the following four options to specify which server context to use:
 
@@ -36,6 +37,7 @@
     # content of .env
     WK_TOKEN="my_webknossos_token"
     WK_URL="…"
+    WK_TIMEOUT="3600"  # in seconds
     ```
 
 4. If nothing else is specified and authentication is needed,
@@ -54,7 +56,7 @@
 from dotenv import load_dotenv
 from rich.prompt import Prompt
 
-from webknossos.client._defaults import DEFAULT_WEBKNOSSOS_URL
+from webknossos.client._defaults import DEFAULT_HTTP_TIMEOUT, DEFAULT_WEBKNOSSOS_URL
 from webknossos.client._generated import Client as GeneratedClient
 
 load_dotenv()
@@ -101,13 +103,14 @@ def _cached_get_datastore_token(context: "_WebknossosContext") -> str:
 def _cached__get_generated_client(
     webknossos_url: str,
     token: Optional[str],
+    timeout: int,
 ) -> GeneratedClient:
     """Generates a client which might contain an x-auth-token header."""
     if token is None:
-        return GeneratedClient(base_url=webknossos_url, timeout=30)
+        return GeneratedClient(base_url=webknossos_url, timeout=timeout)
     else:
         return GeneratedClient(
-            base_url=webknossos_url, headers={"X-Auth-Token": token}, timeout=30
+            base_url=webknossos_url, headers={"X-Auth-Token": token}, timeout=timeout
         )
 
 
@@ -122,6 +125,7 @@ def _clear_all_context_caches() -> None:
 class _WebknossosContext:
     url: str = os.environ.get("WK_URL", default=DEFAULT_WEBKNOSSOS_URL)
     token: Optional[str] = os.environ.get("WK_TOKEN", default=None)
+    timeout: int = int(os.environ.get("WK_TIMEOUT", default=DEFAULT_HTTP_TIMEOUT))
 
     # all properties are cached outside to allow re-usability
     # if same context is instantiated twice
@@ -149,14 +153,16 @@ def datastore_required_token(self) -> str:
 
     @property
     def generated_client(self) -> GeneratedClient:
-        return _cached__get_generated_client(self.url, self.token)
+        return _cached__get_generated_client(self.url, self.token, self.timeout)
 
     @property
     def generated_auth_client(self) -> GeneratedClient:
-        return _cached__get_generated_client(self.url, self.required_token)
+        return _cached__get_generated_client(
+            self.url, self.required_token, self.timeout
+        )
 
     def get_generated_datastore_client(self, datastore_url: str) -> GeneratedClient:
-        return GeneratedClient(base_url=datastore_url, timeout=120)
+        return GeneratedClient(base_url=datastore_url, timeout=self.timeout)
 
 
 _webknossos_context_var: ContextVar[_WebknossosContext] = ContextVar(
@@ -166,10 +172,32 @@ def get_generated_datastore_client(self, datastore_url: str) -> GeneratedClient:
 
 @contextmanager
 def webknossos_context(
-    url: str = DEFAULT_WEBKNOSSOS_URL,
+    url: Optional[str] = None,
     token: Optional[str] = None,
+    timeout: Optional[int] = None,
 ) -> Iterator[None]:
-    context_var_token = _webknossos_context_var.set(_WebknossosContext(url, token))
+    """Returns a new webKnossos server contextmanager. Use with the `with` statement:
+    ```python
+    with webknossos_context(token="my_webknossos_token"):
+       # code that interacts with webknossos
+    ```
+
+    You can specify the following arguments:
+    * `url`, by default [https://webknossos.org](https://www.webknossos.org),
+    * `token`, as displayed on [https://webknossos.org/auth/token](https://webknossos.org/auth/token),
+    * `timeout` to specify a custom network request timeout in seconds, `1800` (30min) by default.
+
+    `url` and `timeout` are taken from the previous context (e.g. environment variables) if not specified.
+    `token` must be set explicitly, it is not available when not specified.
+    """
+
+    if url is None:
+        url = _get_context().url
+    if timeout is None:
+        timeout = _get_context().timeout
+    context_var_token = _webknossos_context_var.set(
+        _WebknossosContext(url, token, timeout)
+    )
     try:
         yield
     finally:

webknossos/webknossos/dataset/dataset.py

@@ -219,7 +219,7 @@ def open(cls, dataset_path: Union[str, PathLike]) -> "Dataset":
     def download(
         cls,
         dataset_name: str,
-        organization_name: str,
+        organization_name: Optional[str] = None,
         bbox: Optional[BoundingBox] = None,
         layers: Optional[List[str]] = None,
         mags: Optional[List[Mag]] = None,