handle http response status code 429

PiperOrigin-RevId: 677896725
chronicle · Sep 23, 2024 · b4e8a3e · b4e8a3e
1 parent 915bdf7
commit b4e8a3e
Show file tree

Hide file tree

Showing 14 changed files with 235 additions and 96 deletions.
diff --git a/tools/rule_manager/chronicle_api/reference_lists/create_reference_list.py b/tools/rule_manager/chronicle_api/reference_lists/create_reference_list.py
@@ -19,6 +19,7 @@
 """
 
 import os
+import time
 from typing import Mapping, Any, Sequence
 
 from google.auth.transport import requests
@@ -30,6 +31,7 @@ def create_reference_list(
     description: str,
     entries: Sequence[str | None],
     syntax_type: str | None = None,
+    max_retries: int = 3,
 ) -> Mapping[str, Any]:
   """Creates a new reference list.
 
@@ -41,6 +43,9 @@ def create_reference_list(
       syntax_type: The syntax type indicating how list entries should be
         validated. Reference:
         https://cloud.google.com/chronicle/docs/reference/rest/v1alpha/projects.locations.instances.referenceLists#ReferenceListSyntaxType
+      max_retries (optional): Maximum number of times to retry HTTP request if
+        certain response codes are returned. For example: HTTP response status
+        code 429 (Too Many Requests)
 
   Returns:
       New reference list.
@@ -51,6 +56,7 @@ def create_reference_list(
   """
   url = f"{os.environ['CHRONICLE_API_BASE_URL']}/{os.environ['CHRONICLE_INSTANCE']}/referenceLists"
   params = {"referenceListId": name}
+  response = None
 
   if len(entries) == 0:  # pylint: disable="g-explicit-length-test"
     # If 'entries' is an empty list, the reference list is empty [{}]
@@ -68,12 +74,20 @@ def create_reference_list(
       "syntax_type": syntax_type,
   }
 
-  response = http_session.request(
-      method="POST", url=url, json=body, params=params
-  )
+  for _ in range(max_retries + 1):
+    response = http_session.request(
+        method="POST", url=url, json=body, params=params
+    )
+
+    if response.status_code >= 400:
+      print(response.text)
+
+    if response.status_code == 429:
+      print("API rate limit exceeded. Sleeping for 60s before retrying")
+      time.sleep(60)
+    else:
+      break
 
-  if response.status_code >= 400:
-    print(response.text)
   response.raise_for_status()
 
   return response.json()
diff --git a/tools/rule_manager/chronicle_api/reference_lists/get_reference_list.py b/tools/rule_manager/chronicle_api/reference_lists/get_reference_list.py
@@ -19,7 +19,8 @@
 """
 
 import os
-from typing import Any, Mapping
+import time
+from typing import Mapping, Any
 
 from google.auth.transport import requests
 
@@ -28,6 +29,7 @@ def get_reference_list(
     http_session: requests.AuthorizedSession,
     resource_name: str,
     view: str | None = "REFERENCE_LIST_VIEW_FULL",
+    max_retries: int = 3,
 ) -> Mapping[str, Any]:
   """Retrieves a reference list.
 
@@ -38,6 +40,9 @@ def get_reference_list(
       view (optional): The scope of fields to populate for the ReferenceList
         being returned. Reference:
         https://cloud.google.com/chronicle/docs/reference/rest/v1alpha/ReferenceListView
+      max_retries (optional): Maximum number of times to retry HTTP request if
+        certain response codes are returned. For example: HTTP response status
+        code 429 (Too Many Requests)
 
   Returns:
       Content and metadata about the requested reference list.
@@ -48,11 +53,20 @@ def get_reference_list(
   """
   url = f"{os.environ['CHRONICLE_API_BASE_URL']}/{resource_name}"
   params = {"view": view}
+  response = None
 
-  response = http_session.request(method="GET", url=url, params=params)
+  for _ in range(max_retries + 1):
+    response = http_session.request(method="GET", url=url, params=params)
 
-  if response.status_code >= 400:
-    print(response.text)
-    response.raise_for_status()
+    if response.status_code >= 400:
+      print(response.text)
+
+    if response.status_code == 429:
+      print("API rate limit exceeded. Sleeping for 60s before retrying")
+      time.sleep(60)
+    else:
+      break
+
+  response.raise_for_status()
 
   return response.json()
diff --git a/tools/rule_manager/chronicle_api/reference_lists/list_reference_lists.py b/tools/rule_manager/chronicle_api/reference_lists/list_reference_lists.py
@@ -19,6 +19,7 @@
 """
 
 import os
+import time
 from typing import Mapping, Any, List, Tuple
 
 from google.auth.transport import requests
@@ -29,6 +30,7 @@ def list_reference_lists(
     page_size: int | None = None,
     page_token: str | None = None,
     view: str | None = "REFERENCE_LIST_VIEW_FULL",
+    max_retries: int = 3,
 ) -> Tuple[List[Mapping[str, Any]], str]:
   """Retrieve a list of reference lists.
 
@@ -43,6 +45,9 @@ def list_reference_lists(
       view (optional): The scope of fields to populate for the ReferenceList
         being returned. Reference:
         https://cloud.google.com/chronicle/docs/reference/rest/v1alpha/ReferenceListView
+      max_retries (optional): Maximum number of times to retry HTTP request if
+        certain response codes are returned. For example: HTTP response status
+        code 429 (Too Many Requests)
 
   Returns:
       List of reference lists and a page token for the next page of reference
@@ -54,12 +59,21 @@ def list_reference_lists(
   """
   url = f"{os.environ['CHRONICLE_API_BASE_URL']}/{os.environ['CHRONICLE_INSTANCE']}/referenceLists"
   params = {"page_size": page_size, "page_token": page_token, "view": view}
+  response = None
 
-  response = http_session.request(method="GET", url=url, params=params)
+  for _ in range(max_retries + 1):
+    response = http_session.request(method="GET", url=url, params=params)
 
-  if response.status_code >= 400:
-    print(response.text)
-    response.raise_for_status()
+    if response.status_code >= 400:
+      print(response.text)
+
+    if response.status_code == 429:
+      print("API rate limit exceeded. Sleeping for 60s before retrying")
+      time.sleep(60)
+    else:
+      break
+
+  response.raise_for_status()
 
   response_json = response.json()
 

diff --git a/tools/rule_manager/chronicle_api/reference_lists/update_reference_list.py b/tools/rule_manager/chronicle_api/reference_lists/update_reference_list.py
@@ -20,6 +20,7 @@
 
 import copy
 import os
+import time
 from typing import Mapping, Any, List
 
 from google.auth.transport import requests
@@ -30,6 +31,7 @@ def update_reference_list(
     resource_name: str,
     updates: Mapping[str, Any],
     update_mask: List[str] | None = None,
+    max_retries: int = 3,
 ) -> Mapping[str, Any]:
   """Updates an existing reference list.
 
@@ -44,6 +46,9 @@ def update_reference_list(
         list. If no update_mask is provided, all non-empty fields will be
         updated. example: An update_mask of ["entries"] will update the entries
         for a reference list.
+      max_retries (optional): Maximum number of times to retry HTTP request if
+        certain response codes are returned. For example: HTTP response status
+        code 429 (Too Many Requests)
 
   Returns:
       New version of the reference list.
@@ -53,6 +58,7 @@ def update_reference_list(
       (response.status_code >= 400).
   """
   url = f"{os.environ['CHRONICLE_API_BASE_URL']}/{resource_name}"
+  response = None
 
   # If no update_mask is provided, all non-empty fields will be updated
   if update_mask is None:
@@ -61,7 +67,7 @@ def update_reference_list(
     params = {"updateMask": update_mask}
 
   if updates.get("entries") is not None:
-    if len(updates.get("entries")) == 0:  # pylint: disable="g-explicit-length-test"
+    if len(updates.get("entries")) == 0:  # pylint: disable=g-explicit-length-test
       # If 'entries' is an empty list, the reference list is empty [{}]
       updates["entries"] = [{}]
     else:
@@ -72,12 +78,20 @@ def update_reference_list(
         reference_list_entries.append({"value": entry.strip()})
       updates["entries"] = copy.deepcopy(reference_list_entries)
 
-  response = http_session.request(
-      method="PATCH", url=url, params=params, json=updates
-  )
+  for _ in range(max_retries + 1):
+    response = http_session.request(
+        method="PATCH", url=url, params=params, json=updates
+    )
+
+    if response.status_code >= 400:
+      print(response.text)
+
+    if response.status_code == 429:
+      print("API rate limit exceeded. Sleeping for 60s before retrying")
+      time.sleep(60)
+    else:
+      break
 
-  if response.status_code >= 400:
-    print(response.text)
   response.raise_for_status()
 
   return response.json()
diff --git a/tools/rule_manager/chronicle_api/rules/create_rule.py b/tools/rule_manager/chronicle_api/rules/create_rule.py
@@ -1,4 +1,4 @@
-# Copyright 2023 Google LLC
+# Copyright 2024 Google LLC
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -19,34 +19,49 @@
 """
 
 import os
-from typing import Mapping, Any
+import time
+from typing import Any, Mapping
 
 from google.auth.transport import requests
 
 
 def create_rule(
-    http_session: requests.AuthorizedSession, rule_text: str
+    http_session: requests.AuthorizedSession,
+    rule_text: str,
+    max_retries: int = 3,
 ) -> Mapping[str, Any]:
   """Creates a new rule.
 
   Args:
     http_session: Authorized session for HTTP requests.
     rule_text: The content of the YARA-L 2.0 rule.
+    max_retries (optional): Maximum number of times to retry HTTP request if
+      certain response codes are returned. For example: HTTP response status
+          code 429 (Too Many Requests)
 
   Returns:
     New rule.
 
   Raises:
     requests.exceptions.HTTPError: HTTP request resulted in an error
-      (response.status_code >= 400).
+    (response.status_code >= 400).
   """
   url = f"{os.environ['CHRONICLE_API_BASE_URL']}/{os.environ['CHRONICLE_INSTANCE']}/rules"
   body = {"text": rule_text}
+  response = None
 
-  response = http_session.request(method="POST", url=url, json=body)
+  for _ in range(max_retries + 1):
+    response = http_session.request(method="POST", url=url, json=body)
+
+    if response.status_code >= 400:
+      print(response.text)
+
+    if response.status_code == 429:
+      print("API rate limit exceeded. Sleeping for 60s before retrying")
+      time.sleep(60)
+    else:
+      break
 
-  if response.status_code >= 400:
-    print(response.text)
   response.raise_for_status()
 
   return response.json()
diff --git a/tools/rule_manager/chronicle_api/rules/get_rule.py b/tools/rule_manager/chronicle_api/rules/get_rule.py
@@ -1,4 +1,4 @@
-# Copyright 2023 Google LLC
+# Copyright 2024 Google LLC
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -19,34 +19,49 @@
 """
 
 import os
+import time
 from typing import Any, Mapping
 
 from google.auth.transport import requests
 
 
 def get_rule(
-    http_session: requests.AuthorizedSession, resource_name: str
+    http_session: requests.AuthorizedSession,
+    resource_name: str,
+    max_retries: int = 3,
 ) -> Mapping[str, Any]:
   """Retrieves a rule.
 
   Args:
     http_session: Authorized session for HTTP requests.
     resource_name: The resource name of the rule to retrieve. Format -
       projects/{project}/locations/{location}/instances/{instance}/rules/{rule_id}
+    max_retries (optional): Maximum number of times to retry HTTP request if
+      certain response codes are returned. For example: HTTP response status
+      code 429 (Too Many Requests)
 
   Returns:
     Content and metadata about the requested rule.
 
   Raises:
     requests.exceptions.HTTPError: HTTP request resulted in an error
-      (response.status_code >= 400).
+    (response.status_code >= 400).
   """
   url = f"{os.environ['CHRONICLE_API_BASE_URL']}/{resource_name}"
+  response = None
 
-  response = http_session.request(method="GET", url=url)
+  for _ in range(max_retries + 1):
+    response = http_session.request(method="GET", url=url)
 
-  if response.status_code >= 400:
-    print(response.text)
-    response.raise_for_status()
+    if response.status_code >= 400:
+      print(response.text)
+
+    if response.status_code == 429:
+      print("API rate limit exceeded. Sleeping for 60s before retrying")
+      time.sleep(60)
+    else:
+      break
+
+  response.raise_for_status()
 
   return response.json()