feat: Check if dataset exists before creation for Ray on Vertex BigQu…

…ery Write PiperOrigin-RevId: 589905230
googleapis · Dec 11, 2023 · 544d6fe · 544d6fe
1 parent 3d8835e
commit 544d6fe
Showing 1 changed file with 6 additions and 6 deletions.
diff --git a/google/cloud/aiplatform/preview/vertex_ray/bigquery_datasource.py b/google/cloud/aiplatform/preview/vertex_ray/bigquery_datasource.py
@@ -250,14 +250,14 @@ def _write_single_block(
         client = bigquery.Client(project=project_id, client_info=bq_info)
         dataset_id = dataset.split(".", 1)[0]
         try:
-            client.create_dataset(f"{project_id}.{dataset_id}", timeout=30)
-            print("[Ray on Vertex AI]: Created dataset", dataset_id)
-        except exceptions.Conflict:
+            client.get_dataset(dataset_id)
             print(
-                "[Ray on Vertex AI]: Dataset",
-                dataset_id,
-                "already exists. The table will be overwritten if it already exists.",
+                f"[Ray on Vertex AI]: Dataset {dataset_id} already exists."
+                + "The table will be overwritten if it already exists."
             )
+        except exceptions.NotFound:
+            client.create_dataset(f"{project_id}.{dataset_id}", timeout=30)
+            print(f"[Ray on Vertex AI]: Created dataset {dataset_id}")
 
         # Delete table if it already exists
         client.delete_table(f"{project_id}.{dataset}", not_found_ok=True)