Skip to content

Commit

Permalink
Batch ops should give access to modified row count (#2069)
Browse files Browse the repository at this point in the history
Add `executePreparedBatch` to supersede `executeAndReturnGeneratedKeys`,
which allows access to the modified row count after execution. Also add
a convenience method to the iterable that allows splitting up the list
of results into per-batch sets.

Addresses the problems described in #2060

Based on a proposal by @doppelrittberger (Markus Ritter)
  • Loading branch information
hgschmie committed Jul 19, 2022
1 parent 3357177 commit ef3b9eb
Show file tree
Hide file tree
Showing 4 changed files with 356 additions and 3 deletions.
157 changes: 157 additions & 0 deletions core/src/main/java/org/jdbi/v3/core/result/BatchResultBearing.java
Original file line number Diff line number Diff line change
@@ -0,0 +1,157 @@
/*
* 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 org.jdbi.v3.core.result;

import java.lang.reflect.Type;
import java.sql.Statement;
import java.util.Map;
import java.util.function.BiConsumer;
import java.util.function.BiFunction;
import java.util.function.Supplier;
import java.util.stream.Collector;
import java.util.stream.Stream;

import org.jdbi.v3.core.generic.GenericType;
import org.jdbi.v3.core.mapper.ColumnMapper;
import org.jdbi.v3.core.mapper.RowMapper;
import org.jdbi.v3.core.mapper.RowViewMapper;
import org.jdbi.v3.core.qualifier.QualifiedType;

/**
* Extends the {@link ResultBearing} class to provide access to the per-batch row modification counts.
*/
public final class BatchResultBearing implements ResultBearing {

private final ResultBearing delegate;
private final Supplier<int[]> modifiedRowCountsSupplier;

public BatchResultBearing(ResultBearing delegate, Supplier<int[]> modifiedRowCountsSupplier) {
this.delegate = delegate;
this.modifiedRowCountsSupplier = modifiedRowCountsSupplier;
}

@Override
public <R> R scanResultSet(ResultSetScanner<R> mapper) {
return delegate.scanResultSet(mapper);
}

@Override
public <T> BatchResultIterable<T> mapTo(Class<T> type) {
return BatchResultIterable.of(delegate.mapTo(type), modifiedRowCountsSupplier);
}

@Override
public <T> BatchResultIterable<T> mapTo(GenericType<T> type) {
return BatchResultIterable.of(delegate.mapTo(type), modifiedRowCountsSupplier);
}

@Override
public BatchResultIterable<?> mapTo(Type type) {
return BatchResultIterable.of(delegate.mapTo(type), modifiedRowCountsSupplier);
}

@Override
public <T> BatchResultIterable<T> mapTo(QualifiedType<T> type) {
return BatchResultIterable.of(delegate.mapTo(type), modifiedRowCountsSupplier);
}

@Override
public <T> BatchResultIterable<T> mapToBean(Class<T> type) {
return BatchResultIterable.of(delegate.mapToBean(type), modifiedRowCountsSupplier);
}

@Override
public BatchResultIterable<Map<String, Object>> mapToMap() {
return BatchResultIterable.of(delegate.mapToMap(), modifiedRowCountsSupplier);
}

@Override
public <T> BatchResultIterable<Map<String, T>> mapToMap(Class<T> valueType) {
return BatchResultIterable.of(delegate.mapToMap(valueType), modifiedRowCountsSupplier);
}

@Override
public <T> BatchResultIterable<Map<String, T>> mapToMap(GenericType<T> valueType) {
return BatchResultIterable.of(delegate.mapToMap(valueType), modifiedRowCountsSupplier);
}

@Override
public <T> BatchResultIterable<T> map(ColumnMapper<T> mapper) {
return BatchResultIterable.of(delegate.map(mapper), modifiedRowCountsSupplier);
}

@Override
public <T> BatchResultIterable<T> map(RowMapper<T> mapper) {
return BatchResultIterable.of(delegate.map(mapper), modifiedRowCountsSupplier);
}

@Override
public <T> BatchResultIterable<T> map(RowViewMapper<T> mapper) {
return BatchResultIterable.of(delegate.map(mapper), modifiedRowCountsSupplier);
}

@Override
public <C, R> Stream<R> reduceRows(RowReducer<C, R> reducer) {
return delegate.reduceRows(reducer);
}

@Override
public <K, V> Stream<V> reduceRows(BiConsumer<Map<K, V>, RowView> accumulator) {
return delegate.reduceRows(accumulator);
}

@Override
public <U> U reduceRows(U seed, BiFunction<U, RowView, U> accumulator) {
return delegate.reduceRows(seed, accumulator);
}

@Override
public <U> U reduceResultSet(U seed, ResultSetAccumulator<U> accumulator) {
return delegate.reduceResultSet(seed, accumulator);
}

@Override
public <A, R> R collectRows(Collector<RowView, A, R> collector) {
return delegate.collectRows(collector);
}

@Override
public <R> R collectInto(Class<R> containerType) {
return delegate.collectInto(containerType);
}

@Override
public <R> R collectInto(GenericType<R> containerType) {
return delegate.collectInto(containerType);
}

@Override
public Object collectInto(Type containerType) {
return delegate.collectInto(containerType);
}

/**
* Returns the mod counts for the executed {@link org.jdbi.v3.core.statement.PreparedBatch}
* Note that some database drivers might return special values like {@link Statement#SUCCESS_NO_INFO}
* or {@link Statement#EXECUTE_FAILED}.
* <br>
* <b>Note that the result is only available after the statement was executed (eg. by calling map())</b>. Calling this method before execution
* will return an empty array.
*
* @return the number of modified rows per batch part for the executed {@link org.jdbi.v3.core.statement.PreparedBatch}.
*/
public int[] modifiedRowCounts() {
return modifiedRowCountsSupplier.get();
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,64 @@
/*
* 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 org.jdbi.v3.core.result;

import java.util.ArrayList;
import java.util.Collections;
import java.util.LinkedList;
import java.util.List;
import java.util.function.Supplier;

/**
* Extend the {@link ResultIterable} for batch operations.
* @param <T>
*/
public interface BatchResultIterable<T> extends ResultIterable<T> {

/**
* Split the results into per-batch sub-lists. Note that this may not be correct if any of the executed batches returned an error code.
*
* @return results in a {@link List} of {@link List}s.
*/
List<List<T>> listPerBatch();

static <U> BatchResultIterable<U> of(ResultIterable<U> delegate, Supplier<int[]> modifiedRowCountsSupplier) {

return new BatchResultIterable<U>() {
@Override
public List<List<U>> listPerBatch() {
List<List<U>> results = new LinkedList<>();
try (ResultIterator<U> iterator = delegate.iterator()) {
for (int modCount : modifiedRowCountsSupplier.get()) {
if (modCount <= 0) {
// error return (SUCCESS_NO_INFO or EXECUTE_FAILED) or empty.
results.add(Collections.emptyList());
} else {
List<U> batchResult = new ArrayList<>(modCount);
for (int i = 0; i < modCount; i++) {
batchResult.add(iterator.next());
}
results.add(batchResult);
}
}
}
return results;
}

@Override
public ResultIterator<U> iterator() {
return delegate.iterator();
}
};
}
}
54 changes: 51 additions & 3 deletions core/src/main/java/org/jdbi/v3/core/statement/PreparedBatch.java
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@
import java.util.Map;
import java.util.NoSuchElementException;
import java.util.Optional;
import java.util.function.Consumer;
import java.util.function.Function;
import java.util.function.Supplier;

Expand All @@ -35,6 +36,7 @@
import org.jdbi.v3.core.argument.internal.NamedArgumentFinderFactory;
import org.jdbi.v3.core.argument.internal.NamedArgumentFinderFactory.PrepareKey;
import org.jdbi.v3.core.qualifier.QualifiedType;
import org.jdbi.v3.core.result.BatchResultBearing;
import org.jdbi.v3.core.result.ResultBearing;
import org.jdbi.v3.core.result.ResultIterator;
import org.jdbi.v3.core.result.ResultProducer;
Expand Down Expand Up @@ -152,10 +154,30 @@ public void close() {
};
}

/**
* Execute the batch and give access to any generated keys returned by the operation.
*
* @param columnNames The column names for generated keys.
* @return A {@link ResultBearing} object that can be used to access the results of the batch.
* @deprecated Use {@link #executePreparedBatch(String...)} which has the same functionality but also returns the per-batch modified row counts.
*/
@Deprecated
public ResultBearing executeAndReturnGeneratedKeys(String... columnNames) {
return execute(returningGeneratedKeys(columnNames));
}

/**
* Execute the batch and give access to any generated keys returned by the operation.
*
* @param columnNames The column names for generated keys.
* @return A {@link BatchResultBearing} object that can be used to access the results of the batch and the per-batch modified row counts.
*/
public BatchResultBearing executePreparedBatch(String... columnNames) {
final ExecutedBatchConsumer executedBatchConsumer = new ExecutedBatchConsumer();
final ResultBearing resultBearing = execute(returningGeneratedKeys(columnNames), executedBatchConsumer);
return new BatchResultBearing(resultBearing, executedBatchConsumer);
}

/**
* Executes the batch, returning the result obtained from the given {@link ResultProducer}.
*
Expand All @@ -164,8 +186,16 @@ public ResultBearing executeAndReturnGeneratedKeys(String... columnNames) {
* @return value returned by the result producer.
*/
public <R> R execute(ResultProducer<R> producer) {
return execute(producer, x -> {});
}

private <R> R execute(ResultProducer<R> producer, Consumer<ExecutedBatch> batchConsumer) {
try {
return producer.produce(() -> internalBatchExecute().stmt, getContext());
return producer.produce(() -> {
ExecutedBatch batch = internalBatchExecute();
batchConsumer.accept(batch);
return batch.stmt;
}, getContext());
} catch (SQLException e) {
try {
close();
Expand Down Expand Up @@ -222,13 +252,13 @@ private ExecutedBatch internalBatchExecute() {
beforeExecution();

try {
final int[] rs = SqlLoggerUtil.wrap(stmt::executeBatch, ctx, getConfig(SqlStatements.class).getSqlLogger());
final int[] modifiedRows = SqlLoggerUtil.wrap(stmt::executeBatch, ctx, getConfig(SqlStatements.class).getSqlLogger());

afterExecution();

ctx.setBinding(new PreparedBinding(ctx));

return new ExecutedBatch(stmt, rs);
return new ExecutedBatch(stmt, modifiedRows);
} catch (SQLException e) {
throw new UnableToExecuteStatementException(Batch.mungeBatchException(e), ctx);
}
Expand Down Expand Up @@ -297,4 +327,22 @@ private static class ExecutedBatch {
this.updateCounts = Arrays.copyOf(updateCounts, updateCounts.length);
}
}

private static final class ExecutedBatchConsumer implements Consumer<ExecutedBatch>, Supplier<int[]> {

private int[] modifiedRowCounts = new int[0];

@Override
public void accept(ExecutedBatch executedBatch) {
// has been copied within the executed batch
this.modifiedRowCounts = executedBatch.updateCounts;
}

@Override
@SuppressWarnings("PMD.MethodReturnsInternalArray")
public int[] get() {
// Array was copied as when the ExecutedBatch was created, so exposing it is fine.
return modifiedRowCounts;
}
}
}
Loading

0 comments on commit ef3b9eb

Please sign in to comment.