feat: add properties to set universe domain and endpoint in bigquery (G…

…oogleCloudPlatform#3158) * feat: add properties to set universe domain and endpoint in bigquery
renovate-bot · Sep 26, 2024 · 9b3c780 · 9b3c780
1 parent 4718eaf
commit 9b3c780
Show file tree

Hide file tree

Showing 4 changed files with 163 additions and 9 deletions.
diff --git a/docs/src/main/asciidoc/bigquery.adoc b/docs/src/main/asciidoc/bigquery.adoc
@@ -38,6 +38,8 @@ The following application properties may be configured with Spring Framework on
 | `spring.cloud.gcp.bigquery.credentials.location` | Credentials file location for authenticating with the Google Cloud BigQuery APIs, if different from the ones in the <<spring-cloud-gcp-core,Spring Framework on Google Cloud Core Module>> | No | Inferred from https://cloud.google.com/docs/authentication/production[Application Default Credentials], typically set by https://cloud.google.com/sdk/gcloud/reference/auth/application-default[`gcloud`].
 | `spring.cloud.gcp.bigquery.jsonWriterBatchSize` | Batch size which will be used by `BigQueryJsonDataWriter` while using https://cloud.google.com/bigquery/docs/write-api[BigQuery Storage Write API]. Note too large or too low values might impact performance. | No | 1000
 | `spring.cloud.gcp.bigquery.threadPoolSize` | The size of thread pool of `ThreadPoolTaskScheduler` which is used by `BigQueryTemplate` | No | 4
+| `spring.cloud.gcp.bigquery.universe-domain` | Universe domain of the Bigquery service. The universe domain is a part of the endpoint which is formatted as ${service}.${universeDomain}:${port} | Relies on client library’s default universe domain which is googleapis.com
+| `spring.cloud.gcp.bigquery.endpoint` | Endpoint of the Bigquery service. Follows the ${service}.${universeDomain}:${port} format for the BigqueryWriteClient otherwise reformats it to `https://${service}.${universeDomain}/` when setting it to Bigquery client. 
 |===========================================================================
 
 ==== BigQuery Client Object

diff --git a/...ain/java/com/google/cloud/spring/autoconfigure/bigquery/GcpBigQueryAutoConfiguration.java b/...ain/java/com/google/cloud/spring/autoconfigure/bigquery/GcpBigQueryAutoConfiguration.java
@@ -27,6 +27,7 @@
 import com.google.cloud.spring.core.GcpProjectIdProvider;
 import com.google.cloud.spring.core.UserAgentHeaderProvider;
 import java.io.IOException;
+import java.net.URISyntaxException;
 import java.util.HashMap;
 import java.util.Map;
 import org.springframework.beans.factory.annotation.Qualifier;
@@ -57,6 +58,9 @@ public class GcpBigQueryAutoConfiguration {
 
   private int threadPoolSize;
 
+  private String universeDomain;
+  private String endpoint;
+
   GcpBigQueryAutoConfiguration(
       GcpBigQueryProperties gcpBigQueryProperties,
       GcpProjectIdProvider projectIdProvider,
@@ -78,6 +82,10 @@ public class GcpBigQueryAutoConfiguration {
     this.jsonWriterBatchSize = gcpBigQueryProperties.getJsonWriterBatchSize();
 
     this.threadPoolSize = getThreadPoolSize(gcpBigQueryProperties.getThreadPoolSize());
+
+    this.universeDomain = gcpBigQueryProperties.getUniverseDomain();
+
+    this.endpoint = gcpBigQueryProperties.getEndpoint();
   }
 
   /**
@@ -96,26 +104,36 @@ private int getThreadPoolSize(int threadPoolSize) {
 
   @Bean
   @ConditionalOnMissingBean
-  public BigQuery bigQuery() throws IOException {
-    BigQueryOptions bigQueryOptions =
+  public BigQuery bigQuery() throws IOException, URISyntaxException {
+    BigQueryOptions.Builder bigQueryOptionsBuilder =
         BigQueryOptions.newBuilder()
             .setProjectId(this.projectId)
             .setCredentials(this.credentialsProvider.getCredentials())
-            .setHeaderProvider(new UserAgentHeaderProvider(GcpBigQueryAutoConfiguration.class))
-            .build();
-    return bigQueryOptions.getService();
+            .setHeaderProvider(new UserAgentHeaderProvider(GcpBigQueryAutoConfiguration.class));
+    if (this.universeDomain != null) {
+      bigQueryOptionsBuilder.setUniverseDomain(this.universeDomain);
+    }
+    if (this.endpoint != null) {
+      bigQueryOptionsBuilder.setHost(resolveToHost(this.endpoint));
+    }
+    return bigQueryOptionsBuilder.build().getService();
   }
 
   @Bean
   @ConditionalOnMissingBean
   public BigQueryWriteClient bigQueryWriteClient() throws IOException {
-    BigQueryWriteSettings bigQueryWriteSettings =
+    BigQueryWriteSettings.Builder bigQueryWriteSettingsBuilder =
         BigQueryWriteSettings.newBuilder()
             .setCredentialsProvider(this.credentialsProvider)
             .setQuotaProjectId(this.projectId)
-            .setHeaderProvider(new UserAgentHeaderProvider(GcpBigQueryAutoConfiguration.class))
-            .build();
-    return BigQueryWriteClient.create(bigQueryWriteSettings);
+            .setHeaderProvider(new UserAgentHeaderProvider(GcpBigQueryAutoConfiguration.class));
+    if (this.universeDomain != null) {
+      bigQueryWriteSettingsBuilder.setUniverseDomain(this.universeDomain);
+    }
+    if (this.endpoint != null) {
+      bigQueryWriteSettingsBuilder.setEndpoint(this.endpoint);
+    }
+    return BigQueryWriteClient.create(bigQueryWriteSettingsBuilder.build());
   }
 
   @Bean
@@ -141,4 +159,12 @@ public BigQueryTemplate bigQueryTemplate(
     return new BigQueryTemplate(
         bigQuery, bigQueryWriteClient, bqInitSettings, bigQueryThreadPoolTaskScheduler);
   }
+
+  private String resolveToHost(String endpoint) throws URISyntaxException {
+    int portIndex = endpoint.indexOf(":");
+    if (portIndex != -1) {
+      return "https://" + endpoint.substring(0, portIndex) + "/";
+    }
+    return "https://" + endpoint + "/";
+  }
 }
diff --git a/...e/src/main/java/com/google/cloud/spring/autoconfigure/bigquery/GcpBigQueryProperties.java b/...e/src/main/java/com/google/cloud/spring/autoconfigure/bigquery/GcpBigQueryProperties.java
@@ -45,6 +45,13 @@ public class GcpBigQueryProperties implements CredentialsSupplier {
   /** The size of thread pool of ThreadPoolTaskScheduler used by GcpBigQueryAutoConfiguration */
   private int threadPoolSize;
 
+  private String universeDomain;
+
+  /**
+   * Endpoint (formatted as `{service}.{universeDomain}:${port}`)
+   */
+  private String endpoint;
+
   public int getJsonWriterBatchSize() {
     return jsonWriterBatchSize;
   }
@@ -80,4 +87,21 @@ public String getDatasetName() {
   public void setDatasetName(String datasetName) {
     this.datasetName = datasetName;
   }
+
+  public String getUniverseDomain() {
+    return universeDomain;
+  }
+
+  public void setUniverseDomain(String universeDomain) {
+    this.universeDomain = universeDomain;
+  }
+
+  public String getEndpoint() {
+    return endpoint;
+  }
+
+  public void setEndpoint(String endpoint) {
+    this.endpoint = endpoint;
+  }
+
 }
diff --git a/...ava/com/google/cloud/spring/autoconfigure/bigquery/GcpBigQueryAutoConfigurationTests.java b/...ava/com/google/cloud/spring/autoconfigure/bigquery/GcpBigQueryAutoConfigurationTests.java
@@ -23,6 +23,7 @@
 import com.google.auth.oauth2.GoogleCredentials;
 import com.google.cloud.bigquery.BigQuery;
 import com.google.cloud.bigquery.BigQueryOptions;
+import com.google.cloud.bigquery.storage.v1.BigQueryWriteClient;
 import com.google.cloud.spring.autoconfigure.core.GcpContextAutoConfiguration;
 import com.google.cloud.spring.bigquery.core.BigQueryTemplate;
 import org.junit.jupiter.api.Test;
@@ -60,6 +61,107 @@ void testSettingBigQueryOptions() {
         });
   }
 
+  @Test
+  void testBigQuery_universeDomain() {
+    this.contextRunner
+        .withPropertyValues("spring.cloud.gcp.bigquery.universe-domain=myUniverseDomain")
+        .run(
+            ctx -> {
+              BigQueryOptions options = ctx.getBean(BigQuery.class).getOptions();
+              assertThat(options.getUniverseDomain()).isEqualTo("myUniverseDomain");
+              assertThat(options.getResolvedApiaryHost("bigquery"))
+                  .isEqualTo("https://bigquery.myUniverseDomain/");
+            });
+  }
+
+  @Test
+  void testBigQuery_noUniverseDomainAndEndpointSet_useClientDefault() {
+    this.contextRunner.run(
+        ctx -> {
+          BigQueryOptions options = ctx.getBean(BigQuery.class).getOptions();
+          assertThat(options.getUniverseDomain()).isNull();
+          assertThat(options.getResolvedApiaryHost("bigquery"))
+              .isEqualTo("https://bigquery.googleapis.com/");
+        });
+  }
+
+  @Test
+  void testBigQuery_endpoint() {
+    this.contextRunner
+        .withPropertyValues("spring.cloud.gcp.bigquery.endpoint=bigquery.example.com:443")
+        .run(
+            ctx -> {
+              BigQueryOptions options = ctx.getBean(BigQuery.class).getOptions();
+              assertThat(options.getResolvedApiaryHost("bigquery"))
+                  .isEqualTo("https://bigquery.example.com/");
+            });
+  }
+
+  @Test
+  void testBigQuery_bothEndpointAndUniverseDomainSet() {
+    this.contextRunner
+        .withPropertyValues("spring.cloud.gcp.bigquery.endpoint=bigquery.example.com:123")
+        .withPropertyValues("spring.cloud.gcp.bigquery.universe-domain=myUniverseDomain")
+        .run(
+            ctx -> {
+              BigQueryOptions options = ctx.getBean(BigQuery.class).getOptions();
+              assertThat(options.getResolvedApiaryHost("bigquery"))
+                  .isEqualTo("https://bigquery.example.com/");
+              assertThat(options.getUniverseDomain()).isEqualTo("myUniverseDomain");
+            });
+  }
+
+  @Test
+  void testBigQueryWrite_universeDomain() {
+    this.contextRunner
+        .withPropertyValues("spring.cloud.gcp.bigquery.universe-domain=myUniverseDomain")
+        .run(
+            ctx -> {
+              BigQueryWriteClient writeClient = ctx.getBean(BigQueryWriteClient.class);
+              assertThat(writeClient.getSettings().getUniverseDomain())
+                  .isEqualTo("myUniverseDomain");
+            });
+  }
+
+  @Test
+  void testBigQueryWrite_endpoint() {
+    this.contextRunner
+        .withPropertyValues(
+            "spring.cloud.gcp.bigquery.endpoint=bigquerystorage.example.com:123")
+        .run(
+            ctx -> {
+              BigQueryWriteClient client = ctx.getBean(BigQueryWriteClient.class);
+              assertThat(client.getSettings().getEndpoint())
+                  .isEqualTo("bigquerystorage.example.com:123");
+            });
+  }
+
+  @Test
+  void testBigQueryWrite_bothUniverseDomainAndEndpointSet() {
+    this.contextRunner
+        .withPropertyValues("spring.cloud.gcp.bigquery.universe-domain=myUniverseDomain")
+        .withPropertyValues(
+            "spring.cloud.gcp.bigquery.endpoint=bigquerystorage.example.com:123")
+        .run(
+            ctx -> {
+              BigQueryWriteClient client = ctx.getBean(BigQueryWriteClient.class);
+              assertThat(client.getSettings().getUniverseDomain()).isEqualTo("myUniverseDomain");
+              assertThat(client.getSettings().getEndpoint())
+                  .isEqualTo("bigquerystorage.example.com:123");
+            });
+  }
+
+  @Test
+  void testBigQueryWrite_noUniverseDomainOrEndpointSet_useClientDefault() {
+    this.contextRunner.run(
+        ctx -> {
+          BigQueryWriteClient client = ctx.getBean(BigQueryWriteClient.class);
+          assertThat(client.getSettings().getUniverseDomain()).isEqualTo("googleapis.com");
+          assertThat(client.getSettings().getEndpoint())
+              .isEqualTo("bigquerystorage.googleapis.com:443");
+        });
+  }
+
   /** Spring Boot config for tests. */
   @AutoConfigurationPackage
   static class TestConfiguration {