Examples with WriteRequest org.elasticsearch.action.support.WriteRequest used on opensource projects

Example 1 with WriteRequest

use of org.elasticsearch.action.support.WriteRequest in project snow-owl by b2ihealthcare.

the class RestHighLevelClientExt method bulk.

static Request bulk(BulkRequest bulkRequest) throws IOException {
    // Bulk API only supports newline delimited JSON or Smile. Before executing
    // the bulk, we need to check that all requests have the same content-type
    // and this content-type is supported by the Bulk API.
    XContentType bulkContentType = null;
    String index = null;
    for (int i = 0; i < bulkRequest.numberOfActions(); i++) {
        DocWriteRequest<?> writeRequest = bulkRequest.requests().get(i);
        index = enforceSameIndex(writeRequest.index(), index);
        // Remove index property, as it will be encoded in the request path
        DocWriteRequest.OpType opType = writeRequest.opType();
        switch(opType) {
            // $FALL-THROUGH$
            case INDEX:
            case CREATE:
                bulkContentType = enforceSameContentType((IndexRequest) writeRequest, bulkContentType);
                ((IndexRequest) writeRequest).index(null);
                break;
            case DELETE:
                ((DeleteRequest) writeRequest).index(null);
                break;
            case UPDATE:
                UpdateRequest updateRequest = (UpdateRequest) writeRequest;
                if (updateRequest.doc() != null) {
                    bulkContentType = enforceSameContentType(updateRequest.doc(), bulkContentType);
                }
                if (updateRequest.upsertRequest() != null) {
                    bulkContentType = enforceSameContentType(updateRequest.upsertRequest(), bulkContentType);
                }
                updateRequest.index(null);
                break;
        }
    }
    if (bulkContentType == null) {
        bulkContentType = XContentType.JSON;
    }
    String endpoint = endpoint(index, "_bulk");
    Request request = new Request(HttpPost.METHOD_NAME, endpoint);
    if (bulkRequest.timeout() != null) {
        request.addParameter("timeout", bulkRequest.timeout().getStringRep());
    }
    if (bulkRequest.getRefreshPolicy() != WriteRequest.RefreshPolicy.NONE) {
        request.addParameter("refresh", bulkRequest.getRefreshPolicy().getValue());
    }
    final byte separator = bulkContentType.xContent().streamSeparator();
    final ContentType requestContentType = RequestConverters.createContentType(bulkContentType);
    ByteArrayOutputStream content = new ByteArrayOutputStream();
    for (DocWriteRequest<?> writeRequest : bulkRequest.requests()) {
        DocWriteRequest.OpType opType = writeRequest.opType();
        try (XContentBuilder metadata = XContentBuilder.builder(bulkContentType.xContent())) {
            metadata.startObject();
            {
                metadata.startObject(opType.getLowercase());
                if (Strings.hasLength(writeRequest.index())) {
                    metadata.field("_index", writeRequest.index());
                }
                if (Strings.hasLength(writeRequest.id())) {
                    metadata.field("_id", writeRequest.id());
                }
                if (Strings.hasLength(writeRequest.routing())) {
                    metadata.field("routing", writeRequest.routing());
                }
                if (writeRequest.version() != Versions.MATCH_ANY) {
                    metadata.field("version", writeRequest.version());
                }
                VersionType versionType = writeRequest.versionType();
                if (versionType != VersionType.INTERNAL) {
                    if (versionType == VersionType.EXTERNAL) {
                        metadata.field("version_type", "external");
                    } else if (versionType == VersionType.EXTERNAL_GTE) {
                        metadata.field("version_type", "external_gte");
                    }
                }
                if (opType == DocWriteRequest.OpType.INDEX || opType == DocWriteRequest.OpType.CREATE) {
                    IndexRequest indexRequest = (IndexRequest) writeRequest;
                    if (Strings.hasLength(indexRequest.getPipeline())) {
                        metadata.field("pipeline", indexRequest.getPipeline());
                    }
                } else if (opType == DocWriteRequest.OpType.UPDATE) {
                    UpdateRequest updateRequest = (UpdateRequest) writeRequest;
                    if (updateRequest.retryOnConflict() > 0) {
                        metadata.field("retry_on_conflict", updateRequest.retryOnConflict());
                    }
                    if (updateRequest.fetchSource() != null) {
                        metadata.field("_source", updateRequest.fetchSource());
                    }
                }
                metadata.endObject();
            }
            metadata.endObject();
            BytesRef metadataSource = BytesReference.bytes(metadata).toBytesRef();
            content.write(metadataSource.bytes, metadataSource.offset, metadataSource.length);
            content.write(separator);
        }
        BytesRef source = null;
        if (opType == DocWriteRequest.OpType.INDEX || opType == DocWriteRequest.OpType.CREATE) {
            IndexRequest indexRequest = (IndexRequest) writeRequest;
            BytesReference indexSource = indexRequest.source();
            XContentType indexXContentType = indexRequest.getContentType();
            try (XContentParser parser = XContentHelper.createParser(NamedXContentRegistry.EMPTY, LoggingDeprecationHandler.INSTANCE, indexSource, indexXContentType)) {
                try (XContentBuilder builder = XContentBuilder.builder(bulkContentType.xContent())) {
                    builder.copyCurrentStructure(parser);
                    source = BytesReference.bytes(builder).toBytesRef();
                }
            }
        } else if (opType == DocWriteRequest.OpType.UPDATE) {
            source = XContentHelper.toXContent((UpdateRequest) writeRequest, bulkContentType, false).toBytesRef();
        }
        if (source != null) {
            content.write(source.bytes, source.offset, source.length);
            content.write(separator);
        }
    }
    HttpEntity entity = new ByteArrayEntity(content.toByteArray(), 0, content.size(), requestContentType);
    request.setEntity(entity);
    return request;
}