Skip to content

Commit b0b4bc5

Browse files
NathanQingyangXuNathanQingyangXu
committed
implement batch processing for non-native query scenario.
1 parent 19803e1 commit b0b4bc5

File tree

10 files changed

+623
-190
lines changed

10 files changed

+623
-190
lines changed

‎src/integrationTest/java/com/mongodb/hibernate/jdbc/MongoPreparedStatementIntegrationTests.java

Lines changed: 148 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,6 @@
1616

1717
package com.mongodb.hibernate.jdbc;
1818

19-
import static com.mongodb.hibernate.internal.MongoAssertions.assertNotNull;
2019
import static org.junit.jupiter.api.Assertions.assertEquals;
2120

2221
import java.sql.Connection;
@@ -27,8 +26,9 @@
2726
import org.bson.BsonDocument;
2827
import org.hibernate.Session;
2928
import org.hibernate.SessionFactory;
29+
import org.hibernate.cfg.AvailableSettings;
3030
import org.hibernate.cfg.Configuration;
31-
import org.jspecify.annotations.Nullable;
31+
import org.jspecify.annotations.NullUnmarked;
3232
import org.junit.jupiter.api.AfterAll;
3333
import org.junit.jupiter.api.AfterEach;
3434
import org.junit.jupiter.api.BeforeAll;
@@ -37,11 +37,11 @@
3737
import org.junit.jupiter.params.ParameterizedTest;
3838
import org.junit.jupiter.params.provider.ValueSource;
3939

