Introduce Stream variant methods for SqlQuery

Closes GH-34474

Signed-off-by: Yanming Zhou <[email protected]>

Author: quaff

spring-jdbc/src/main/java/org/springframework/jdbc/object/SqlQuery.java

@@ -18,13 +18,16 @@
 
 import java.util.List;
 import java.util.Map;
+import java.util.function.BiFunction;
+import java.util.stream.Stream;
 
 import javax.sql.DataSource;
 
 import org.jspecify.annotations.Nullable;
 
 import org.springframework.dao.DataAccessException;
 import org.springframework.dao.support.DataAccessUtils;
+import org.springframework.jdbc.core.PreparedStatementCreator;
 import org.springframework.jdbc.core.RowMapper;
 import org.springframework.jdbc.core.namedparam.MapSqlParameterSource;
 import org.springframework.jdbc.core.namedparam.NamedParameterUtils;
@@ -52,6 +55,7 @@
  * @author Rod Johnson
  * @author Juergen Hoeller
  * @author Thomas Risberg
+ * @author Yanming Zhou
  * @param <T> the result type
  * @see SqlUpdate
  */
@@ -94,6 +98,23 @@ public List<T> execute(Object @Nullable [] params, @Nullable Map<?, ?> context)
 		return getJdbcTemplate().query(newPreparedStatementCreator(params), rowMapper);
 	}
 
+	/**
+	 * Central stream method. All un-named parameter execution goes through this method.
+	 * @param params parameters, similar to JDO query parameters.
+	 * Primitive parameters must be represented by their Object wrapper type.
+	 * The ordering of parameters is significant.
+	 * @param context the contextual information passed to the {@code mapRow}
+	 * callback method. The JDBC operation itself doesn't rely on this parameter,
+	 * but it can be useful for creating the objects of the result list.
+	 * @return a result Stream of objects, one per row of the ResultSet. Normally all these
+	 * will be of the same class, although it is possible to use different types.
+	 */
+	public Stream<T> stream(Object @Nullable [] params, @Nullable Map<?, ?> context) throws DataAccessException {
+		validateParameters(params);
+		RowMapper<T> rowMapper = newRowMapper(params, context);
+		return getJdbcTemplate().queryForStream(newPreparedStatementCreator(params), rowMapper);
+	}
+
 	/**
 	 * Convenient method to execute without context.
 	 * @param params parameters for the query. Primitive parameters must
@@ -104,6 +125,16 @@ public List<T> execute(Object... params) throws DataAccessException {
 		return execute(params, null);
 	}
 
+	/**
+	 * Convenient method to stream without context.
+	 * @param params parameters for the query. Primitive parameters must
+	 * be represented by their Object wrapper type. The ordering of parameters is
+	 * significant.
+	 */
+	public Stream<T> stream(Object... params) throws DataAccessException {
+		return stream(params, null);
+	}
+
 	/**
 	 * Convenient method to execute without parameters.
 	 * @param context the contextual information for object creation
@@ -112,13 +143,28 @@ public List<T> execute(Map<?, ?> context) throws DataAccessException {
 		return execute((Object[]) null, context);
 	}
 
+	/**
+	 * Convenient method to stream without parameters.
+	 * @param context the contextual information for object creation
+	 */
+	public Stream<T> stream(Map<?, ?> context) throws DataAccessException {
+		return stream(null, context);
+	}
+
 	/**
 	 * Convenient method to execute without parameters nor context.
 	 */
 	public List<T> execute() throws DataAccessException {
 		return execute((Object[]) null, null);
 	}
 
+	/**
+	 * Convenient method to stream without parameters nor context.
+	 */
+	public Stream<T> stream() throws DataAccessException {
+		return stream(null, null);
+	}
+
 	/**
 	 * Convenient method to execute with a single int parameter and context.
 	 * @param p1 single int parameter
@@ -202,13 +248,23 @@ public List<T> execute(String p1) throws DataAccessException {
 	 * will be of the same class, although it is possible to use different types.
 	 */
 	public List<T> executeByNamedParam(Map<String, ?> paramMap, @Nullable Map<?, ?> context) throws DataAccessException {
-		validateNamedParameters(paramMap);
-		ParsedSql parsedSql = getParsedSql();
-		MapSqlParameterSource paramSource = new MapSqlParameterSource(paramMap);
-		String sqlToUse = NamedParameterUtils.substituteNamedParameters(parsedSql, paramSource);
-		@Nullable Object[] params = NamedParameterUtils.buildValueArray(parsedSql, paramSource, getDeclaredParameters());
-		RowMapper<T> rowMapper = newRowMapper(params, context);
-		return getJdbcTemplate().query(newPreparedStatementCreator(sqlToUse, params), rowMapper);
+		return queryByNamedParam(paramMap, context, getJdbcTemplate()::query);
+	}
+
+	/**
+	 * Central stream method. All named parameter execution goes through this method.
+	 * @param paramMap parameters associated with the name specified while declaring
+	 * the SqlParameters. Primitive parameters must be represented by their Object wrapper
+	 * type. The ordering of parameters is not significant since they are supplied in a
+	 * SqlParameterMap which is an implementation of the Map interface.
+	 * @param context the contextual information passed to the {@code mapRow}
+	 * callback method. The JDBC operation itself doesn't rely on this parameter,
+	 * but it can be useful for creating the objects of the result list.
+	 * @return a Stream of objects, one per row of the ResultSet. Normally all these
+	 * will be of the same class, although it is possible to use different types.
+	 */
+	public Stream<T> streamByNamedParam(Map<String, ?> paramMap, @Nullable Map<?, ?> context) throws DataAccessException {
+		return queryByNamedParam(paramMap, context, getJdbcTemplate()::queryForStream);
 	}
 
 	/**
@@ -221,6 +277,15 @@ public List<T> executeByNamedParam(Map<String, ? extends @Nullable Object> param
 		return executeByNamedParam(paramMap, null);
 	}
 
+	/**
+	 * Convenient method to stream without context.
+	 * @param paramMap parameters associated with the name specified while declaring
+	 * the SqlParameters. Primitive parameters must be represented by their Object wrapper
+	 * type. The ordering of parameters is not significant.
+	 */
+	public Stream<T> streamByNamedParam(Map<String, ? extends @Nullable Object> paramMap) throws DataAccessException {
+		return streamByNamedParam(paramMap, null);
+	}
 
 	/**
 	 * Generic object finder method, used by all other {@code findObject} methods.
@@ -342,4 +407,14 @@ public List<T> executeByNamedParam(Map<String, ? extends @Nullable Object> param
 	 */
 	protected abstract RowMapper<T> newRowMapper(@Nullable Object @Nullable [] parameters, @Nullable Map<?, ?> context);
 
