feat(command): projection subset for findOneAndUpdate/findOneAndReplace

docs/COMPATIBILITY.md

@@ -29,8 +29,8 @@ Certification context:
 | `bulkWrite` | Partial | Ordered mode only (`ordered=true`); supports `insertOne/updateOne/updateMany/deleteOne/deleteMany/replaceOne` and stops on first write error |
 | `countDocuments` | Partial | Filter + skip/limit + hint/collation/readConcern shape validation |
 | `replaceOne` | Partial | Rewrites to single replacement `update` path (`multi=false`) |
-| `findOneAndUpdate` | Partial | Rewrites to `findAndModify`; operator updates only |
-| `findOneAndReplace` | Partial | Rewrites to `findAndModify`; replacement updates only |
+| `findOneAndUpdate` | Partial | Rewrites to `findAndModify`; operator updates only; supports projection include/exclude subset (including `_id` override) |
+| `findOneAndReplace` | Partial | Rewrites to `findAndModify`; replacement updates only; supports projection include/exclude subset (including `_id` override) |
 | `commitTransaction`, `abortTransaction` | Supported | Session/txn envelope supported |
 
 ## Query Operators

src/main/java/org/jongodb/command/FindAndModifyCommandHandler.java

@@ -1,7 +1,9 @@
 package org.jongodb.command;
 
 import java.util.ArrayList;
+import java.util.LinkedHashSet;
 import java.util.List;
+import java.util.Set;
 import org.bson.BsonBoolean;
 import org.bson.BsonDocument;
 import org.bson.BsonDouble;
@@ -62,6 +64,22 @@ public BsonDocument handle(final BsonDocument command) {
             sort = sortValue.asDocument();
         }
 
