Set custom http timeouts (#155)

Author: GianlucaFicarelli

src/entitysdk/config.py

@@ -33,5 +33,34 @@ class Settings(BaseSettings):
         ),
     ] = "https://www.openbraininstitute.org/api/entitycore"
 
+    connect_timeout: Annotated[
+        float,
+        Field(
+            alias="ENTITYSDK_CONNECT_TIMEOUT",
+            description="Maximum time to wait until a connection is established, in seconds.",
+        ),
+    ] = 5
+    read_timeout: Annotated[
+        float,
+        Field(
+            alias="ENTITYSDK_READ_TIMEOUT",
+            description="Maximum time to wait for a chunk of data to be received, in seconds.",
+        ),
+    ] = 30
+    write_timeout: Annotated[
+        float,
+        Field(
+            alias="ENTITYSDK_WRITE_TIMEOUT",
+            description="Maximum time to wait for a chunk of data to be sent, in seconds.",
+        ),
+    ] = 30
+    pool_timeout: Annotated[
+        float,
+        Field(
+            alias="ENTITYSDK_POOL_TIMEOUT",
+            description="Maximum time to acquire a connection from the pool, in seconds.",
+        ),
+    ] = 5
+
 
 settings = Settings()

src/entitysdk/core.py

@@ -11,6 +11,7 @@
 
 from entitysdk import serdes
 from entitysdk.common import ProjectContext
+from entitysdk.config import settings
 from entitysdk.exception import EntitySDKError
 from entitysdk.models.asset import (
     Asset,
@@ -216,7 +217,7 @@ def delete_entity(
     entity_type: type[Identifiable],
     token: str,
     http_client: httpx.Client | None = None,
-):
+) -> None:
     """Delete entity."""
     make_db_api_request(
         url=url,
@@ -310,16 +311,22 @@ def upload_asset_directory(
     js = response.json()
 
     def upload(to_upload):
+        upload_client = http_client or httpx.Client()
         failed = {}
         for path, url in to_upload.items():
             with open(paths[Path(path)], "rb") as fd:
                 try:
-                    response = http_client.request(
+                    response = upload_client.request(
                         method="PUT",
                         url=url,
                         content=fd,
                         follow_redirects=True,
-                        timeout=20,
+                        timeout=httpx.Timeout(
+                            connect=settings.connect_timeout,
+                            read=settings.read_timeout,
+                            write=settings.write_timeout,
+                            pool=settings.pool_timeout,
+                        ),
                     )
                 except httpx.HTTPError:
                     L.exception("Upload failed, will retry again")

src/entitysdk/util.py

@@ -46,6 +46,12 @@ def make_db_api_request(
             data=data,
             params=parameters,
             follow_redirects=True,
+            timeout=httpx.Timeout(
+                connect=settings.connect_timeout,
+                read=settings.read_timeout,
+                write=settings.write_timeout,
+                pool=settings.pool_timeout,
+            ),
         )
     except httpx.RequestError as e:
         raise EntitySDKError(f"Request error: {e}") from e