40+
@NullUnmarked
4041
class MongoPreparedStatementIntegrationTests {
42+
private static SessionFactory sessionFactory;
4143

42-
private static @Nullable SessionFactory sessionFactory;
43-
44-
private @Nullable Session session;
44+
private Session session;
4545

4646
@BeforeAll
4747
static void beforeAll() {
@@ -57,7 +57,18 @@ static void afterAll() {
5757

5858
@BeforeEach
5959
void setUp() {
60-
session = assertNotNull(sessionFactory).openSession();
60+
session = sessionFactory.openSession();
61+
session.doWork(conn -> {
62+
conn.createStatement()
63+
.executeUpdate(
64+
"""
65+
{
66+
delete: "books",
67+
deletes: [
68+
{ q: {}, limit: 0 }
69+
]
70+
}""");
71+
});
6172
}
6273

6374
@AfterEach
@@ -70,21 +81,6 @@ void tearDown() {
7081
@Nested
7182
class ExecuteUpdateTests {
7283

73-
@BeforeEach
74-
void setUp() {
75-
assertNotNull(session).doWork(conn -> {
76-
conn.createStatement()
77-
.executeUpdate(
78-
"""
79-
{
80-
delete: "books",
81-
deletes: [
82-
{ q: {}, limit: 0 }
83-
]
84-
}""");
85-
});
86-
}
87-
8884
private static final String INSERT_MQL =
8985
"""
9086
{
@@ -181,7 +177,7 @@ void testUpdate(boolean autoCommit) {
181177
}
182178

183179
private void prepareData() {
184-
assertNotNull(session).doWork(connection -> {
180+
session.doWork(connection -> {
185181
connection.setAutoCommit(true);
186182
var statement = connection.createStatement();
187183
statement.executeUpdate(INSERT_MQL);
@@ -193,7 +189,7 @@ private void assertExecuteUpdate(
193189
boolean autoCommit,
194190
int expectedUpdatedRowCount,
195191
Set<? extends BsonDocument> expectedDocuments) {
196-
assertNotNull(session).doWork(connection -> {
192+
session.doWork(connection -> {
197193
connection.setAutoCommit(autoCommit);
198194
try (var pstmt = pstmtProvider.apply(connection)) {
199195
try {
@@ -212,4 +208,133 @@ private void assertExecuteUpdate(
212208
});
213209
}
214210
}
211+
212+
@Nested
213+
class BatchTests {
214+
private static final int BATCH_SIZE = 2;
215+
216+
private static SessionFactory batchableSessionFactory;
217+
218+
private Session batchableSession;
219+
220+
private static final String MQL =
221+
"""
222+
{
223+
insert: "books",
224+
documents: [
225+
{
226+
_id: { $undefined: true },
227+
title: { $undefined: true }
228+
}
229+
]
230+
}""";
231+
232+
@BeforeAll
233+
static void beforeAll() {
234+
batchableSessionFactory = new Configuration()
235+
.setProperty(AvailableSettings.STATEMENT_BATCH_SIZE, BATCH_SIZE)
236+
.buildSessionFactory();
237+
}
238+
239+
@AfterAll
240+
static void afterAll() {
241+
if (batchableSessionFactory != null) {
242+
batchableSessionFactory.close();
243+
}
244+
}
245+
246+
@BeforeEach
247+
void setUp() {
248+
batchableSession = batchableSessionFactory.openSession();
249+
}
250+
251+
@AfterEach
252+
void tearDown() {
253+
if (batchableSession != null) {
254+
batchableSession.close();
255+
}
256+
}
257+
258+
@ParameterizedTest
259+
@ValueSource(booleans = {true, false})
260+
void testExecuteBatch(boolean autoCommit) {
261+
batchableSession.doWork(connection -> {
262+
connection.setAutoCommit(autoCommit);
263+
try (var pstmt = connection.prepareStatement(MQL)) {
264+
try {
265+
pstmt.setInt(1, 1);
266+
pstmt.setString(2, "War and Peace");
267+
pstmt.addBatch();
268+
269+
pstmt.setInt(1, 2);
270+
pstmt.setString(2, "Anna Karenina");
271+
pstmt.addBatch();
272+
273+
pstmt.executeBatch();
274+
275+
pstmt.setInt(1, 3);
276+
pstmt.setString(2, "Crime and Punishment");
277+
pstmt.addBatch();
278+
279+
pstmt.setInt(1, 4);
280+
pstmt.setString(2, "Notes from Underground");
281+
pstmt.addBatch();
282+
283+
pstmt.executeBatch();
284+
285+
pstmt.setInt(1, 5);
286+
pstmt.setString(2, "Fathers and Sons");
287+
288+
pstmt.addBatch();
289+
290+
pstmt.executeBatch();
291+
} finally {
292+
if (!autoCommit) {
293+
connection.commit();
294+
}
295+
pstmt.clearBatch();
296+
}
297+
298+
var expectedDocuments = Set.of(
299+
BsonDocument.parse(
300+
"""
301+
{
302+
_id: 1,
303+
title: "War and Peace"
304+
}"""),
305+
BsonDocument.parse(
306+
"""
307+
{
308+
_id: 2,
309+
title: "Anna Karenina"
310+
}"""),
311+
BsonDocument.parse(
312+
"""
313+
{
314+
_id: 3,
315+
title: "Crime and Punishment"
316+
}"""),
317+
BsonDocument.parse(
318+
"""
319+
{
320+
_id: 4,
321+
title: "Notes from Underground"
322+
}"""),
323+
BsonDocument.parse(
324+
"""
325+
{
326+
_id: 5,
327+
title: "Fathers and Sons"
328+
}"""));
329+
330+
var realDocuments = ((MongoPreparedStatement) pstmt)
331+
.getMongoDatabase()
332+
.getCollection("books", BsonDocument.class)
333+
.find()
334+
.into(new HashSet<>());
335+
assertEquals(expectedDocuments, realDocuments);
336+
}
337+
});
338+
}
339+
}
215340
}

‎src/main/java/com/mongodb/hibernate/internal/ConfigurationHelper.java

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,47 @@
1+
/*
2+
* Copyright 2025-present MongoDB, Inc.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package com.mongodb.hibernate.internal;
18+
19+
import java.util.Map;
20+
import org.hibernate.internal.util.config.ConfigurationException;
21+
22+
/** This class is not part of the public API and may be removed or changed at any time */
23+
public class ConfigurationHelper {
24+
25+
/**
26+
* Get the config value as an int
27+
*
28+
* @param name The config setting name.
29+
* @param values The map of config values
30+
* @param defaultValue The default value to use if not found
31+
* @return The value.
32+
*/
33+
public static int getInt(String name, Map<String, Object> values, int defaultValue) {
34+
Object value = values.get(name);
35+
if (value == null) {
36+
return defaultValue;
37+
}
38+
if (value instanceof Integer) {
39+
return (Integer) value;
40+
}
41+
if (value instanceof String) {
42+
return Integer.parseInt((String) value);
43+
}
44+
throw new ConfigurationException("Could not determine how to handle configuration value [name=" + name
45+
+ ", value=" + value + "(" + value.getClass().getName() + ")] as int");
46+
}
47+
}

‎src/main/java/com/mongodb/hibernate/jdbc/MongoBatchPreparedStatement.java

Lines changed: 0 additions & 111 deletions
This file was deleted.

0 commit comments

Comments
 (0)