Fix MongoStatement.getFieldNamesFromProjectStage

Author: stIncMale

src/main/java/com/mongodb/hibernate/jdbc/MongoStatement.java

@@ -89,7 +89,7 @@ ResultSet executeQueryCommand(BsonDocument command) throws SQLException {
     @VisibleForTesting(otherwise = PRIVATE)
     static List<String> getFieldNamesFromProjectStage(BsonDocument projectStage) {
         var fieldNames = projectStage.entrySet().stream()
-                .filter(field -> trueOrOne(field.getValue()))
+                .filter(MongoStatement::isIncludeProjectSpecification)
                 .map(Map.Entry::getKey)
                 .collect(toCollection(ArrayList::new));
         if (!projectStage.containsKey(ID_FIELD_NAME)) {
@@ -99,9 +99,19 @@ static List<String> getFieldNamesFromProjectStage(BsonDocument projectStage) {
         return fieldNames;
     }
 
-    private static boolean trueOrOne(BsonValue value) {
-        return (value.isBoolean() && value.asBoolean().getValue())
+    private static boolean isIncludeProjectSpecification(Map.Entry<String, BsonValue> specification) {
+        var value = specification.getValue();
+        if (value.isString() || value.isDocument()) {
+            throw new FeatureNotSupportedException("Expressions are not supported in `$project` specifications");
+        }
+        var include = (value.isBoolean() && value.asBoolean().getValue())
                 || (value.isNumber() && value.asNumber().intValue() == 1);
+        if (!include && !specification.getKey().equals(ID_FIELD_NAME)) {
+            throw new RuntimeException(format(
+                    "Exclusions are not allowed in `$project` specifications, except for the [%s] field",
+                    ID_FIELD_NAME));
+        }
+        return include;
     }
 
     @Override

src/test/java/com/mongodb/hibernate/jdbc/MongoStatementTests.java

@@ -183,15 +183,23 @@ void testExecuteBatch() throws SQLException {
 
     @Test
     void testGetProjectStageFieldNames() {
-        BiConsumer<String, List<String>> asserter = (projectStage, expectedFieldNames) -> assertEquals(
+        BiConsumer<String, List<String>> successAsserter = (projectStage, expectedFieldNames) -> assertEquals(
                 expectedFieldNames, MongoStatement.getFieldNamesFromProjectStage(BsonDocument.parse(projectStage)));
+        BiConsumer<String, String> failureAsserter = (projectStage, expectedMessageFragment) -> {
+            Throwable e = assertThrows(
+                    RuntimeException.class,
+                    () -> MongoStatement.getFieldNamesFromProjectStage(BsonDocument.parse(projectStage)));
+            assertThat(e.getMessage()).contains(expectedMessageFragment);
+        };
         assertAll(
-                () -> asserter.accept("{title: 1, publishYear: 1}", List.of("title", "publishYear", "_id")),
-                () -> asserter.accept("{title: 1, publishYear: 0}", List.of("title", "_id")),
-                () -> asserter.accept("{title: 1, publishYear: false}", List.of("title", "_id")),
-                () -> asserter.accept("{title: 1, _id: 0}", List.of("title")),
-                () -> asserter.accept("{title: 1, _id: false}", List.of("title")),
-                () -> asserter.accept("{_id: 1, title: 1}", List.of("_id", "title")));
+                () -> successAsserter.accept("{title: 1, publishYear: 1}", List.of("title", "publishYear", "_id")),
+                () -> successAsserter.accept("{title: 1, _id: 0}", List.of("title")),
+                () -> successAsserter.accept("{title: 1, _id: false}", List.of("title")),
+                () -> successAsserter.accept("{_id: 1, title: 1}", List.of("_id", "title")),
+                () -> failureAsserter.accept("{title: 1, publishYear: 0}", "Exclusions are not allowed"),
+                () -> failureAsserter.accept("{title: 1, publishYear: false}", "Exclusions are not allowed"),
+                () -> failureAsserter.accept("{title: '$fieldName'}", "Expressions are not supported"),
+                () -> failureAsserter.accept("{title: {fieldName: 'fieldValue'}}", "Expressions are not supported"));
     }
 
     @Nested