Code review comment: Ok I'll revert to having 2 separate classes and a base class 👍 

Code before review:
public boolean start() throws IOException {
restClient = source.spec.getConnectionConfiguration().createClient();
String query = source.spec.getQuery() != null ? source.spec.getQuery().get() : null;
if (query == null) {
query = "{\"query\": { \"match_all\": {} }}";
}
if ((source.backendVersion >= 5) && source.numSlices != null && source.numSlices > 1) {
String sliceQuery =
String.format("\"slice\": {\"id\": %s,\"max\": %s}", source.sliceId, source.numSlices);
query = query.replaceFirst("\\{", "{" + sliceQuery + ",");
}
String endPoint =
String.format(
"/%s/%s/_search",
source.spec.getConnectionConfiguration().getIndex(),
source.spec.getConnectionConfiguration().getType());
Map<String, String> params = new HashMap<>();
params.put("scroll", source.spec.getScrollKeepalive());
if (source.backendVersion == 2) {
params.put("size", String.valueOf(source.spec.getBatchSize()));
if (source.shardPreference != null) {
params.put("preference", "_shards:" + source.shardPreference);
}
}
HttpEntity queryEntity = new NStringEntity(query, ContentType.APPLICATION_JSON);
Request request = new Request("GET", endPoint);
request.addParameters(params);
request.setEntity(queryEntity);
Response response = restClient.performRequest(request);
JsonNode searchResult = parseResponse(response.getEntity());
updateScrollId(searchResult);
return readNextBatchAndReturnFirstDocument(searchResult);
}
private void updateScrollId(JsonNode searchResult) {
scrollId = searchResult.path("_scroll_id").asText();
}
@Override
public boolean advance() throws IOException {
if (batchIterator.hasNext()) {
current = batchIterator.next();
return true;
} else {
String requestBody =
String.format(
"{\"scroll\" : \"%s\",\"scroll_id\" : \"%s\"}",
source.spec.getScrollKeepalive(), scrollId);
HttpEntity scrollEntity = new NStringEntity(requestBody, ContentType.APPLICATION_JSON);
Request request = new Request("GET", "/_search/scroll");
request.addParameters(Collections.emptyMap());
request.setEntity(scrollEntity);
Response response = restClient.performRequest(request);
JsonNode searchResult = parseResponse(response.getEntity());
updateScrollId(searchResult);
return readNextBatchAndReturnFirstDocument(searchResult);
}
}
private boolean readNextBatchAndReturnFirstDocument(JsonNode searchResult) {
JsonNode hits = searchResult.path("hits").path("hits");
if (hits.size() == 0) {
current = null;
batchIterator = null;
return false;
}
List<String> batch = new ArrayList<>();
boolean withMetadata = source.spec.isWithMetadata();
for (JsonNode hit : hits) {
if (withMetadata) {
batch.add(hit.toString());
} else {
String document = hit.path("_source").toString();
batch.add(document);
}
}
batchIterator = batch.listIterator();
current = batchIterator.next();
return true;
}
@Override
public String getCurrent() throws NoSuchElementException {
if (current == null) {
throw new NoSuchElementException();
}
return current;
}
@Override
public void close() throws IOException {
String requestBody = String.format("{\"scroll_id\" : [\"%s\"]}", scrollId);
HttpEntity entity = new NStringEntity(requestBody, ContentType.APPLICATION_JSON);
try {
Request request = new Request("DELETE", "/_search/scroll");
request.addParameters(Collections.emptyMap());
request.setEntity(entity);
restClient.performRequest(request);
} finally {
if (restClient != null) {
restClient.close();
}
}
}
@Override
public BoundedSource<String> getCurrentSource() {
return source;
}
}
/**
* A POJO encapsulating a configuration for retry behavior when issuing requests to ES. A retry
* will be attempted until the maxAttempts or maxDuration is exceeded, whichever comes first, for
* 429 TOO_MANY_REQUESTS error.
*/
@AutoValue
public abstract static class RetryConfiguration implements Serializable {
@VisibleForTesting
static final RetryPredicate DEFAULT_RETRY_PREDICATE = new DefaultRetryPredicate();
abstract int getMaxAttempts();
abstract Duration getMaxDuration();
abstract RetryPredicate getRetryPredicate();
abstract Builder builder();
@AutoValue.Builder
abstract static class Builder {
abstract ElasticsearchIO.RetryConfiguration.Builder setMaxAttempts(int maxAttempts);
abstract ElasticsearchIO.RetryConfiguration.Builder setMaxDuration(Duration maxDuration);
abstract ElasticsearchIO.RetryConfiguration.Builder setRetryPredicate(
RetryPredicate retryPredicate);
abstract ElasticsearchIO.RetryConfiguration build();
}
/**
* Creates RetryConfiguration for {@link ElasticsearchIO} with provided maxAttempts,
* maxDurations and exponential backoff based retries.
*
* @param maxAttempts max number of attempts.
* @param maxDuration maximum duration for retries.
* @return {@link RetryConfiguration} object with provided settings.
*/
public static RetryConfiguration create(int maxAttempts, Duration maxDuration) {
checkArgument(maxAttempts > 0, "maxAttempts must be greater than 0");
checkArgument(
maxDuration != null && maxDuration.isLongerThan(Duration.ZERO),
"maxDuration must be greater than 0");
return new AutoValue_ElasticsearchIO_RetryConfiguration.Builder()
.setMaxAttempts(maxAttempts)
.setMaxDuration(maxDuration)
.setRetryPredicate(DEFAULT_RETRY_PREDICATE)
.build();
}
@VisibleForTesting
RetryConfiguration withRetryPredicate(RetryPredicate predicate) {
checkArgument(predicate != null, "predicate must be provided");
return builder().setRetryPredicate(predicate).build();
}
/**
* An interface used to control if we retry the Elasticsearch call when a {@link Response} is
* obtained. If {@link RetryPredicate
* the requests to the Elasticsearch server if the {@link RetryConfiguration} permits it.
*/
@FunctionalInterface
interface RetryPredicate extends Predicate<HttpEntity>, Serializable {}
/**
* This is the default predicate used to test if a failed ES operation should be retried. A
* retry will be attempted until the maxAttempts or maxDuration is exceeded, whichever comes
* first, for TOO_MANY_REQUESTS(429) error.
*/
@VisibleForTesting
static class DefaultRetryPredicate implements RetryPredicate {
private int errorCode;
DefaultRetryPredicate(int code) {
this.errorCode = code;
}
DefaultRetryPredicate() {
this(429);
}
/** Returns true if the response has the error code for any mutation. */
private static boolean errorCodePresent(HttpEntity responseEntity, int errorCode) {
try {
JsonNode json = parseResponse(responseEntity);
if (json.path("errors").asBoolean()) {
for (JsonNode item : json.path("items")) {
if (item.findValue("status").asInt() == errorCode) {
return true;
}
}
}
} catch (IOException e) {
LOG.warn("Could not extract error codes from responseEntity {}", responseEntity);
}
return false;
}
@Override
public boolean test(HttpEntity responseEntity) {
return errorCodePresent(responseEntity, errorCode);
}
}
}
/** A {@link PTransform} converting docs to their Bulk API counterparts. */
@AutoValue
public abstract static class DocToBulk
extends PTransform<PCollection<String>, PCollection<String>> {
private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();
private static final int DEFAULT_RETRY_ON_CONFLICT = 5;
static {
SimpleModule module = new SimpleModule();
module.addSerializer(DocumentMetadata.class, new DocumentMetadataSerializer());
OBJECT_MAPPER.registerModule(module);
}
abstract @Nullable ConnectionConfiguration getConnectionConfiguration();
abstract Write.@Nullable FieldValueExtractFn getIdFn();
abstract Write.@Nullable FieldValueExtractFn getIndexFn();
abstract Write.@Nullable FieldValueExtractFn getRoutingFn();
abstract Write.@Nullable FieldValueExtractFn getTypeFn();
abstract Write.@Nullable FieldValueExtractFn getDocVersionFn();
abstract @Nullable String getDocVersionType();
abstract @Nullable String getUpsertScript();
abstract @Nullable Boolean getUsePartialUpdate();
abstract Write.@Nullable BooleanFieldValueExtractFn getIsDeleteFn();
abstract @Nullable Integer getBackendVersion();
abstract Builder builder();
@AutoValue.Builder
abstract static class Builder {
abstract Builder setConnectionConfiguration(ConnectionConfiguration connectionConfiguration);
abstract Builder setIdFn(Write.FieldValueExtractFn idFunction);
abstract Builder setIndexFn(Write.FieldValueExtractFn indexFn);
abstract Builder setRoutingFn(Write.FieldValueExtractFn routingFunction);
abstract Builder setTypeFn(Write.FieldValueExtractFn typeFn);
abstract Builder setDocVersionFn(Write.FieldValueExtractFn docVersionFn);
abstract Builder setDocVersionType(String docVersionType);
abstract Builder setIsDeleteFn(Write.BooleanFieldValueExtractFn isDeleteFn);
abstract Builder setUsePartialUpdate(Boolean usePartialUpdate);
abstract Builder setUpsertScript(String source);
abstract Builder setBackendVersion(Integer assumedBackendVersion);
abstract DocToBulk build();
}
/**
* Provide the Elasticsearch connection configuration object. Only required if
* withBackendVersion was not used i.e. getBackendVersion() returns null.
*
* @param connectionConfiguration the Elasticsearch {@link ConnectionConfiguration} object
* @return the {@link DocToBulk} with connection configuration set
*/
public DocToBulk withConnectionConfiguration(ConnectionConfiguration connectionConfiguration) {
checkArgument(connectionConfiguration != null, "connectionConfiguration can not be null");
return builder().setConnectionConfiguration(connectionConfiguration).build();
}
/**
* Provide a function to extract the id from the document. This id will be used as the document
* id in Elasticsearch. Should the function throw an Exception then the batch will fail and the
* exception propagated.
*
* @param idFn to extract the document ID
* @return the {@link DocToBulk} with the function set
*/
public DocToBulk withIdFn(Write.FieldValueExtractFn idFn) {
checkArgument(idFn != null, "idFn must not be null");
return builder().setIdFn(idFn).build();
}
/**
* Provide a function to extract the target index from the document allowing for dynamic
* document routing. Should the function throw an Exception then the batch will fail and the
* exception propagated.
*
* @param indexFn to extract the destination index from
* @return the {@link DocToBulk} with the function set
*/
public DocToBulk withIndexFn(Write.FieldValueExtractFn indexFn) {
checkArgument(indexFn != null, "indexFn must not be null");
return builder().setIndexFn(indexFn).build();
}
/**
* Provide a function to extract the target routing from the document allowing for dynamic
* document routing. Should the function throw an Exception then the batch will fail and the
* exception propagated.
*
* @param routingFn to extract the destination index from
* @return the {@link DocToBulk} with the function set
*/
public DocToBulk withRoutingFn(Write.FieldValueExtractFn routingFn) {
checkArgument(routingFn != null, "routingFn must not be null");
return builder().setRoutingFn(routingFn).build();
}
/**
* Provide a function to extract the target type from the document allowing for dynamic document
* routing. Should the function throw an Exception then the batch will fail and the exception
* propagated. Users are encouraged to consider carefully if multipe types are a sensible model
* <a
* href="https:
* discussed in this blog</a>.
*
* @param typeFn to extract the destination index from
* @return the {@link DocToBulk} with the function set
*/
public DocToBulk withTypeFn(Write.FieldValueExtractFn typeFn) {
checkArgument(typeFn != null, "typeFn must not be null");
return builder().setTypeFn(typeFn).build();
}
/**
* Provide an instruction to control whether partial updates or inserts (default) are issued to
* Elasticsearch.
*
* @param usePartialUpdate set to true to issue partial updates
* @return the {@link DocToBulk} with the partial update control set
*/
public DocToBulk withUsePartialUpdate(boolean usePartialUpdate) {
return builder().setUsePartialUpdate(usePartialUpdate).build();
}
/**
* Whether to use scripted updates and what script to use.
*
* @param source set to the value of the script source, painless lang
* @return the {@link DocToBulk} with the scripted updates set
*/
public DocToBulk withUpsertScript(String source) {
return builder().setUsePartialUpdate(false).setUpsertScript(source).build();
}
/**
* Provide a function to extract the doc version from the document. This version number will be
* used as the document version in Elasticsearch. Should the function throw an Exception then
* the batch will fail and the exception propagated. Incompatible with update operations and
* should only be used with withUsePartialUpdate(false)
*
* @param docVersionFn to extract the document version
* @return the {@link DocToBulk} with the function set
*/
public DocToBulk withDocVersionFn(Write.FieldValueExtractFn docVersionFn) {
checkArgument(docVersionFn != null, "docVersionFn must not be null");
return builder().setDocVersionFn(docVersionFn).build();
}
/**
* Provide a function to extract the target operation either upsert or delete from the document
* fields allowing dynamic bulk operation decision. While using withIsDeleteFn, it should be
* taken care that the document's id extraction is defined using the withIdFn function or else
* IllegalArgumentException is thrown. Should the function throw an Exception then the batch
* will fail and the exception propagated.
*
* @param isDeleteFn set to true for deleting the specific document
* @return the {@link Write} with the function set
*/
public DocToBulk withIsDeleteFn(Write.BooleanFieldValueExtractFn isDeleteFn) {
checkArgument(isDeleteFn != null, "deleteFn is required");
return builder().setIsDeleteFn(isDeleteFn).build();
}
/**
* Provide a function to extract the doc version from the document. This version number will be
* used as the document version in Elasticsearch. Should the function throw an Exception then
* the batch will fail and the exception propagated. Incompatible with update operations and
* should only be used with withUsePartialUpdate(false)
*
* @param docVersionType the version type to use, one of {@value ElasticsearchIO
* @return the {@link DocToBulk} with the doc version type set
*/
public DocToBulk withDocVersionType(String docVersionType) {
checkArgument(
VERSION_TYPES.contains(docVersionType),
"docVersionType must be one of " + "%s",
String.join(", ", VERSION_TYPES));
return builder().setDocVersionType(docVersionType).build();
}
/**
* Use to set explicitly which version of Elasticsearch the destination cluster is running.
* Providing this hint means there is no need for setting {@link
* DocToBulk
*
* @param backendVersion the major version number of the version of Elasticsearch being run in
*     the cluster where documents will be indexed.
* @return the {@link DocToBulk} with the Elasticsearch major version number set
*/
public DocToBulk withBackendVersion(int backendVersion) {
checkArgument(
VALID_CLUSTER_VERSIONS.contains(backendVersion),
"Backend version may only be one of " + "%s",
String.join(", ", VERSION_TYPES));
return builder().setBackendVersion(backendVersion).build();
}
@Override
public PCollection<String> expand(PCollection<String> docs) {
ConnectionConfiguration connectionConfiguration = getConnectionConfiguration();
Integer backendVersion = getBackendVersion();
Write.FieldValueExtractFn idFn = getIdFn();
Write.BooleanFieldValueExtractFn isDeleteFn = getIsDeleteFn();
checkState(
(backendVersion != null || connectionConfiguration != null),
"withBackendVersion() or withConnectionConfiguration() is required");
checkArgument(
isDeleteFn == null || idFn != null,
"Id needs to be specified by withIdFn for delete operation");
return docs.apply(ParDo.of(new DocToBulkFn(this)));
}
private static class DocumentMetadata implements Serializable {
final String index;
final String type;
final String id;
final Integer retryOnConflict;
final String routing;
final Integer backendVersion;
final String version;
final String versionType;
DocumentMetadata(
String index,
String type,
String id,
Integer retryOnConflict,
String routing,
Integer backendVersion,
String version,
String versionType) {
this.index = index;
this.id = id;
this.type = type;
this.retryOnConflict = retryOnConflict;
this.routing = routing;
this.backendVersion = backendVersion;
this.version = version;
this.versionType = versionType;
}
}
private static class DocumentMetadataSerializer extends StdSerializer<DocumentMetadata> {
private DocumentMetadataSerializer() {
super(DocumentMetadata.class);
}
@Override
public void serialize(DocumentMetadata value, JsonGenerator gen, SerializerProvider provider)
throws IOException {
gen.writeStartObject();
if (value.index != null) {
gen.writeStringField("_index", value.index);
}
if (value.type != null) {
gen.writeStringField("_type", value.type);
}
if (value.id != null) {
gen.writeStringField("_id", value.id);
}
if (value.routing != null) {
gen.writeStringField("routing", value.routing);
}
if (value.retryOnConflict != null && value.backendVersion <= 6) {
gen.writeNumberField("_retry_on_conflict", value.retryOnConflict);
}
if (value.retryOnConflict != null && value.backendVersion >= 7) {
gen.writeNumberField("retry_on_conflict", value.retryOnConflict);
}
if (value.version != null) {
gen.writeStringField("version", value.version);
}
if (value.versionType != null) {
gen.writeStringField("version_type", value.versionType);
}
gen.writeEndObject();
}
}
@VisibleForTesting
static String createBulkApiEntity(DocToBulk spec, String document, int backendVersion)
throws IOException {
String documentMetadata = "{}";
boolean isDelete = false;
if (spec.getIndexFn() != null || spec.getTypeFn() != null || spec.getIdFn() != null) {
JsonNode parsedDocument = OBJECT_MAPPER.readTree(document);
documentMetadata = getDocumentMetadata(spec, parsedDocument, backendVersion);
if (spec.getIsDeleteFn() != null) {
isDelete = spec.getIsDeleteFn().apply(parsedDocument);
}
}
if (isDelete) {
return String.format("{ \"delete\" : %s }%n", documentMetadata);
} else {
if (spec.getUsePartialUpdate()) {
return String.format(
"{ \"update\" : %s }%n{ \"doc\" : %s, " + "\"doc_as_upsert\" : true }%n",
documentMetadata, document);
} else if (spec.getUpsertScript() != null) {
return String.format(
"{ \"update\" : %s }%n{ \"script\" : {\"source\": \"%s\", "
+ "\"params\": %s}, \"upsert\" : %s }%n",
documentMetadata, spec.getUpsertScript(), document, document);
} else {
return String.format("{ \"index\" : %s }%n%s%n", documentMetadata, document);
}
}
}
private static String lowerCaseOrNull(String input) {
return input == null ? null : input.toLowerCase();
}
/**
* Extracts the components that comprise the document address from the document using the {@link
* Write.FieldValueExtractFn} configured. This allows any or all of the index, type and document
* id to be controlled on a per document basis. If none are provided then an empty default of
* {@code {}} is returned. Sanitization of the index is performed, automatically lower-casing
* the value as required by Elasticsearch.
*
* @param parsedDocument the json from which the index, type and id may be extracted
* @return the document address as JSON or the default
* @throws IOException if the document cannot be parsed as JSON
*/
private static String getDocumentMetadata(
DocToBulk spec, JsonNode parsedDocument, int backendVersion) throws IOException {
DocumentMetadata metadata =
new DocumentMetadata(
spec.getIndexFn() != null
? lowerCaseOrNull(spec.getIndexFn().apply(parsedDocument))
: null,
spec.getTypeFn() != null ? spec.getTypeFn().apply(parsedDocument) : null,
spec.getIdFn() != null ? spec.getIdFn().apply(parsedDocument) : null,
(spec.getUsePartialUpdate()
|| (spec.getUpsertScript() != null && !spec.getUpsertScript().isEmpty()))
? DEFAULT_RETRY_ON_CONFLICT
: null,
spec.getRoutingFn() != null ? spec.getRoutingFn().apply(parsedDocument) : null,
backendVersion,
spec.getDocVersionFn() != null ? spec.getDocVersionFn().apply(parsedDocument) : null,
spec.getDocVersionType());
return OBJECT_MAPPER.writeValueAsString(metadata);
}
/** {@link DoFn} to for the {@link DocToBulk} transform. */
@VisibleForTesting
static class DocToBulkFn extends DoFn<String, String> {
private final DocToBulk spec;
private int backendVersion;
public DocToBulkFn(DocToBulk spec) {
this.spec = spec;
}
@Setup
public void setup() throws IOException {
ConnectionConfiguration connectionConfiguration = spec.getConnectionConfiguration();
if (spec.getBackendVersion() == null) {
backendVersion = ElasticsearchIO.getBackendVersion(connectionConfiguration);
} else {
backendVersion = spec.getBackendVersion();
}
}
@ProcessElement
public void processElement(ProcessContext c) throws IOException {
c.output(createBulkApiEntity(spec, c.element(), backendVersion));
}
}
}
/**
* A {@link PTransform} convenience wrapper for doing both document to bulk API serialization as
* well as batching those Bulk API entities and writing them to an Elasticsearch cluster. This
* class is effectively a thin proxy for DocToBulk->BulkIO all-in-one for convenience and backward
* compatibility.
*/
@AutoValue
public abstract static class Write extends PTransform<PCollection<String>, PDone> {
public interface FieldValueExtractFn extends SerializableFunction<JsonNode, String> {}
public interface BooleanFieldValueExtractFn extends SerializableFunction<JsonNode, Boolean> {}
public abstract DocToBulk getDocToBulk();
public abstract BulkIO getBulkIO();
abstract Builder writeBuilder();
@AutoValue.Builder
abstract static class Builder {
abstract Builder setDocToBulk(DocToBulk docToBulk);
abstract Builder setBulkIO(BulkIO bulkIO);
abstract Write build();
}
/** Refer to {@link DocToBulk
public Write withIdFn(FieldValueExtractFn idFn) {
return writeBuilder().setDocToBulk(getDocToBulk().withIdFn(idFn)).build();
}
/** Refer to {@link DocToBulk
public Write withIndexFn(FieldValueExtractFn indexFn) {
return writeBuilder().setDocToBulk(getDocToBulk().withIndexFn(indexFn)).build();
}
/** Refer to {@link DocToBulk
public Write withRoutingFn(FieldValueExtractFn routingFn) {
return writeBuilder().setDocToBulk(getDocToBulk().withRoutingFn(routingFn)).build();
}
/** Refer to {@link DocToBulk
public Write withTypeFn(FieldValueExtractFn typeFn) {
return writeBuilder().setDocToBulk(getDocToBulk().withTypeFn(typeFn)).build();
}
/** Refer to {@link DocToBulk
public Write withDocVersionFn(FieldValueExtractFn docVersionFn) {
return writeBuilder().setDocToBulk(getDocToBulk().withDocVersionFn(docVersionFn)).build();
}
/** Refer to {@link DocToBulk
public Write withDocVersionType(String docVersionType) {
return writeBuilder().setDocToBulk(getDocToBulk().withDocVersionType(docVersionType)).build();
}
/** Refer to {@link DocToBulk
public Write withUsePartialUpdate(boolean usePartialUpdate) {
return writeBuilder()
.setDocToBulk(getDocToBulk().withUsePartialUpdate(usePartialUpdate))
.build();
}
/** Refer to {@link DocToBulk
public Write withUpsertScript(String source) {
return writeBuilder().setDocToBulk(getDocToBulk().withUpsertScript(source)).build();
}
/** Refer to {@link DocToBulk
public Write withBackendVersion(int backendVersion) {
return writeBuilder().setDocToBulk(getDocToBulk().withBackendVersion(backendVersion)).build();
}
/** Refer to {@link DocToBulk
public Write withIsDeleteFn(Write.BooleanFieldValueExtractFn isDeleteFn) {
return writeBuilder().setDocToBulk(getDocToBulk().withIsDeleteFn(isDeleteFn)).build();
}
/** Refer to {@link BulkIO
public Write withConnectionConfiguration(ConnectionConfiguration connectionConfiguration) {
checkArgument(connectionConfiguration != null, "connectionConfiguration can not be null");
return writeBuilder()
.setDocToBulk(getDocToBulk().withConnectionConfiguration(connectionConfiguration))
.setBulkIO(getBulkIO().withConnectionConfiguration(connectionConfiguration))
.build();
}
/** Refer to {@link BulkIO
public Write withMaxBatchSize(long batchSize) {
return writeBuilder().setBulkIO(getBulkIO().withMaxBatchSize(batchSize)).build();
}
/** Refer to {@link BulkIO
public Write withMaxBatchSizeBytes(long batchSizeBytes) {
return writeBuilder().setBulkIO(getBulkIO().withMaxBatchSizeBytes(batchSizeBytes)).build();
}
/** Refer to {@link BulkIO
public Write withRetryConfiguration(RetryConfiguration retryConfiguration) {
return writeBuilder()
.setBulkIO(getBulkIO().withRetryConfiguration(retryConfiguration))
.build();
}
/** Refer to {@link BulkIO
public Write withIgnoreVersionConflicts(boolean ignoreVersionConflicts) {
return writeBuilder()
.setBulkIO(getBulkIO().withIgnoreVersionConflicts(ignoreVersionConflicts))
.build();
}
/** Refer to {@link BulkIO
public Write withUseStatefulBatches(boolean useStatefulBatches) {
return writeBuilder()
.setBulkIO(getBulkIO().withUseStatefulBatches(useStatefulBatches))
.build();
}
/** Refer to {@link BulkIO
public Write withMaxBufferingDuration(Duration maxBufferingDuration) {
return writeBuilder()
.setBulkIO(getBulkIO().withMaxBufferingDuration(maxBufferingDuration))
.build();
}
/** Refer to {@link BulkIO
public Write withMaxParallelRquestsPerWindow(int maxParallelRquestsPerWindow) {
return writeBuilder()
.setBulkIO(getBulkIO().withMaxParallelRequestsPerWindow(maxParallelRquestsPerWindow))
.build();
}
/** Refer to {@link BulkIO
public Write withAllowableResponseErrors(@Nullable Set<String> allowableResponseErrors) {
if (allowableResponseErrors == null) {
allowableResponseErrors = new HashSet<>();
}
return writeBuilder()
.setBulkIO(getBulkIO().withAllowableResponseErrors(allowableResponseErrors))
.build();
}
@Override
public PDone expand(PCollection<String> input) {
input.apply(getDocToBulk()).apply(getBulkIO());
return PDone.in(input.getPipeline());
}
}
/** A {@link PTransform} writing data to Elasticsearch. */
@AutoValue
public abstract static class BulkIO extends PTransform<PCollection<String>, PDone> {
@VisibleForTesting
static final String RETRY_ATTEMPT_LOG = "Error writing to Elasticsearch. Retry attempt[%d]";
@VisibleForTesting
static final String RETRY_FAILED_LOG =
"Error writing to ES after %d attempt(s). No more attempts allowed";
abstract @Nullable ConnectionConfiguration getConnectionConfiguration();
abstract long getMaxBatchSize();
abstract long getMaxBatchSizeBytes();
abstract @Nullable Duration getMaxBufferingDuration();
abstract boolean getUseStatefulBatches();
abstract int getMaxParallelRequestsPerWindow();
abstract @Nullable RetryConfiguration getRetryConfiguration();
abstract @Nullable Set<String> getAllowedResponseErrors();
abstract Builder builder();
@AutoValue.Builder
abstract static class Builder {
abstract Builder setConnectionConfiguration(ConnectionConfiguration connectionConfiguration);
abstract Builder setMaxBatchSize(long maxBatchSize);
abstract Builder setMaxBatchSizeBytes(long maxBatchSizeBytes);
abstract Builder setRetryConfiguration(RetryConfiguration retryConfiguration);
abstract Builder setAllowedResponseErrors(Set<String> allowedResponseErrors);
abstract Builder setMaxBufferingDuration(Duration maxBufferingDuration);
abstract Builder setUseStatefulBatches(boolean useStatefulBatches);
abstract Builder setMaxParallelRequestsPerWindow(int maxParallelRequestsPerWindow);
abstract BulkIO build();
}
/**
* Provide the Elasticsearch connection configuration object.
*
* @param connectionConfiguration the Elasticsearch {@link ConnectionConfiguration} object
* @return the {@link BulkIO} with connection configuration set
*/
public BulkIO withConnectionConfiguration(ConnectionConfiguration connectionConfiguration) {
checkArgument(connectionConfiguration != null, "connectionConfiguration can not be null");
return builder().setConnectionConfiguration(connectionConfiguration).build();
}
/**
* Provide a maximum size in number of documents for the batch see bulk API
* (https:
* docs (like Elasticsearch bulk size advice). See
* https:
* execution engine, size of bundles may vary, this sets the maximum size. Change this if you
* need to have smaller ElasticSearch bulks.
*
* @param batchSize maximum batch size in number of documents
* @return the {@link BulkIO} with connection batch size set
*/
public BulkIO withMaxBatchSize(long batchSize) {
checkArgument(batchSize > 0, "batchSize must be > 0, but was %s", batchSize);
return builder().setMaxBatchSize(batchSize).build();
}
/**
* Provide a maximum size in bytes for the batch see bulk API
* (https:
* (like Elasticsearch bulk size advice). See
* https:
* execution engine, size of bundles may vary, this sets the maximum size. Change this if you
* need to have smaller ElasticSearch bulks.
*
* @param batchSizeBytes maximum batch size in bytes
* @return the {@link BulkIO} with connection batch size in bytes set
*/
public BulkIO withMaxBatchSizeBytes(long batchSizeBytes) {
checkArgument(batchSizeBytes > 0, "batchSizeBytes must be > 0, but was %s", batchSizeBytes);
return builder().setMaxBatchSizeBytes(batchSizeBytes).build();
}
/**
* Provides configuration to retry a failed batch call to Elasticsearch. A batch is considered
* as failed if the underlying {@link RestClient} surfaces 429 HTTP status code as error for one
* or more of the items in the {@link Response}. Users should consider that retrying might
* compound the underlying problem which caused the initial failure. Users should also be aware
* that once retrying is exhausted the error is surfaced to the runner which <em>may</em> then
* opt to retry the current bundle in entirety or abort if the max number of retries of the
* runner is completed. Retrying uses an exponential backoff algorithm, with minimum backoff of
* 5 seconds and then surfacing the error once the maximum number of retries or maximum
* configuration duration is exceeded.
*
* <p>Example use:
*
* <pre>{@code
* ElasticsearchIO.write()
*   .withRetryConfiguration(ElasticsearchIO.RetryConfiguration.create(10, Duration.standardMinutes(3))
*   ...
* }</pre>
*
* @param retryConfiguration the rules which govern the retry behavior
* @return the {@link BulkIO} with retrying configured
*/
public BulkIO withRetryConfiguration(RetryConfiguration retryConfiguration) {
checkArgument(retryConfiguration != null, "retryConfiguration is required");
return builder().setRetryConfiguration(retryConfiguration).build();
}
/**
* Whether or not to suppress version conflict errors in a Bulk API response. This can be useful
* if your use case involves using external version types.
*
* @param ignoreVersionConflicts true to suppress version conflicts, false to surface version
*     conflict errors.
* @return the {@link BulkIO} with version conflict handling configured
*/
public BulkIO withIgnoreVersionConflicts(boolean ignoreVersionConflicts) {
Set<String> allowedResponseErrors = getAllowedResponseErrors();
if (allowedResponseErrors == null) {
allowedResponseErrors = new HashSet<>();
}
if (ignoreVersionConflicts) {
allowedResponseErrors.add(VERSION_CONFLICT_ERROR);
}
return builder().setAllowedResponseErrors(allowedResponseErrors).build();
}
/**
* Provide a set of textual error types which can be contained in Bulk API response
* items[].error.type field. Any element in @param allowableResponseErrorTypes will suppress
* errors of the same type in Bulk responses.
*
* <p>See also
* https:
*
* @param allowableResponseErrorTypes
* @return the {@link BulkIO} with allowable response errors set
*/
public BulkIO withAllowableResponseErrors(@Nullable Set<String> allowableResponseErrorTypes) {
if (allowableResponseErrorTypes == null) {
allowableResponseErrorTypes = new HashSet<>();
}
return builder().setAllowedResponseErrors(allowableResponseErrorTypes).build();
}
/**
* If using {@link BulkIO
* time before buffered elements are emitted to Elasticsearch as a Bulk API request. If this
* config is not set, Bulk requests will not be issued until {@link BulkIO
* number of documents have been buffered. This may result in higher latency in particular if
* your max batch size is set to a large value and your pipeline input is low volume.
*
* @param maxBufferingDuration the maximum duration to wait before sending any buffered
*     documents to Elasticsearch, regardless of maxBatchSize.
* @return the {@link BulkIO} with maximum buffering duration set
*/
public BulkIO withMaxBufferingDuration(Duration maxBufferingDuration) {
LOG.warn(
"Use of withMaxBufferingDuration requires withUseStatefulBatches(true). "
+ "Setting that automatically.");
return builder()
.setUseStatefulBatches(true)
.setMaxBufferingDuration(maxBufferingDuration)
.build();
}
/**
* Whether or not to use Stateful Processing to ensure bulk requests have the desired number of
* entities i.e. as close to the maxBatchSize as possible. By default without this feature
* enabled, Bulk requests will not contain more than maxBatchSize entities, but the lower bound
* of batch size is determined by Beam Runner bundle sizes, which may be as few as 1.
*
* @param useStatefulBatches true enables the use of Stateful Processing to ensure that batches
*     are as close to the maxBatchSize as possible.
* @return the {@link BulkIO} with Stateful Processing enabled or disabled
*/
public BulkIO withUseStatefulBatches(boolean useStatefulBatches) {
return builder().setUseStatefulBatches(useStatefulBatches).build();
}
/**
* When using {@link BulkIO
* batches are maintained per-key-per-window. If data is globally windowed and this
* configuration is set to 1, there will only ever be 1 request in flight. Having only a single
* request in flight can be beneficial for ensuring an Elasticsearch cluster is not overwhelmed
* by parallel requests, but may not work for all use cases. If this number is less than the
* number of maximum workers in your pipeline, the IO work may not be distributed across all
* workers.
*
* @param maxParallelRequestsPerWindow the maximum number of parallel bulk requests for a window
*     of data
* @return the {@link BulkIO} with maximum parallel bulk requests per window set
*/
public BulkIO withMaxParallelRequestsPerWindow(int maxParallelRequestsPerWindow) {
checkArgument(
maxParallelRequestsPerWindow > 0, "parameter value must be positive " + "a integer");
return builder().setMaxParallelRequestsPerWindow(maxParallelRequestsPerWindow).build();
}
@Override
public PDone expand(PCollection<String> input) {
ConnectionConfiguration connectionConfiguration = getConnectionConfiguration();
checkState(connectionConfiguration != null, "withConnectionConfiguration() is required");
if (getUseStatefulBatches()) {
GroupIntoBatches<Integer, String> groupIntoBatches =
GroupIntoBatches.ofSize(getMaxBatchSize());
if (getMaxBufferingDuration() != null) {
groupIntoBatches = groupIntoBatches.withMaxBufferingDuration(getMaxBufferingDuration());
}
input
.apply(ParDo.of(new AssignShardFn<>(getMaxParallelRequestsPerWindow())))
.apply(groupIntoBatches)
.apply(
"Remove key no longer needed",
MapElements.into(TypeDescriptors.iterables(TypeDescriptors.strings()))
.via(KV::getValue))
.apply(ParDo.of(new BulkIOFn(this)));
} else {
input
.apply(
"Make elements iterable",
MapElements.into(TypeDescriptors.iterables(TypeDescriptors.strings()))
.via(Collections::singletonList))
.apply(ParDo.of(new BulkIOFn(this)));
}
return PDone.in(input.getPipeline());
}
/** {@link DoFn} to for the {@link BulkIO} transform. */
@VisibleForTesting
static class BulkIOFn extends DoFn<Iterable<String>, Void> {
private static final Duration RETRY_INITIAL_BACKOFF = Duration.standardSeconds(5);
private transient FluentBackoff retryBackoff;
protected BulkIO spec;
private transient RestClient restClient;
protected ArrayList<String> batch;
long currentBatchSizeBytes;
@VisibleForTesting
BulkIOFn(BulkIO bulkSpec) {
this.spec = bulkSpec;
}
@Setup
public void setup() throws IOException {
ConnectionConfiguration connectionConfiguration = spec.getConnectionConfiguration();
restClient = connectionConfiguration.createClient();
retryBackoff =
FluentBackoff.DEFAULT.withMaxRetries(0).withInitialBackoff(RETRY_INITIAL_BACKOFF);
if (spec.getRetryConfiguration() != null) {
retryBackoff =
FluentBackoff.DEFAULT
.withInitialBackoff(RETRY_INITIAL_BACKOFF)
.withMaxRetries(spec.getRetryConfiguration().getMaxAttempts() - 1)
.withMaxCumulativeBackoff(spec.getRetryConfiguration().getMaxDuration());
}
}
@StartBundle
public void startBundle(StartBundleContext context) {
batch = new ArrayList<>();
currentBatchSizeBytes = 0;
}
@FinishBundle
public void finishBundle(FinishBundleContext context)
throws IOException, InterruptedException {
flushBatch();
}
@ProcessElement
public void processElement(@Element @NonNull Iterable<String> bulkApiEntities)
throws Exception {
for (String bulkApiEntity : bulkApiEntities) {
addAndMaybeFlush(bulkApiEntity);
}
}
protected void addAndMaybeFlush(String bulkApiEntity)
throws IOException, InterruptedException {
batch.add(bulkApiEntity);
currentBatchSizeBytes += bulkApiEntity.getBytes(StandardCharsets.UTF_8).length;
if (batch.size() >= spec.getMaxBatchSize()
|| currentBatchSizeBytes >= spec.getMaxBatchSizeBytes()) {
flushBatch();
}
}
private void flushBatch() throws IOException, InterruptedException {
if (batch.isEmpty()) {
return;
}
LOG.info(
"ElasticsearchIO batch size: {}, batch size bytes: {}",
batch.size(),
currentBatchSizeBytes);
StringBuilder bulkRequest = new StringBuilder();
for (String json : batch) {
bulkRequest.append(json);
}
batch.clear();
currentBatchSizeBytes = 0L;
Response response = null;
HttpEntity responseEntity = null;
String endPoint = spec.getConnectionConfiguration().getBulkEndPoint();
HttpEntity requestBody =
new NStringEntity(bulkRequest.toString(), ContentType.APPLICATION_JSON);
try {
Request request = new Request("POST", endPoint);
request.addParameters(Collections.emptyMap());
request.setEntity(requestBody);
response = restClient.performRequest(request);
responseEntity = new BufferedHttpEntity(response.getEntity());
} catch (java.io.IOException ex) {
if (spec.getRetryConfiguration() == null) {
throw ex;
}
LOG.error("Caught ES timeout, retrying", ex);
}
if (spec.getRetryConfiguration() != null
&& (response == null
|| responseEntity == null
|| spec.getRetryConfiguration().getRetryPredicate().test(responseEntity))) {
if (responseEntity != null
&& spec.getRetryConfiguration().getRetryPredicate().test(responseEntity)) {
LOG.warn("ES Cluster is responding with HTP 429 - TOO_MANY_REQUESTS.");
}
responseEntity = handleRetry("POST", endPoint, Collections.emptyMap(), requestBody);
}
checkForErrors(responseEntity, spec.getAllowedResponseErrors());
}
/** retry request based on retry configuration policy. */
private HttpEntity handleRetry(
String method, String endpoint, Map<String, String> params, HttpEntity requestBody)
throws IOException, InterruptedException {
Response response;
HttpEntity responseEntity;
Sleeper sleeper = Sleeper.DEFAULT;
BackOff backoff = retryBackoff.backoff();
int attempt = 0;
while (BackOffUtils.next(sleeper, backoff)) {
LOG.warn(String.format(RETRY_ATTEMPT_LOG, ++attempt));
try {
Request request = new Request(method, endpoint);
request.addParameters(params);
request.setEntity(requestBody);
response = restClient.performRequest(request);
responseEntity = new BufferedHttpEntity(response.getEntity());
} catch (java.io.IOException ex) {
LOG.error("Caught ES timeout, retrying", ex);
continue;
}
if (!Objects.requireNonNull(spec.getRetryConfiguration())
.getRetryPredicate()
.test(responseEntity)) {
return responseEntity;
} else {
LOG.warn("ES Cluster is responding with HTP 429 - TOO_MANY_REQUESTS.");
}
}
throw new IOException(String.format(RETRY_FAILED_LOG, attempt));
}
@Teardown
public void closeClient() throws IOException {
if (restClient != null) {
restClient.close();
}
}
}
}
static int getBackendVersion(ConnectionConfiguration connectionConfiguration) {
try (RestClient restClient = connectionConfiguration.createClient()) {
Request request = new Request("GET", "");
Response response = restClient.performRequest(request);
JsonNode jsonNode = parseResponse(response.getEntity());
int backendVersion =
Integer.parseInt(jsonNode.path("version").path("number").asText().substring(0, 1));
checkArgument(
(VALID_CLUSTER_VERSIONS.contains(backendVersion)),
"The Elasticsearch version to connect to is %s.x. "
+ "This version of the ElasticsearchIO is only compatible with "
+ "Elasticsearch v7.x, v6.x, v5.x and v2.x",
backendVersion);
return backendVersion;
} catch (IOException e) {
throw new IllegalArgumentException("Cannot get Elasticsearch version", e);
}
}
}

