feat(ingestion/bigquery): Add ability to filter GCP project ingestion based on project labels

metadata-ingestion/setup.py

@@ -181,6 +181,7 @@
  "google-cloud-logging<=3.5.0",
  "google-cloud-bigquery",
  "google-cloud-datacatalog>=1.5.0",
+ "google-cloud-resource-manager",
  "more-itertools>=8.12.0",
  "sqlalchemy-bigquery>=1.4.1",
 }

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

@@ -113,8 +113,9 @@ def __init__(self, ctx: PipelineContext, config: BigQueryV2Config):
  BigqueryTableIdentifier._BQ_SHARDED_TABLE_SUFFIX = ""
 
  self.bigquery_data_dictionary = BigQuerySchemaApi(
- self.report.schema_api_perf,
- self.config.get_bigquery_client(),
+ report=BigQueryV2Report().schema_api_perf,
+ projects_client=config.get_projects_client(),
+ client=config.get_bigquery_client(),
  )
  if self.config.extract_policy_tags_from_catalog:
  self.bigquery_data_dictionary.datacatalog_client = (
@@ -257,14 +258,37 @@ def get_workunits_internal(self) -> Iterable[MetadataWorkUnit]:
 
  def _get_projects(self) -> List[BigqueryProject]:
  logger.info("Getting projects")
+
  if self.config.project_ids or self.config.project_id:
  project_ids = self.config.project_ids or [self.config.project_id] # type: ignore
  return [
  BigqueryProject(id=project_id, name=project_id)
  for project_id in project_ids
  ]
- else:
- return list(self._query_project_list())
+
+ if self.config.project_labels:
+ return list(self._query_project_list_from_labels())
+
+ return list(self._query_project_list())
+
+ def _query_project_list_from_labels(self) -> Iterable[BigqueryProject]:
+ projects = self.bigquery_data_dictionary.get_projects_with_labels(
+ self.config.project_labels
+ )
+
+ if not projects: # Report failure on exception and if empty list is returned
+ self.report.report_failure(
+ "metadata-extraction",
+ "Get projects didn't return any project with any of the specified label(s). "
+ "Maybe resourcemanager.projects.list permission is missing for the service account. "
+ "You can assign predefined roles/bigquery.metadataViewer role to your service account.",
+ )
+
+ for project in projects:
+ if self.config.project_id_pattern.allowed(project.id):
+ yield project
+ else:
+ self.report.report_dropped(project.id)
 
  def _query_project_list(self) -> Iterable[BigqueryProject]:
  try:

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

@@ -3,7 +3,7 @@
 from datetime import timedelta
 from typing import Any, Dict, List, Optional, Union
 
-from google.cloud import bigquery, datacatalog_v1
+from google.cloud import bigquery, datacatalog_v1, resourcemanager_v3
 from google.cloud.logging_v2.client import Client as GCPLoggingClient
 from pydantic import Field, PositiveInt, PrivateAttr, root_validator, validator
 
@@ -74,6 +74,9 @@ def get_bigquery_client(self) -> bigquery.Client:
  client_options = self.extra_client_options
  return bigquery.Client(self.project_on_behalf, **client_options)
 
+ def get_projects_client(self) -> resourcemanager_v3.ProjectsClient:
+ return resourcemanager_v3.ProjectsClient()
+
  def get_policy_tag_manager_client(self) -> datacatalog_v1.PolicyTagManagerClient:
  return datacatalog_v1.PolicyTagManagerClient()
 
@@ -214,6 +217,13 @@ def have_table_data_read_permission(self) -> bool:
  "Overrides `project_id_pattern`."
  ),
  )
+ project_labels: List[str] = Field(
+ default_factory=list,
+ description=(
+ "Ingests projects with the specified labels. Use this property if you want to specify what "
+ "projects to ingest based on project-level labels."
+ ),
+ )
 
  storage_project_id: None = Field(default=None, hidden_from_docs=True)
 

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

@@ -31,6 +31,7 @@ class BigQuerySchemaApiPerfReport(Report):
  num_get_snapshots_for_dataset_api_requests: int = 0
 
  list_projects: PerfTimer = field(default_factory=PerfTimer)
+ list_projects_with_labels: PerfTimer = field(default_factory=PerfTimer)
  list_datasets: PerfTimer = field(default_factory=PerfTimer)
 
  get_columns_for_dataset_sec: float = 0

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

@@ -5,7 +5,7 @@
 from typing import Any, Dict, Iterable, Iterator, List, Optional
 
 from google.api_core import retry
-from google.cloud import bigquery, datacatalog_v1
+from google.cloud import bigquery, datacatalog_v1, resourcemanager_v3
 from google.cloud.bigquery.table import (
  RowIterator,
  TableListItem,
@@ -144,9 +144,11 @@ def __init__(
  self,
  report: BigQuerySchemaApiPerfReport,
  client: bigquery.Client,
+ projects_client: resourcemanager_v3.ProjectsClient,
  datacatalog_client: Optional[datacatalog_v1.PolicyTagManagerClient] = None,
  ) -> None:
  self.bq_client = client
+ self.projects_client = projects_client
  self.report = report
  self.datacatalog_client = datacatalog_client
 
@@ -175,7 +177,7 @@ def _should_retry(exc: BaseException) -> bool:
  # '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.
+ # to restrict the request calls inside list_project but issue still occurred.
  projects_iterator = self.bq_client.list_projects(
  max_results=max_results_per_page,
  page_token=page_token,
@@ -202,6 +204,26 @@ def _should_retry(exc: BaseException) -> bool:
  return []
  return projects
 
+ def get_projects_with_labels(self, labels: List[str]) -> List[BigqueryProject]:
+ with self.report.list_projects_with_labels:
+ try:
+ projects = []
+ labels_query = " OR ".join([f"labels.{label}" for label in labels])
+ for project in self.projects_client.search_projects(query=labels_query):
+ projects.append(
+ BigqueryProject(
+ id=project.project_id, name=project.display_name
+ )
+ )
+
+ return projects
+
+ except Exception as e:
+ logger.error(
+ f"Error getting projects with labels: {labels}. {e}", exc_info=True
+ )
+ return []
+
  def get_datasets_for_project_id(
  self, project_id: str, maxResults: Optional[int] = None
  ) -> List[BigqueryDataset]:

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

@@ -96,7 +96,9 @@ def metadata_read_capability_test(
  client: bigquery.Client = config.get_bigquery_client()
  assert client
  bigquery_data_dictionary = BigQuerySchemaApi(
- BigQueryV2Report().schema_api_perf, client
+ report=BigQueryV2Report().schema_api_perf,
+ projects_client=config.get_projects_client(),
+ client=client,
  )
  result = bigquery_data_dictionary.get_datasets_for_project_id(
  project_id, 10

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

@@ -479,7 +479,9 @@ def lineage_via_catalog_lineage_api(
  lineage_client: lineage_v1.LineageClient = lineage_v1.LineageClient()
 
  data_dictionary = BigQuerySchemaApi(
- self.report.schema_api_perf, self.config.get_bigquery_client()
+ self.report.schema_api_perf,
+ self.config.get_bigquery_client(),
+ self.config.get_projects_client(),
  )
 
  # Filtering datasets