Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: add properties to set universe domain and endpoint in bigquery #3158

Merged
merged 12 commits into from
Sep 26, 2024
2 changes: 2 additions & 0 deletions docs/src/main/asciidoc/bigquery.adoc
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
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 + "/";
}
}
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;
}

}
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");
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Please verify that it's intended to have getEndpoint not return the URI scheme, but getResolvedApiaryHost does return the URI scheme.

LGTM

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, returning endpoint with this pattern is intended for Bigquerystorage which will otherwise throw an exception if it doesn't follow a specific convention. For example, passing in bigquerystorage.example.com without the port will result in:

Caused by: org.springframework.beans.BeanInstantiationException: Failed to instantiate [com.google.cloud.bigquery.storage.v1.BigQueryWriteClient]: Factory method 'bigQueryWriteClient' threw exception with message: invalid endpoint, expecting "<host>:<port>"
	at org.springframework.beans.factory.support.SimpleInstantiationStrategy.instantiate(SimpleInstantiationStrategy.java:178)
	at org.springframework.beans.factory.support.ConstructorResolver.instantiate(ConstructorResolver.java:644)
	... 31 more
Caused by: java.lang.IllegalArgumentException: invalid endpoint, expecting "<host>:<port>"

On the other hand bigquery accepts either accepts user-provided endpoint as-is if one is provided or builds the host in the https://service.universeDomain format if only the universe domain is provided:
https://github.com/googleapis/sdk-platform-java/blob/1a988df22f7e3d15ce6b121bf26897c59ab468e4/java-core/google-cloud-core/src/main/java/com/google/cloud/ServiceOptions.java#L868 but doesn't contain an explicit check to verify this schema. And in this PR, we reformat the provided endpoint to follow the same pattern that getResolvedApiaryHost() defaults to.

});
}

@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
Loading