package org.vertexium.elasticsearch7.bulk; import com.google.common.collect.ImmutableMap; import org.elasticsearch.action.bulk.BulkRequestBuilder; import org.elasticsearch.action.update.UpdateRequest; import org.elasticsearch.action.update.UpdateRequestBuilder; import org.elasticsearch.client.Client; import org.elasticsearch.script.Script; import org.elasticsearch.script.ScriptType; import org.vertexium.ElementLocation; import org.vertexium.VertexiumException; import org.vertexium.VertexiumObjectId; import org.vertexium.elasticsearch7.Elasticsearch7SearchIndex; import java.util.*; import java.util.stream.Collectors; import static org.vertexium.elasticsearch7.bulk.BulkUtils.*; public class BulkUpdateItem extends BulkItem<UpdateItem> { private final ElementLocation sourceElementLocation; private final Map<String, String> source = new HashMap<>(); private final Map<String, Set<Object>> fieldsToSet = new HashMap<>(); private final Set<String> fieldsToRemove = new HashSet<>(); private final Map<String, String> fieldsToRename = new HashMap<>(); private final Set<String> additionalVisibilities = new HashSet<>(); private final Set<String> additionalVisibilitiesToDelete = new HashSet<>(); private boolean updateOnly = true; private Integer size; public BulkUpdateItem( String indexName, String type, String documentId, VertexiumObjectId vertexiumObjectId, ElementLocation sourceElementLocation ) { super(indexName, type, documentId, vertexiumObjectId); this.sourceElementLocation = sourceElementLocation; } @SuppressWarnings({"rawtypes", "unchecked"}) @Override public void add(UpdateItem item) { super.add(item); size = null; for (Map.Entry<String, Object> itemEntry : item.getFieldsToSet().entrySet()) { Object itemValue = itemEntry.getValue(); fieldsToSet.compute(itemEntry.getKey(), (key, existingValue) -> { if (existingValue == null) { if (itemValue instanceof Collection) { return new HashSet<>((Collection<?>) itemValue); } else { Set newValue = new HashSet<>(); newValue.add(itemValue); return newValue; } } else { if (itemValue instanceof Collection) { existingValue.addAll((Collection) itemValue); } else { existingValue.add(itemValue); } return existingValue; } }); } for (Map.Entry<String, String> itemEntry : item.getFieldsToRename().entrySet()) { String itemValue = itemEntry.getValue(); fieldsToRename.compute(itemEntry.getKey(), (key, existingValue) -> { if (existingValue == null) { return itemValue; } else if (existingValue.equals(itemValue)) { return itemValue; } else { throw new VertexiumException("Changing the same property to two different visibilities in the same batch is not allowed: " + itemEntry.getKey()); } }); } source.putAll(item.getSource()); fieldsToRemove.addAll(item.getFieldsToRemove()); additionalVisibilities.addAll(item.getAdditionalVisibilities()); additionalVisibilitiesToDelete.addAll(item.getAdditionalVisibilitiesToDelete()); if (!item.isExistingElement()) { updateOnly = false; } } @SuppressWarnings({"unchecked", "rawtypes"}) @Override public void addToBulkRequest(Client client, BulkRequestBuilder bulkRequestBuilder) { UpdateRequestBuilder updateRequestBuilder = client .prepareUpdate(getIndexName(), getType(), getDocumentId()); if (!updateOnly) { updateRequestBuilder = updateRequestBuilder .setScriptedUpsert(true) .setUpsert((Map<String, Object>) (Map) source); } UpdateRequest updateRequest = updateRequestBuilder .setScript(new Script( ScriptType.STORED, null, "updateFieldsOnDocumentScript", ImmutableMap.of( "fieldsToSet", fieldsToSet.entrySet().stream() .collect(Collectors.toMap( Map.Entry::getKey, entry -> new ArrayList<>(entry.getValue()) )), "fieldsToRemove", new ArrayList<>(fieldsToRemove), "fieldsToRename", fieldsToRename, "additionalVisibilities", new ArrayList<>(additionalVisibilities), "additionalVisibilitiesToDelete", new ArrayList<>(additionalVisibilitiesToDelete) ) )) .setRetryOnConflict(Elasticsearch7SearchIndex.MAX_RETRIES) .request(); bulkRequestBuilder.add(updateRequest); } @Override public int getSize() { if (size == null) { size = getIndexName().length() + getType().length() + getDocumentId().length() + calculateSizeOfId(getVertexiumObjectId()) + calculateSizeOfMap(source) + calculateSizeOfMap(fieldsToSet) + calculateSizeOfCollection(fieldsToRemove) + calculateSizeOfMap(fieldsToRename) + calculateSizeOfCollection(additionalVisibilities) + calculateSizeOfCollection(additionalVisibilitiesToDelete); } return size; } public ElementLocation getSourceElementLocation() { return sourceElementLocation; } }