Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
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
Original file line number Diff line number Diff line change
Expand Up @@ -17,28 +17,19 @@
import io.micrometer.core.instrument.Counter;
import io.micrometer.core.instrument.DistributionSummary;
import io.micrometer.core.instrument.Timer;
import org.apache.commons.lang3.StringUtils;
import org.opensearch.dataprepper.http.BaseHttpService;
import org.opensearch.dataprepper.http.codec.MultiLineJsonCodec;
import org.opensearch.dataprepper.metrics.PluginMetrics;
import org.opensearch.dataprepper.model.buffer.Buffer;
import org.opensearch.dataprepper.model.event.Event;
import org.opensearch.dataprepper.model.event.EventType;
import org.opensearch.dataprepper.model.event.JacksonEvent;
import org.opensearch.dataprepper.model.opensearch.OpenSearchBulkActions;
import org.opensearch.dataprepper.model.record.Record;
import org.opensearch.dataprepper.plugins.source.opensearchapi.model.BulkAPIEventMetadataKeyAttributes;
import org.opensearch.dataprepper.plugins.source.opensearchapi.model.BulkAPIRequestParams;
import org.opensearch.dataprepper.plugins.source.opensearchapi.model.BulkActionAndMetadataObject;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.io.ByteArrayInputStream;
import java.time.Instant;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Iterator;
import java.util.List;
import java.util.Map;

