Limit/Offset/Fetch MQL translation

src/main/java/com/mongodb/hibernate/internal/translate/AstVisitorValueDescriptor.java

@@ -20,7 +20,6 @@
 import static com.mongodb.hibernate.internal.MongoAssertions.fail;
 
 import com.mongodb.hibernate.internal.translate.mongoast.AstValue;
-import com.mongodb.hibernate.internal.translate.mongoast.command.AstCommand;
 import com.mongodb.hibernate.internal.translate.mongoast.command.aggregate.AstProjectStageSpecification;
 import com.mongodb.hibernate.internal.translate.mongoast.command.aggregate.AstSortField;
 import com.mongodb.hibernate.internal.translate.mongoast.filter.AstFilter;
@@ -34,13 +33,15 @@
 @SuppressWarnings("UnusedTypeParameter")
 final class AstVisitorValueDescriptor<T> {
 
-    static final AstVisitorValueDescriptor<AstCommand> COLLECTION_MUTATION = new AstVisitorValueDescriptor<>();
-    static final AstVisitorValueDescriptor<AstCommand> COLLECTION_AGGREGATE = new AstVisitorValueDescriptor<>();
+    static final AstVisitorValueDescriptor<ModelMutationMqlTranslator.Result> MUTATION_RESULT =
+            new AstVisitorValueDescriptor<>();
+    static final AstVisitorValueDescriptor<SelectMqlTranslator.Result> SELECT_RESULT =
+            new AstVisitorValueDescriptor<>();
 
     static final AstVisitorValueDescriptor<String> COLLECTION_NAME = new AstVisitorValueDescriptor<>();
 
     static final AstVisitorValueDescriptor<String> FIELD_PATH = new AstVisitorValueDescriptor<>();
-    static final AstVisitorValueDescriptor<AstValue> FIELD_VALUE = new AstVisitorValueDescriptor<>();
+    static final AstVisitorValueDescriptor<AstValue> VALUE = new AstVisitorValueDescriptor<>();
 
     static final AstVisitorValueDescriptor<List<AstProjectStageSpecification>> PROJECT_STAGE_SPECIFICATIONS =
             new AstVisitorValueDescriptor<>();

src/main/java/com/mongodb/hibernate/internal/translate/ModelMutationMqlTranslator.java

@@ -16,11 +16,16 @@
 
 package com.mongodb.hibernate.internal.translate;
 
+import static com.mongodb.hibernate.internal.MongoAssertions.assertNotNull;
 import static com.mongodb.hibernate.internal.MongoAssertions.assertNull;
-import static com.mongodb.hibernate.internal.translate.AstVisitorValueDescriptor.COLLECTION_MUTATION;
+import static com.mongodb.hibernate.internal.translate.AstVisitorValueDescriptor.MUTATION_RESULT;
+import static java.util.Collections.emptyList;
 
+import com.mongodb.hibernate.internal.translate.mongoast.command.AstCommand;
+import java.util.List;
 import org.hibernate.engine.spi.SessionFactoryImplementor;
 import org.hibernate.query.spi.QueryOptions;
+import org.hibernate.sql.exec.spi.JdbcParameterBinder;
 import org.hibernate.sql.exec.spi.JdbcParameterBindings;
 import org.hibernate.sql.model.ast.TableMutation;
 import org.hibernate.sql.model.internal.TableUpdateNoSet;
@@ -39,15 +44,38 @@ final class ModelMutationMqlTranslator<O extends JdbcMutationOperation> extends
     @Override
     public O translate(@Nullable JdbcParameterBindings jdbcParameterBindings, QueryOptions queryOptions) {
         assertNull(jdbcParameterBindings);
-        checkQueryOptionsSupportability(queryOptions);
+        applyQueryOptions(queryOptions);
 
-        String mql;
+        Result result;
         if ((TableMutation<?>) tableMutation instanceof TableUpdateNoSet) {
-            mql = "";
+            result = Result.empty();
         } else {
-            var mutationCommand = acceptAndYield(tableMutation, COLLECTION_MUTATION);
-            mql = renderMongoAstNode(mutationCommand);
+            result = acceptAndYield(tableMutation, MUTATION_RESULT);
+        }
+        return result.createJdbcMutationOperation(tableMutation);
+    }
+
+    static final class Result {
+        private final @Nullable AstCommand command;
+
+        private final List<JdbcParameterBinder> parameterBinders;
+
+        private Result(@Nullable AstCommand command, List<JdbcParameterBinder> parameterBinders) {
+            this.command = command;
+            this.parameterBinders = parameterBinders;
+        }
+
+        static Result create(AstCommand command, List<JdbcParameterBinder> parameterBinders) {
+            return new Result(assertNotNull(command), parameterBinders);
+        }
+
+        private static Result empty() {
+            return new Result(null, emptyList());
+        }
+
+        private <O extends JdbcMutationOperation> O createJdbcMutationOperation(TableMutation<O> tableMutation) {
+            var mql = command == null ? "" : renderMongoAstNode(command);
+            return tableMutation.createMutationOperation(mql, parameterBinders);
         }
-        return tableMutation.createMutationOperation(mql, getParameterBinders());
     }
 }

