From 726eaae5cd08178d58741ff10c468d16bba27bc0 Mon Sep 17 00:00:00 2001 From: Vincent Koc Date: Sun, 5 Jun 2022 20:39:26 +1000 Subject: [PATCH] Update dbt.py --- airflow/providers/dbt/cloud/hooks/dbt.py | 22 ++++++++-------------- 1 file changed, 8 insertions(+), 14 deletions(-) diff --git a/airflow/providers/dbt/cloud/hooks/dbt.py b/airflow/providers/dbt/cloud/hooks/dbt.py index 13d10dbc6283..d88c0053d318 100644 --- 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(