[PLUGIN-1787] Ignore 409 Conflict during retry of bucket creation in BigQuery source plugin

src/main/java/io/cdap/plugin/gcp/bigquery/source/BigQuerySourceUtils.java

@@ -87,30 +87,43 @@ public static String getOrCreateBucket(Configuration configuration,
       // By default, this option is false, meaning the job can not delete the bucket. So enable it only when bucket name
       // is not provided.
       configuration.setBoolean("fs.gs.bucket.delete.enable", true);
-      GCPUtils.createBucket(storage, bucket, dataset.getLocation(), cmekKeyName);
+      createBucket(storage, bucket, dataset, cmekKeyName);
     } else if (storage != null && storage.get(bucket) == null) {
-      try {
-        GCPUtils.createBucket(storage, bucket, dataset.getLocation(), cmekKeyName);
-      } catch (StorageException e) {
-        if (e.getCode() == 409) {
-          // A conflict means the bucket already exists
-          // This most likely means multiple stages in the same pipeline are trying to create the same bucket.
-          // Ignore this and move on, since all that matters is that the bucket exists.
-          return bucket;
-        }
-        String errorMessage = String.format("Unable to create Cloud Storage bucket '%s' in the same "
-                + "location ('%s') as BigQuery dataset '%s'. " + "Please use a bucket "
-                + "that is in the same location as the dataset. For more details, see %s",
-            bucket, dataset.getLocation(), dataset.getDatasetId().getDataset(),
-            GCPUtils.GCS_SUPPORTED_DOC_URL);
-        throw ErrorUtils.getProgramFailureException(new ErrorCategory(ErrorCategory.ErrorCategoryEnum.PLUGIN),
-          errorMessage, e.getMessage(), ErrorType.USER, true, ErrorCodeType.HTTP,
-            String.valueOf(e.getCode()), GCPUtils.GCS_SUPPORTED_DOC_URL, e);
-      }
+      createBucket(storage, bucket, dataset, cmekKeyName);
     }
     return bucket;
   }
 
+  private static void createBucket(@Nullable Storage storage,
+                                   String bucket,
+                                   Dataset dataset,
+                                   @Nullable CryptoKeyName cmekKeyName) {
+    if (storage == null) {
+      return;
+    }
+    try {
+      GCPUtils.createBucket(storage, bucket, dataset.getLocation(), cmekKeyName);
+    } catch (StorageException e) {
+      if (e.getCode() == 409) {
+        // A conflict means the bucket already exists.
+        // This most likely means multiple stages in the same pipeline are trying to create the same bucket,
+        // or a retry occurred after a successful bucket creation.
+        // Ignore this and move on, since all that matters is that the bucket exists.
+        LOG.debug("Bucket '{}' already exists, ignoring 409 Conflict: {}", bucket, e.getMessage());
+        return;
+      }
+      String datasetName = dataset.getDatasetId() != null ? dataset.getDatasetId().getDataset() : "";
+      String errorMessage = String.format("Unable to create Cloud Storage bucket '%s' in the same "
+              + "location ('%s') as BigQuery dataset '%s'. " + "Please use a bucket "
+              + "that is in the same location as the dataset. For more details, see %s",
+          bucket, dataset.getLocation(), datasetName,
+          GCPUtils.GCS_SUPPORTED_DOC_URL);
+      throw ErrorUtils.getProgramFailureException(new ErrorCategory(ErrorCategory.ErrorCategoryEnum.PLUGIN),
+        errorMessage, e.getMessage(), ErrorType.USER, true, ErrorCodeType.HTTP,
+          String.valueOf(e.getCode()), GCPUtils.GCS_SUPPORTED_DOC_URL, e);
+    }
+  }
+
   /**
    * Sets up service account credentials into supplied Hadoop configuration.
    *

src/test/java/io/cdap/plugin/gcp/bigquery/source/BigQuerySourceUtilsTest.java

@@ -17,15 +17,18 @@
 package io.cdap.plugin.gcp.bigquery.source;
 
 import com.google.cloud.bigquery.Dataset;
+import com.google.cloud.bigquery.DatasetId;
 import com.google.cloud.storage.Bucket;
+import com.google.cloud.storage.BucketInfo;
 import com.google.cloud.storage.Storage;
+import com.google.cloud.storage.StorageException;
+import io.cdap.cdap.api.exception.ProgramFailureException;
 import org.apache.hadoop.conf.Configuration;
 import org.junit.Assert;
 import org.junit.Test;
 import org.mockito.Mockito;
 
 import java.io.IOException;
-import java.lang.reflect.Field;
 
 public class BigQuerySourceUtilsTest {
 
@@ -50,4 +53,39 @@ public void getOrCreateBucket() throws IllegalAccessException, NoSuchFieldExcept
                                                            "some-path", null);
     Assert.assertEquals("a-bucket", bucket2);
   }
+
+  @Test
+  public void getOrCreateBucket_nullBucketConflict409_returnsBucketAndEnablesDelete() throws IOException {
+    Configuration configuration = new Configuration();
+    BigQuerySourceConfig config = BigQuerySourceConfig.builder().build();
+    Dataset ds = Mockito.mock(Dataset.class);
+    Storage st = Mockito.mock(Storage.class);
+    Mockito.when(st.create(Mockito.any(BucketInfo.class)))
+      .thenThrow(new StorageException(409, "Your previous request to create the named bucket succeeded "
+        + "and you already own it."));
+
+    String bucket = BigQuerySourceUtils.getOrCreateBucket(configuration, st, config.getBucket(), ds,
+                                                         "some-path", null);
+
+    Assert.assertEquals("bq-source-bucket-some-path", bucket);
+    Assert.assertTrue(configuration.getBoolean("fs.gs.bucket.delete.enable", false));
+  }
+
+  @Test
+  public void getOrCreateBucket_nullBucketNon409Error_throwsProgramFailureException() {
+    Configuration configuration = new Configuration();
+    BigQuerySourceConfig config = BigQuerySourceConfig.builder().build();
+    Dataset ds = Mockito.mock(Dataset.class);
+    DatasetId datasetId = DatasetId.of("project", "dataset");
+    Mockito.when(ds.getDatasetId()).thenReturn(datasetId);
+    Mockito.when(ds.getLocation()).thenReturn("US");
+    Storage st = Mockito.mock(Storage.class);
+    Mockito.when(st.create(Mockito.any(BucketInfo.class)))
+      .thenThrow(new StorageException(403, "Access denied"));
+
+    ProgramFailureException exception = Assert.assertThrows(ProgramFailureException.class, () ->
+      BigQuerySourceUtils.getOrCreateBucket(configuration, st, config.getBucket(), ds, "some-path", null));
+
+    Assert.assertTrue(exception.getMessage().contains("Access denied"));
+  }
 }