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.

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

Add withDefaultHeaders to connection configuration for ElasticsearchIO #25024

Merged
merged 3 commits into from
Jan 18, 2023
Merged

Add withDefaultHeaders to connection configuration for ElasticsearchIO #25024

Changes from 1 commit
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
6f87445
Add withDefaultHeaders to connection configuration for ElasticsearchIO
Jan 16, 2023
de21be6
Address PR comments, run spotless
Jan 17, 2023
a05a9c0
Fix copy/paste error in Javadoc, add example for usage of withDefault…
Jan 18, 2023
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
29 changes: 28 additions & 1 deletion .../io/elasticsearch/src/main/java/org/apache/beam/sdk/io/elasticsearch/ElasticsearchIO.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -137,7 +137,8 @@
* }</pre>
*
* <p>The connection configuration also accepts optional configuration: {@code withUsername()},
* {@code withPassword()}, {@code withApiKey()} and {@code withBearerToken()}.
* {@code withPassword()}, {@code withApiKey()}, {@code withBearerToken()} and
* {@code withDefaultHeaders()}.
*
* <p>You can also specify a query on the {@code read()} using {@code withQuery()}.
*
Expand Down Expand Up @@ -326,6 +327,8 @@ public abstract static class ConnectionConfiguration implements Serializable {

public abstract @Nullable String getBearerToken();

public abstract @Nullable List<Header> getDefaultHeaders();

public abstract @Nullable String getKeystorePath();

public abstract @Nullable String getKeystorePassword();
Expand Down Expand Up @@ -354,6 +357,8 @@ abstract static class Builder {

abstract Builder setBearerToken(String bearerToken);

abstract Builder setDefaultHeaders(List<Header> defaultHeaders);

abstract Builder setKeystorePath(String keystorePath);

abstract Builder setKeystorePassword(String password);
Expand Down Expand Up @@ -502,6 +507,8 @@ public ConnectionConfiguration withPassword(String password) {

/**
* If Elasticsearch authentication is enabled, provide an API key.
* Be aware that you can only use one of {@Code withApiToken()}, {@code withBearerToken()} and
* {@code withDefaultHeaders} at the same time, as they will override eachother.
*
* @param apiKey the API key used to authenticate to Elasticsearch
* @return a {@link ConnectionConfiguration} describes a connection configuration to
Expand All @@ -514,6 +521,8 @@ public ConnectionConfiguration withApiKey(String apiKey) {

/**
* If Elasticsearch authentication is enabled, provide a bearer token.
* Be aware that you can only use one of {@Code withApiToken()}, {@code withBearerToken()} and
* {@code withDefaultHeaders} at the same time, as they will override eachother.
Copy link
Member

Choose a reason for hiding this comment

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

Is there a reasonable way to add Precondition checks to try to alert the user when these mutually exclusive fields have been used in conjunction?

Copy link
Author

Choose a reason for hiding this comment

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

There is but I didn't add it because the same issue was already lurking with just withApiKey and withBearerToken. I added checks to all three methods now to check that only one is used at a time - see my last update.

*
* @param bearerToken the bearer token used to authenticate to Elasticsearch
* @return a {@link ConnectionConfiguration} describes a connection configuration to
Expand All @@ -524,6 +533,21 @@ public ConnectionConfiguration withBearerToken(String bearerToken) {
return builder().setBearerToken(bearerToken).build();
}

/**
* For authentication or custom requirements, provide a set if default headers for the client.
* Be aware that you can only use one of {@code withApiToken()}, {@code withBearerToken()} and
* {@code withDefaultHeaders} at the same time, as they will override eachother.
*
* @param defaultHeaders the headers to add to outgoing requests
* @return a {@link ConnectionConfiguration} describes a connection configuration to
* Elasticsearch.
*/
public ConnectionConfiguration withDefaultHeaders(Header [] defaultHeaders) {
checkArgument(defaultHeaders != null, "defaultHeaders can not be null");
checkArgument(defaultHeaders.length > 0, "defaultHeaders can not be empty");
return builder().setDefaultHeaders(Arrays.asList(defaultHeaders)).build();
}

/**
* If Elasticsearch uses SSL/TLS with mutual authentication (via shield), provide the keystore
* containing the client key.
Expand Down Expand Up @@ -640,6 +664,9 @@ RestClient createClient() throws IOException {
restClientBuilder.setDefaultHeaders(
new Header[] {new BasicHeader("Authorization", "Bearer " + getBearerToken())});
}
if (getDefaultHeaders() != null) {
restClientBuilder.setDefaultHeaders(getDefaultHeaders().toArray(Header[]::new));
}

restClientBuilder.setHttpClientConfigCallback(
httpClientBuilder -> {
Expand Down