[Data] Check if BigQuery dataset exists before creation (#41630)

We found issue that user may not grant permission to create Google BigQuery dataset when using Ray Data `write_bigquery()` API. So the `client.create_dataset` call would fail with permission error: ``` Access Denied: Project ...: User does not have bigquery.datasets.create permission in project ... ``` So here we change to check if BigQuery dataset already exists, if not then trying to create it. Signed-off-by: Cheng Su <[email protected]>
ray-project · Dec 6, 2023 · 9516630 · 9516630
1 parent e92f015
commit 9516630
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/python/ray/data/datasource/bigquery_datasink.py b/python/ray/data/datasource/bigquery_datasink.py
@@ -45,13 +45,14 @@ def on_write_start(self) -> None:
  client = bigquery.Client(project=self.project_id)
  dataset_id = self.dataset.split(".", 1)[0]
  try:
- client.create_dataset(f"{self.project_id}.{dataset_id}", timeout=30)
- logger.info("Created dataset " + dataset_id)
- except exceptions.Conflict:
+ client.get_dataset(dataset_id)
  logger.info(
  f"Dataset {dataset_id} already exists. "
  "The table will be overwritten if it already exists."
  )
+ except exceptions.NotFound:
+ client.create_dataset(f"{self.project_id}.{dataset_id}", timeout=30)
+ logger.info("Created dataset " + dataset_id)
 
  # Delete table if it already exists
  client.delete_table(f"{self.project_id}.{self.dataset}", not_found_ok=True)