Configurable Timeout (#138)

GitOrigin-RevId: 08db966ce345de9f68e0b3f08d3eb05946391b94
descarteslabs · Nov 20, 2023 · c66a6dd · c66a6dd
1 parent 645ef5a
commit c66a6dd
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 23 deletions.
diff --git a/descarteslabs/vector/common.py b/descarteslabs/vector/common.py
@@ -4,8 +4,12 @@
 import geopandas as gpd
 import pandas as pd
 
+from . import __version__
+
 API_HOST = os.getenv("VECTOR_API_HOST", "https://vector.descarteslabs.com")
 TYPES = (gpd.GeoDataFrame, pd.DataFrame)
+VECTOR_TIMEOUT = int(os.environ.get("VECTOR_TIMEOUT", "600"))
+USERAGENT = f"dl-vector/{__version__}"
 
 
 def get_token() -> str:

diff --git a/descarteslabs/vector/features.py b/descarteslabs/vector/features.py
@@ -10,13 +10,9 @@
 import requests
 from descarteslabs.utils import Properties
 
-from . import __version__
-from .common import API_HOST, TYPES, get_token
+from .common import API_HOST, TYPES, USERAGENT, VECTOR_TIMEOUT, get_token
 from .util import backoff_wrapper, check_response, response_to_dataframe
 
-REQUEST_TIMEOUT = 600
-USERAGENT = f"dl-vector/{__version__}"
-
 
 class Statistic(str, Enum):
     """
@@ -69,7 +65,7 @@ def add(
             "User-Agent": USERAGENT,
         },
         files=files,
-        timeout=REQUEST_TIMEOUT,
+        timeout=VECTOR_TIMEOUT,
     )
 
     check_response(response, "add feature")
@@ -113,7 +109,7 @@ def query(
             "aoi": aoi,
             "columns": columns,
         },
-        timeout=REQUEST_TIMEOUT,
+        timeout=VECTOR_TIMEOUT,
     )
     check_response(response, "query feature")
 
@@ -148,7 +144,7 @@ def _join(params: dict) -> Union[gpd.GeoDataFrame, pd.DataFrame]:
         f"{API_HOST}/products/features/join",
         headers={"Authorization": get_token(), "User-Agent": USERAGENT},
         json=params,
-        timeout=REQUEST_TIMEOUT,
+        timeout=VECTOR_TIMEOUT,
     )
     check_response(response, "join feature")
 
@@ -298,7 +294,7 @@ def get(product_id: str, feature_id: str) -> Union[gpd.GeoDataFrame, pd.DataFram
         f"{API_HOST}/products/{product_id}/features/{feature_id}",
         headers={"Authorization": get_token(), "User-Agent": USERAGENT},
         params={"format": "Parquet"},
-        timeout=REQUEST_TIMEOUT,
+        timeout=VECTOR_TIMEOUT,
     )
 
     check_response(response, "get feature")
@@ -351,7 +347,7 @@ def update(
             "User-Agent": USERAGENT,
         },
         files=files,
-        timeout=REQUEST_TIMEOUT,
+        timeout=VECTOR_TIMEOUT,
     )
 
     check_response(response, "update feature")
@@ -404,7 +400,7 @@ def aggregate(
             "aoi": aoi,
             "columns": columns,
         },
-        timeout=REQUEST_TIMEOUT,
+        timeout=VECTOR_TIMEOUT,
     )
     check_response(response, "aggregate feature")
 
@@ -427,7 +423,7 @@ def delete(product_id: str, feature_id: str):
     response = requests.delete(
         f"{API_HOST}/products/{product_id}/features/{feature_id}",
         headers={"Authorization": get_token(), "User-Agent": USERAGENT},
-        timeout=REQUEST_TIMEOUT,
+        timeout=VECTOR_TIMEOUT,
     )
 
     check_response(response, "delete feature")
diff --git a/descarteslabs/vector/products.py b/descarteslabs/vector/products.py
@@ -2,15 +2,11 @@
 
 import requests
 
-from . import __version__
-from .common import API_HOST, get_token
+from .common import API_HOST, USERAGENT, VECTOR_TIMEOUT, get_token
 from .models import GenericFeatureBaseModel, VectorBaseModel
 from .util import backoff_wrapper, check_response
 from .vector_exceptions import ClientException
 
-REQUEST_TIMEOUT = 120
-USERAGENT = f"dl-vector/{__version__}"
-
 
 def _check_tags(tags: Union[List[str], None] = None):
     if tags:
@@ -101,7 +97,7 @@ def create(
         f"{API_HOST}/products/",
         headers={"Authorization": get_token(), "User-Agent": USERAGENT},
         json=request_json,
-        timeout=REQUEST_TIMEOUT,
+        timeout=VECTOR_TIMEOUT,
     )
     check_response(response, "create product")
     return response.json()
@@ -128,13 +124,13 @@ def list(tags: Union[List[str], None] = None) -> List[dict]:
             f"{API_HOST}/products/",
             headers={"Authorization": get_token(), "User-Agent": USERAGENT},
             params={"tags": ",".join(tags)},
-            timeout=REQUEST_TIMEOUT,
+            timeout=VECTOR_TIMEOUT,
         )
     else:
         response = requests.get(
             f"{API_HOST}/products/",
             headers={"Authorization": get_token(), "User-Agent": USERAGENT},
-            timeout=REQUEST_TIMEOUT,
+            timeout=VECTOR_TIMEOUT,
         )
     check_response(response, "list products")
     return response.json()
@@ -157,7 +153,7 @@ def get(product_id: str) -> dict:
     response = requests.get(
         f"{API_HOST}/products/{product_id}",
         headers={"Authorization": get_token(), "User-Agent": USERAGENT},
-        timeout=REQUEST_TIMEOUT,
+        timeout=VECTOR_TIMEOUT,
     )
     check_response(response, "get product")
     return response.json()
@@ -214,7 +210,7 @@ def update(
                 "owners": owners,
             },
         ),
-        timeout=REQUEST_TIMEOUT,
+        timeout=VECTOR_TIMEOUT,
     )
     check_response(response, "update product")
     return response.json()
@@ -237,6 +233,6 @@ def delete(product_id: str) -> None:
     response = requests.delete(
         f"{API_HOST}/products/{product_id}",
         headers={"Authorization": get_token(), "User-Agent": USERAGENT},
-        timeout=REQUEST_TIMEOUT,
+        timeout=VECTOR_TIMEOUT,
     )
     check_response(response, "delete product")