DATAJDBC-356 - Support for large datasets with stream

spring-data-jdbc/src/main/java/org/springframework/data/jdbc/repository/support/JdbcOpenRowSetTemplate.java

@@ -0,0 +1,66 @@
+/*
+ * Copyright 2018-2019 the original author or authors.
+ *
+ * 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
+ *
+ *      https://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.springframework.data.jdbc.repository.support;
+
+import java.sql.Connection;
+import java.sql.PreparedStatement;
+import java.sql.ResultSet;
+import java.sql.SQLException;
+
+import org.springframework.jdbc.InvalidResultSetAccessException;
+import org.springframework.jdbc.core.JdbcOperations;
+import org.springframework.jdbc.core.PreparedStatementCreator;
+import org.springframework.jdbc.core.namedparam.NamedParameterJdbcTemplate;
+import org.springframework.jdbc.core.namedparam.SqlParameterSource;
+import org.springframework.jdbc.datasource.DataSourceUtils;
+import org.springframework.jdbc.support.JdbcUtils;
+import org.springframework.util.Assert;
+
+/**
+ * An extended {@link NamedParameterJdbcTemplate} to return an {@link JdbcOpenSqlRowSet}.
+ *
+ * @author Marcos Vinicius da Silva
+ */
+class JdbcOpenRowSetTemplate extends NamedParameterJdbcTemplate {
+
+    JdbcOpenRowSetTemplate(JdbcOperations classicJdbcTemplate) {
+        super(classicJdbcTemplate);
+    }
+
+    JdbcOpenSqlRowSet queryForOpenCursorRowSet(String sql, SqlParameterSource paramSource)  {
+        Assert.state(this.getJdbcTemplate().getDataSource() != null, "No DataSource set");
+
+        Connection connection = DataSourceUtils.getConnection(this.getJdbcTemplate().getDataSource());
+        PreparedStatementCreator preparedStatementCreator = this.getPreparedStatementCreator(sql, paramSource);
+        PreparedStatement preparedStatement = null;
+        ResultSet resultSet = null;
+        try {
+            preparedStatement = preparedStatementCreator.createPreparedStatement(connection);
+            if (this.getJdbcTemplate().getFetchSize() >= 0) {
+                preparedStatement.setFetchSize(this.getJdbcTemplate().getFetchSize());
+            }
+            resultSet = preparedStatement.executeQuery();
+
+            return new JdbcOpenSqlRowSet(this.getJdbcTemplate().getDataSource(), resultSet);
+        } catch (SQLException e) {
+            JdbcUtils.closeResultSet(resultSet);
+            JdbcUtils.closeStatement(preparedStatement);
+            DataSourceUtils.releaseConnection(connection, this.getJdbcTemplate().getDataSource());
+
+            throw new InvalidResultSetAccessException(e);
+        }
+    }
+}

spring-data-jdbc/src/main/java/org/springframework/data/jdbc/repository/support/JdbcOpenSqlRowSet.java

@@ -0,0 +1,75 @@
+/*
+ * Copyright 2018-2019 the original author or authors.
+ *
+ * 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
+ *
+ *      https://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.springframework.data.jdbc.repository.support;
+
+import java.sql.Connection;
+import java.sql.ResultSet;
+import java.sql.SQLException;
+import java.sql.Statement;
+import java.util.concurrent.atomic.AtomicBoolean;
+
+import javax.sql.DataSource;
+
+import org.springframework.jdbc.InvalidResultSetAccessException;
+import org.springframework.jdbc.datasource.DataSourceUtils;
+import org.springframework.jdbc.support.JdbcUtils;
+import org.springframework.jdbc.support.rowset.ResultSetWrappingSqlRowSet;
+import org.springframework.jdbc.support.rowset.SqlRowSet;
+
+/**
+ * An implementation of Spring's {@link SqlRowSet} interface, wrapping a
+ * <b>connected</b> {@link ResultSet}. This implementation also keeps track of
+ * the {@link DataSource} that originated the {@link ResultSet} and releases
+ * the resources when {@link JdbcOpenSqlRowSet#close()} is called.
+ *
+ * @author Marcos Vinicius da Silva
+ */
+class JdbcOpenSqlRowSet extends ResultSetWrappingSqlRowSet implements AutoCloseable {
+    private final DataSource dataSource;
+    private final AtomicBoolean closed = new AtomicBoolean(false);
+
+    /**
+     * Create a new JdbcOpenSqlRowSet for the given ResultSet.
+     *
+     * @param resultSet  a <b>connected</b> ResultSet to wrap. The client code is responsible
+     *                   for closing the resources
+     * @param dataSource the dataSource that originated the ResultSet connection
+     * @throws InvalidResultSetAccessException if extracting
+     *                                         the ResultSetMetaData failed
+     * @see java.sql.ResultSet#getMetaData
+     * @see JdbcOpenSqlRowSet
+     */
+    public JdbcOpenSqlRowSet(DataSource dataSource, ResultSet resultSet) {
+        super(resultSet);
+        this.dataSource = dataSource;
+    }
+
+    @Override
+    public void close() {
+        if (this.closed.compareAndSet(false, true)) {
+            try {
+                final Statement statement = getResultSet().getStatement();
+                final Connection connection = statement.getConnection();
+
+                JdbcUtils.closeResultSet(getResultSet());
+                JdbcUtils.closeStatement(statement);
+                DataSourceUtils.releaseConnection(connection, dataSource);
+            } catch (SQLException e) {
+                throw new InvalidResultSetAccessException(e);
+            }
+        }
+    }
+}

