Skip to content

Commit

Permalink
feat: add properties to set universe domain and endpoint in bigquery (#…
Browse files Browse the repository at this point in the history 
…3158)

* feat: add properties to set universe domain and endpoint in bigquery
  • Loading branch information
@mpeddada1
mpeddada1 authored Sep 26, 2024
1 parent 4718eaf commit 9b3c780
Show file tree
Hide file tree
Showing 4 changed files with 163 additions and 9 deletions.
2 changes: 2 additions & 0 deletions docs/src/main/asciidoc/bigquery.adoc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
44 changes: 35 additions & 9 deletions ...ain/java/com/google/cloud/spring/autoconfigure/bigquery/GcpBigQueryAutoConfiguration.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -57,6 +58,9 @@ public class GcpBigQueryAutoConfiguration {

private int threadPoolSize;

private String universeDomain;
private String endpoint;

GcpBigQueryAutoConfiguration(
GcpBigQueryProperties gcpBigQueryProperties,
GcpProjectIdProvider projectIdProvider,
Expand All @@ -78,6 +82,10 @@ public class GcpBigQueryAutoConfiguration {
this.jsonWriterBatchSize = gcpBigQueryProperties.getJsonWriterBatchSize();

this.threadPoolSize = getThreadPoolSize(gcpBigQueryProperties.getThreadPoolSize());

this.universeDomain = gcpBigQueryProperties.getUniverseDomain();

this.endpoint = gcpBigQueryProperties.getEndpoint();
}

/**
Expand All @@ -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
Expand All @@ -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 + "/";
}
}
24 changes: 24 additions & 0 deletions ...e/src/main/java/com/google/cloud/spring/autoconfigure/bigquery/GcpBigQueryProperties.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}
Expand Down Expand Up @@ -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;
}

}
102 changes: 102 additions & 0 deletions ...ava/com/google/cloud/spring/autoconfigure/bigquery/GcpBigQueryAutoConfigurationTests.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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 {
Expand Down

0 comments on commit 9b3c780

Please sign in to comment.