src/main/java/com/mongodb/hibernate/internal/translate/SelectMqlTranslator.java

@@ -16,28 +16,33 @@
 
 package com.mongodb.hibernate.internal.translate;
 
-import static com.mongodb.hibernate.internal.translate.AstVisitorValueDescriptor.COLLECTION_AGGREGATE;
+import static com.mongodb.hibernate.internal.translate.AstVisitorValueDescriptor.SELECT_RESULT;
+import static java.lang.Integer.MAX_VALUE;
+import static java.util.Collections.emptyMap;
 import static org.hibernate.sql.ast.SqlTreePrinter.logSqlAst;
+import static org.hibernate.sql.exec.spi.JdbcLockStrategy.NONE;
 
+import com.mongodb.hibernate.internal.translate.mongoast.command.AstCommand;
+import java.util.List;
+import java.util.Set;
 import org.hibernate.engine.spi.SessionFactoryImplementor;
 import org.hibernate.query.spi.QueryOptions;
 import org.hibernate.sql.ast.tree.Statement;
+import org.hibernate.sql.ast.tree.expression.JdbcParameter;
 import org.hibernate.sql.ast.tree.select.SelectStatement;
 import org.hibernate.sql.exec.spi.JdbcOperationQuerySelect;
+import org.hibernate.sql.exec.spi.JdbcParameterBinder;
 import org.hibernate.sql.exec.spi.JdbcParameterBindings;
 import org.hibernate.sql.results.jdbc.spi.JdbcValuesMappingProducerProvider;
 import org.jspecify.annotations.Nullable;
 
 final class SelectMqlTranslator extends AbstractMqlTranslator<JdbcOperationQuerySelect> {
 
     private final SelectStatement selectStatement;
-    private final JdbcValuesMappingProducerProvider jdbcValuesMappingProducerProvider;
 
     SelectMqlTranslator(SessionFactoryImplementor sessionFactory, SelectStatement selectStatement) {
         super(sessionFactory);
         this.selectStatement = selectStatement;
-        jdbcValuesMappingProducerProvider =
-                sessionFactory.getServiceRegistry().requireService(JdbcValuesMappingProducerProvider.class);
     }
 
     @Override
@@ -47,16 +52,51 @@ public JdbcOperationQuerySelect translate(
         logSqlAst(selectStatement);
 
         checkJdbcParameterBindingsSupportability(jdbcParameterBindings);
-        checkQueryOptionsSupportability(queryOptions);
+        applyQueryOptions(queryOptions);
 
-        var aggregateCommand = acceptAndYield((Statement) selectStatement, COLLECTION_AGGREGATE);
-        var jdbcValuesMappingProducer =
-                jdbcValuesMappingProducerProvider.buildMappingProducer(selectStatement, getSessionFactory());
+        var result = acceptAndYield((Statement) selectStatement, SELECT_RESULT);
+        return result.createJdbcOperationQuerySelect(selectStatement, getSessionFactory());
+    }
+
+    static final class Result {
+        private final AstCommand command;
+        private final List<JdbcParameterBinder> parameterBinders;
+        private final Set<String> affectedTableNames;
+        private final @Nullable JdbcParameter offsetParameter;
+        private final @Nullable JdbcParameter limitParameter;
+
+        Result(
+                AstCommand command,
+                List<JdbcParameterBinder> parameterBinders,
+                Set<String> affectedTableNames,
+                @Nullable JdbcParameter offsetParameter,
+                @Nullable JdbcParameter limitParameter) {
+            this.command = command;
+            this.parameterBinders = parameterBinders;
+            this.affectedTableNames = affectedTableNames;
+            this.offsetParameter = offsetParameter;
+            this.limitParameter = limitParameter;
+        }
 
-        return new JdbcOperationQuerySelect(
-                renderMongoAstNode(aggregateCommand),
-                getParameterBinders(),
-                jdbcValuesMappingProducer,
-                getAffectedTableNames());
+        private JdbcOperationQuerySelect createJdbcOperationQuerySelect(
+                SelectStatement selectStatement, SessionFactoryImplementor sessionFactory) {
+            var jdbcValuesMappingProducerProvider =
+                    sessionFactory.getServiceRegistry().requireService(JdbcValuesMappingProducerProvider.class);
+            var jdbcValuesMappingProducer =
+                    jdbcValuesMappingProducerProvider.buildMappingProducer(selectStatement, sessionFactory);
+            return new JdbcOperationQuerySelect(
+                    renderMongoAstNode(command),
+                    parameterBinders,
+                    jdbcValuesMappingProducer,
+                    affectedTableNames,
+                    0,
+                    MAX_VALUE,
+                    emptyMap(),
+                    NONE,
+                    // The following parameters are provided for query plan cache purposes.
+                    // Not setting them could result in reusing the wrong query plan and subsequently the wrong MQL.
+                    offsetParameter,
+                    limitParameter);
+        }
     }
 }