/*
* OpenSearch API Service class is responsible for handling bulk API requests.
Expand All @@ -56,8 +47,7 @@ public class OpenSearchAPIService implements BaseHttpService {

private static final Logger LOG = LoggerFactory.getLogger(OpenSearchAPIService.class);

// TODO: support other data-types as request body, e.g. json_lines, msgpack
private final MultiLineJsonCodec jsonCodec = new MultiLineJsonCodec();
private final OpenSearchBulkByteDecoder bulkByteDecoder = new OpenSearchBulkByteDecoder();
private final Buffer<Record<Event>> buffer;
private final int bufferWriteTimeoutInMillis;
private final Counter requestsReceivedCounter;
Expand Down Expand Up @@ -110,21 +100,13 @@ private HttpResponse processBulkRequest(final ServiceRequestContext serviceReque
}

final HttpData content = aggregatedHttpRequest.content();
List<Map<String, Object>> bulkRequestPayloadList;

// parse the request payload
try {
bulkRequestPayloadList = jsonCodec.parse(content);
} catch (IOException e) {
LOG.error("Failed to parse the request of size {} due to: {}", content.length(), e.getMessage());
throw new IOException("Bad request data format.", e.getCause());
}

try {
if (buffer.isByteBuffer()) {
buffer.writeBytes(content.array(), null, bufferWriteTimeoutInMillis);
} else {
List<Record<Event>> records = generateEventsFromBulkRequest(bulkRequestPayloadList, bulkAPIRequestParams);
List<Record<Event>> records = new ArrayList<>();
bulkByteDecoder.parse(new ByteArrayInputStream(content.array()), Instant.now(), bulkAPIRequestParams, records::add);
buffer.writeAll(records, bufferWriteTimeoutInMillis);
}
} catch (Exception e) {
Expand All @@ -134,72 +116,4 @@ private HttpResponse processBulkRequest(final ServiceRequestContext serviceReque
successRequestsCounter.increment();
return HttpResponse.of(HttpStatus.OK);
}

private boolean isValidBulkAction(Map<String, Object> actionMap) {
return Arrays.stream(OpenSearchBulkActions.values())
.anyMatch(bulkAction -> actionMap.containsKey(bulkAction.toString()));
}

private List<Record<Event>> generateEventsFromBulkRequest(final List<Map<String, Object>> bulkRequestPayloadList, final BulkAPIRequestParams bulkAPIRequestParams) throws Exception {
List<Record<Event>> records = new ArrayList<>();
Iterator<Map<String, Object>> bulkRequestPayloadListIterator = bulkRequestPayloadList.iterator();

while (bulkRequestPayloadListIterator.hasNext()) {
Map<String, Object> actionMetadataRow = bulkRequestPayloadListIterator.next();
if (!isValidBulkAction(actionMetadataRow)) {
throw new IOException("Invalid request data.");
}

BulkActionAndMetadataObject bulkActionAndMetadataObject = new BulkActionAndMetadataObject(actionMetadataRow);
final boolean isDeleteAction = bulkActionAndMetadataObject.getAction().equals(OpenSearchBulkActions.DELETE.toString());
Map<String, Object> documentDataObject = null;
if (!isDeleteAction) {
if (!bulkRequestPayloadListIterator.hasNext()) {
throw new IOException("Invalid request data.");
}
documentDataObject = bulkRequestPayloadListIterator.next();
// Performing another validation check to make sure that the doc row is not a valid action row
if (isValidBulkAction(documentDataObject)) {
throw new IOException("Invalid request data.");
}
}
final JacksonEvent event = createBulkRequestActionEvent(bulkActionAndMetadataObject, bulkAPIRequestParams, documentDataObject);
records.add(new Record<>(event));
}

return records;
}

private JacksonEvent createBulkRequestActionEvent(
final BulkActionAndMetadataObject bulkActionAndMetadataObject,
final BulkAPIRequestParams bulkAPIRequestParams, Map<String, Object> optionalDocumentData) {

final JacksonEvent.Builder eventBuilder = JacksonEvent.builder().withEventType(EventType.DOCUMENT.toString());
if (optionalDocumentData != null) {
eventBuilder.withData(optionalDocumentData);
}
final JacksonEvent event = eventBuilder.build();

final String index = !StringUtils.isEmpty(bulkAPIRequestParams.getIndex()) ? bulkAPIRequestParams.getIndex() : bulkActionAndMetadataObject.getIndex();

event.getMetadata().setAttribute(BulkAPIEventMetadataKeyAttributes.BULK_API_EVENT_METADATA_ATTRIBUTE_ACTION, bulkActionAndMetadataObject.getAction());
event.getMetadata().setAttribute(BulkAPIEventMetadataKeyAttributes.BULK_API_EVENT_METADATA_ATTRIBUTE_INDEX, index);

String docId = bulkActionAndMetadataObject.getDocId();
if (!StringUtils.isBlank(docId) && !StringUtils.isEmpty(docId)) {
event.getMetadata().setAttribute(BulkAPIEventMetadataKeyAttributes.BULK_API_EVENT_METADATA_ATTRIBUTE_ID, docId);
}

String pipeline = bulkAPIRequestParams.getPipeline();
if (!StringUtils.isBlank(pipeline) && !StringUtils.isEmpty(pipeline)) {
event.getMetadata().setAttribute(BulkAPIEventMetadataKeyAttributes.BULK_API_EVENT_METADATA_ATTRIBUTE_PIPELINE, pipeline);
}

String routing = bulkAPIRequestParams.getRouting();
if (!StringUtils.isBlank(routing) && !StringUtils.isEmpty(routing)) {
event.getMetadata().setAttribute(BulkAPIEventMetadataKeyAttributes.BULK_API_EVENT_METADATA_ATTRIBUTE_ROUTING, routing);
}

return event;
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@
import org.opensearch.dataprepper.model.annotations.DataPrepperPlugin;
import org.opensearch.dataprepper.model.annotations.DataPrepperPluginConstructor;
import org.opensearch.dataprepper.model.buffer.Buffer;
import org.opensearch.dataprepper.model.codec.ByteDecoder;
import org.opensearch.dataprepper.model.configuration.PipelineDescription;
import org.opensearch.dataprepper.model.event.Event;
import org.opensearch.dataprepper.model.plugin.PluginFactory;
Expand Down Expand Up @@ -38,4 +39,9 @@ public BaseHttpService getHttpService(final int bufferWriteTimeoutInMillis, fina
public String getHttpHealthCheckPath() {
return HTTP_HEALTH_CHECK_PATH;
}

@Override
public ByteDecoder getDecoder() {
return new OpenSearchBulkByteDecoder();
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,137 @@
/*
* Copyright OpenSearch Contributors
* SPDX-License-Identifier: Apache-2.0
*
* The OpenSearch Contributors require contributions made to
* this file be licensed under the Apache-2.0 license or a
* compatible open source license.
*
*/

package org.opensearch.dataprepper.plugins.source.opensearchapi;

import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.opensearch.dataprepper.model.codec.ByteDecoder;
import org.opensearch.dataprepper.model.event.Event;
import org.opensearch.dataprepper.model.event.EventType;
import org.opensearch.dataprepper.model.event.JacksonEvent;
import org.opensearch.dataprepper.model.opensearch.OpenSearchBulkActions;
import org.opensearch.dataprepper.model.record.Record;
import org.opensearch.dataprepper.plugins.source.opensearchapi.model.BulkAPIEventMetadataKeyAttributes;
import org.opensearch.dataprepper.plugins.source.opensearchapi.model.BulkAPIRequestParams;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.nio.charset.StandardCharsets;
import java.time.Instant;
import java.util.Arrays;
import java.util.Map;
import java.util.function.Consumer;

