-
Notifications
You must be signed in to change notification settings - Fork 12
Introduce batch support. #136
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from 28 commits
4b65eb2
5b17d58
3f1ebcb
ded2bb7
92d9c41
b78cb77
a7f7792
7372d56
330de0f
5c81496
e4d7afd
183aafc
9577603
76acbc5
39df511
11a2b07
510bc2d
6faa9cb
b27ef0e
7c8a702
51103be
4317e2f
2b5be78
40c75cf
91a88f1
9435785
46d4682
afd8853
18a18f7
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Large diffs are not rendered by default.
| 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 | ||||
|---|---|---|---|---|---|---|
|
|
@@ -28,6 +28,7 @@ | |||||
| import com.mongodb.hibernate.internal.type.ObjectIdJdbcType; | ||||||
| import java.math.BigDecimal; | ||||||
| import java.sql.Array; | ||||||
| import java.sql.BatchUpdateException; | ||||||
| import java.sql.JDBCType; | ||||||
| import java.sql.PreparedStatement; | ||||||
| import java.sql.ResultSet; | ||||||
|
|
@@ -45,18 +46,22 @@ | |||||
| import org.bson.BsonType; | ||||||
| import org.bson.BsonValue; | ||||||
| import org.bson.types.ObjectId; | ||||||
| import org.jspecify.annotations.Nullable; | ||||||
|
|
||||||
| final class MongoPreparedStatement extends MongoStatement implements PreparedStatementAdapter { | ||||||
|
|
||||||
| private final BsonDocument command; | ||||||
|
|
||||||
| private final List<BsonDocument> commandBatch; | ||||||
| private final List<ParameterValueSetter> parameterValueSetters; | ||||||
|
|
||||||
| private static final int @Nullable [] NULL_UPDATE_COUNTS = null; | ||||||
|
|
||||||
| MongoPreparedStatement( | ||||||
| MongoDatabase mongoDatabase, ClientSession clientSession, MongoConnection mongoConnection, String mql) | ||||||
| throws SQLSyntaxErrorException { | ||||||
| super(mongoDatabase, clientSession, mongoConnection); | ||||||
| this.command = MongoStatement.parse(mql); | ||||||
| this.commandBatch = new ArrayList<>(); | ||||||
| this.parameterValueSetters = new ArrayList<>(); | ||||||
| parseParameters(command, parameterValueSetters); | ||||||
| } | ||||||
|
|
@@ -66,15 +71,16 @@ public ResultSet executeQuery() throws SQLException { | |||||
| checkClosed(); | ||||||
| closeLastOpenResultSet(); | ||||||
| checkAllParametersSet(); | ||||||
| return executeQueryCommand(command); | ||||||
| return executeQuery(command); | ||||||
| } | ||||||
|
|
||||||
| @Override | ||||||
| public int executeUpdate() throws SQLException { | ||||||
| checkClosed(); | ||||||
| closeLastOpenResultSet(); | ||||||
| checkAllParametersSet(); | ||||||
| return executeUpdateCommand(command); | ||||||
| checkSupportedUpdateCommand(command); | ||||||
| return executeUpdate(command); | ||||||
| } | ||||||
|
|
||||||
| private void checkAllParametersSet() throws SQLException { | ||||||
|
|
@@ -179,7 +185,47 @@ 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()); | ||||||
|
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. [optional] This check may not help us at all, as batching parameterless commands is hardly practical, but neither will it hurt, and it makes it clear that we have to clone specifically because of the parameter value setters.
Suggested change
|
||||||
| } | ||||||
|
|
||||||
| @Override | ||||||
| public void clearBatch() throws SQLException { | ||||||
| checkClosed(); | ||||||
| commandBatch.clear(); | ||||||
| } | ||||||
|
|
||||||
| @Override | ||||||
| public int[] executeBatch() throws SQLException { | ||||||
| checkClosed(); | ||||||
| try { | ||||||
| closeLastOpenResultSet(); | ||||||
| if (commandBatch.isEmpty()) { | ||||||
| return EMPTY_BATCH_RESULT; | ||||||
| } | ||||||
| checkSupportedBatchCommand(commandBatch.get(0)); | ||||||
| return executeBatch(commandBatch); | ||||||
| } finally { | ||||||
| commandBatch.clear(); | ||||||
| } | ||||||
| } | ||||||
|
|
||||||
| /** @throws BatchUpdateException if any of the commands in the batch attempts to return a result set. */ | ||||||
| private void checkSupportedBatchCommand(BsonDocument command) | ||||||
vbabanin marked this conversation as resolved.
Outdated
Show resolved
Hide resolved
|
||||||
| throws SQLFeatureNotSupportedException, BatchUpdateException, SQLSyntaxErrorException { | ||||||
| var commandDescription = getCommandDescription(command); | ||||||
| if (commandDescription.returnsResultSet()) { | ||||||
| throw new BatchUpdateException( | ||||||
| "Commands returning result set are not allowed. Received command: %s" | ||||||
| .formatted(commandDescription.getCommandName()), | ||||||
| NULL_SQL_STATE, | ||||||
| NO_ERROR_CODE, | ||||||
| NULL_UPDATE_COUNTS); | ||||||
|
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Why do we use |
||||||
| } | ||||||
| if (!commandDescription.isUpdate()) { | ||||||
| throw new SQLFeatureNotSupportedException( | ||||||
| "Unsupported command for batch operation: %s".formatted(commandDescription.getCommandName())); | ||||||
| } | ||||||
| } | ||||||
|
|
||||||
| @Override | ||||||
|
|
||||||
Uh oh!
There was an error while loading. Please reload this page.