spring-data-jdbc/src/main/java/org/springframework/data/jdbc/repository/support/JdbcOpenSqlRowSetIterator.java

@@ -0,0 +1,61 @@
+/*
+ * Copyright 2018-2019 the original author or authors.
+ *
+ * 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
+ *
+ *      https://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.springframework.data.jdbc.repository.support;
+
+import java.sql.SQLException;
+import java.util.Iterator;
+
+import org.springframework.jdbc.InvalidResultSetAccessException;
+import org.springframework.jdbc.core.ResultSetExtractor;
+import org.springframework.jdbc.core.RowMapper;
+
+/**
+ * An iterator implementation that wraps a {@link JdbcOpenSqlRowSet} and applies
+ * an {@link RowMapper} or {@link ResultSetExtractor} to each row.
+ *
+ * @author Marcos Vinicius da Silva
+ */
+class JdbcOpenSqlRowSetIterator<T> implements Iterator<T> {
+    private final JdbcOpenSqlRowSet openCursorSqlRowSet;
+    private final RowMapper<T> rowMapper;
+    private final ResultSetExtractor<T> resultSetExtractor;
+    private int rowIndex = 1;
+
+    public JdbcOpenSqlRowSetIterator(JdbcOpenSqlRowSet openCursorSqlRowSet,
+                                     RowMapper<T> rowMapper, ResultSetExtractor<T> resultSetExtractor) {
+        this.openCursorSqlRowSet = openCursorSqlRowSet;
+        this.rowMapper = rowMapper;
+        this.resultSetExtractor = resultSetExtractor;
+    }
+
+    @Override
+    public boolean hasNext() {
+            return openCursorSqlRowSet.next();
+    }
+
+    @Override
+    public T next() {
+        try {
+            if (rowMapper != null) {
+                return rowMapper.mapRow(openCursorSqlRowSet.getResultSet(), rowIndex++);
+            } else {
+                return resultSetExtractor.extractData(openCursorSqlRowSet.getResultSet());
+            }
+        } catch (SQLException e) {
+            throw new InvalidResultSetAccessException(e);
+        }
+    }
+}

spring-data-jdbc/src/main/java/org/springframework/data/jdbc/repository/support/JdbcRepositoryQuery.java

@@ -17,6 +17,11 @@
 
 import java.lang.reflect.Constructor;
 import java.util.List;
+import java.util.Spliterator;
+import java.util.Spliterators;
+import java.util.function.Supplier;
+import java.util.stream.Stream;
+import java.util.stream.StreamSupport;
 
 import org.springframework.beans.BeanUtils;
 import org.springframework.context.ApplicationEventPublisher;