public class OpenSearchBulkByteDecoder implements ByteDecoder {

private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();
private static final TypeReference<Map<String, Object>> MAP_TYPE = new TypeReference<>() {};

@Override
public void parse(InputStream inputStream, Instant timeReceived, Consumer<Record<Event>> eventConsumer) throws IOException {
parse(inputStream, timeReceived, null, eventConsumer);
}

public void parse(InputStream inputStream, Instant timeReceived, BulkAPIRequestParams requestParams, Consumer<Record<Event>> eventConsumer) throws IOException {
final BufferedReader reader = new BufferedReader(new InputStreamReader(inputStream, StandardCharsets.UTF_8));

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

Wrap the reader in try-with-resources

String line;
while ((line = reader.readLine()) != null) {

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

This code should be the same as the existing code in OpenSearchAPIService for handling the response for non-byte buffers. The current structure of this class looks reusable. I think we might be able to refactor OpenSearchAPIService to use this in processBulkRequest.

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

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

@dlvenable Refactored OpenSearchAPIService.processBulkRequest to delegate to OpenSearchBulkByteDecoder for the non-byte-buffer path, so there's a single copy of the parsing logic. Let me know if you had a different structure in mind.

if (line.isBlank()) continue;

Map<String, Object> actionLine = OBJECT_MAPPER.readValue(line, MAP_TYPE);
String action = extractAction(actionLine);
if (action == null) continue;

@SuppressWarnings("unchecked")
Map<String, Object> metadata = (Map<String, Object>) actionLine.get(action);
boolean isDelete = OpenSearchBulkActions.DELETE.toString().equals(action);

Map<String, Object> documentData = null;
if (!isDelete) {
String docLine = reader.readLine();
if (docLine == null || docLine.isBlank()) continue;
documentData = OBJECT_MAPPER.readValue(docLine, MAP_TYPE);
}

JacksonEvent.Builder builder = JacksonEvent.builder()
.withEventType(EventType.DOCUMENT.toString());
if (documentData != null) {
builder.withData(documentData);
}
if (timeReceived != null) {
builder.withTimeReceived(timeReceived);
}
JacksonEvent event = builder.build();

event.getMetadata().setAttribute(
BulkAPIEventMetadataKeyAttributes.BULK_API_EVENT_METADATA_ATTRIBUTE_ACTION, action);

// Set index: action line takes priority, then URL-level fallback

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

This is a breaking change, but I think it is good. It is the way OpenSearch works. I'm going to add the breaking change label to this PR because of this. No need to make any changes. I'm just pointing it out.

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

Added the label @dlvenable

String index = getMetadataValue(metadata, "_index");
if (isBlank(index) && requestParams != null && !isBlank(requestParams.getIndex())) {
index = requestParams.getIndex();
}
if (!isBlank(index)) {
event.getMetadata().setAttribute(BulkAPIEventMetadataKeyAttributes.BULK_API_EVENT_METADATA_ATTRIBUTE_INDEX, index);
}

// Set id
setIfPresent(event, metadata, "_id", BulkAPIEventMetadataKeyAttributes.BULK_API_EVENT_METADATA_ATTRIBUTE_ID);

// Set routing: action line takes priority, then URL-level fallback
String routing = getMetadataValue(metadata, "routing");
if (isBlank(routing) && requestParams != null && !isBlank(requestParams.getRouting())) {
routing = requestParams.getRouting();
}
if (!isBlank(routing)) {
event.getMetadata().setAttribute(BulkAPIEventMetadataKeyAttributes.BULK_API_EVENT_METADATA_ATTRIBUTE_ROUTING, routing);
}

// Set pipeline: action line takes priority, then URL-level fallback
String pipeline = getMetadataValue(metadata, "pipeline");
if (isBlank(pipeline) && requestParams != null && !isBlank(requestParams.getPipeline())) {
pipeline = requestParams.getPipeline();
}
if (!isBlank(pipeline)) {
event.getMetadata().setAttribute(BulkAPIEventMetadataKeyAttributes.BULK_API_EVENT_METADATA_ATTRIBUTE_PIPELINE, pipeline);
}

eventConsumer.accept(new Record<>(event));
}
}

private String extractAction(Map<String, Object> actionLine) {
return Arrays.stream(OpenSearchBulkActions.values())
.map(OpenSearchBulkActions::toString)
.filter(actionLine::containsKey)
.findFirst()
.orElse(null);
}

private String getMetadataValue(Map<String, Object> metadata, String key) {
if (metadata == null) return null;
Object value = metadata.get(key);
return value != null ? value.toString() : null;
}

private void setIfPresent(JacksonEvent event, Map<String, Object> metadata, String key, String attribute) {
if (metadata == null) return;
Object value = metadata.get(key);
if (value != null) {
event.getMetadata().setAttribute(attribute, value.toString());
}
}

private boolean isBlank(String value) {
return value == null || value.isBlank();
}
}
Loading
Loading