src/main/java/com/mongodb/hibernate/internal/translate/mongoast/command/aggregate/AstLimitStage.java

@@ -0,0 +1,32 @@
+/*
+ * Copyright 2025-present MongoDB, Inc.
+ *
+ * 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
+ *
+ *   http://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 com.mongodb.hibernate.internal.translate.mongoast.command.aggregate;
+
+import com.mongodb.hibernate.internal.translate.mongoast.AstValue;
+import org.bson.BsonWriter;
+
+public record AstLimitStage(AstValue value) implements AstStage {
+    @Override
+    public void render(BsonWriter writer) {
+        writer.writeStartDocument();
+        {
+            writer.writeName("$limit");
+            value.render(writer);
+        }
+        writer.writeEndDocument();
+    }
+}

src/main/java/com/mongodb/hibernate/internal/translate/mongoast/command/aggregate/AstSkipStage.java

@@ -0,0 +1,32 @@
+/*
+ * Copyright 2025-present MongoDB, Inc.
+ *
+ * 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
+ *
+ *   http://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 com.mongodb.hibernate.internal.translate.mongoast.command.aggregate;
+
+import com.mongodb.hibernate.internal.translate.mongoast.AstValue;
+import org.bson.BsonWriter;
+
+public record AstSkipStage(AstValue value) implements AstStage {
+    @Override
+    public void render(BsonWriter writer) {
+        writer.writeStartDocument();
+        {
+            writer.writeName("$skip");
+            value.render(writer);
+        }
+        writer.writeEndDocument();
+    }
+}

src/test/java/com/mongodb/hibernate/internal/translate/AstVisitorValueDescriptorTests.java

@@ -24,6 +24,6 @@ class AstVisitorValueDescriptorTests {
 
     @Test
     void testToString() {
-        assertEquals("COLLECTION_MUTATION", AstVisitorValueDescriptor.COLLECTION_MUTATION.toString());
+        assertEquals("MUTATION_RESULT", AstVisitorValueDescriptor.MUTATION_RESULT.toString());
     }
 }

src/test/java/com/mongodb/hibernate/internal/translate/AstVisitorValueHolderTests.java

@@ -16,8 +16,9 @@
 
 package com.mongodb.hibernate.internal.translate;
 
-import static com.mongodb.hibernate.internal.translate.AstVisitorValueDescriptor.COLLECTION_MUTATION;
-import static com.mongodb.hibernate.internal.translate.AstVisitorValueDescriptor.FIELD_VALUE;
+import static com.mongodb.hibernate.internal.translate.AstVisitorValueDescriptor.MUTATION_RESULT;
+import static com.mongodb.hibernate.internal.translate.AstVisitorValueDescriptor.VALUE;
+import static java.util.Collections.emptyList;
 import static org.junit.jupiter.api.Assertions.assertSame;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 
@@ -45,9 +46,9 @@ void beforeEach() {
     void testSimpleUsage() {
 
         var value = new AstLiteralValue(new BsonString("field_value"));
-        Runnable valueYielder = () -> astVisitorValueHolder.yield(FIELD_VALUE, value);
+        Runnable valueYielder = () -> astVisitorValueHolder.yield(VALUE, value);
 
-        var valueGotten = astVisitorValueHolder.execute(FIELD_VALUE, valueYielder);
+        var valueGotten = astVisitorValueHolder.execute(VALUE, valueYielder);
 
         assertSame(value, valueGotten);
     }
@@ -57,42 +58,44 @@ void testRecursiveUsage() {
 
         Runnable tableInserter = () -> {
             Runnable fieldValueYielder = () -> {
-                astVisitorValueHolder.yield(FIELD_VALUE, AstParameterMarker.INSTANCE);
+                astVisitorValueHolder.yield(VALUE, AstParameterMarker.INSTANCE);
             };
-            var fieldValue = astVisitorValueHolder.execute(FIELD_VALUE, fieldValueYielder);
+            var fieldValue = astVisitorValueHolder.execute(VALUE, fieldValueYielder);
             AstElement astElement = new AstElement("province", fieldValue);
             astVisitorValueHolder.yield(
-                    COLLECTION_MUTATION, new AstInsertCommand("city", new AstDocument(List.of(astElement))));
+                    MUTATION_RESULT,
+                    ModelMutationMqlTranslator.Result.create(
+                            new AstInsertCommand("city", new AstDocument(List.of(astElement))), emptyList()));
         };
 
-        astVisitorValueHolder.execute(COLLECTION_MUTATION, tableInserter);
+        astVisitorValueHolder.execute(MUTATION_RESULT, tableInserter);
     }
 
     @Test
     @DisplayName("Exception is thrown when holder is not empty when setting value")
     void testHolderNotEmptyWhenSetting() {
 
         Runnable valueYielder = () -> {
-            astVisitorValueHolder.yield(FIELD_VALUE, new AstLiteralValue(new BsonString("value1")));
-            astVisitorValueHolder.yield(FIELD_VALUE, new AstLiteralValue(new BsonString("value2")));
+            astVisitorValueHolder.yield(VALUE, new AstLiteralValue(new BsonString("value1")));
+            astVisitorValueHolder.yield(VALUE, new AstLiteralValue(new BsonString("value2")));
         };
 
-        assertThrows(Error.class, () -> astVisitorValueHolder.execute(FIELD_VALUE, valueYielder));
+        assertThrows(Error.class, () -> astVisitorValueHolder.execute(VALUE, valueYielder));
     }
 
     @Test
     @DisplayName("Exception is thrown when holder is expecting a descriptor different from that of actual data")
     void testHolderExpectingDifferentDescriptor() {
 
         Runnable valueYielder =
-                () -> astVisitorValueHolder.yield(FIELD_VALUE, new AstLiteralValue(new BsonString("some_value")));
+                () -> astVisitorValueHolder.yield(VALUE, new AstLiteralValue(new BsonString("some_value")));
 
-        assertThrows(Error.class, () -> astVisitorValueHolder.execute(COLLECTION_MUTATION, valueYielder));
+        assertThrows(Error.class, () -> astVisitorValueHolder.execute(MUTATION_RESULT, valueYielder));
     }
 
     @Test
     @DisplayName("Exception is thrown when no value is yielded")
     void testHolderStillEmpty() {
-        assertThrows(Error.class, () -> astVisitorValueHolder.execute(FIELD_VALUE, () -> {}));
+        assertThrows(Error.class, () -> astVisitorValueHolder.execute(VALUE, () -> {}));
     }
 }