[ASOC-2886] Add retries when sending graphql requests

finite_state_sdk/__init__.py

@@ -5,6 +5,12 @@
 import time
 from warnings import warn
 import finite_state_sdk.queries as queries
+from tenacity import retry, retry_if_exception, stop_after_attempt, wait_fixed
+from finite_state_sdk.utils import (
+    BreakoutException,
+    is_mutation,
+    is_not_breakout_exception,
+)
 
 API_URL = 'https://platform.finitestate.io/api/v1/graphql'
 AUDIENCE = "https://platform.finitestate.io/api/v1/graphql"
@@ -1968,6 +1974,7 @@ def search_sbom(token, organization_context, name=None, version=None, asset_vers
     return records
 
 
+@retry(stop=stop_after_attempt(5), wait=wait_fixed(5), retry=retry_if_exception(is_not_breakout_exception))
 def send_graphql_query(token, organization_context, query, variables=None):
     """
     Send a GraphQL query to the API
@@ -1989,26 +1996,27 @@ def send_graphql_query(token, organization_context, query, variables=None):
         dict: Response JSON
     """
     headers = {
-        'Content-Type': 'application/json',
-        'Authorization': f'Bearer {token}',
-        'Organization-Context': organization_context
-    }
-    data = {
-        'query': query,
-        'variables': variables
+        "Content-Type": "application/json",
+        "Authorization": f"Bearer {token}",
+        "Organization-Context": organization_context,
     }
+    data = {"query": query, "variables": variables}
 
     response = requests.post(API_URL, headers=headers, json=data)
-
     if response.status_code == 200:
         thejson = response.json()
 
         if "errors" in thejson:
-            raise Exception(f"Error: {thejson['errors']}")
+            # Raise a BreakoutException for GraphQL errors
+            raise BreakoutException(f"Error: {thejson['errors']}")
 
         return thejson
     else:
-        raise Exception(f"Error: {response.status_code} - {response.text}")
+        is_mutation_operation = is_mutation(query)
+        if is_mutation_operation:
+            raise BreakoutException(f"Error: {response.status_code} - {response.text}")
+        else:
+            raise Exception(f"Error: {response.status_code} - {response.text}")
 
 
 def update_finding_statuses(token, organization_context, user_id=None, finding_ids=None, status=None,

finite_state_sdk/utils.py

@@ -0,0 +1,40 @@
+from gql import gql
+from graphql.language.ast import OperationDefinitionNode, OperationType
+
+
+class BreakoutException(Exception):
+    """Exception raised for errors in the BreakoutException."""
+
+    pass
+
+
+def is_not_breakout_exception(exception):
+    """Check if the exception is not a BreakoutException."""
+    return not isinstance(exception, BreakoutException)
+
+
+def is_mutation(query_string):
+    """
+    Check if the provided GraphQL query string contains any mutations.
+
+    Args:
+        query_string (str): The GraphQL query string.
+
+    Returns:
+        bool: True if there is a mutation, False otherwise.
+    """
+    operation_types = determine_operation_types(query_string)
+    return OperationType.MUTATION in operation_types
+
+
+def determine_operation_types(query_string):
+    # Parse the query string
+    query_doc = gql(query_string)
+    operation_types = []
+
+    # Check the type of the first operation in the document
+    for definition in query_doc.definitions:
+        if isinstance(definition, OperationDefinitionNode):
+            operation_types.append(definition.operation)
+
+    return operation_types