BigQuery: Fixes When Table,Dataset and Job is not found (#4389)

* Fixes when found return null value * added comment and modified methods name * format * added more test cases to cover code coverage
googleapis · Feb 25, 2019 · 786ce15 · 786ce15
1 parent c8e0f2f
commit 786ce15
Show file tree

Hide file tree

Showing 3 changed files with 88 additions and 0 deletions.
diff --git a/...d-clients/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java b/...d-clients/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java
@@ -18,6 +18,7 @@
 
 import static com.google.cloud.RetryHelper.runWithRetries;
 import static com.google.common.base.Preconditions.checkArgument;
+import static java.net.HttpURLConnection.HTTP_NOT_FOUND;
 
 import com.google.api.core.InternalApi;
 import com.google.api.gax.paging.Page;
@@ -293,6 +294,9 @@ public com.google.api.services.bigquery.model.Dataset call() {
               getOptions().getRetrySettings(),
               EXCEPTION_HANDLER,
               getOptions().getClock());
+      if (getOptions().getThrowNotFound() && answer == null) {
+        throw new BigQueryException(HTTP_NOT_FOUND, "Dataset not found");
+      }
       return answer == null ? null : Dataset.fromPb(this, answer);
     } catch (RetryHelper.RetryHelperException e) {
       throw BigQueryException.translateAndThrow(e);
@@ -482,6 +486,9 @@ public com.google.api.services.bigquery.model.Table call() {
               getOptions().getRetrySettings(),
               EXCEPTION_HANDLER,
               getOptions().getClock());
+      if (getOptions().getThrowNotFound() && answer == null) {
+        throw new BigQueryException(HTTP_NOT_FOUND, "Table not found");
+      }
       return answer == null ? null : Table.fromPb(this, answer);
     } catch (RetryHelper.RetryHelperException e) {
       throw BigQueryException.translateAndThrow(e);
@@ -702,6 +709,9 @@ public com.google.api.services.bigquery.model.Job call() {
               getOptions().getRetrySettings(),
               EXCEPTION_HANDLER,
               getOptions().getClock());
+      if (getOptions().getThrowNotFound() && answer == null) {
+        throw new BigQueryException(HTTP_NOT_FOUND, "Job not found");
+      }
       return answer == null ? null : Job.fromPb(this, answer);
     } catch (RetryHelper.RetryHelperException e) {
       throw BigQueryException.translateAndThrow(e);

diff --git a/...lients/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryOptions.java b/...lients/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryOptions.java
@@ -34,6 +34,8 @@ public class BigQueryOptions extends ServiceOptions<BigQuery, BigQueryOptions> {
   private static final Set<String> SCOPES = ImmutableSet.of(BIGQUERY_SCOPE);
   private static final long serialVersionUID = -2437598817433266049L;
   private final String location;
+  // set the option ThrowNotFound when you want to throw the exception when the value not found
+  private boolean setThrowNotFound;
 
   public static class DefaultBigQueryFactory implements BigQueryFactory {
 
@@ -125,6 +127,14 @@ public String getLocation() {
     return location;
   }
 
+  public void setThrowNotFound(boolean setThrowNotFound) {
+    this.setThrowNotFound = setThrowNotFound;
+  }
+
+  public boolean getThrowNotFound() {
+    return setThrowNotFound;
+  }
+
   @SuppressWarnings("unchecked")
   @Override
   public Builder toBuilder() {

diff --git a/...ients/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/BigQueryImplTest.java b/...ients/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/BigQueryImplTest.java
@@ -366,6 +366,29 @@ public void testGetDataset() {
         new Dataset(bigquery, new DatasetInfo.BuilderImpl(DATASET_INFO_WITH_PROJECT)), dataset);
   }
 
+  @Test
+  public void testGetDatasetNotFoundWhenThrowIsDisabled() {
+    EasyMock.expect(bigqueryRpcMock.getDataset(PROJECT, DATASET, EMPTY_RPC_OPTIONS))
+        .andReturn(DATASET_INFO_WITH_PROJECT.toPb());
+    EasyMock.replay(bigqueryRpcMock);
+    options.setThrowNotFound(false);
+    bigquery = options.getService();
+    Dataset dataset = bigquery.getDataset(DATASET);
+    assertEquals(
+        new Dataset(bigquery, new DatasetInfo.BuilderImpl(DATASET_INFO_WITH_PROJECT)), dataset);
+  }
+
+  @Test
+  public void testGetDatasetNotFoundWhenThrowIsEnabled() {
+    EasyMock.expect(bigqueryRpcMock.getDataset(PROJECT, "dataset-not-found", EMPTY_RPC_OPTIONS))
+        .andThrow(new BigQueryException(404, "Dataset not found"));
+    EasyMock.replay(bigqueryRpcMock);
+    options.setThrowNotFound(true);
+    bigquery = options.getService();
+    thrown.expect(BigQueryException.class);
+    bigquery.getDataset("dataset-not-found");
+  }
+
   @Test
   public void testGetDatasetFromDatasetId() {
     EasyMock.expect(bigqueryRpcMock.getDataset(PROJECT, DATASET, EMPTY_RPC_OPTIONS))
@@ -603,6 +626,29 @@ public void testGetTable() {
     assertEquals(new Table(bigquery, new TableInfo.BuilderImpl(TABLE_INFO_WITH_PROJECT)), table);
   }
 
+  @Test
+  public void testGetTableNotFoundWhenThrowIsDisabled() {
+    EasyMock.expect(bigqueryRpcMock.getTable(PROJECT, DATASET, TABLE, EMPTY_RPC_OPTIONS))
+        .andReturn(TABLE_INFO_WITH_PROJECT.toPb());
+    EasyMock.replay(bigqueryRpcMock);
+    options.setThrowNotFound(false);
+    bigquery = options.getService();
+    Table table = bigquery.getTable(DATASET, TABLE);
+    assertEquals(new Table(bigquery, new TableInfo.BuilderImpl(TABLE_INFO_WITH_PROJECT)), table);
+  }
+
+  @Test
+  public void testGetTableNotFoundWhenThrowIsEnabled() {
+    EasyMock.expect(
+            bigqueryRpcMock.getTable(PROJECT, DATASET, "table-not-found", EMPTY_RPC_OPTIONS))
+        .andThrow(new BigQueryException(404, "Table not found"));
+    EasyMock.replay(bigqueryRpcMock);
+    options.setThrowNotFound(true);
+    bigquery = options.getService();
+    thrown.expect(BigQueryException.class);
+    bigquery.getTable(DATASET, "table-not-found");
+  }
+
   @Test
   public void testGetTableFromTableId() {
     EasyMock.expect(bigqueryRpcMock.getTable(PROJECT, DATASET, TABLE, EMPTY_RPC_OPTIONS))
@@ -1201,6 +1247,28 @@ public void testGetJobWithLocation() {
     assertEquals(new Job(bigquery, new JobInfo.BuilderImpl(COMPLETE_COPY_JOB)), job);
   }
 
+  @Test
+  public void testGetJobNotFoundWhenThrowIsDisabled() {
+    EasyMock.expect(bigqueryRpcMock.getJob(PROJECT, JOB, null, EMPTY_RPC_OPTIONS))
+        .andReturn(COMPLETE_COPY_JOB.toPb());
+    EasyMock.replay(bigqueryRpcMock);
+    options.setThrowNotFound(false);
+    bigquery = options.getService();
+    Job job = bigquery.getJob(JOB);
+    assertEquals(new Job(bigquery, new JobInfo.BuilderImpl(COMPLETE_COPY_JOB)), job);
+  }
+
+  @Test
+  public void testGetJobNotFoundWhenThrowIsEnabled() {
+    EasyMock.expect(bigqueryRpcMock.getJob(PROJECT, "job-not-found", null, EMPTY_RPC_OPTIONS))
+        .andThrow(new BigQueryException(404, "Job not found"));
+    EasyMock.replay(bigqueryRpcMock);
+    options.setThrowNotFound(true);
+    bigquery = options.getService();
+    thrown.expect(BigQueryException.class);
+    bigquery.getJob("job-not-found");
+  }
+
   @Test
   public void testGetJobFromJobId() {
     EasyMock.expect(bigqueryRpcMock.getJob(PROJECT, JOB, null, EMPTY_RPC_OPTIONS))