Browse Source
Extract common code into BulkOperationsSupport. Reorder methods. Add missing verifyComplete to tests. See #2821 Original pull request: #4342pull/4373/head
Mark Paluch
3 years ago
8 changed files with 425 additions and 504 deletions
@ -0,0 +1,221 @@ |
|||||||
|
/* |
||||||
|
* Copyright 2023 the original author or authors. |
||||||
|
* |
||||||
|
* Licensed under the Apache License, Version 2.0 (the "License"); |
||||||
|
* you may not use this file except in compliance with the License. |
||||||
|
* You may obtain a copy of the License at |
||||||
|
* |
||||||
|
* https://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
* |
||||||
|
* Unless required by applicable law or agreed to in writing, software |
||||||
|
* distributed under the License is distributed on an "AS IS" BASIS, |
||||||
|
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
||||||
|
* See the License for the specific language governing permissions and |
||||||
|
* limitations under the License. |
||||||
|
*/ |
||||||
|
package org.springframework.data.mongodb.core; |
||||||
|
|
||||||
|
import java.util.ArrayList; |
||||||
|
import java.util.List; |
||||||
|
import java.util.Optional; |
||||||
|
|
||||||
|
import org.bson.Document; |
||||||
|
import org.bson.conversions.Bson; |
||||||
|
import org.springframework.context.ApplicationEvent; |
||||||
|
import org.springframework.data.mapping.PersistentEntity; |
||||||
|
import org.springframework.data.mongodb.core.BulkOperations.BulkMode; |
||||||
|
import org.springframework.data.mongodb.core.aggregation.AggregationOperationContext; |
||||||
|
import org.springframework.data.mongodb.core.aggregation.AggregationUpdate; |
||||||
|
import org.springframework.data.mongodb.core.aggregation.RelaxedTypeBasedAggregationOperationContext; |
||||||
|
import org.springframework.data.mongodb.core.convert.QueryMapper; |
||||||
|
import org.springframework.data.mongodb.core.convert.UpdateMapper; |
||||||
|
import org.springframework.data.mongodb.core.mapping.MongoPersistentEntity; |
||||||
|
import org.springframework.data.mongodb.core.mapping.event.AfterSaveEvent; |
||||||
|
import org.springframework.data.mongodb.core.mapping.event.BeforeSaveEvent; |
||||||
|
import org.springframework.data.mongodb.core.query.Collation; |
||||||
|
import org.springframework.data.mongodb.core.query.Query; |
||||||
|
import org.springframework.data.mongodb.core.query.Update; |
||||||
|
import org.springframework.data.mongodb.core.query.UpdateDefinition; |
||||||
|
import org.springframework.data.mongodb.core.query.UpdateDefinition.ArrayFilter; |
||||||
|
import org.springframework.util.Assert; |
||||||
|
|
||||||
|
import com.mongodb.client.model.BulkWriteOptions; |
||||||
|
import com.mongodb.client.model.DeleteManyModel; |
||||||
|
import com.mongodb.client.model.DeleteOneModel; |
||||||
|
import com.mongodb.client.model.InsertOneModel; |
||||||
|
import com.mongodb.client.model.ReplaceOneModel; |
||||||
|
import com.mongodb.client.model.UpdateManyModel; |
||||||
|
import com.mongodb.client.model.UpdateOneModel; |
||||||
|
import com.mongodb.client.model.UpdateOptions; |
||||||
|
import com.mongodb.client.model.WriteModel; |
||||||
|
|
||||||
|
/** |
||||||
|
* Support class for bulk operations. |
||||||
|
* |
||||||
|
* @author Mark Paluch |
||||||
|
* @since 4.1 |
||||||
|
*/ |
||||||
|
abstract class BulkOperationsSupport { |
||||||
|
|
||||||
|
private final String collectionName; |
||||||
|
|
||||||
|
BulkOperationsSupport(String collectionName) { |
||||||
|
|
||||||
|
Assert.hasText(collectionName, "CollectionName must not be null nor empty"); |
||||||
|
|
||||||
|
this.collectionName = collectionName; |
||||||
|
} |
||||||
|
|
||||||
|
/** |
||||||
|
* Emit a {@link BeforeSaveEvent}. |
||||||
|
* |
||||||
|
* @param holder |
||||||
|
*/ |
||||||
|
void maybeEmitBeforeSaveEvent(SourceAwareWriteModelHolder holder) { |
||||||
|
|
||||||
|
if (holder.model() instanceof InsertOneModel) { |
||||||
|
|
||||||
|
Document target = ((InsertOneModel<Document>) holder.model()).getDocument(); |
||||||
|
maybeEmitEvent(new BeforeSaveEvent<>(holder.source(), target, collectionName)); |
||||||
|
} else if (holder.model() instanceof ReplaceOneModel) { |
||||||
|
|
||||||
|
Document target = ((ReplaceOneModel<Document>) holder.model()).getReplacement(); |
||||||
|
maybeEmitEvent(new BeforeSaveEvent<>(holder.source(), target, collectionName)); |
||||||
|
} |
||||||
|
} |
||||||
|
|
||||||
|
/** |
||||||
|
* Emit a {@link AfterSaveEvent}. |
||||||
|
* |
||||||
|
* @param holder |
||||||
|
*/ |
||||||
|
void maybeEmitAfterSaveEvent(SourceAwareWriteModelHolder holder) { |
||||||
|
|
||||||
|
if (holder.model() instanceof InsertOneModel) { |
||||||
|
|
||||||
|
Document target = ((InsertOneModel<Document>) holder.model()).getDocument(); |
||||||
|
maybeEmitEvent(new AfterSaveEvent<>(holder.source(), target, collectionName)); |
||||||
|
} else if (holder.model() instanceof ReplaceOneModel) { |
||||||
|
|
||||||
|
Document target = ((ReplaceOneModel<Document>) holder.model()).getReplacement(); |
||||||
|
maybeEmitEvent(new AfterSaveEvent<>(holder.source(), target, collectionName)); |
||||||
|
} |
||||||
|
} |
||||||
|
|
||||||
|
WriteModel<Document> mapWriteModel(Object source, WriteModel<Document> writeModel) { |
||||||
|
|
||||||
|
if (writeModel instanceof UpdateOneModel<Document> model) { |
||||||
|
|
||||||
|
if (source instanceof AggregationUpdate aggregationUpdate) { |
||||||
|
|
||||||
|
List<Document> pipeline = mapUpdatePipeline(aggregationUpdate); |
||||||
|
return new UpdateOneModel<>(getMappedQuery(model.getFilter()), pipeline, model.getOptions()); |
||||||
|
} |
||||||
|
|
||||||
|
return new UpdateOneModel<>(getMappedQuery(model.getFilter()), getMappedUpdate(model.getUpdate()), |
||||||
|
model.getOptions()); |
||||||
|
} |
||||||
|
|
||||||
|
if (writeModel instanceof UpdateManyModel<Document> model) { |
||||||
|
|
||||||
|
if (source instanceof AggregationUpdate aggregationUpdate) { |
||||||
|
|
||||||
|
List<Document> pipeline = mapUpdatePipeline(aggregationUpdate); |
||||||
|
return new UpdateManyModel<>(getMappedQuery(model.getFilter()), pipeline, model.getOptions()); |
||||||
|
} |
||||||
|
|
||||||
|
return new UpdateManyModel<>(getMappedQuery(model.getFilter()), getMappedUpdate(model.getUpdate()), |
||||||
|
model.getOptions()); |
||||||
|
} |
||||||
|
|
||||||
|
if (writeModel instanceof DeleteOneModel<Document> model) { |
||||||
|
return new DeleteOneModel<>(getMappedQuery(model.getFilter()), model.getOptions()); |
||||||
|
} |
||||||
|
|
||||||
|
if (writeModel instanceof DeleteManyModel<Document> model) { |
||||||
|
return new DeleteManyModel<>(getMappedQuery(model.getFilter()), model.getOptions()); |
||||||
|
} |
||||||
|
|
||||||
|
return writeModel; |
||||||
|
} |
||||||
|
|
||||||
|
private List<Document> mapUpdatePipeline(AggregationUpdate source) { |
||||||
|
|
||||||
|
Class<?> type = entity().isPresent() ? entity().map(PersistentEntity::getType).get() : Object.class; |
||||||
|
AggregationOperationContext context = new RelaxedTypeBasedAggregationOperationContext(type, |
||||||
|
updateMapper().getMappingContext(), queryMapper()); |
||||||
|
|
||||||
|
return new AggregationUtil(queryMapper(), queryMapper().getMappingContext()).createPipeline(source, context); |
||||||
|
} |
||||||
|
|
||||||
|
/** |
||||||
|
* Emit a {@link ApplicationEvent} if event multicasting is enabled. |
||||||
|
* |
||||||
|
* @param event |
||||||
|
*/ |
||||||
|
protected abstract void maybeEmitEvent(ApplicationEvent event); |
||||||
|
|
||||||
|
/** |
||||||
|
* @return the {@link UpdateMapper} to use. |
||||||
|
*/ |
||||||
|
protected abstract UpdateMapper updateMapper(); |
||||||
|
|
||||||
|
/** |
||||||
|
* @return the {@link QueryMapper} to use. |
||||||
|
*/ |
||||||
|
protected abstract QueryMapper queryMapper(); |
||||||
|
|
||||||
|
/** |
||||||
|
* @return the associated {@link PersistentEntity}. Can be {@link Optional#empty()}. |
||||||
|
*/ |
||||||
|
protected abstract Optional<? extends MongoPersistentEntity<?>> entity(); |
||||||
|
|
||||||
|
protected Bson getMappedUpdate(Bson update) { |
||||||
|
return updateMapper().getMappedObject(update, entity()); |
||||||
|
} |
||||||
|
|
||||||
|
protected Bson getMappedQuery(Bson query) { |
||||||
|
return queryMapper().getMappedObject(query, entity()); |
||||||
|
} |
||||||
|
|
||||||
|
protected static BulkWriteOptions getBulkWriteOptions(BulkMode bulkMode) { |
||||||
|
|
||||||
|
BulkWriteOptions options = new BulkWriteOptions(); |
||||||
|
|
||||||
|
return switch (bulkMode) { |
||||||
|
case ORDERED -> options.ordered(true); |
||||||
|
case UNORDERED -> options.ordered(false); |
||||||
|
}; |
||||||
|
} |
||||||
|
|
||||||
|
/** |
||||||
|
* @param filterQuery The {@link Query} to read a potential {@link Collation} from. Must not be {@literal null}. |
||||||
|
* @param update The {@link Update} to apply |
||||||
|
* @param upsert flag to indicate if document should be upserted. |
||||||
|
* @return new instance of {@link UpdateOptions}. |
||||||
|
*/ |
||||||
|
protected static UpdateOptions computeUpdateOptions(Query filterQuery, UpdateDefinition update, boolean upsert) { |
||||||
|
|
||||||
|
UpdateOptions options = new UpdateOptions(); |
||||||
|
options.upsert(upsert); |
||||||
|
|
||||||
|
if (update.hasArrayFilters()) { |
||||||
|
List<Document> list = new ArrayList<>(update.getArrayFilters().size()); |
||||||
|
for (ArrayFilter arrayFilter : update.getArrayFilters()) { |
||||||
|
list.add(arrayFilter.asDocument()); |
||||||
|
} |
||||||
|
options.arrayFilters(list); |
||||||
|
} |
||||||
|
|
||||||
|
filterQuery.getCollation().map(Collation::toMongoCollation).ifPresent(options::collation); |
||||||
|
return options; |
||||||
|
} |
||||||
|
|
||||||
|
/** |
||||||
|
* Value object chaining together an actual source with its {@link WriteModel} representation. |
||||||
|
* |
||||||
|
* @author Christoph Strobl |
||||||
|
*/ |
||||||
|
record SourceAwareWriteModelHolder(Object source, WriteModel<Document> model) { |
||||||
|
} |
||||||
|
} |
||||||
Loading…
Reference in new issue