+        final BsonValue fieldsValue = command.containsKey("fields")
+                ? command.get("fields")
+                : command.get("projection");
+        final ProjectionSpec projectionSpec;
+        if (fieldsValue == null) {
+            projectionSpec = ProjectionSpec.none();
+        } else if (!fieldsValue.isDocument()) {
+            return CommandErrors.typeMismatch("fields must be a document");
+        } else {
+            try {
+                projectionSpec = parseProjectionSpec(fieldsValue.asDocument());
+            } catch (final IllegalArgumentException exception) {
+                return CommandErrors.badValue(exception.getMessage());
+            }
+        }
+
         final BsonValue updateValue = command.get("update");
         if (!remove && (updateValue == null || !updateValue.isDocument())) {
             return CommandErrors.typeMismatch("update must be a document");
@@ -79,9 +97,9 @@ public BsonDocument handle(final BsonDocument command) {
 
         try {
             if (remove) {
-                return handleRemove(database, collection, selected);
+                return handleRemove(database, collection, selected, projectionSpec);
             }
-            return handleUpdate(database, collection, query, update, selected, upsert, returnNew);
+            return handleUpdate(database, collection, query, update, selected, upsert, returnNew, projectionSpec);
         } catch (final DuplicateKeyException exception) {
             return CommandErrors.duplicateKey(exception.getMessage());
         } catch (final IllegalArgumentException exception) {
@@ -92,7 +110,8 @@ public BsonDocument handle(final BsonDocument command) {
     private BsonDocument handleRemove(
             final String database,
             final String collection,
-            final BsonDocument selected) {
+            final BsonDocument selected,
+            final ProjectionSpec projectionSpec) {
         if (selected == null) {
             return successResponse(0, false, null, null);
         }
@@ -101,7 +120,7 @@ private BsonDocument handleRemove(
                 database,
                 collection,
                 List.of(new CommandStore.DeleteRequest(singleDocumentFilter(selected), 1)));
-        final BsonDocument value = deleted > 0 ? selected : null;
+        final BsonDocument value = deleted > 0 ? applyProjection(selected, projectionSpec) : null;
         return successResponse(deleted > 0 ? 1 : 0, false, null, value);
     }
 
@@ -112,14 +131,16 @@ private BsonDocument handleUpdate(
             final BsonDocument update,
             final BsonDocument selected,
             final boolean upsert,
-            final boolean returnNew) {
+            final boolean returnNew,
+            final ProjectionSpec projectionSpec) {
         if (selected != null) {
             final BsonDocument oneFilter = singleDocumentFilter(selected);
             final CommandStore.UpdateResult result = store.update(
                     database,
                     collection,
                     List.of(new CommandStore.UpdateRequest(oneFilter, update, false, false)));
-            final BsonDocument value = returnNew ? firstMatch(database, collection, oneFilter) : selected;
+            final BsonDocument selectedValue = returnNew ? firstMatch(database, collection, oneFilter) : selected;
+            final BsonDocument value = applyProjection(selectedValue, projectionSpec);
             return successResponse(result.matchedCount() > 0 ? 1 : 0, result.matchedCount() > 0, null, value);
         }
 
@@ -141,6 +162,7 @@ private BsonDocument handleUpdate(
             } else {
                 value = firstMatch(database, collection, query);
             }
+            value = applyProjection(value, projectionSpec);
         }
 
         return successResponse(upsertedId == null ? 0 : 1, false, upsertedId, value);
@@ -357,5 +379,158 @@ private static Long readIntegralLong(final BsonValue value) {
         return null;
     }
 
+    private static ProjectionSpec parseProjectionSpec(final BsonDocument projectionDocument) {
+        if (projectionDocument.isEmpty()) {
+            return ProjectionSpec.none();
+        }
+
+        final Set<String> includePaths = new LinkedHashSet<>();
+        final Set<String> excludePaths = new LinkedHashSet<>();
+        boolean includeId = true;
+
+        for (final String field : projectionDocument.keySet()) {
+            final ProjectionFlag flag = normalizeProjectionFlag(projectionDocument.get(field));
+            if ("_id".equals(field)) {
+                includeId = flag.include();
+                continue;
+            }
+            if (flag.include()) {
+                includePaths.add(field);
+            } else {
+                excludePaths.add(field);
+            }
+        }
+
+        if (!includePaths.isEmpty() && !excludePaths.isEmpty()) {
+            throw new IllegalArgumentException("projection cannot mix inclusion and exclusion except for _id");
+        }
+        if (!includePaths.isEmpty()) {
+            return ProjectionSpec.include(includePaths, includeId);
+        }
+        if (!excludePaths.isEmpty()) {
+            return ProjectionSpec.exclude(excludePaths, includeId);
+        }
+        return ProjectionSpec.none(includeId);
+    }
+
+    private static ProjectionFlag normalizeProjectionFlag(final BsonValue value) {
+        if (value == null) {
+            throw new IllegalArgumentException("projection values must be 0 or 1");
+        }
+        if (value.isBoolean()) {
+            return new ProjectionFlag(value.asBoolean().getValue());
+        }
+        final Long parsed = readIntegralLong(value);
+        if (parsed == null || (parsed != 0L && parsed != 1L)) {
+            throw new IllegalArgumentException("projection values must be 0 or 1");
+        }
+        return new ProjectionFlag(parsed == 1L);
+    }
+
+    private static BsonDocument applyProjection(final BsonDocument value, final ProjectionSpec projectionSpec) {
+        if (value == null) {
+            return null;
+        }
+        if (projectionSpec.noProjection()) {
+            return value;
+        }
+
+        if (projectionSpec.includeMode()) {
+            final BsonDocument projected = new BsonDocument();
+            if (projectionSpec.includeId()) {
+                final BsonValue id = value.get("_id");
+                if (id != null) {
+                    projected.put("_id", id);
+                }
+            }
+            for (final String includePath : projectionSpec.paths()) {
+                final BsonValue projectedValue = resolveProjectionPath(value, includePath);
+                if (projectedValue != null) {
+                    setProjectionPath(projected, includePath, projectedValue);
+                }
+            }
+            if (!projectionSpec.includeId()) {
+                projected.remove("_id");
+            }
+            return projected;
+        }
+
+        final BsonDocument projected = value.clone();
+        for (final String excludedPath : projectionSpec.paths()) {
+            removeProjectionPath(projected, excludedPath);
+        }
+        if (!projectionSpec.includeId()) {
+            projected.remove("_id");
+        }
+        return projected;
+    }
+
+    private static BsonValue resolveProjectionPath(final BsonDocument source, final String path) {
+        if (path == null || path.isEmpty()) {
+            return null;
+        }
+        final String[] segments = path.split("\\.");
+        BsonValue current = source;
+        for (final String segment : segments) {
+            if (current == null || !current.isDocument()) {
+                return null;
+            }
+            current = current.asDocument().get(segment);
+        }
+        return current;
+    }
+
+    private static void setProjectionPath(final BsonDocument target, final String path, final BsonValue value) {
+        final String[] segments = path.split("\\.");
+        BsonDocument current = target;
+        for (int index = 0; index < segments.length - 1; index++) {
+            final String segment = segments[index];
+            final BsonValue existing = current.get(segment);
+            if (existing != null && !existing.isDocument()) {
+                return;
+            }
+            if (existing == null) {
+                final BsonDocument child = new BsonDocument();
+                current.put(segment, child);
+                current = child;
+            } else {
+                current = existing.asDocument();
+            }
+        }
+        current.put(segments[segments.length - 1], value);
+    }
+
+    private static void removeProjectionPath(final BsonDocument target, final String path) {
+        final String[] segments = path.split("\\.");
+        BsonDocument current = target;
+        for (int index = 0; index < segments.length - 1; index++) {
+            final BsonValue next = current.get(segments[index]);
+            if (next == null || !next.isDocument()) {
+                return;
+            }
+            current = next.asDocument();
+        }
+        current.remove(segments[segments.length - 1]);
+    }
+
+    private record ProjectionFlag(boolean include) {}
+    private record ProjectionSpec(Set<String> paths, boolean includeMode, boolean includeId, boolean noProjection) {
+        private static ProjectionSpec include(final Set<String> paths, final boolean includeId) {
+            return new ProjectionSpec(Set.copyOf(paths), true, includeId, false);
+        }
+
+        private static ProjectionSpec exclude(final Set<String> paths, final boolean includeId) {
+            return new ProjectionSpec(Set.copyOf(paths), false, includeId, false);
+        }
+
+        private static ProjectionSpec none() {
+            return none(true);
+        }
+
+        private static ProjectionSpec none(final boolean includeId) {
+            return new ProjectionSpec(Set.of(), false, includeId, true);
+        }
+    }
+
     private record SortKey(String field, int direction) {}
 }

src/main/java/org/jongodb/command/FindOneAndReplaceCommandHandler.java

@@ -63,6 +63,16 @@ public BsonDocument handle(final BsonDocument command) {
             sort = sortValue.asDocument();
         }
 
+        final BsonValue projectionValue = command.get("projection");
+        final BsonDocument projection;
+        if (projectionValue == null) {
+            projection = null;
+        } else if (!projectionValue.isDocument()) {
+            return CommandErrors.typeMismatch("projection must be a document");
+        } else {
+            projection = projectionValue.asDocument();
+        }
+
         final BsonValue upsertValue = command.get("upsert");
         final boolean upsert;
         if (upsertValue == null) {
@@ -89,6 +99,9 @@ public BsonDocument handle(final BsonDocument command) {
         if (sort != null) {
             translatedFindAndModify.append("sort", sort);
         }
+        if (projection != null) {
+            translatedFindAndModify.append("fields", projection);
+        }
 
         appendIfPresent(command, translatedFindAndModify, "hint");
         appendIfPresent(command, translatedFindAndModify, "collation");

src/main/java/org/jongodb/command/FindOneAndUpdateCommandHandler.java

@@ -65,6 +65,16 @@ public BsonDocument handle(final BsonDocument command) {
             sort = sortValue.asDocument();
         }
 
+        final BsonValue projectionValue = command.get("projection");
+        final BsonDocument projection;
+        if (projectionValue == null) {
+            projection = null;
+        } else if (!projectionValue.isDocument()) {
+            return CommandErrors.typeMismatch("projection must be a document");
+        } else {
+            projection = projectionValue.asDocument();
+        }
+
         final BsonValue upsertValue = command.get("upsert");
         final boolean upsert;
         if (upsertValue == null) {
@@ -99,6 +109,9 @@ public BsonDocument handle(final BsonDocument command) {
         if (sort != null) {
             translatedFindAndModify.append("sort", sort);
         }
+        if (projection != null) {
+            translatedFindAndModify.append("fields", projection);
+        }
 
         appendIfPresent(command, translatedFindAndModify, "hint");
         appendIfPresent(command, translatedFindAndModify, "collation");

src/test/java/org/jongodb/command/CommandDispatcherE2ETest.java

@@ -704,6 +704,34 @@ void findOneAndUpdateCommandSupportsBeforeAndAfterSemantics() {
                 afterResponse.getDocument("value").getString("name").getValue());
     }
 
+    @Test
+    void findOneAndUpdateCommandSupportsProjectionWithNestedInclude() {
+        final CommandDispatcher dispatcher = new CommandDispatcher(new EngineBackedCommandStore(new InMemoryEngineStore()));
+        dispatcher.dispatch(BsonDocument.parse(
+                "{\"insert\":\"users\",\"$db\":\"app\",\"documents\":[{\"_id\":1,\"name\":\"before\",\"email\":\"a@example.com\",\"profile\":{\"city\":\"seoul\",\"zip\":\"123\"}}]}"));
+
+        final BsonDocument response = dispatcher.dispatch(BsonDocument.parse(
+                "{\"findOneAndUpdate\":\"users\",\"$db\":\"app\",\"filter\":{\"_id\":1},\"update\":{\"$set\":{\"name\":\"after\"}},\"returnDocument\":\"after\",\"projection\":{\"name\":1,\"profile.city\":1,\"_id\":0}}"));
+        assertEquals(1.0, response.get("ok").asNumber().doubleValue());
+        final BsonDocument value = response.getDocument("value");
+        assertTrue(!value.containsKey("_id"));
+        assertEquals("after", value.getString("name").getValue());
+        assertTrue(!value.containsKey("email"));
+        assertEquals("seoul", value.getDocument("profile").getString("city").getValue());
+        assertTrue(!value.getDocument("profile").containsKey("zip"));
+    }
+
+    @Test
+    void findOneAndUpdateCommandRejectsMixedProjectionModes() {
+        final CommandDispatcher dispatcher = new CommandDispatcher(new EngineBackedCommandStore(new InMemoryEngineStore()));
+        dispatcher.dispatch(BsonDocument.parse(
+                "{\"insert\":\"users\",\"$db\":\"app\",\"documents\":[{\"_id\":1,\"name\":\"before\",\"email\":\"a@example.com\"}]}"));
+
+        final BsonDocument response = dispatcher.dispatch(BsonDocument.parse(
+                "{\"findOneAndUpdate\":\"users\",\"$db\":\"app\",\"filter\":{\"_id\":1},\"update\":{\"$set\":{\"name\":\"after\"}},\"projection\":{\"name\":1,\"email\":0}}"));
+        assertCommandError(response, "BadValue");
+    }
+
     @Test
     void findOneAndUpdateCommandRejectsInvalidPayloadShapes() {
         final CommandDispatcher dispatcher = new CommandDispatcher(new RecordingStore());
@@ -784,6 +812,23 @@ void findOneAndReplaceCommandSupportsUpsertWithReturnAfter() {
                 response.getDocument("value").getString("name").getValue());
     }
 
+    @Test
+    void findOneAndReplaceCommandSupportsProjectionExcludePaths() {
+        final CommandDispatcher dispatcher = new CommandDispatcher(new EngineBackedCommandStore(new InMemoryEngineStore()));
+        dispatcher.dispatch(BsonDocument.parse(
+                "{\"insert\":\"users\",\"$db\":\"app\",\"documents\":[{\"_id\":1,\"name\":\"before\",\"email\":\"a@example.com\",\"profile\":{\"city\":\"seoul\",\"zip\":\"123\"}}]}"));
+
+        final BsonDocument response = dispatcher.dispatch(BsonDocument.parse(
+                "{\"findOneAndReplace\":\"users\",\"$db\":\"app\",\"filter\":{\"_id\":1},\"replacement\":{\"name\":\"after\",\"email\":\"b@example.com\",\"profile\":{\"city\":\"busan\",\"zip\":\"999\"}},\"returnDocument\":\"after\",\"projection\":{\"email\":0,\"profile.zip\":0}}"));
+        assertEquals(1.0, response.get("ok").asNumber().doubleValue());
+        final BsonDocument value = response.getDocument("value");
+        assertEquals(1, value.getInt32("_id").getValue());
+        assertEquals("after", value.getString("name").getValue());
+        assertTrue(!value.containsKey("email"));
+        assertEquals("busan", value.getDocument("profile").getString("city").getValue());
+        assertTrue(!value.getDocument("profile").containsKey("zip"));
+    }
+
     @Test
     void updateCommandRejectsInvalidPayloadShapes() {
         final RecordingStore store = new RecordingStore();