From e40a53940653062ff5d4c164c4033c9b8d88bc69 Mon Sep 17 00:00:00 2001 From: Robin Bygrave Date: Thu, 18 Mar 2021 17:02:41 +1300 Subject: [PATCH] #2200 - ENH: Add DtoQuery findEach with batch consumer - findEach(int batchSize, Consumer> consumer) --- .../src/main/java/io/ebean/DtoQuery.java | 11 ++++ .../io/ebeaninternal/api/SpiEbeanServer.java | 5 ++ .../server/core/DefaultServer.java | 11 ++++ .../server/core/DtoQueryRequest.java | 5 ++ .../server/query/DtoQueryEngine.java | 21 +++++++ .../server/querydefn/DefaultDtoQuery.java | 5 ++ .../src/test/java/io/ebean/DtoQueryTest.java | 62 +++++++++++++++++++ .../ebeaninternal/api/TDSpiEbeanServer.java | 4 ++ .../java/org/tests/model/basic/EBasicLog.java | 4 ++ .../org/tests/model/basic/EBasicWithLog.java | 2 +- 10 files changed, 129 insertions(+), 1 deletion(-) diff --git a/ebean-api/src/main/java/io/ebean/DtoQuery.java b/ebean-api/src/main/java/io/ebean/DtoQuery.java index bb2fa43ab..33f439b93 100644 --- a/ebean-api/src/main/java/io/ebean/DtoQuery.java +++ b/ebean-api/src/main/java/io/ebean/DtoQuery.java @@ -53,6 +53,17 @@ public interface DtoQuery { */ void findEach(Consumer consumer); + /** + * Execute the query iterating the results and batching them for the consumer. + *

