BigQuery : Fix Location configurable at BigQueryOptions (#4329)

* Fix Location configurable from BigQueryOptions * modified code * modified code and add test case * removed unused location
googleapis · Jan 25, 2019 · f4bac2f · f4bac2f
1 parent b92a48b
commit f4bac2f
Show file tree

Hide file tree

Showing 4 changed files with 84 additions and 3 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
@@ -679,7 +679,13 @@ public Job getJob(String jobId, JobOption... options) {
   @Override
   public Job getJob(JobId jobId, JobOption... options) {
     final Map<BigQueryRpc.Option, ?> 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<Boolean>() {
@@ -784,7 +796,13 @@ private static QueryResponse getQueryResults(
       JobId jobId,
       final BigQueryOptions serviceOptions,
       final Map<BigQueryRpc.Option, ?> 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/...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
@@ -33,6 +33,7 @@ public class BigQueryOptions extends ServiceOptions<BigQuery, BigQueryOptions> {
   private static final String BIGQUERY_SCOPE = "https://www.googleapis.com/auth/bigquery";
   private static final Set<String> 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<BigQuery, BigQueryOptions, 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<BigQuery, BigQueryOptions> {
@@ -112,6 +121,10 @@ protected BigQueryRpc getBigQueryRpcV2() {
     return (BigQueryRpc) getRpc();
   }
 
+  public String getLocation() {
+    return location;
+  }
+
   @SuppressWarnings("unchecked")
   @Override
   public Builder toBuilder() {

diff --git a/...le-cloud-clients/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/JobId.java b/...le-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/...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
@@ -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();