@@ -103,9 +108,16 @@ private QueryExecutor<Object> createExecutor(JdbcQueryMethod queryMethod, @Nulla
 		if (queryMethod.isModifyingQuery()) {
 			return createModifyingQueryExecutor(query);
 		}
-		if (queryMethod.isCollectionQuery() || queryMethod.isStreamQuery()) {
+
+		if (queryMethod.isStreamQuery()) {
+			QueryExecutor<Object> innerExecutor = createRowMapperOpenQueryStreamExecutor(query, rowMapper, extractor);
+			return createOpenStreamQueryExecutor(innerExecutor);
+		}
+
+		if (queryMethod.isCollectionQuery()) {
 			QueryExecutor<Object> innerExecutor = extractor != null ? createResultSetExtractorQueryExecutor(query, extractor)
 					: createListRowMapperQueryExecutor(query, rowMapper);
+
 			return createCollectionQueryExecutor(innerExecutor);
 		}
 
@@ -170,6 +182,19 @@ private QueryExecutor<Object> createModifyingQueryExecutor(String query) {
 		};
 	}
 
+	private QueryExecutor<Object> createOpenStreamQueryExecutor(QueryExecutor<Object> executor) {
+		return parameters -> {
+
+			Stream<?> result = (Stream<?>) executor.execute(parameters);
+
+			Assert.notNull(result, "A stream valued result must never be null.");
+
+			return result.peek(element -> {
+				publishAfterLoad(element);
+			});
+		};
+	}
+
 	private QueryExecutor<Object> createListRowMapperQueryExecutor(String query, RowMapper<?> rowMapper) {
 		return parameters -> operations.query(query, parameters, rowMapper);
 	}
@@ -183,6 +208,21 @@ private QueryExecutor<Object> createResultSetExtractorQueryExecutor(String query
 		return parameters -> operations.query(query, parameters, resultSetExtractor);
 	}
 
+	private QueryExecutor<Object> createRowMapperOpenQueryStreamExecutor(String query, RowMapper rowMapper, ResultSetExtractor extractor) {
+		JdbcOpenRowSetTemplate openResultSetNamedParameterJdbcTemplate =
+				new JdbcOpenRowSetTemplate(operations.getJdbcOperations());
+		return parameters -> {
+			final JdbcOpenSqlRowSet rowSet =
+					openResultSetNamedParameterJdbcTemplate.queryForOpenCursorRowSet(query, parameters);
+
+			final Spliterator<Object> spliterator = Spliterators
+					.spliteratorUnknownSize(new JdbcOpenSqlRowSetIterator<Object>(rowSet, rowMapper, extractor), Spliterator.IMMUTABLE);
+			final Supplier<Spliterator<Object>> supplier = () -> spliterator;
+			return StreamSupport.stream(supplier, Spliterator.IMMUTABLE, false)
+                    .onClose(rowSet::close);
+		};
+	}
+
 	/*
 	 * (non-Javadoc)
 	 * @see org.springframework.data.repository.query.RepositoryQuery#getQueryMethod()

spring-data-jdbc/src/test/java/org/springframework/data/jdbc/repository/query/QueryAnnotationHsqlIntegrationTests.java

@@ -15,7 +15,8 @@
  */
 package org.springframework.data.jdbc.repository.query;
 
-import static org.assertj.core.api.Assertions.*;
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatExceptionOfType;
 
 import lombok.Value;
 
@@ -35,8 +36,11 @@
 import org.springframework.context.annotation.Import;
 import org.springframework.dao.DataAccessException;
 import org.springframework.data.annotation.Id;
+import org.springframework.data.annotation.Transient;
 import org.springframework.data.jdbc.repository.config.EnableJdbcRepositories;
 import org.springframework.data.jdbc.testing.TestConfiguration;
+import org.springframework.data.relational.core.mapping.event.AfterLoadCallback;
+import org.springframework.data.relational.core.mapping.event.AfterLoadEvent;
 import org.springframework.data.repository.CrudRepository;
 import org.springframework.data.repository.query.Param;
 import org.springframework.lang.Nullable;
@@ -163,11 +167,26 @@ public void executeCustomQueryWithReturnTypeIsStream() {
 		repository.save(dummyEntity("a"));
 		repository.save(dummyEntity("b"));
 
-		Stream<DummyEntity> entities = repository.findAllWithReturnTypeIsStream();
+		try (Stream<DummyEntity> entities = repository.findAllWithReturnTypeIsStream()) {
 
-		assertThat(entities) //
-				.extracting(e -> e.name) //
-				.containsExactlyInAnyOrder("a", "b");
+			assertThat(entities) //
+					.extracting(e -> e.name) //
+					.containsExactlyInAnyOrder("a", "b");
+		}
+
+	}
+
+	@Test // DATAJDBC-356
+	public void executeCustomQueryWithReturnTypeIsStreamAndPublishEvents() {
+
+		repository.save(dummyEntity("a"));
+		repository.save(dummyEntity("b"));
+
+		try (Stream<DummyEntity> entities = repository.findAllWithReturnTypeIsStream()) {
+
+			assertThat(entities) //
+					.allMatch(dummyEntity -> dummyEntity.loaded);
+		}
 
 	}
 
@@ -283,13 +302,24 @@ static class Config {
 		Class<?> testClass() {
 			return QueryAnnotationHsqlIntegrationTests.class;
 		}
+
+		@Bean
+		AfterLoadCallback<DummyEntity> dummyEntityAfterLoadCallback() {
+			return (dummy) -> {
+				dummy.loaded = true;
+				return dummy;
+			};
+		}
 	}
 
 	private static class DummyEntity {
 
 		@Id Long id;
 
 		String name;
+
+		@Transient
+		boolean loaded = false;
 	}
 
 	private interface DummyEntityRepository extends CrudRepository<DummyEntity, Long> {