Merge pull request #355 from ynput/bugfix/graphql-infinite-pagination-loop

GraphQl: Prevent infinite query loop on unexpected responses

Author: iLLiCiTiT

ayon_api/graphql.py

@@ -353,6 +353,25 @@ def parse_result(
         for child in self._children:
             child.parse_result(data, output, progress_data)
 
+    def _query_data(self, con: ServerAPI) -> dict[str, Any]:
+        """Send single query to server and return 'data' of the response."""
+        query_str = self.calculate_query()
+        variables = self.get_variables_values()
+        response = con.query_graphql(query_str, variables)
+        if response.errors:
+            raise GraphQlQueryFailed(response.errors, query_str, variables)
+
+        data = response.data.get("data")
+        if data is None:
+            # Parsing 'None' would not change pagination state and the same
+            #   query would be sent again in an infinite loop.
+            raise GraphQlQueryError(
+                f"GraphQl query '{self._name}' response does not contain"
+                f" 'data'. Response: {str(response.data)[:1000]}"
+                f"\nQuery:\n{query_str}\nVariables: {variables}"
+            )
+        return data
+
     def query(self, con: ServerAPI) -> dict[str, Any]:
         """Do a query from server.
 
@@ -366,15 +385,8 @@ def query(self, con: ServerAPI) -> dict[str, Any]:
         progress_data = {}
         output = {}
         while self.need_query:
-            query_str = self.calculate_query()
-            variables = self.get_variables_values()
-            response = con.query_graphql(
-                query_str,
-                variables
-            )
-            if response.errors:
-                raise GraphQlQueryFailed(response.errors, query_str, variables)
-            self.parse_result(response.data["data"], output, progress_data)
+            data = self._query_data(con)
+            self.parse_result(data, output, progress_data)
 
         return output
 
@@ -394,30 +406,16 @@ def continuous_query(
         if self.has_multiple_edge_fields:
             output = {}
             while self.need_query:
-                query_str = self.calculate_query()
-                variables = self.get_variables_values()
-
-                response = con.query_graphql(query_str, variables)
-                if response.errors:
-                    raise GraphQlQueryFailed(
-                        response.errors, query_str, variables
-                    )
-                self.parse_result(response.data["data"], output, progress_data)
+                data = self._query_data(con)
+                self.parse_result(data, output, progress_data)
 
             yield output
 
         else:
             while self.need_query:
                 output = {}
-                query_str = self.calculate_query()
-                variables = self.get_variables_values()
-                response = con.query_graphql(query_str, variables)
-                if response.errors:
-                    raise GraphQlQueryFailed(
-                        response.errors, query_str, variables
-                    )
-
-                self.parse_result(response.data["data"], output, progress_data)
+                data = self._query_data(con)
+                self.parse_result(data, output, progress_data)
 
                 yield output
 
@@ -945,6 +943,14 @@ def parse_result(
                 change_cursor = False
 
         if change_cursor and self._need_query:
+            if new_cursor is None:
+                # Without cursor the pagination would start from beginning
+                raise GraphQlQueryError(
+                    f"Field '{self.path}' reported another page without"
+                    " a cursor. Stopped pagination after"
+                    f" {self._fetched_counter} items."
+                )
+
             if new_cursor == self._cursor:
                 raise GraphQlQueryError(
                     "Cursor didn't change during pagination."

tests/test_download_resume.py

@@ -51,10 +51,3 @@ def get_func(url, **kwargs):
     content, progress = _download(con, get_func)
     assert content == CONTENT
     assert progress.transferred_size == len(CONTENT)
-
-
-def test_download_without_content_length(con):
-    content, _ = _download(
-        con, lambda url, **kwargs: FakeResponse(200, CONTENT)
-    )
-    assert content == CONTENT

tests/test_graphql_infinite_loop.py

@@ -0,0 +1,46 @@
+"""Unexpected GraphQl responses must not cause infinite loop.
+
+Does not require running AYON server.
+"""
+import pytest
+
+from ayon_api.exceptions import GraphQlQueryError
+from ayon_api.graphql_queries import events_graphql_query
+from ayon_api.utils import SortOrder
+
+from .graphql_fake_server import FakeResponse
+
+
+class _ScriptedServer:
+    """Return prepared responses, fail if queried too many times."""
+    def __init__(self, responses):
+        self._responses = responses
+        self.queries = []
+
+    def query_graphql(self, query_str, variables):
+        self.queries.append(query_str)
+        if len(self.queries) > 10:
+            raise RuntimeError("Infinite query loop")
+        idx = min(len(self.queries), len(self._responses)) - 1
+        return FakeResponse(self._responses[idx])
+
+
+def test_null_data_raises_instead_of_infinite_loop():
+    server = _ScriptedServer([{"data": None}])
+    query = events_graphql_query({"id"}, SortOrder.ascending)
+    with pytest.raises(GraphQlQueryError, match="does not contain 'data'"):
+        query.query(server)
+
+
+def test_missing_cursor_stops_pagination():
+    def page(ids, end_cursor):
+        return {"data": {"events": {
+            "edges": [{"node": {"id": id_}} for id_ in ids],
+            "pageInfo": {"endCursor": end_cursor, "hasNextPage": True},
+        }}}
+
+    server = _ScriptedServer([page(["e0"], "c0"), page([], None)])
+    query = events_graphql_query({"id"}, SortOrder.ascending)
+
+    with pytest.raises(GraphQlQueryError, match="page without a cursor"):
+        query.query(server)