simple DML translation

src/integrationTest/java/com/mongodb/hibernate/query/select/AbstractSelectionQueryIntegrationTests.java → src/integrationTest/java/com/mongodb/hibernate/query/AbstractQueryIntegrationTests.java

@@ -14,18 +14,19 @@
  * limitations under the License.
  */
 
-package com.mongodb.hibernate.query.select;
+package com.mongodb.hibernate.query;
 
 import static com.mongodb.hibernate.MongoTestAssertions.assertIterableEq;
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
 
+import com.mongodb.client.MongoCollection;
 import com.mongodb.hibernate.TestCommandListener;
 import com.mongodb.hibernate.junit.MongoExtension;
-import java.util.List;
 import java.util.function.Consumer;
 import org.assertj.core.api.InstanceOfAssertFactories;
 import org.bson.BsonDocument;
+import org.hibernate.query.MutationQuery;
 import org.hibernate.query.SelectionQuery;
 import org.hibernate.testing.orm.junit.ServiceRegistryScope;
 import org.hibernate.testing.orm.junit.ServiceRegistryScopeAware;
@@ -36,17 +37,17 @@
 
 @SessionFactory(exportSchema = false)
 @ExtendWith(MongoExtension.class)
-abstract class AbstractSelectionQueryIntegrationTests implements SessionFactoryScopeAware, ServiceRegistryScopeAware {
+public abstract class AbstractQueryIntegrationTests implements SessionFactoryScopeAware, ServiceRegistryScopeAware {
 
     private SessionFactoryScope sessionFactoryScope;
 
     private TestCommandListener testCommandListener;
 
-    SessionFactoryScope getSessionFactoryScope() {
+    protected SessionFactoryScope getSessionFactoryScope() {
         return sessionFactoryScope;
     }
 
-    TestCommandListener getTestCommandListener() {
+    protected TestCommandListener getTestCommandListener() {
         return testCommandListener;
     }
 
@@ -60,12 +61,12 @@ public void injectServiceRegistryScope(ServiceRegistryScope serviceRegistryScope
         this.testCommandListener = serviceRegistryScope.getRegistry().requireService(TestCommandListener.class);
     }
 
-    <T> void assertSelectionQuery(
+    protected <T> void assertSelectionQuery(
             String hql,
             Class<T> resultType,
             Consumer<SelectionQuery<T>> queryPostProcessor,
             String expectedMql,
-            List<T> expectedResultList) {
+            Iterable<T> expectedResultList) {
         assertSelectionQuery(
                 hql,
                 resultType,
@@ -74,16 +75,17 @@ <T> void assertSelectionQuery(
                 resultList -> assertIterableEq(expectedResultList, resultList));
     }
 
-    <T> void assertSelectionQuery(String hql, Class<T> resultType, String expectedMql, List<T> expectedResultList) {
+    protected <T> void assertSelectionQuery(
+            String hql, Class<T> resultType, String expectedMql, Iterable<T> expectedResultList) {
         assertSelectionQuery(hql, resultType, null, expectedMql, expectedResultList);
     }
 
-    <T> void assertSelectionQuery(
+    protected <T> void assertSelectionQuery(
             String hql,
             Class<T> resultType,
             Consumer<SelectionQuery<T>> queryPostProcessor,
             String expectedMql,
-            Consumer<List<T>> resultListVerifier) {
+            Consumer<Iterable<? extends T>> resultListVerifier) {
         sessionFactoryScope.inTransaction(session -> {
             var selectionQuery = session.createSelectionQuery(hql, resultType);
             if (queryPostProcessor != null) {
@@ -97,12 +99,12 @@ <T> void assertSelectionQuery(
         });
     }
 
-    <T> void assertSelectionQuery(
-            String hql, Class<T> resultType, String expectedMql, Consumer<List<T>> resultListVerifier) {
+    protected <T> void assertSelectionQuery(
+            String hql, Class<T> resultType, String expectedMql, Consumer<Iterable<? extends T>> resultListVerifier) {
         assertSelectionQuery(hql, resultType, null, expectedMql, resultListVerifier);
     }
 
-    <T> void assertSelectQueryFailure(
+    protected <T> void assertSelectQueryFailure(
             String hql,
             Class<T> resultType,
             Consumer<SelectionQuery<T>> queryPostProcessor,
@@ -120,7 +122,7 @@ <T> void assertSelectQueryFailure(
                 .hasMessage(expectedExceptionMessage, expectedExceptionMessageParameters));
     }
 
-    <T> void assertSelectQueryFailure(
+    protected <T> void assertSelectQueryFailure(
             String hql,
             Class<T> resultType,
             Class<? extends Exception> expectedExceptionType,
@@ -135,12 +137,31 @@ <T> void assertSelectQueryFailure(
                 expectedExceptionMessageParameters);
     }
 
-    void assertActualCommand(BsonDocument expectedCommand) {
+    protected void assertActualCommand(BsonDocument expectedCommand) {
         var capturedCommands = testCommandListener.getStartedCommands();
 
         assertThat(capturedCommands)
                 .singleElement()
                 .asInstanceOf(InstanceOfAssertFactories.MAP)
                 .containsAllEntriesOf(expectedCommand);
     }
+
+    protected void assertMutationQuery(
+            String hql,
+            Consumer<MutationQuery> queryPostProcessor,
+            int expectedMutationCount,
+            String expectedMql,
+            MongoCollection<BsonDocument> collection,
+            Iterable<? extends BsonDocument> expectedDocuments) {
+        sessionFactoryScope.inTransaction(session -> {
+            var query = session.createMutationQuery(hql);
+            if (queryPostProcessor != null) {
+                queryPostProcessor.accept(query);
+            }
+            var mutationCount = query.executeUpdate();
+            assertActualCommand(BsonDocument.parse(expectedMql));
+            assertThat(mutationCount).isEqualTo(expectedMutationCount);
+        });
+        assertThat(collection.find()).containsExactlyElementsOf(expectedDocuments);
+    }
 }

src/integrationTest/java/com/mongodb/hibernate/query/select/Book.java → src/integrationTest/java/com/mongodb/hibernate/query/Book.java

@@ -14,30 +14,32 @@
  * limitations under the License.
  */
 
-package com.mongodb.hibernate.query.select;
+package com.mongodb.hibernate.query;
 
 import jakarta.persistence.Entity;
 import jakarta.persistence.Id;
 import jakarta.persistence.Table;
 import java.math.BigDecimal;
 
 @Entity(name = "Book")
-@Table(name = "books")
-class Book {
+@Table(name = Book.COLLECTION)
+public class Book {
+    public static final String COLLECTION = "books";
+
     @Id
-    int id;
+    public int id;
 
     // TODO-HIBERNATE-48 dummy values are set for currently null value is not supported
-    String title = "";
-    Boolean outOfStock = false;
-    Integer publishYear = 0;
-    Long isbn13 = 0L;
-    Double discount = 0.0;
-    BigDecimal price = new BigDecimal("0.0");
+    public String title = "";
+    public Boolean outOfStock = false;
+    public Integer publishYear = 0;
+    public Long isbn13 = 0L;
+    public Double discount = 0.0;
+    public BigDecimal price = new BigDecimal("0.0");
 
-    Book() {}
+    public Book() {}
 
-    Book(int id, String title, Integer publishYear, Boolean outOfStock) {
+    public Book(int id, String title, Integer publishYear, Boolean outOfStock) {
         this.id = id;
         this.title = title;
         this.publishYear = publishYear;

src/integrationTest/java/com/mongodb/hibernate/query/mutation/AbstractMutationQueryIntegrationTests.java

@@ -0,0 +1,101 @@
+/*
+ * 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.query.mutation;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.hibernate.cfg.JdbcSettings.DIALECT;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.spy;
+
+import com.mongodb.hibernate.dialect.MongoDialect;
+import com.mongodb.hibernate.query.AbstractQueryIntegrationTests;
+import org.hibernate.dialect.Dialect;
+import org.hibernate.engine.jdbc.dialect.spi.DialectResolutionInfo;
+import org.hibernate.engine.spi.SessionFactoryImplementor;
+import org.hibernate.sql.ast.SqlAstTranslator;
+import org.hibernate.sql.ast.SqlAstTranslatorFactory;
+import org.hibernate.sql.ast.tree.MutationStatement;
+import org.hibernate.sql.ast.tree.select.SelectStatement;
+import org.hibernate.sql.exec.spi.AbstractJdbcOperationQuery;
+import org.hibernate.sql.exec.spi.JdbcOperationQueryMutation;
+import org.hibernate.sql.exec.spi.JdbcOperationQuerySelect;
+import org.hibernate.sql.model.ast.TableMutation;
+import org.hibernate.sql.model.jdbc.JdbcMutationOperation;
+import org.hibernate.testing.orm.junit.ServiceRegistry;
+import org.hibernate.testing.orm.junit.Setting;
+import org.mockito.stubbing.Answer;
+
+@ServiceRegistry(
+        settings =
+                @Setting(
+                        name = DIALECT,
+                        value =
+                                "com.mongodb.hibernate.query.mutation.AbstractMutationQueryIntegrationTests$MutationTranslateResultAwareDialect"))
+public class AbstractMutationQueryIntegrationTests extends AbstractQueryIntegrationTests {
+
+    protected void assertExpectedAffectedCollections(String... expectedAffectedfCollections) {
+        assertThat(((MutationTranslateResultAwareDialect) getSessionFactoryScope()
+                                .getSessionFactory()
+                                .getJdbcServices()
+                                .getDialect())
+                        .capturedTranslateResult.getAffectedTableNames())
+                .containsExactlyInAnyOrder(expectedAffectedfCollections);
+    }
+
+    public static final class MutationTranslateResultAwareDialect extends Dialect {
+        private final Dialect delegate;
+        private AbstractJdbcOperationQuery capturedTranslateResult;
+
+        public MutationTranslateResultAwareDialect(DialectResolutionInfo info) {
+            super(info);
+            delegate = new MongoDialect(info);
+        }
+
+        @Override
+        public SqlAstTranslatorFactory getSqlAstTranslatorFactory() {
+            return new SqlAstTranslatorFactory() {
+                @Override
+                public SqlAstTranslator<JdbcOperationQuerySelect> buildSelectTranslator(
+                        SessionFactoryImplementor sessionFactory, SelectStatement statement) {
+                    return delegate.getSqlAstTranslatorFactory().buildSelectTranslator(sessionFactory, statement);
+                }
+
+                @Override
+                public SqlAstTranslator<? extends JdbcOperationQueryMutation> buildMutationTranslator(
+                        SessionFactoryImplementor sessionFactory, MutationStatement statement) {
+                    var originalTranslator =
+                            delegate.getSqlAstTranslatorFactory().buildMutationTranslator(sessionFactory, statement);
+                    var translatorSpy = spy(originalTranslator);
+                    doAnswer((Answer<AbstractJdbcOperationQuery>) invocation -> {
+                                capturedTranslateResult = (AbstractJdbcOperationQuery) invocation.callRealMethod();
+                                return capturedTranslateResult;
+                            })
+                            .when(translatorSpy)
+                            .translate(any(), any());
+                    return translatorSpy;
+                }
+
+                @Override
+                public <O extends JdbcMutationOperation> SqlAstTranslator<O> buildModelMutationTranslator(
+                        TableMutation<O> mutation, SessionFactoryImplementor sessionFactory) {
+                    return delegate.getSqlAstTranslatorFactory().buildModelMutationTranslator(mutation, sessionFactory);
+                }
+            };
+        }
+    }
+}