fix(bigquery/ingestor): handle quota exceeded for project.list requests

metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/bigquery_schema.py

@@ -154,37 +154,51 @@ def get_query_result(self, query: str) -> RowIterator:
         resp = self.bq_client.query(query)
         return resp.result()
 
-    def get_projects(self) -> List[BigqueryProject]:
+    def get_projects(self, max_results_per_page: int = 100) -> List[BigqueryProject]:
         def _should_retry(exc: BaseException) -> bool:
             logger.debug(
                 f"Exception occured for project.list api. Reason: {exc}. Retrying api request..."
             )
             self.report.num_list_projects_retry_request += 1
             return True
 
+        page_token = None
+        projects: List[BigqueryProject] = []
         with self.report.list_projects:
-            try:
-                # Bigquery API has limit in calling project.list request i.e. 2 request per second.
-                # https://cloud.google.com/bigquery/quotas#api_request_quotas
-                # Whenever this limit reached an exception occur with msg
-                # 'Quota exceeded: Your user exceeded quota for concurrent project.lists requests.'
-                # Hence, added the api request retry of 15 min.
-                # We already tried adding rate_limit externally, proving max_result and page_size
-                # to restrict the request calls inside list_project but issue still occured.
-                projects_iterator = self.bq_client.list_projects(
-                    retry=retry.Retry(
-                        predicate=_should_retry, initial=10, maximum=180, timeout=900
+            while True:
+                try:
+                    # Bigquery API has limit in calling project.list request i.e. 2 request per second.
+                    # https://cloud.google.com/bigquery/quotas#api_request_quotas
+                    # Whenever this limit reached an exception occur with msg
+                    # 'Quota exceeded: Your user exceeded quota for concurrent project.lists requests.'
+                    # Hence, added the api request retry of 15 min.
+                    # We already tried adding rate_limit externally, proving max_result and page_size
+                    # to restrict the request calls inside list_project but issue still occured.
+                    projects_iterator = self.bq_client.list_projects(
+                        max_results=max_results_per_page,
+                        page_token=page_token,
+                        timeout=900,
+                        retry=retry.Retry(
+                            predicate=_should_retry,
+                            initial=10,
+                            maximum=180,
+                            multiplier=4,
+                            timeout=900,
+                        ),
                     )
-                )
-                projects: List[BigqueryProject] = [
-                    BigqueryProject(id=p.project_id, name=p.friendly_name)
-                    for p in projects_iterator
-                ]
-                self.report.num_list_projects = len(projects)
-                return projects
-            except Exception as e:
-                logger.error(f"Error getting projects. {e}", exc_info=True)
-                return []
+                    _projects: List[BigqueryProject] = [
+                        BigqueryProject(id=p.project_id, name=p.friendly_name)
+                        for p in projects_iterator
+                    ]
+                    projects.extend(_projects)
+                    self.report.num_list_projects = len(projects)
+                    page_token = projects_iterator.next_page_token
+                    if not page_token:
+                        break
+                except Exception as e:
+                    logger.error(f"Error getting projects. {e}", exc_info=True)
+                    return []
+        return projects
 
     def get_datasets_for_project_id(
         self, project_id: str, maxResults: Optional[int] = None

metadata-ingestion/tests/unit/test_bigquery_source.py

@@ -273,24 +273,35 @@ def test_get_projects_with_single_project_id(get_bq_client_mock):
 def test_get_projects_by_list(get_bq_client_mock):
     client_mock = MagicMock()
     get_bq_client_mock.return_value = client_mock
-    client_mock.list_projects.return_value = [
-        SimpleNamespace(
-            project_id="test-1",
-            friendly_name="one",
-        ),
-        SimpleNamespace(
-            project_id="test-2",
-            friendly_name="two",
-        ),
-    ]
+
+    first_page = MagicMock()
+    first_page.__iter__.return_value = iter(
+        [
+            SimpleNamespace(project_id="test-1", friendly_name="one"),
+            SimpleNamespace(project_id="test-2", friendly_name="two"),
+        ]
+    )
+    first_page.next_page_token = "token1"
+
+    second_page = MagicMock()
+    second_page.__iter__.return_value = iter(
+        [
+            SimpleNamespace(project_id="test-3", friendly_name="three"),
+            SimpleNamespace(project_id="test-4", friendly_name="four"),
+        ]
+    )
+    second_page.next_page_token = None
+    client_mock.list_projects.side_effect = [first_page, second_page]
 
     config = BigQueryV2Config.parse_obj({})
     source = BigqueryV2Source(config=config, ctx=PipelineContext(run_id="test1"))
     assert source._get_projects() == [
         BigqueryProject("test-1", "one"),
         BigqueryProject("test-2", "two"),
+        BigqueryProject("test-3", "three"),
+        BigqueryProject("test-4", "four"),
     ]
-    assert client_mock.list_projects.call_count == 1
+    assert client_mock.list_projects.call_count == 2
 
 
 @patch.object(BigQuerySchemaApi, "get_projects")