Line under review: MapElements.into(TypeDescriptors.iterables(TypeDescriptors.strings()))

Code after addressing the review:
public boolean start() throws IOException {
restClient = source.spec.getConnectionConfiguration().createClient();
String query = source.spec.getQuery() != null ? source.spec.getQuery().get() : null;
if (query == null) {
query = "{\"query\": { \"match_all\": {} }}";
}
if ((source.backendVersion >= 5) && source.numSlices != null && source.numSlices > 1) {
String sliceQuery =
String.format("\"slice\": {\"id\": %s,\"max\": %s}", source.sliceId, source.numSlices);
query = query.replaceFirst("\\{", "{" + sliceQuery + ",");
}
String endPoint =
String.format(
"/%s/%s/_search",
source.spec.getConnectionConfiguration().getIndex(),
source.spec.getConnectionConfiguration().getType());
Map<String, String> params = new HashMap<>();
params.put("scroll", source.spec.getScrollKeepalive());
if (source.backendVersion == 2) {
params.put("size", String.valueOf(source.spec.getBatchSize()));
if (source.shardPreference != null) {
params.put("preference", "_shards:" + source.shardPreference);
}
}
HttpEntity queryEntity = new NStringEntity(query, ContentType.APPLICATION_JSON);
Request request = new Request("GET", endPoint);
request.addParameters(params);
request.setEntity(queryEntity);
Response response = restClient.performRequest(request);
JsonNode searchResult = parseResponse(response.getEntity());
updateScrollId(searchResult);
return readNextBatchAndReturnFirstDocument(searchResult);
}
private void updateScrollId(JsonNode searchResult) {
scrollId = searchResult.path("_scroll_id").asText();
}
@Override
public boolean advance() throws IOException {
if (batchIterator.hasNext()) {
current = batchIterator.next();
return true;
} else {
String requestBody =
String.format(
"{\"scroll\" : \"%s\",\"scroll_id\" : \"%s\"}",
source.spec.getScrollKeepalive(), scrollId);
HttpEntity scrollEntity = new NStringEntity(requestBody, ContentType.APPLICATION_JSON);
Request request = new Request("GET", "/_search/scroll");
request.addParameters(Collections.emptyMap());
request.setEntity(scrollEntity);
Response response = restClient.performRequest(request);
JsonNode searchResult = parseResponse(response.getEntity());
updateScrollId(searchResult);
return readNextBatchAndReturnFirstDocument(searchResult);
}
}
private boolean readNextBatchAndReturnFirstDocument(JsonNode searchResult) {
JsonNode hits = searchResult.path("hits").path("hits");
if (hits.size() == 0) {
current = null;
batchIterator = null;
return false;
}
List<String> batch = new ArrayList<>();
boolean withMetadata = source.spec.isWithMetadata();
for (JsonNode hit : hits) {
if (withMetadata) {
batch.add(hit.toString());
} else {
String document = hit.path("_source").toString();
batch.add(document);
}
}
batchIterator = batch.listIterator();
current = batchIterator.next();
return true;
}
@Override
public String getCurrent() throws NoSuchElementException {
if (current == null) {
throw new NoSuchElementException();
}
return current;
}
@Override
public void close() throws IOException {
String requestBody = String.format("{\"scroll_id\" : [\"%s\"]}", scrollId);
HttpEntity entity = new NStringEntity(requestBody, ContentType.APPLICATION_JSON);
try {
Request request = new Request("DELETE", "/_search/scroll");
request.addParameters(Collections.emptyMap());
request.setEntity(entity);
restClient.performRequest(request);
} finally {
if (restClient != null) {
restClient.close();
}
}
}
@Override
public BoundedSource<String> getCurrentSource() {
return source;
}
}
/**
* A POJO encapsulating a configuration for retry behavior when issuing requests to ES. A retry
* will be attempted until the maxAttempts or maxDuration is exceeded, whichever comes first, for
* 429 TOO_MANY_REQUESTS error.
*/
@AutoValue
public abstract static class RetryConfiguration implements Serializable {
@VisibleForTesting
static final RetryPredicate DEFAULT_RETRY_PREDICATE = new DefaultRetryPredicate();
abstract int getMaxAttempts();
abstract Duration getMaxDuration();
abstract RetryPredicate getRetryPredicate();
abstract Builder builder();
@AutoValue.Builder
abstract static class Builder {
abstract ElasticsearchIO.RetryConfiguration.Builder setMaxAttempts(int maxAttempts);
abstract ElasticsearchIO.RetryConfiguration.Builder setMaxDuration(Duration maxDuration);
abstract ElasticsearchIO.RetryConfiguration.Builder setRetryPredicate(
RetryPredicate retryPredicate);
abstract ElasticsearchIO.RetryConfiguration build();
}
/**
* Creates RetryConfiguration for {@link ElasticsearchIO} with provided maxAttempts,
* maxDurations and exponential backoff based retries.
*
* @param maxAttempts max number of attempts.
* @param maxDuration maximum duration for retries.
* @return {@link RetryConfiguration} object with provided settings.
*/
public static RetryConfiguration create(int maxAttempts, Duration maxDuration) {
checkArgument(maxAttempts > 0, "maxAttempts must be greater than 0");
checkArgument(
maxDuration != null && maxDuration.isLongerThan(Duration.ZERO),
"maxDuration must be greater than 0");
return new AutoValue_ElasticsearchIO_RetryConfiguration.Builder()
.setMaxAttempts(maxAttempts)
.setMaxDuration(maxDuration)
.setRetryPredicate(DEFAULT_RETRY_PREDICATE)
.build();
}
@VisibleForTesting
RetryConfiguration withRetryPredicate(RetryPredicate predicate) {
checkArgument(predicate != null, "predicate must be provided");
return builder().setRetryPredicate(predicate).build();
}
/**
* An interface used to control if we retry the Elasticsearch call when a {@link Response} is
* obtained. If {@link RetryPredicate
* the requests to the Elasticsearch server if the {@link RetryConfiguration} permits it.
*/
@FunctionalInterface
interface RetryPredicate extends Predicate<HttpEntity>, Serializable {}
/**
* This is the default predicate used to test if a failed ES operation should be retried. A
* retry will be attempted until the maxAttempts or maxDuration is exceeded, whichever comes
* first, for TOO_MANY_REQUESTS(429) error.
*/
@VisibleForTesting
static class DefaultRetryPredicate implements RetryPredicate {
private int errorCode;
DefaultRetryPredicate(int code) {
this.errorCode = code;
}
DefaultRetryPredicate() {
this(429);
}
/** Returns true if the response has the error code for any mutation. */
private static boolean errorCodePresent(HttpEntity responseEntity, int errorCode) {
try {
JsonNode json = parseResponse(responseEntity);
if (json.path("errors").asBoolean()) {
for (JsonNode item : json.path("items")) {
if (item.findValue("status").asInt() == errorCode) {
return true;
}
}
}
} catch (IOException e) {
LOG.warn("Could not extract error codes from responseEntity {}", responseEntity);
}
return false;
}
@Override
public boolean test(HttpEntity responseEntity) {
return errorCodePresent(responseEntity, errorCode);
}
}
}
/** A {@link PTransform} converting docs to their Bulk API counterparts. */
@AutoValue
public abstract static class DocToBulk
extends PTransform<PCollection<String>, PCollection<String>> {
private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();
private static final int DEFAULT_RETRY_ON_CONFLICT = 5;
static {
SimpleModule module = new SimpleModule();
module.addSerializer(DocumentMetadata.class, new DocumentMetadataSerializer());
OBJECT_MAPPER.registerModule(module);
}
abstract @Nullable ConnectionConfiguration getConnectionConfiguration();
abstract Write.@Nullable FieldValueExtractFn getIdFn();
abstract Write.@Nullable FieldValueExtractFn getIndexFn();
abstract Write.@Nullable FieldValueExtractFn getRoutingFn();
abstract Write.@Nullable FieldValueExtractFn getTypeFn();
abstract Write.@Nullable FieldValueExtractFn getDocVersionFn();
abstract @Nullable String getDocVersionType();
abstract @Nullable String getUpsertScript();
abstract @Nullable Boolean getUsePartialUpdate();
abstract Write.@Nullable BooleanFieldValueExtractFn getIsDeleteFn();
abstract @Nullable Integer getBackendVersion();
abstract Builder builder();
@AutoValue.Builder
abstract static class Builder {
abstract Builder setConnectionConfiguration(ConnectionConfiguration connectionConfiguration);
abstract Builder setIdFn(Write.FieldValueExtractFn idFunction);
abstract Builder setIndexFn(Write.FieldValueExtractFn indexFn);
abstract Builder setRoutingFn(Write.FieldValueExtractFn routingFunction);
abstract Builder setTypeFn(Write.FieldValueExtractFn typeFn);
abstract Builder setDocVersionFn(Write.FieldValueExtractFn docVersionFn);
abstract Builder setDocVersionType(String docVersionType);
abstract Builder setIsDeleteFn(Write.BooleanFieldValueExtractFn isDeleteFn);
abstract Builder setUsePartialUpdate(Boolean usePartialUpdate);
abstract Builder setUpsertScript(String source);
abstract Builder setBackendVersion(Integer assumedBackendVersion);
abstract DocToBulk build();
}
/**
* Provide the Elasticsearch connection configuration object. Only required if
* withBackendVersion was not used i.e. getBackendVersion() returns null.
*
* @param connectionConfiguration the Elasticsearch {@link ConnectionConfiguration} object
* @return the {@link DocToBulk} with connection configuration set
*/
public DocToBulk withConnectionConfiguration(ConnectionConfiguration connectionConfiguration) {
checkArgument(connectionConfiguration != null, "connectionConfiguration can not be null");
return builder().setConnectionConfiguration(connectionConfiguration).build();
}
/**
* Provide a function to extract the id from the document. This id will be used as the document
* id in Elasticsearch. Should the function throw an Exception then the batch will fail and the
* exception propagated.
*
* @param idFn to extract the document ID
* @return the {@link DocToBulk} with the function set
*/
public DocToBulk withIdFn(Write.FieldValueExtractFn idFn) {
checkArgument(idFn != null, "idFn must not be null");
return builder().setIdFn(idFn).build();
}
/**
* Provide a function to extract the target index from the document allowing for dynamic
* document routing. Should the function throw an Exception then the batch will fail and the
* exception propagated.
*
* @param indexFn to extract the destination index from
* @return the {@link DocToBulk} with the function set
*/
public DocToBulk withIndexFn(Write.FieldValueExtractFn indexFn) {
checkArgument(indexFn != null, "indexFn must not be null");
return builder().setIndexFn(indexFn).build();
}
/**
* Provide a function to extract the target routing from the document allowing for dynamic
* document routing. Should the function throw an Exception then the batch will fail and the
* exception propagated.
*
* @param routingFn to extract the destination index from
* @return the {@link DocToBulk} with the function set
*/
public DocToBulk withRoutingFn(Write.FieldValueExtractFn routingFn) {
checkArgument(routingFn != null, "routingFn must not be null");
return builder().setRoutingFn(routingFn).build();
}
/**
* Provide a function to extract the target type from the document allowing for dynamic document
* routing. Should the function throw an Exception then the batch will fail and the exception
* propagated. Users are encouraged to consider carefully if multipe types are a sensible model
* <a
* href="https:
* discussed in this blog</a>.
*
* @param typeFn to extract the destination index from
* @return the {@link DocToBulk} with the function set
*/
public DocToBulk withTypeFn(Write.FieldValueExtractFn typeFn) {
checkArgument(typeFn != null, "typeFn must not be null");
return builder().setTypeFn(typeFn).build();
}
/**
* Provide an instruction to control whether partial updates or inserts (default) are issued to
* Elasticsearch.
*
* @param usePartialUpdate set to true to issue partial updates
* @return the {@link DocToBulk} with the partial update control set
*/
public DocToBulk withUsePartialUpdate(boolean usePartialUpdate) {
return builder().setUsePartialUpdate(usePartialUpdate).build();
}
/**
* Whether to use scripted updates and what script to use.
*
* @param source set to the value of the script source, painless lang
* @return the {@link DocToBulk} with the scripted updates set
*/
public DocToBulk withUpsertScript(String source) {
if (getBackendVersion() == null || getBackendVersion() == 2) {
LOG.warn("Painless scripts are not supported on Elasticsearch clusters before version 5.0");
}
return builder().setUsePartialUpdate(false).setUpsertScript(source).build();
}
/**
* Provide a function to extract the doc version from the document. This version number will be
* used as the document version in Elasticsearch. Should the function throw an Exception then
* the batch will fail and the exception propagated. Incompatible with update operations and
* should only be used with withUsePartialUpdate(false)
*
* @param docVersionFn to extract the document version
* @return the {@link DocToBulk} with the function set
*/
public DocToBulk withDocVersionFn(Write.FieldValueExtractFn docVersionFn) {
checkArgument(docVersionFn != null, "docVersionFn must not be null");
return builder().setDocVersionFn(docVersionFn).build();
}
/**
* Provide a function to extract the target operation either upsert or delete from the document
* fields allowing dynamic bulk operation decision. While using withIsDeleteFn, it should be
* taken care that the document's id extraction is defined using the withIdFn function or else
* IllegalArgumentException is thrown. Should the function throw an Exception then the batch
* will fail and the exception propagated.
*
* @param isDeleteFn set to true for deleting the specific document
* @return the {@link Write} with the function set
*/
public DocToBulk withIsDeleteFn(Write.BooleanFieldValueExtractFn isDeleteFn) {
checkArgument(isDeleteFn != null, "deleteFn is required");
return builder().setIsDeleteFn(isDeleteFn).build();
}
/**
* Provide a function to extract the doc version from the document. This version number will be
* used as the document version in Elasticsearch. Should the function throw an Exception then
* the batch will fail and the exception propagated. Incompatible with update operations and
* should only be used with withUsePartialUpdate(false)
*
* @param docVersionType the version type to use, one of {@value VERSION_TYPES}
* @return the {@link DocToBulk} with the doc version type set
*/
public DocToBulk withDocVersionType(String docVersionType) {
checkArgument(
VERSION_TYPES.contains(docVersionType),
"docVersionType must be one of " + "%s",
String.join(", ", VERSION_TYPES));
return builder().setDocVersionType(docVersionType).build();
}
/**
* Use to set explicitly which version of Elasticsearch the destination cluster is running.
* Providing this hint means there is no need for setting {@link
* DocToBulk
*
* <p>Note: if the value of @param backendVersion differs from the version the destination
* cluster is running, behavior is undefined and likely to yield errors.
*
* @param backendVersion the major version number of the version of Elasticsearch being run in
*     the cluster where documents will be indexed.
* @return the {@link DocToBulk} with the Elasticsearch major version number set
*/
public DocToBulk withBackendVersion(int backendVersion) {
checkArgument(
VALID_CLUSTER_VERSIONS.contains(backendVersion),
"Backend version may only be one of " + "%s",
String.join(", ", VERSION_TYPES));
return builder().setBackendVersion(backendVersion).build();
}
@Override
public PCollection<String> expand(PCollection<String> docs) {
ConnectionConfiguration connectionConfiguration = getConnectionConfiguration();
Integer backendVersion = getBackendVersion();
Write.FieldValueExtractFn idFn = getIdFn();
Write.BooleanFieldValueExtractFn isDeleteFn = getIsDeleteFn();
checkState(
(backendVersion != null || connectionConfiguration != null),
"withBackendVersion() or withConnectionConfiguration() is required");
checkArgument(
isDeleteFn == null || idFn != null,
"Id needs to be specified by withIdFn for delete operation");
return docs.apply(ParDo.of(new DocToBulkFn(this)));
}
private static class DocumentMetadata implements Serializable {
final String index;
final String type;
final String id;
final Integer retryOnConflict;
final String routing;
final Integer backendVersion;
final String version;
final String versionType;
DocumentMetadata(
String index,
String type,
String id,
Integer retryOnConflict,
String routing,
Integer backendVersion,
String version,
String versionType) {
this.index = index;
this.id = id;
this.type = type;
this.retryOnConflict = retryOnConflict;
this.routing = routing;
this.backendVersion = backendVersion;
this.version = version;
this.versionType = versionType;
}
}
private static class DocumentMetadataSerializer extends StdSerializer<DocumentMetadata> {
private DocumentMetadataSerializer() {
super(DocumentMetadata.class);
}
@Override
public void serialize(DocumentMetadata value, JsonGenerator gen, SerializerProvider provider)
throws IOException {
gen.writeStartObject();
if (value.index != null) {
gen.writeStringField("_index", value.index);
}
if (value.type != null) {
gen.writeStringField("_type", value.type);
}
if (value.id != null) {
gen.writeStringField("_id", value.id);
}
if (value.routing != null) {
gen.writeStringField("routing", value.routing);
}
if (value.retryOnConflict != null && value.backendVersion <= 6) {
gen.writeNumberField("_retry_on_conflict", value.retryOnConflict);
}
if (value.retryOnConflict != null && value.backendVersion >= 7) {
gen.writeNumberField("retry_on_conflict", value.retryOnConflict);
}
if (value.version != null) {
gen.writeStringField("version", value.version);
}
if (value.versionType != null) {
gen.writeStringField("version_type", value.versionType);
}
gen.writeEndObject();
}
}
@VisibleForTesting
static String createBulkApiEntity(DocToBulk spec, String document, int backendVersion)
throws IOException {
String documentMetadata = "{}";
boolean isDelete = false;
if (spec.getIndexFn() != null
|| spec.getTypeFn() != null
|| spec.getIdFn() != null
|| spec.getRoutingFn() != null) {
JsonNode parsedDocument = OBJECT_MAPPER.readTree(document);
documentMetadata = getDocumentMetadata(spec, parsedDocument, backendVersion);
if (spec.getIsDeleteFn() != null) {
isDelete = spec.getIsDeleteFn().apply(parsedDocument);
}
}
if (isDelete) {
return String.format("{ \"delete\" : %s }%n", documentMetadata);
} else {
if (spec.getUsePartialUpdate()) {
return String.format(
"{ \"update\" : %s }%n{ \"doc\" : %s, " + "\"doc_as_upsert\" : true }%n",
documentMetadata, document);
} else if (spec.getUpsertScript() != null) {
return String.format(
"{ \"update\" : %s }%n{ \"script\" : {\"source\": \"%s\", "
+ "\"params\": %s}, \"upsert\" : %s, \"scripted_upsert\": true}%n",
documentMetadata, spec.getUpsertScript(), document, document);
} else {
return String.format("{ \"index\" : %s }%n%s%n", documentMetadata, document);
}
}
}
private static String lowerCaseOrNull(String input) {
return input == null ? null : input.toLowerCase();
}
/**
* Extracts the components that comprise the document address from the document using the {@link
* Write.FieldValueExtractFn} configured. This allows any or all of the index, type and document
* id to be controlled on a per document basis. If none are provided then an empty default of
* {@code {}} is returned. Sanitization of the index is performed, automatically lower-casing
* the value as required by Elasticsearch.
*
* @param parsedDocument the json from which the index, type and id may be extracted
* @return the document address as JSON or the default
* @throws IOException if the document cannot be parsed as JSON
*/
private static String getDocumentMetadata(
DocToBulk spec, JsonNode parsedDocument, int backendVersion) throws IOException {
DocumentMetadata metadata =
new DocumentMetadata(
spec.getIndexFn() != null
? lowerCaseOrNull(spec.getIndexFn().apply(parsedDocument))
: null,
spec.getTypeFn() != null ? spec.getTypeFn().apply(parsedDocument) : null,
spec.getIdFn() != null ? spec.getIdFn().apply(parsedDocument) : null,
(spec.getUsePartialUpdate()
|| (spec.getUpsertScript() != null && !spec.getUpsertScript().isEmpty()))
? DEFAULT_RETRY_ON_CONFLICT
: null,
spec.getRoutingFn() != null ? spec.getRoutingFn().apply(parsedDocument) : null,
backendVersion,
spec.getDocVersionFn() != null ? spec.getDocVersionFn().apply(parsedDocument) : null,
spec.getDocVersionType());
return OBJECT_MAPPER.writeValueAsString(metadata);
}
/** {@link DoFn} to for the {@link DocToBulk} transform. */
@VisibleForTesting
static class DocToBulkFn extends DoFn<String, String> {
private final DocToBulk spec;
private int backendVersion;
public DocToBulkFn(DocToBulk spec) {
this.spec = spec;
}
@Setup
public void setup() throws IOException {
if (spec.getBackendVersion() != null) {
backendVersion = spec.getBackendVersion();
} else {
backendVersion = ElasticsearchIO.getBackendVersion(spec.getConnectionConfiguration());
}
}
@ProcessElement
public void processElement(ProcessContext c) throws IOException {
c.output(createBulkApiEntity(spec, c.element(), backendVersion));
}
}
}
/**
* A {@link PTransform} writing data to Elasticsearch.
*
* <p>This {@link PTransform} acts as a convenience wrapper for doing both document to bulk API
* serialization as well as batching those Bulk API entities and writing them to an Elasticsearch
* cluster. This class is effectively a thin proxy for DocToBulk->BulkIO all-in-one for
* convenience and backward compatibility.
*/
public static class Write extends PTransform<PCollection<String>, PDone> {
public interface FieldValueExtractFn extends SerializableFunction<JsonNode, String> {}
public interface BooleanFieldValueExtractFn extends SerializableFunction<JsonNode, Boolean> {}
private DocToBulk docToBulk =
new AutoValue_ElasticsearchIO_DocToBulk.Builder()
.setUsePartialUpdate(false)
.build();
private BulkIO bulkIO =
new AutoValue_ElasticsearchIO_BulkIO.Builder()
.setMaxBatchSize(1000L)
.setMaxBatchSizeBytes(5L * 1024L * 1024L)
.setUseStatefulBatches(false)
.setMaxParallelRequestsPerWindow(1)
.build();
public DocToBulk getDocToBulk() {
return docToBulk;
}
public BulkIO getBulkIO() {
return bulkIO;
}
/** Refer to {@link DocToBulk
public Write withIdFn(FieldValueExtractFn idFn) {
docToBulk = docToBulk.withIdFn(idFn);
return this;
}
/** Refer to {@link DocToBulk
public Write withIndexFn(FieldValueExtractFn indexFn) {
docToBulk = docToBulk.withIndexFn(indexFn);
return this;
}
/** Refer to {@link DocToBulk
public Write withRoutingFn(FieldValueExtractFn routingFn) {
docToBulk = docToBulk.withRoutingFn(routingFn);
return this;
}
/** Refer to {@link DocToBulk
public Write withTypeFn(FieldValueExtractFn typeFn) {
docToBulk = docToBulk.withTypeFn(typeFn);
return this;
}
/** Refer to {@link DocToBulk
public Write withDocVersionFn(FieldValueExtractFn docVersionFn) {
docToBulk = docToBulk.withDocVersionFn(docVersionFn);
return this;
}
/** Refer to {@link DocToBulk
public Write withDocVersionType(String docVersionType) {
docToBulk = docToBulk.withDocVersionType(docVersionType);
return this;
}
/** Refer to {@link DocToBulk
public Write withUsePartialUpdate(boolean usePartialUpdate) {
docToBulk = docToBulk.withUsePartialUpdate(usePartialUpdate);
return this;
}
/** Refer to {@link DocToBulk
public Write withUpsertScript(String source) {
docToBulk = docToBulk.withUpsertScript(source);
return this;
}
/** Refer to {@link DocToBulk
public Write withBackendVersion(int backendVersion) {
docToBulk = docToBulk.withBackendVersion(backendVersion);
return this;
}
/** Refer to {@link DocToBulk
public Write withIsDeleteFn(Write.BooleanFieldValueExtractFn isDeleteFn) {
docToBulk = docToBulk.withIsDeleteFn(isDeleteFn);
return this;
}
/** Refer to {@link BulkIO
public Write withConnectionConfiguration(ConnectionConfiguration connectionConfiguration) {
checkArgument(connectionConfiguration != null, "connectionConfiguration can not be null");
docToBulk = docToBulk.withConnectionConfiguration(connectionConfiguration);
bulkIO = bulkIO.withConnectionConfiguration(connectionConfiguration);
return this;
}
/** Refer to {@link BulkIO
public Write withMaxBatchSize(long batchSize) {
bulkIO = bulkIO.withMaxBatchSize(batchSize);
return this;
}
/** Refer to {@link BulkIO
public Write withMaxBatchSizeBytes(long batchSizeBytes) {
bulkIO = bulkIO.withMaxBatchSizeBytes(batchSizeBytes);
return this;
}
/** Refer to {@link BulkIO
public Write withRetryConfiguration(RetryConfiguration retryConfiguration) {
bulkIO = bulkIO.withRetryConfiguration(retryConfiguration);
return this;
}
/** Refer to {@link BulkIO
public Write withIgnoreVersionConflicts(boolean ignoreVersionConflicts) {
bulkIO = bulkIO.withIgnoreVersionConflicts(ignoreVersionConflicts);
return this;
}
/** Refer to {@link BulkIO
public Write withUseStatefulBatches(boolean useStatefulBatches) {
bulkIO = bulkIO.withUseStatefulBatches(useStatefulBatches);
return this;
}
/** Refer to {@link BulkIO
public Write withMaxBufferingDuration(Duration maxBufferingDuration) {
bulkIO = bulkIO.withMaxBufferingDuration(maxBufferingDuration);
return this;
}
/** Refer to {@link BulkIO
public Write withMaxParallelRequestsPerWindow(int maxParallelRequestsPerWindow) {
bulkIO = bulkIO.withMaxParallelRequestsPerWindow(maxParallelRequestsPerWindow);
return this;
}
/** Refer to {@link BulkIO
public Write withAllowableResponseErrors(@Nullable Set<String> allowableResponseErrors) {
if (allowableResponseErrors == null) {
allowableResponseErrors = new HashSet<>();
}
bulkIO = bulkIO.withAllowableResponseErrors(allowableResponseErrors);
return this;
}
@Override
public PDone expand(PCollection<String> input) {
return input.apply(docToBulk).apply(bulkIO);
}
}
/**
* A {@link PTransform} writing Bulk API entities created by {@link ElasticsearchIO.DocToBulk} to
* an Elasticsearch cluster. Typically, using {@link ElasticsearchIO.Write} is preferred, whereas
* using {@link ElasticsearchIO.DocToBulk} and BulkIO separately is for advanced use cases such as
* mirroring data to multiple clusters or data lakes without recomputation.
*/
@AutoValue
public abstract static class BulkIO extends PTransform<PCollection<String>, PDone> {
@VisibleForTesting
static final String RETRY_ATTEMPT_LOG = "Error writing to Elasticsearch. Retry attempt[%d]";
@VisibleForTesting
static final String RETRY_FAILED_LOG =
"Error writing to ES after %d attempt(s). No more attempts allowed";
abstract @Nullable ConnectionConfiguration getConnectionConfiguration();
abstract long getMaxBatchSize();
abstract long getMaxBatchSizeBytes();
abstract @Nullable Duration getMaxBufferingDuration();
abstract boolean getUseStatefulBatches();
abstract int getMaxParallelRequestsPerWindow();
abstract @Nullable RetryConfiguration getRetryConfiguration();
abstract @Nullable Set<String> getAllowedResponseErrors();
abstract Builder builder();
@AutoValue.Builder
abstract static class Builder {
abstract Builder setConnectionConfiguration(ConnectionConfiguration connectionConfiguration);
abstract Builder setMaxBatchSize(long maxBatchSize);
abstract Builder setMaxBatchSizeBytes(long maxBatchSizeBytes);
abstract Builder setRetryConfiguration(RetryConfiguration retryConfiguration);
abstract Builder setAllowedResponseErrors(Set<String> allowedResponseErrors);
abstract Builder setMaxBufferingDuration(Duration maxBufferingDuration);
abstract Builder setUseStatefulBatches(boolean useStatefulBatches);
abstract Builder setMaxParallelRequestsPerWindow(int maxParallelRequestsPerWindow);
abstract BulkIO build();
}
/**
* Provide the Elasticsearch connection configuration object.
*
* @param connectionConfiguration the Elasticsearch {@link ConnectionConfiguration} object
* @return the {@link BulkIO} with connection configuration set
*/
public BulkIO withConnectionConfiguration(ConnectionConfiguration connectionConfiguration) {
checkArgument(connectionConfiguration != null, "connectionConfiguration can not be null");
return builder().setConnectionConfiguration(connectionConfiguration).build();
}
/**
* Provide a maximum size in number of documents for the batch see bulk API
* (https:
* docs (like Elasticsearch bulk size advice). See
* https:
* execution engine, size of bundles may vary, this sets the maximum size. Change this if you
* need to have smaller ElasticSearch bulks.
*
* @param batchSize maximum batch size in number of documents
* @return the {@link BulkIO} with connection batch size set
*/
public BulkIO withMaxBatchSize(long batchSize) {
checkArgument(batchSize > 0, "batchSize must be > 0, but was %s", batchSize);
return builder().setMaxBatchSize(batchSize).build();
}
/**
* Provide a maximum size in bytes for the batch see bulk API
* (https:
* (like Elasticsearch bulk size advice). See
* https:
* execution engine, size of bundles may vary, this sets the maximum size. Change this if you
* need to have smaller ElasticSearch bulks.
*
* @param batchSizeBytes maximum batch size in bytes
* @return the {@link BulkIO} with connection batch size in bytes set
*/
public BulkIO withMaxBatchSizeBytes(long batchSizeBytes) {
checkArgument(batchSizeBytes > 0, "batchSizeBytes must be > 0, but was %s", batchSizeBytes);
return builder().setMaxBatchSizeBytes(batchSizeBytes).build();
}
/**
* Provides configuration to retry a failed batch call to Elasticsearch. A batch is considered
* as failed if the underlying {@link RestClient} surfaces 429 HTTP status code as error for one
* or more of the items in the {@link Response}. Users should consider that retrying might
* compound the underlying problem which caused the initial failure. Users should also be aware
* that once retrying is exhausted the error is surfaced to the runner which <em>may</em> then
* opt to retry the current bundle in entirety or abort if the max number of retries of the
* runner is completed. Retrying uses an exponential backoff algorithm, with minimum backoff of
* 5 seconds and then surfacing the error once the maximum number of retries or maximum
* configuration duration is exceeded.
*
* <p>Example use:
*
* <pre>{@code
* ElasticsearchIO.write()
*   .withRetryConfiguration(ElasticsearchIO.RetryConfiguration.create(10, Duration.standardMinutes(3))
*   ...
* }</pre>
*
* @param retryConfiguration the rules which govern the retry behavior
* @return the {@link BulkIO} with retrying configured
*/
public BulkIO withRetryConfiguration(RetryConfiguration retryConfiguration) {
checkArgument(retryConfiguration != null, "retryConfiguration is required");
return builder().setRetryConfiguration(retryConfiguration).build();
}
/**
* Whether or not to suppress version conflict errors in a Bulk API response. This can be useful
* if your use case involves using external version types.
*
* @param ignoreVersionConflicts true to suppress version conflicts, false to surface version
*     conflict errors.
* @return the {@link BulkIO} with version conflict handling configured
*/
public BulkIO withIgnoreVersionConflicts(boolean ignoreVersionConflicts) {
Set<String> allowedResponseErrors = getAllowedResponseErrors();
if (allowedResponseErrors == null) {
allowedResponseErrors = new HashSet<>();
}
if (ignoreVersionConflicts) {
allowedResponseErrors.add(VERSION_CONFLICT_ERROR);
}
return builder().setAllowedResponseErrors(allowedResponseErrors).build();
}
/**
* Provide a set of textual error types which can be contained in Bulk API response
* items[].error.type field. Any element in @param allowableResponseErrorTypes will suppress
* errors of the same type in Bulk responses.
*
* <p>See also
* https:
*
* @param allowableResponseErrorTypes
* @return the {@link BulkIO} with allowable response errors set
*/
public BulkIO withAllowableResponseErrors(@Nullable Set<String> allowableResponseErrorTypes) {
if (allowableResponseErrorTypes == null) {
allowableResponseErrorTypes = new HashSet<>();
}
return builder().setAllowedResponseErrors(allowableResponseErrorTypes).build();
}
/**
* If using {@link BulkIO
* time before buffered elements are emitted to Elasticsearch as a Bulk API request. If this
* config is not set, Bulk requests will not be issued until {@link BulkIO
* number of documents have been buffered. This may result in higher latency in particular if
* your max batch size is set to a large value and your pipeline input is low volume.
*
* @param maxBufferingDuration the maximum duration to wait before sending any buffered
*     documents to Elasticsearch, regardless of maxBatchSize.
* @return the {@link BulkIO} with maximum buffering duration set
*/
public BulkIO withMaxBufferingDuration(Duration maxBufferingDuration) {
LOG.warn(
"Use of withMaxBufferingDuration requires withUseStatefulBatches(true). "
+ "Setting that automatically.");
return builder()
.setUseStatefulBatches(true)
.setMaxBufferingDuration(maxBufferingDuration)
.build();
}
/**
* Whether or not to use Stateful Processing to ensure bulk requests have the desired number of
* entities i.e. as close to the maxBatchSize as possible. By default without this feature
* enabled, Bulk requests will not contain more than maxBatchSize entities, but the lower bound
* of batch size is determined by Beam Runner bundle sizes, which may be as few as 1.
*
* @param useStatefulBatches true enables the use of Stateful Processing to ensure that batches
*     are as close to the maxBatchSize as possible.
* @return the {@link BulkIO} with Stateful Processing enabled or disabled
*/
public BulkIO withUseStatefulBatches(boolean useStatefulBatches) {
return builder().setUseStatefulBatches(useStatefulBatches).build();
}
/**
* When using {@link BulkIO
* batches are maintained per-key-per-window. BE AWARE that low values for @param
* maxParallelRequestsPerWindow, in particular if the input data has a finite number of windows,
* can reduce parallelism greatly. If data is globally windowed and @param
* maxParallelRequestsPerWindow is set to 1,there will only ever be 1 request in flight. Having
* only a single request in flight can be beneficial for ensuring an Elasticsearch cluster is
* not overwhelmed by parallel requests,but may not work for all use cases. If this number is
* less than the number of maximum workers in your pipeline, the IO work will result in a
* sub-distribution of the last write step with most of the runners.
*
* @param maxParallelRequestsPerWindow the maximum number of parallel bulk requests for a window
*     of data
* @return the {@link BulkIO} with maximum parallel bulk requests per window set
*/
public BulkIO withMaxParallelRequestsPerWindow(int maxParallelRequestsPerWindow) {
checkArgument(
maxParallelRequestsPerWindow > 0, "parameter value must be positive " + "a integer");
return builder().setMaxParallelRequestsPerWindow(maxParallelRequestsPerWindow).build();
}
/**
* Creates batches of documents using Stateful Processing based on user configurable settings of
* withMaxBufferingDuration and withMaxParallelRequestsPerWindow.
*
* <p>Mostly exists for testability of withMaxParallelRequestsPerWindow.
*/
@VisibleForTesting
static class StatefulBatching
extends PTransform<PCollection<String>, PCollection<KV<Integer, Iterable<String>>>> {
final BulkIO spec;
private StatefulBatching(BulkIO bulkSpec) {
spec = bulkSpec;
}
public static StatefulBatching fromSpec(BulkIO spec) {
return new StatefulBatching(spec);
}
@Override
public PCollection<KV<Integer, Iterable<String>>> expand(PCollection<String> input) {
GroupIntoBatches<Integer, String> groupIntoBatches =
GroupIntoBatches.ofSize(spec.getMaxBatchSize());
if (spec.getMaxBufferingDuration() != null) {
groupIntoBatches =
groupIntoBatches.withMaxBufferingDuration(spec.getMaxBufferingDuration());
}
return input
.apply(ParDo.of(new Reshuffle.AssignShardFn<>(spec.getMaxParallelRequestsPerWindow())))
.apply(groupIntoBatches);
}
}
@Override
public PDone expand(PCollection<String> input) {
ConnectionConfiguration connectionConfiguration = getConnectionConfiguration();
checkState(connectionConfiguration != null, "withConnectionConfiguration() is required");
if (getUseStatefulBatches()) {
input.apply(StatefulBatching.fromSpec(this)).apply(ParDo.of(new BulkIOStatefulFn(this)));
} else {
input.apply(ParDo.of(new BulkIOBundleFn(this)));
}
return PDone.in(input.getPipeline());
}
static class BulkIOBundleFn extends BulkIOBaseFn<String> {
@VisibleForTesting
BulkIOBundleFn(BulkIO bulkSpec) {
super(bulkSpec);
}
@ProcessElement
public void processElement(ProcessContext context) throws Exception {
String bulkApiEntity = context.element();
addAndMaybeFlush(bulkApiEntity);
}
}
/*
Intended for use in conjunction with {@link GroupIntoBatches}
*/
static class BulkIOStatefulFn extends BulkIOBaseFn<KV<Integer, Iterable<String>>> {
@VisibleForTesting
BulkIOStatefulFn(BulkIO bulkSpec) {
super(bulkSpec);
}
@ProcessElement
public void processElement(ProcessContext context) throws Exception {
Iterable<String> bulkApiEntities = context.element().getValue();
for (String bulkApiEntity : bulkApiEntities) {
addAndMaybeFlush(bulkApiEntity);
}
}
}
/** {@link DoFn} to for the {@link BulkIO} transform. */
@VisibleForTesting
private abstract static class BulkIOBaseFn<T> extends DoFn<T, Void> {
private static final Duration RETRY_INITIAL_BACKOFF = Duration.standardSeconds(5);
private transient FluentBackoff retryBackoff;
private BulkIO spec;
private transient RestClient restClient;
private ArrayList<String> batch;
long currentBatchSizeBytes;
protected BulkIOBaseFn(BulkIO bulkSpec) {
this.spec = bulkSpec;
}
@Setup
public void setup() throws IOException {
ConnectionConfiguration connectionConfiguration = spec.getConnectionConfiguration();
restClient = connectionConfiguration.createClient();
retryBackoff =
FluentBackoff.DEFAULT.withMaxRetries(0).withInitialBackoff(RETRY_INITIAL_BACKOFF);
if (spec.getRetryConfiguration() != null) {
retryBackoff =
FluentBackoff.DEFAULT
.withInitialBackoff(RETRY_INITIAL_BACKOFF)
.withMaxRetries(spec.getRetryConfiguration().getMaxAttempts() - 1)
.withMaxCumulativeBackoff(spec.getRetryConfiguration().getMaxDuration());
}
}
@StartBundle
public void startBundle(StartBundleContext context) {
batch = new ArrayList<>();
currentBatchSizeBytes = 0;
}
@FinishBundle
public void finishBundle(FinishBundleContext context)
throws IOException, InterruptedException {
flushBatch();
}
protected void addAndMaybeFlush(String bulkApiEntity)
throws IOException, InterruptedException {
batch.add(bulkApiEntity);
currentBatchSizeBytes += bulkApiEntity.getBytes(StandardCharsets.UTF_8).length;
if (batch.size() >= spec.getMaxBatchSize()
|| currentBatchSizeBytes >= spec.getMaxBatchSizeBytes()) {
flushBatch();
}
}
private boolean isRetryableClientException(Throwable t) {
return t.getCause() instanceof ConnectTimeoutException
|| t.getCause() instanceof SocketTimeoutException
|| t.getCause() instanceof ConnectionClosedException
|| t.getCause() instanceof ConnectException;
}
private void flushBatch() throws IOException, InterruptedException {
if (batch.isEmpty()) {
return;
}
LOG.info(
"ElasticsearchIO batch size: {}, batch size bytes: {}",
batch.size(),
currentBatchSizeBytes);
StringBuilder bulkRequest = new StringBuilder();
for (String json : batch) {
bulkRequest.append(json);
}
batch.clear();
currentBatchSizeBytes = 0L;
Response response = null;
HttpEntity responseEntity = null;
String endPoint = spec.getConnectionConfiguration().getBulkEndPoint();
HttpEntity requestBody =
new NStringEntity(bulkRequest.toString(), ContentType.APPLICATION_JSON);
try {
Request request = new Request("POST", endPoint);
request.addParameters(Collections.emptyMap());
request.setEntity(requestBody);
response = restClient.performRequest(request);
responseEntity = new BufferedHttpEntity(response.getEntity());
} catch (java.io.IOException ex) {
if (spec.getRetryConfiguration() == null || !isRetryableClientException(ex)) {
throw ex;
}
LOG.error("Caught ES timeout, retrying", ex);
}
if (spec.getRetryConfiguration() != null
&& (response == null
|| responseEntity == null
|| spec.getRetryConfiguration().getRetryPredicate().test(responseEntity))) {
if (responseEntity != null
&& spec.getRetryConfiguration().getRetryPredicate().test(responseEntity)) {
LOG.warn("ES Cluster is responding with HTP 429 - TOO_MANY_REQUESTS.");
}
responseEntity = handleRetry("POST", endPoint, Collections.emptyMap(), requestBody);
}
checkForErrors(responseEntity, spec.getAllowedResponseErrors());
}
/** retry request based on retry configuration policy. */
private HttpEntity handleRetry(
String method, String endpoint, Map<String, String> params, HttpEntity requestBody)
throws IOException, InterruptedException {
Response response;
HttpEntity responseEntity = null;
Sleeper sleeper = Sleeper.DEFAULT;
BackOff backoff = retryBackoff.backoff();
int attempt = 0;
while (BackOffUtils.next(sleeper, backoff)) {
LOG.warn(String.format(RETRY_ATTEMPT_LOG, ++attempt));
try {
Request request = new Request(method, endpoint);
request.addParameters(params);
request.setEntity(requestBody);
response = restClient.performRequest(request);
responseEntity = new BufferedHttpEntity(response.getEntity());
} catch (java.io.IOException ex) {
if (isRetryableClientException(ex)) {
LOG.error("Caught ES timeout, retrying", ex);
continue;
}
}
if (!Objects.requireNonNull(spec.getRetryConfiguration())
.getRetryPredicate()
.test(responseEntity)) {
return responseEntity;
} else {
LOG.warn("ES Cluster is responding with HTP 429 - TOO_MANY_REQUESTS.");
}
}
throw new IOException(String.format(RETRY_FAILED_LOG, attempt));
}
@Teardown
public void closeClient() throws IOException {
if (restClient != null) {
restClient.close();
}
}
}
}
static int getBackendVersion(ConnectionConfiguration connectionConfiguration) {
try (RestClient restClient = connectionConfiguration.createClient()) {
Request request = new Request("GET", "");
Response response = restClient.performRequest(request);
JsonNode jsonNode = parseResponse(response.getEntity());
int backendVersion =
Integer.parseInt(jsonNode.path("version").path("number").asText().substring(0, 1));
checkArgument(
(VALID_CLUSTER_VERSIONS.contains(backendVersion)),
"The Elasticsearch version to connect to is %s.x. "
+ "This version of the ElasticsearchIO is only compatible with "
+ "Elasticsearch v7.x, v6.x, v5.x and v2.x",
backendVersion);
return backendVersion;
} catch (IOException e) {
throw new IllegalArgumentException("Cannot get Elasticsearch version", e);
}
}
}