+	private <R> R queryByNamedParam(Map<String, ?> paramMap, @Nullable Map<?, ?> context, BiFunction<PreparedStatementCreator, RowMapper<T>, R> queryFunction) {
+		validateNamedParameters(paramMap);
+		ParsedSql parsedSql = getParsedSql();
+		MapSqlParameterSource paramSource = new MapSqlParameterSource(paramMap);
+		String sqlToUse = NamedParameterUtils.substituteNamedParameters(parsedSql, paramSource);
+		@Nullable Object[] params = NamedParameterUtils.buildValueArray(parsedSql, paramSource, getDeclaredParameters());
+		RowMapper<T> rowMapper = newRowMapper(params, context);
+		return queryFunction.apply(newPreparedStatementCreator(sqlToUse, params), rowMapper);
+	}
+
 }

spring-jdbc/src/test/java/org/springframework/jdbc/object/SqlQueryTests.java

@@ -26,6 +26,7 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.stream.Stream;
 
 import javax.sql.DataSource;
 
@@ -52,6 +53,7 @@
  * @author Trevor Cook
  * @author Thomas Risberg
  * @author Juergen Hoeller
+ * @author Yanming Zhou
  */
 class SqlQueryTests {
 
@@ -125,6 +127,72 @@ protected Integer mapRow(ResultSet rs, int rownum, Object @Nullable [] params, @
 		verify(preparedStatement).close();
 	}
 
+	@Test
+	void testStreamWithoutParams() throws SQLException {
+		given(resultSet.next()).willReturn(true, false);
+		given(resultSet.getInt(1)).willReturn(1);
+
+		SqlQuery<Integer> query = new MappingSqlQueryWithParameters<>() {
+			@Override
+			protected Integer mapRow(ResultSet rs, int rownum, Object @Nullable [] params, @Nullable Map<? ,?> context)
+					throws SQLException {
+				assertThat(params).as("params were null").isNull();
+				assertThat(context).as("context was null").isNull();
+				return rs.getInt(1);
+			}
+		};
+		query.setDataSource(dataSource);
+		query.setSql(SELECT_ID);
+		query.compile();
+		try (Stream<Integer> stream = query.stream()) {
+			List<Integer> list = stream.toList();
+			assertThat(list).containsExactly(1);
+		}
+		verify(connection).prepareStatement(SELECT_ID);
+		verify(resultSet).close();
+		verify(preparedStatement).close();
+	}
+
+	@Test
+	void testStreamByNamedParam() throws SQLException {
+		given(resultSet.next()).willReturn(true, false);
+		given(resultSet.getInt("id")).willReturn(1);
+		given(resultSet.getString("forename")).willReturn("rod");
+		given(connection.prepareStatement(SELECT_ID_FORENAME_NAMED_PARAMETERS_PARSED,
+				ResultSet.TYPE_SCROLL_SENSITIVE, ResultSet.CONCUR_READ_ONLY)
+		).willReturn(preparedStatement);
+
+		SqlQuery<Customer> query = new MappingSqlQueryWithParameters<>() {
+			@Override
+			protected Customer mapRow(ResultSet rs, int rownum, Object @Nullable [] params, @Nullable Map<? ,?> context)
+					throws SQLException {
+				assertThat(params).as("params were not null").isNotNull();
+				assertThat(context).as("context was null").isNull();
+				Customer cust = new Customer();
+				cust.setId(rs.getInt(COLUMN_NAMES[0]));
+				cust.setForename(rs.getString(COLUMN_NAMES[1]));
+				return cust;
+			}
+		};
+		query.declareParameter(new SqlParameter("id", Types.NUMERIC));
+		query.declareParameter(new SqlParameter("country", Types.VARCHAR));
+		query.setDataSource(dataSource);
+		query.setSql(SELECT_ID_FORENAME_NAMED_PARAMETERS);
+		query.compile();
+		try (Stream<Customer> stream = query.streamByNamedParam(Map.of("id", 1, "country", "UK"))) {
+			List<Customer> list = stream.toList();
+			assertThat(list).hasSize(1);
+			Customer customer = list.get(0);
+			assertThat(customer.getId()).isEqualTo(1);
+			assertThat(customer.getForename()).isEqualTo("rod");
+		}
+		verify(connection).prepareStatement(SELECT_ID_FORENAME_NAMED_PARAMETERS_PARSED);
+		verify(preparedStatement).setObject(1, 1, Types.NUMERIC);
+		verify(preparedStatement).setString(2, "UK");
+		verify(resultSet).close();
+		verify(preparedStatement).close();
+	}
+
 	@Test
 	void testQueryWithoutEnoughParams() {
 		MappingSqlQuery<Integer> query = new MappingSqlQuery<>() {