Update dbt.py (#24218)

apache · Jun 6, 2022 · 5e6997e · 5e6997e
1 parent 603c555
commit 5e6997e
Showing 1 changed file with 8 additions and 14 deletions.
diff --git a/airflow/providers/dbt/cloud/hooks/dbt.py b/airflow/providers/dbt/cloud/hooks/dbt.py
@@ -168,28 +168,22 @@ def get_conn(self, *args, **kwargs) -> Session:
         return session
 
     def _paginate(self, endpoint: str, payload: Optional[Dict[str, Any]] = None) -> List[Response]:
-        results = []
         response = self.run(endpoint=endpoint, data=payload)
         resp_json = response.json()
         limit = resp_json["extra"]["filters"]["limit"]
         num_total_results = resp_json["extra"]["pagination"]["total_count"]
         num_current_results = resp_json["extra"]["pagination"]["count"]
-        results.append(response)
-
-        if not num_current_results == num_total_results:
+        results = [response]
+        if num_current_results != num_total_results:
             _paginate_payload = payload.copy() if payload else {}
             _paginate_payload["offset"] = limit
 
-            while True:
-                if num_current_results < num_total_results:
-                    response = self.run(endpoint=endpoint, data=_paginate_payload)
-                    resp_json = response.json()
-                    results.append(response)
-                    num_current_results += resp_json["extra"]["pagination"]["count"]
-                    _paginate_payload["offset"] += limit
-                else:
-                    break
-
+            while not num_current_results >= num_total_results:
+                response = self.run(endpoint=endpoint, data=_paginate_payload)
+                resp_json = response.json()
+                results.append(response)
+                num_current_results += resp_json["extra"]["pagination"]["count"]
+                _paginate_payload["offset"] += limit
         return results
 
     def _run_and_get_response(