Makes retaining of the raw document optional/selectable 🎛️

scireum · Apr 16, 2024 · b85243a · b85243a
1 parent 81f3d0a
commit b85243a
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 9 deletions.
diff --git a/src/main/java/sirius/db/mongo/Mango.java b/src/main/java/sirius/db/mongo/Mango.java
@@ -246,25 +246,28 @@ protected <E extends MongoEntity> Optional<E> findEntity(Object id,
                         mongo.find(entityDescriptor.getRealm());
         return finder.where(MongoEntity.ID, id.toString())
                      .singleIn(entityDescriptor.getRelationName())
-                     .map(doc -> make(entityDescriptor, doc));
+                     .map(doc -> make(entityDescriptor, doc, false));
     }
 
     /**
      * Creates a new entity for the given descriptor based on the given doc.
      *
-     * @param descriptor the descriptor of the entity to create
-     * @param doc        the document to read the values from
-     * @param <E>        the effective type of the generated entity
+     * @param descriptor        the descriptor of the entity to create
+     * @param doc               the document to read the values from
+     * @param retainRawDocument whether to retain the raw database document in the entity
+     * @param <E>               the effective type of the generated entity
      * @return the generated entity
      */
     @SuppressWarnings("unchecked")
-    public static <E extends MongoEntity> E make(EntityDescriptor descriptor, Doc doc) {
+    public static <E extends MongoEntity> E make(EntityDescriptor descriptor, Doc doc, boolean retainRawDocument) {
         try {
             E result = (E) descriptor.make(Mango.class,
                                            null,
                                            key -> doc.getUnderlyingObject().containsKey(key) ? doc.get(key) : null);
 
-            result.setMongoDocument(doc);
+            if (retainRawDocument) {
+                result.setMongoDocument(doc);
+            }
 
             if (descriptor.isVersioned()) {
                 result.setVersion(doc.get(VERSION).asInt(0));

diff --git a/src/main/java/sirius/db/mongo/MongoQuery.java b/src/main/java/sirius/db/mongo/MongoQuery.java
@@ -45,6 +45,8 @@ public class MongoQuery<E extends MongoEntity> extends Query<MongoQuery<E>, E, M
 
     private List<MongoFacet> facets;
 
+    private boolean retainRawDocuments = false;
+
     @Part
     private static Mango mango;
 
@@ -150,12 +152,23 @@ public MongoQuery<E> markLongRunning() {
         return this;
     }
 
+    /**
+     * Marks the query to retain the raw database documents.
+     *
+     * @return the query itself for fluent method calls
+     */
+    public MongoQuery<E> retainingRawDocuments() {
+        retainRawDocuments = true;
+        return this;
+    }
+
     @Override
     protected void doIterate(Predicate<E> resultHandler) {
         if (forceFail) {
             return;
         }
-        finder.eachIn(descriptor.getRelationName(), doc -> resultHandler.test(Mango.make(descriptor, doc)));
+        finder.eachIn(descriptor.getRelationName(),
+                      doc -> resultHandler.test(Mango.make(descriptor, doc, retainRawDocuments)));
     }
 
     @Override
@@ -206,7 +219,7 @@ protected Iterator<E> pullNextBlock() {
             if (lastId != null) {
                 query.where(QueryBuilder.FILTERS.gt(MongoEntity.ID, lastId));
             }
-            query.allIn(relation, doc -> buffer.add(Mango.make(descriptor, doc)));
+            query.allIn(relation, doc -> buffer.add(Mango.make(descriptor, doc, retainRawDocuments)));
 
             if (!buffer.isEmpty()) {
                 lastId = buffer.getLast().getId();
@@ -273,7 +286,7 @@ public List<E> randomList() {
         }
 
         finder.sample(descriptor.getRelationName(), doc -> {
-            result.add(Mango.make(descriptor, doc));
+            result.add(Mango.make(descriptor, doc, retainRawDocuments));
             failOnOverflow(result);
             return true;
         });