diff --git a/google-cloud-clients/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java b/google-cloud-clients/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java index 09376fe13154..173e0ca2ca7b 100644 --- a/google-cloud-clients/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java +++ b/google-cloud-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/google-cloud-clients/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryOptions.java b/google-cloud-clients/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryOptions.java index fecf5005f157..ae4db388f80d 100644 --- a/google-cloud-clients/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryOptions.java +++ b/google-cloud-clients/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryOptions.java @@ -34,6 +34,8 @@ public class BigQueryOptions extends ServiceOptions { private static final Set 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/google-cloud-clients/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/BigQueryImplTest.java b/google-cloud-clients/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/BigQueryImplTest.java index 6d3348cae0a6..8ee3bac25e55 100644 --- a/google-cloud-clients/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/BigQueryImplTest.java +++ b/google-cloud-clients/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))