From f4bac2ff79670ccf5d13a692f84ea006767125a5 Mon Sep 17 00:00:00 2001 From: Praful Makani Date: Fri, 25 Jan 2019 23:56:47 +0530 Subject: [PATCH] BigQuery : Fix Location configurable at BigQueryOptions (#4329) * Fix Location configurable from BigQueryOptions * modified code * modified code and add test case * removed unused location --- .../google/cloud/bigquery/BigQueryImpl.java | 24 ++++++++-- .../cloud/bigquery/BigQueryOptions.java | 13 ++++++ .../java/com/google/cloud/bigquery/JobId.java | 4 ++ .../cloud/bigquery/BigQueryImplTest.java | 46 +++++++++++++++++++ 4 files changed, 84 insertions(+), 3 deletions(-) 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 e30b0ffd2156..09376fe13154 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 @@ -679,7 +679,13 @@ public Job getJob(String jobId, JobOption... options) { @Override public Job getJob(JobId jobId, JobOption... options) { final Map optionsMap = optionMap(options); - final JobId completeJobId = jobId.setProjectId(getOptions().getProjectId()); + final JobId completeJobId = + jobId + .setProjectId(getOptions().getProjectId()) + .setLocation( + jobId.getLocation() == null && getOptions().getLocation() != null + ? getOptions().getLocation() + : jobId.getLocation()); try { com.google.api.services.bigquery.model.Job answer = runWithRetries( @@ -742,7 +748,13 @@ public boolean cancel(String jobId) { @Override public boolean cancel(JobId jobId) { - final JobId completeJobId = jobId.setProjectId(getOptions().getProjectId()); + final JobId completeJobId = + jobId + .setProjectId(getOptions().getProjectId()) + .setLocation( + jobId.getLocation() == null && getOptions().getLocation() != null + ? getOptions().getLocation() + : jobId.getLocation()); try { return runWithRetries( new Callable() { @@ -784,7 +796,13 @@ private static QueryResponse getQueryResults( JobId jobId, final BigQueryOptions serviceOptions, final Map optionsMap) { - final JobId completeJobId = jobId.setProjectId(serviceOptions.getProjectId()); + final JobId completeJobId = + jobId + .setProjectId(serviceOptions.getProjectId()) + .setLocation( + jobId.getLocation() == null && serviceOptions.getLocation() != null + ? serviceOptions.getLocation() + : jobId.getLocation()); try { GetQueryResultsResponse results = runWithRetries( 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 080bb57521f3..fecf5005f157 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 @@ -33,6 +33,7 @@ public class BigQueryOptions extends ServiceOptions { private static final String BIGQUERY_SCOPE = "https://www.googleapis.com/auth/bigquery"; private static final Set SCOPES = ImmutableSet.of(BIGQUERY_SCOPE); private static final long serialVersionUID = -2437598817433266049L; + private final String location; public static class DefaultBigQueryFactory implements BigQueryFactory { @@ -56,6 +57,8 @@ public ServiceRpc create(BigQueryOptions options) { public static class Builder extends ServiceOptions.Builder { + private String location; + private Builder() {} private Builder(BigQueryOptions options) { @@ -71,6 +74,11 @@ public Builder setTransportOptions(TransportOptions transportOptions) { return super.setTransportOptions(transportOptions); } + public Builder setLocation(String location) { + this.location = location; + return this; + } + @Override public BigQueryOptions build() { return new BigQueryOptions(this); @@ -79,6 +87,7 @@ public BigQueryOptions build() { private BigQueryOptions(Builder builder) { super(BigQueryFactory.class, BigQueryRpcFactory.class, builder, new BigQueryDefaults()); + this.location = builder.location; } private static class BigQueryDefaults implements ServiceDefaults { @@ -112,6 +121,10 @@ protected BigQueryRpc getBigQueryRpcV2() { return (BigQueryRpc) getRpc(); } + public String getLocation() { + return location; + } + @SuppressWarnings("unchecked") @Override public Builder toBuilder() { diff --git a/google-cloud-clients/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/JobId.java b/google-cloud-clients/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/JobId.java index 254e69647e1a..05f80bbf376f 100644 --- a/google-cloud-clients/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/JobId.java +++ b/google-cloud-clients/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/JobId.java @@ -101,6 +101,10 @@ JobId setProjectId(String projectId) { return getProject() != null ? this : toBuilder().setProject(projectId).build(); } + JobId setLocation(String location) { + return getLocation() != null ? this : toBuilder().setLocation(location).build(); + } + JobReference toPb() { return new JobReference() .setProjectId(getProject()) 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 9b67445acd13..6d3348cae0a6 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 @@ -66,6 +66,7 @@ public class BigQueryImplTest { private static final String PROJECT = "project"; + private static final String LOCATION = "US"; private static final String OTHER_PROJECT = "otherProject"; private static final String DATASET = "dataset"; private static final String TABLE = "table"; @@ -293,6 +294,16 @@ private BigQueryOptions createBigQueryOptionsForProject( .build(); } + private BigQueryOptions createBigQueryOptionsForProjectWithLocation( + String project, BigQueryRpcFactory rpcFactory) { + return BigQueryOptions.newBuilder() + .setProjectId(project) + .setLocation(LOCATION) + .setServiceRpcFactory(rpcFactory) + .setRetrySettings(ServiceOptions.getNoRetrySettings()) + .build(); + } + @Before public void setUp() { rpcFactoryMock = EasyMock.createMock(BigQueryRpcFactory.class); @@ -1179,6 +1190,17 @@ public void testGetJob() { assertEquals(new Job(bigquery, new JobInfo.BuilderImpl(COMPLETE_COPY_JOB)), job); } + @Test + public void testGetJobWithLocation() { + EasyMock.expect(bigqueryRpcMock.getJob(PROJECT, JOB, LOCATION, EMPTY_RPC_OPTIONS)) + .andReturn(COMPLETE_COPY_JOB.toPb()); + EasyMock.replay(bigqueryRpcMock); + BigQueryOptions options = createBigQueryOptionsForProjectWithLocation(PROJECT, rpcFactoryMock); + bigquery = options.getService(); + Job job = bigquery.getJob(JOB); + assertEquals(new Job(bigquery, new JobInfo.BuilderImpl(COMPLETE_COPY_JOB)), job); + } + @Test public void testGetJobFromJobId() { EasyMock.expect(bigqueryRpcMock.getJob(PROJECT, JOB, null, EMPTY_RPC_OPTIONS)) @@ -1189,6 +1211,17 @@ public void testGetJobFromJobId() { assertEquals(new Job(bigquery, new JobInfo.BuilderImpl(COMPLETE_COPY_JOB)), job); } + @Test + public void testGetJobFromJobIdWithLocation() { + EasyMock.expect(bigqueryRpcMock.getJob(PROJECT, JOB, LOCATION, EMPTY_RPC_OPTIONS)) + .andReturn(COMPLETE_COPY_JOB.toPb()); + EasyMock.replay(bigqueryRpcMock); + BigQueryOptions options = createBigQueryOptionsForProjectWithLocation(PROJECT, rpcFactoryMock); + bigquery = options.getService(); + Job job = bigquery.getJob(JobId.of(JOB)); + assertEquals(new Job(bigquery, new JobInfo.BuilderImpl(COMPLETE_COPY_JOB)), job); + } + @Test public void testGetJobFromJobIdWithProject() { JobId jobId = JobId.of(OTHER_PROJECT, JOB); @@ -1201,6 +1234,19 @@ public void testGetJobFromJobIdWithProject() { assertEquals(new Job(bigquery, new JobInfo.BuilderImpl(jobInfo)), job); } + @Test + public void testGetJobFromJobIdWithProjectWithLocation() { + JobId jobId = JobId.of(OTHER_PROJECT, JOB); + JobInfo jobInfo = COPY_JOB.setProjectId(OTHER_PROJECT); + EasyMock.expect(bigqueryRpcMock.getJob(OTHER_PROJECT, JOB, LOCATION, EMPTY_RPC_OPTIONS)) + .andReturn(jobInfo.toPb()); + EasyMock.replay(bigqueryRpcMock); + BigQueryOptions options = createBigQueryOptionsForProjectWithLocation(PROJECT, rpcFactoryMock); + bigquery = options.getService(); + Job job = bigquery.getJob(jobId); + assertEquals(new Job(bigquery, new JobInfo.BuilderImpl(jobInfo)), job); + } + @Test public void testListJobs() { bigquery = options.getService();