+ * This runs like findEach streaming results from the database but just collects the results + * into batches to pass to the consumer. + * + * @param batch The number of dto beans to collect before given them to the consumer + * @param consumer The consumer to process the batch of DTO beans + */ + void findEach(int batch, Consumer> consumer); + /** * Execute the query iterating a row at a time with the ability to stop consuming part way through. *

diff --git a/ebean-core/src/main/java/io/ebeaninternal/api/SpiEbeanServer.java b/ebean-core/src/main/java/io/ebeaninternal/api/SpiEbeanServer.java index e9f26b2b2..deb9c281c 100644 --- a/ebean-core/src/main/java/io/ebeaninternal/api/SpiEbeanServer.java +++ b/ebean-core/src/main/java/io/ebeaninternal/api/SpiEbeanServer.java @@ -264,6 +264,11 @@ public interface SpiEbeanServer extends ExtendedServer, EbeanServer, BeanCollect */ void findDtoEach(SpiDtoQuery query, Consumer consumer); + /** + * DTO findEach batch query. + */ + void findDtoEach(SpiDtoQuery query, int batch, Consumer> consumer); + /** * DTO findEachWhile query. */ diff --git a/ebean-core/src/main/java/io/ebeaninternal/server/core/DefaultServer.java b/ebean-core/src/main/java/io/ebeaninternal/server/core/DefaultServer.java index a51bff04d..34a094301 100644 --- a/ebean-core/src/main/java/io/ebeaninternal/server/core/DefaultServer.java +++ b/ebean-core/src/main/java/io/ebeaninternal/server/core/DefaultServer.java @@ -1597,6 +1597,17 @@ public final class DefaultServer implements SpiServer, SpiEbeanServer { } } + @Override + public void findDtoEach(SpiDtoQuery query, int batch, Consumer> consumer) { + DtoQueryRequest request = new DtoQueryRequest<>(this, dtoQueryEngine, query); + try { + request.initTransIfRequired(); + request.findEach(batch, consumer); + } finally { + request.endTransIfRequired(); + } + } + @Override public void findDtoEachWhile(SpiDtoQuery query, Predicate consumer) { DtoQueryRequest request = new DtoQueryRequest<>(this, dtoQueryEngine, query); diff --git a/ebean-core/src/main/java/io/ebeaninternal/server/core/DtoQueryRequest.java b/ebean-core/src/main/java/io/ebeaninternal/server/core/DtoQueryRequest.java index 0dd1857c3..120e98a27 100644 --- a/ebean-core/src/main/java/io/ebeaninternal/server/core/DtoQueryRequest.java +++ b/ebean-core/src/main/java/io/ebeaninternal/server/core/DtoQueryRequest.java @@ -92,6 +92,11 @@ public final class DtoQueryRequest extends AbstractSqlQueryRequest { queryEngine.findEach(this, consumer); } + public void findEach(int batch, Consumer> consumer) { + flushJdbcBatchOnQuery(); + queryEngine.findEach(this, batch, consumer); + } + public void findEachWhile(Predicate consumer) { flushJdbcBatchOnQuery(); queryEngine.findEachWhile(this, consumer); diff --git a/ebean-core/src/main/java/io/ebeaninternal/server/query/DtoQueryEngine.java b/ebean-core/src/main/java/io/ebeaninternal/server/query/DtoQueryEngine.java index 911bcd511..bbd111f1c 100644 --- a/ebean-core/src/main/java/io/ebeaninternal/server/query/DtoQueryEngine.java +++ b/ebean-core/src/main/java/io/ebeaninternal/server/query/DtoQueryEngine.java @@ -42,7 +42,28 @@ public class DtoQueryEngine { } } catch (Exception e) { throw new PersistenceException(errMsg(e.getMessage(), request.getSql()), e); + } finally { + request.close(); + } + } + public void findEach(DtoQueryRequest request, int batchSize, Consumer> consumer) { + try { + List buffer = new ArrayList<>(); + request.executeSql(binder, SpiQuery.Type.ITERATE); + while (request.next()) { + buffer.add(request.readNextBean()); + if (buffer.size() >= batchSize) { + consumer.accept(buffer); + buffer.clear(); + } + } + if (!buffer.isEmpty()) { + // consume the remainder + consumer.accept(buffer); + } + } catch (Exception e) { + throw new PersistenceException(errMsg(e.getMessage(), request.getSql()), e); } finally { request.close(); } diff --git a/ebean-core/src/main/java/io/ebeaninternal/server/querydefn/DefaultDtoQuery.java b/ebean-core/src/main/java/io/ebeaninternal/server/querydefn/DefaultDtoQuery.java index 12d5545e0..d5877572a 100644 --- a/ebean-core/src/main/java/io/ebeaninternal/server/querydefn/DefaultDtoQuery.java +++ b/ebean-core/src/main/java/io/ebeaninternal/server/querydefn/DefaultDtoQuery.java @@ -102,6 +102,11 @@ public class DefaultDtoQuery implements SpiDtoQuery { server.findDtoEach(this, consumer); } + @Override + public void findEach(int batch, Consumer> consumer) { + server.findDtoEach(this, batch, consumer); + } + @Override public void findEachWhile(Predicate consumer) { server.findDtoEachWhile(this, consumer); diff --git a/ebean-core/src/test/java/io/ebean/DtoQueryTest.java b/ebean-core/src/test/java/io/ebean/DtoQueryTest.java index 3a2c8e35e..eea1c6da3 100644 --- a/ebean-core/src/test/java/io/ebean/DtoQueryTest.java +++ b/ebean-core/src/test/java/io/ebean/DtoQueryTest.java @@ -10,11 +10,13 @@ import org.junit.Test; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import org.tests.model.basic.Customer; +import org.tests.model.basic.EBasicLog; import org.tests.model.basic.ResetBasicData; import java.util.Arrays; import java.util.List; import java.util.Optional; +import java.util.concurrent.atomic.AtomicInteger; import static org.assertj.core.api.Assertions.assertThat; @@ -22,6 +24,9 @@ public class DtoQueryTest extends BaseTestCase { private static final Logger log = LoggerFactory.getLogger(DtoQueryTest.class); + private final AtomicInteger batchCount = new AtomicInteger(); + private final AtomicInteger rowCount = new AtomicInteger(); + @Test public void dto_findList_constructorMatch() { @@ -78,6 +83,63 @@ public class DtoQueryTest extends BaseTestCase { assertSql(sql.get(0)).contains("select id, name from o_customer where id > ?"); } + private void resetFindEachCounts() { + batchCount.set(0); + rowCount.set(0); + } + + @Test + public void dto_findEachBatch() { + seedData(); // 15 rows inserted to fetch + + resetFindEachCounts(); + findEachWithBatch(5); + assertThat(batchCount.get()).isEqualTo(3); + assertThat(rowCount.get()).isEqualTo(15); + + resetFindEachCounts(); + findEachWithBatch(10); + assertThat(batchCount.get()).isEqualTo(2); + assertThat(rowCount.get()).isEqualTo(15); + + resetFindEachCounts(); + findEachWithBatch(14); + assertThat(batchCount.get()).isEqualTo(2); + assertThat(rowCount.get()).isEqualTo(15); + + resetFindEachCounts(); + findEachWithBatch(15); + assertThat(batchCount.get()).isEqualTo(1); + assertThat(rowCount.get()).isEqualTo(15); + + resetFindEachCounts(); + findEachWithBatch(16); + assertThat(batchCount.get()).isEqualTo(1); + assertThat(rowCount.get()).isEqualTo(15); + + resetFindEachCounts(); + findEachWithBatch(20); + assertThat(batchCount.get()).isEqualTo(1); + assertThat(rowCount.get()).isEqualTo(15); + } + + private void findEachWithBatch(int batchSize) { + server().findDto(DCust.class, "select id, name from e_basic_log where name like ?") + .setParameter("dtoFindEachBatch%") + .findEach(batchSize, batch -> { + int batchId = batchCount.incrementAndGet(); + int rows = rowCount.addAndGet(batch.size()); + log.info("batch {} rows {}", batchId, rows); + }); + } + + private void seedData() { + for (int i = 0; i < 15; i++) { + EBasicLog log = new EBasicLog("dtoFindEachBatch "+i); + DB.save(log); + } + } + @Test public void dto_findOneEmpty() { diff --git a/ebean-core/src/test/java/io/ebeaninternal/api/TDSpiEbeanServer.java b/ebean-core/src/test/java/io/ebeaninternal/api/TDSpiEbeanServer.java index 09e75825b..03ed3362f 100644 --- a/ebean-core/src/test/java/io/ebeaninternal/api/TDSpiEbeanServer.java +++ b/ebean-core/src/test/java/io/ebeaninternal/api/TDSpiEbeanServer.java @@ -445,6 +445,10 @@ public class TDSpiEbeanServer implements SpiEbeanServer { public void findDtoEach(SpiDtoQuery query, Consumer consumer) { } + @Override + public void findDtoEach(SpiDtoQuery query, int batch, Consumer> consumer) { + } + @Override public void findDtoEachWhile(SpiDtoQuery query, Predicate consumer) { } diff --git a/ebean-core/src/test/java/org/tests/model/basic/EBasicLog.java b/ebean-core/src/test/java/org/tests/model/basic/EBasicLog.java index e0de4ce92..160fff74a 100644 --- a/ebean-core/src/test/java/org/tests/model/basic/EBasicLog.java +++ b/ebean-core/src/test/java/org/tests/model/basic/EBasicLog.java @@ -13,6 +13,10 @@ public class EBasicLog { String name; + public EBasicLog(String name) { + this.name = name; + } + public Long getId() { return id; } diff --git a/ebean-core/src/test/java/org/tests/model/basic/EBasicWithLog.java b/ebean-core/src/test/java/org/tests/model/basic/EBasicWithLog.java index 743796e52..3ff7634f0 100644 --- a/ebean-core/src/test/java/org/tests/model/basic/EBasicWithLog.java +++ b/ebean-core/src/test/java/org/tests/model/basic/EBasicWithLog.java @@ -116,7 +116,7 @@ public class EBasicWithLog { } private void writeLog(String title) { - EBasicLog log = new EBasicLog(); + EBasicLog log = new EBasicLog(name); log.setName(title); DB.save(log); }