Check type of parameters timeout_seconds

ranking-agent · Aug 8, 2023 · 855ddf5 · 855ddf5
1 parent abcc851
commit 855ddf5
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 11 deletions.
diff --git a/strider/server.py b/strider/server.py
@@ -331,9 +331,9 @@ async def sync_query(
     try:
         LOGGER.info(f"[{qid}] Starting sync query")
         # get max timeout
-        timeout = max(
-            max_process_time, query_dict.get("parameters", {}).get("timeout_seconds", 0)
-        )
+        timeout_seconds = (query_dict.get("parameters") or {}).get("timeout_seconds")
+        timeout_seconds = timeout_seconds if type(timeout_seconds) is int else 0
+        timeout = max(max_process_time, timeout_seconds)
         query_results = await asyncio.wait_for(lookup(query_dict, qid), timeout=timeout)
     except asyncio.TimeoutError:
         LOGGER.error(f"[{qid}] Sync query cancelled due to timeout.")
@@ -448,7 +448,7 @@ async def lookup(
     logger.setLevel(level_number)
     logger.addHandler(log_handler)
 
-    parameters = query_dict.get("parameters", {})
+    parameters = query_dict.get("parameters") or {}
 
     fetcher = Fetcher(logger, parameters)
 
@@ -527,9 +527,9 @@ async def async_lookup(
     query_results = {}
     try:
         # get max timeout
-        timeout = max(
-            max_process_time, query_dict.get("parameters", {}).get("timeout_seconds", 0)
-        )
+        timeout_seconds = (query_dict.get("parameters") or {}).get("timeout_seconds")
+        timeout_seconds = timeout_seconds if type(timeout_seconds) is int else 0
+        timeout = max(max_process_time, timeout_seconds)
         query_results = await asyncio.wait_for(lookup(query_dict, qid), timeout=timeout)
     except asyncio.TimeoutError:
         LOGGER.error(f"[{qid}]: Process cancelled due to timeout.")
@@ -560,10 +560,11 @@ async def single_lookup(query_key):
         query_result = {}
         try:
             # get max timeout
-            timeout = max(
-                max_process_time,
-                queries[query_key].get("parameters", {}).get("timeout_seconds", 0),
+            timeout_seconds = (queries[query_key].get("parameters") or {}).get(
+                "timeout_seconds"
             )
+            timeout_seconds = timeout_seconds if type(timeout_seconds) is int else 0
+            timeout = max(max_process_time, timeout_seconds)
             query_result = await asyncio.wait_for(
                 lookup(queries[query_key], qid), timeout=timeout
             )

diff --git a/strider/throttle.py b/strider/throttle.py
@@ -213,7 +213,10 @@ async def process_batch(
                         ),
                     )
                 )
-                kp_timeout = self.parameters.get("timeout_seconds", settings.kp_timeout)
+                kp_timeout = self.parameters.get("timeout_seconds")
+                kp_timeout = (
+                    kp_timeout if type(kp_timeout) is int else settings.kp_timeout
+                )
                 async with httpx.AsyncClient(timeout=kp_timeout) as client:
                     response = await client.post(
                         self.url,