Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view

Large diffs are not rendered by default.

Original file line number Diff line number Diff line change
Expand Up @@ -128,7 +128,7 @@ protected <T> void assertSelectionQuery(
}
var resultList = selectionQuery.getResultList();

assertActualCommand(BsonDocument.parse(expectedMql));
assertActualCommandsInOrder(BsonDocument.parse(expectedMql));

resultListVerifier.accept(resultList);

Expand Down Expand Up @@ -178,13 +178,13 @@ protected void assertSelectQueryFailure(
expectedExceptionMessageParameters);
}

protected void assertActualCommand(BsonDocument expectedCommand) {
protected void assertActualCommandsInOrder(BsonDocument... expectedCommands) {
var capturedCommands = testCommandListener.getStartedCommands();

assertThat(capturedCommands)
.singleElement()
.asInstanceOf(InstanceOfAssertFactories.MAP)
.containsAllEntriesOf(expectedCommand);
assertThat(capturedCommands).hasSize(expectedCommands.length);
for (int i = 0; i < expectedCommands.length; i++) {
BsonDocument actual = capturedCommands.get(i);
assertThat(actual).asInstanceOf(InstanceOfAssertFactories.MAP).containsAllEntriesOf(expectedCommands[i]);
}
}

protected void assertMutationQuery(
Expand All @@ -201,7 +201,7 @@ protected void assertMutationQuery(
queryPostProcessor.accept(query);
}
var mutationCount = query.executeUpdate();
assertActualCommand(BsonDocument.parse(expectedMql));
assertActualCommandsInOrder(BsonDocument.parse(expectedMql));
assertThat(mutationCount).isEqualTo(expectedMutationCount);
});
assertThat(collection.find()).containsExactlyElementsOf(expectedDocuments);
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,200 @@
/*
* 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.bson.RawBsonDocument.parse;

import com.mongodb.client.MongoCollection;
import com.mongodb.hibernate.junit.InjectMongoCollection;
import com.mongodb.hibernate.query.AbstractQueryIntegrationTests;
import jakarta.persistence.Entity;
import jakarta.persistence.Id;
import jakarta.persistence.Table;
import java.util.List;
import org.bson.BsonDocument;
import org.hibernate.cfg.AvailableSettings;
import org.hibernate.engine.spi.SessionImplementor;
import org.hibernate.testing.orm.junit.DomainModel;
import org.hibernate.testing.orm.junit.ServiceRegistry;
import org.hibernate.testing.orm.junit.Setting;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;

@DomainModel(annotatedClasses = BatchUpdateIntegrationTests.Item.class)
@ServiceRegistry(settings = @Setting(name = AvailableSettings.STATEMENT_BATCH_SIZE, value = "3"))
class BatchUpdateIntegrationTests extends AbstractQueryIntegrationTests {

private static final String COLLECTION_NAME = "items";
private static final int ENTITIES_TO_PERSIST_COUNT = 5;

@InjectMongoCollection(COLLECTION_NAME)
private static MongoCollection<BsonDocument> collection;

@BeforeEach
void beforeEach() {
getTestCommandListener().clear();
}

@Test
void testBatchInsert() {
getSessionFactoryScope().inTransaction(session -> {
for (int i = 1; i <= ENTITIES_TO_PERSIST_COUNT; i++) {
session.persist(new Item(i, String.valueOf(i)));
}
session.flush();
assertActualCommandsInOrder(
parse(
"""
{
"insert": "items",
"ordered": true,
"documents": [
{ "_id": 1, "string": "1"},
{ "_id": 2, "string": "2"},
{ "_id": 3, "string": "3"}
]
}
"""),
parse(
"""
{
"insert": "items",
"ordered": true,
"documents": [
{ "_id": 4, "string": "4"},
{ "_id": 5, "string": "5"}
]
}
"""));
});

assertThat(collection.find())
.containsExactlyElementsOf(List.of(
BsonDocument.parse("{ _id: 1, string: '1' }"),
BsonDocument.parse("{ _id: 2, string: '2' }"),
BsonDocument.parse("{ _id: 3, string: '3' }"),
BsonDocument.parse("{ _id: 4, string: '4' }"),
BsonDocument.parse("{ _id: 5, string: '5' }")));
}

@Test
void testBatchUpdate() {
getSessionFactoryScope().inTransaction(session -> {
insertTestData(session);
for (int i = 1; i <= ENTITIES_TO_PERSIST_COUNT; i++) {
Item item = session.find(Item.class, i);
item.string = "u" + i;
}
session.flush();
assertActualCommandsInOrder(
parse(
"""
{
"update": "items",
"ordered": true,
"updates": [
{ "q": { "_id": { "$eq": 1 } }, "u": { "$set": { "string": "u1" } }, "multi": true },
{ "q": { "_id": { "$eq": 2 } }, "u": { "$set": { "string": "u2" } }, "multi": true },
{ "q": { "_id": { "$eq": 3 } }, "u": { "$set": { "string": "u3" } }, "multi": true }
]
}
"""),
parse(
"""
{
"update": "items",
"ordered": true,
"updates": [
{ "q": { "_id": { "$eq": 4 } }, "u": { "$set": { "string": "u4" } }, "multi": true },
{ "q": { "_id": { "$eq": 5 } }, "u": { "$set": { "string": "u5" } }, "multi": true }
]
}
"""));
});

assertThat(collection.find())
.containsExactlyElementsOf(java.util.List.of(
BsonDocument.parse("{ _id: 1, string: 'u1' }"),
BsonDocument.parse("{ _id: 2, string: 'u2' }"),
BsonDocument.parse("{ _id: 3, string: 'u3' }"),
BsonDocument.parse("{ _id: 4, string: 'u4' }"),
BsonDocument.parse("{ _id: 5, string: 'u5' }")));
}

@Test
void testBatchDelete() {
getSessionFactoryScope().inTransaction(session -> {
insertTestData(session);
for (int i = 1; i <= ENTITIES_TO_PERSIST_COUNT; i++) {
var item = session.find(Item.class, i);
session.remove(item);
}
session.flush();
assertActualCommandsInOrder(
parse(
"""
{
"delete": "items",
"ordered": true,
"deletes": [
{"q": {"_id": {"$eq": 1}}, "limit": 0},
{"q": {"_id": {"$eq": 2}}, "limit": 0},
{"q": {"_id": {"$eq": 3}}, "limit": 0}
]
}
"""),
parse(
"""
{
"delete": "items",
"ordered": true,
"deletes": [
{"q": {"_id": {"$eq": 4}}, "limit": 0},
{"q": {"_id": {"$eq": 5}}, "limit": 0}
]
}
"""));
});

assertThat(collection.find()).isEmpty();
}

private void insertTestData(final SessionImplementor session) {
for (int i = 1; i <= 5; i++) {
session.persist(new Item(i, String.valueOf(i)));
}
session.flush();
getTestCommandListener().clear();
}

@Entity
@Table(name = COLLECTION_NAME)
static class Item {
@Id
int id;

String string;

Item() {}

Item(int id, String string) {
this.id = id;
this.string = string;
}
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -614,7 +614,7 @@ private void setQueryOptionsAndQuery(
query.getResultList();
if (expectedMql != null) {
var expectedCommand = BsonDocument.parse(expectedMql);
assertActualCommand(expectedCommand);
assertActualCommandsInOrder(expectedCommand);
}
}
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -35,10 +35,12 @@
import java.sql.SQLException;
import java.sql.SQLFeatureNotSupportedException;
import java.sql.SQLSyntaxErrorException;
import java.sql.Statement;
import java.sql.Time;
import java.sql.Timestamp;
import java.sql.Types;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Calendar;
import java.util.List;
import java.util.Set;
Expand All @@ -52,15 +54,16 @@
final class MongoPreparedStatement extends MongoStatement implements PreparedStatementAdapter {

private final BsonDocument command;

private final List<BsonDocument> commandBatch;
private final List<ParameterValueSetter> parameterValueSetters;

MongoPreparedStatement(
MongoDatabase mongoDatabase, ClientSession clientSession, MongoConnection mongoConnection, String mql)
throws SQLSyntaxErrorException {
super(mongoDatabase, clientSession, mongoConnection);
this.command = MongoStatement.parse(mql);
this.parameterValueSetters = new ArrayList<>();
command = MongoStatement.parse(mql);
commandBatch = new ArrayList<>();
parameterValueSetters = new ArrayList<>();
parseParameters(command, parameterValueSetters);
}

Expand All @@ -77,6 +80,7 @@ public int executeUpdate() throws SQLException {
checkClosed();
closeLastOpenResultSet();
checkAllParametersSet();
checkSupportedUpdateCommand(command);
return executeUpdateCommand(command);
}

Expand Down Expand Up @@ -200,7 +204,33 @@ public void setObject(int parameterIndex, Object x, int targetSqlType) throws SQ
@Override
public void addBatch() throws SQLException {
checkClosed();
throw new SQLFeatureNotSupportedException("TODO-HIBERNATE-35 https://jira.mongodb.org/browse/HIBERNATE-35");
checkAllParametersSet();
commandBatch.add(command.clone());
}

@Override
public void clearBatch() throws SQLException {
checkClosed();
commandBatch.clear();
}

@Override
public int[] executeBatch() throws SQLException {
checkClosed();
closeLastOpenResultSet();
if (commandBatch.isEmpty()) {
return EMPTY_BATCH_RESULT;
}
checkSupportedBatchCommand(commandBatch.get(0));
try {
executeBulkWrite(commandBatch, ExecutionType.BATCH);
var updateCounts = new int[commandBatch.size()];
// We cannot determine the actual number of rows affected for each command in the batch.
Arrays.fill(updateCounts, Statement.SUCCESS_NO_INFO);
return updateCounts;
} finally {
commandBatch.clear();
}
}

@Override
Expand Down
Loading