Skip to content

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

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.

Sign up for GitHub

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

Merged
merged 12 commits into from
Sep 26, 2024
Merged

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

Show file tree
Hide file tree
Changes from 9 commits
Commits
Show all changes
12 commits
Select commit Hold shift + click to select a range
929844e
feat: add properties to set universe domain and endpoint in bigquery
mpeddada1 Aug 27, 2024
bf6887b
fix test
mpeddada1 Aug 27, 2024
e86aa9e
introduce host configuration for BigqueryOptions
mpeddada1 Sep 16, 2024
7e1f734
verify with getResolvedApiaryHost
mpeddada1 Sep 16, 2024
9419b7b
add documentation for properties
mpeddada1 Sep 16, 2024
8364de1
add global properties for endpoint and universeDomaine; resolve to host
mpeddada1 Sep 23, 2024
a8b8ad2
fix checkstyle
mpeddada1 Sep 23, 2024
8e7b4bf
fix comment
mpeddada1 Sep 23, 2024
83f0fff
make test names consistent
mpeddada1 Sep 23, 2024
c3ebe7c
simplify reformatting logic
mpeddada1 Sep 25, 2024
4b34ff3
reformat
mpeddada1 Sep 25, 2024
2309417
add documentation
mpeddada1 Sep 25, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
49 changes: 40 additions & 9 deletions ...ain/java/com/google/cloud/spring/autoconfigure/bigquery/GcpBigQueryAutoConfiguration.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,8 @@
import com.google.cloud.spring.core.GcpProjectIdProvider;
import com.google.cloud.spring.core.UserAgentHeaderProvider;
import java.io.IOException;
import java.net.URI;
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 +59,9 @@ public class GcpBigQueryAutoConfiguration {

private int threadPoolSize;

private String universeDomain;
private String endpoint;

GcpBigQueryAutoConfiguration(
GcpBigQueryProperties gcpBigQueryProperties,
GcpProjectIdProvider projectIdProvider,
Expand All @@ -78,6 +83,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 +105,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 +160,16 @@ public BigQueryTemplate bigQueryTemplate(
return new BigQueryTemplate(
bigQuery, bigQueryWriteClient, bqInitSettings, bigQueryThreadPoolTaskScheduler);
}

private String resolveToHost(String endpoint) throws URISyntaxException {
if (!endpoint.contains("://")) {
endpoint = "https://" + endpoint;
}
URI uri = new URI(endpoint);

// Construct the new URL with https:// and no port
Copy link
Contributor

Choose a reason for hiding this comment

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

Can we just do String manipulation is host is basically endpoint without ports? Is there any edge cases can not be handled by String manipulation?

Copy link
Contributor Author

Choose a reason for hiding this comment

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

Done. Simplified to use String manipulation.

String newUrl = new URI("https", uri.getUserInfo(), uri.getHost(), -1,
uri.getPath(), uri.getQuery(), uri.getFragment()).toString();
return newUrl;
}
}
24 changes: 24 additions & 0 deletions ...e/src/main/java/com/google/cloud/spring/autoconfigure/bigquery/GcpBigQueryProperties.java
View file
Open in desktop
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
Open in desktop
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