Improving the usage of streams on SQL fetches.

Author: staudtMarius

CHANGELOG.md

@@ -50,6 +50,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Removing `uuid` as required column from input and result time series [#826](https://github.com/ie3-institute/PowerSystemDataModel/issues/826)
 - Removing the support for the old csv format that was marked `deprecated` back in version `1.1.0` [#795](https://github.com/ie3-institute/PowerSystemDataModel/issues/795)
 - BREAKING: Updating PowerSystemUtils dependency to 2.2 [#1006](https://github.com/ie3-institute/PowerSystemDataModel/issues/1006)
+- Improving usage of streams on sql fetches [#827](https://github.com/ie3-institute/PowerSystemDataModel/issues/827)
 
 ## [4.1.0] - 2023-11-02
 

src/main/java/edu/ie3/datamodel/io/connectors/SqlConnector.java

@@ -9,6 +9,8 @@
 import edu.ie3.util.TimeUtil;
 import java.sql.*;
 import java.util.*;
+import java.util.stream.Stream;
+import java.util.stream.StreamSupport;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -66,7 +68,9 @@ public ResultSet executeQuery(Statement stmt, String query) throws SQLException
    */
   public int executeUpdate(String updateQuery) {
     try (Statement stmt = getConnection().createStatement()) {
-      return stmt.executeUpdate(updateQuery);
+      int res = stmt.executeUpdate(updateQuery);
+      getConnection().commit();
+      return res;
     } catch (SQLException e) {
       log.error(String.format("Error at execution of query \"%1.127s\": ", updateQuery), e);
       return -1;
@@ -85,7 +89,8 @@ public Connection getConnection() throws SQLException {
   }
 
   /**
-   * Establishes and returns a database connection
+   * Establishes and returns a database connection. The {@link Connection#getAutoCommit()} is set to
+   * {@code false}.
    *
    * @param reuseConnection should the connection be used again, if it is still valid? If not, a new
    *     connection will be established
@@ -98,6 +103,7 @@ public Connection getConnection(boolean reuseConnection) throws SQLException {
         if (connection != null) connection.close();
 
         connection = DriverManager.getConnection(jdbcUrl, connectionProps);
+        connection.setAutoCommit(false);
       } catch (SQLException e) {
         throw new SQLException("Could not establish connection: ", e);
       }
@@ -115,21 +121,47 @@ public void shutdown() {
   }
 
   /**
-   * Extracts all field to value maps from the ResultSet, one for each row
+   * Method to execute a {@link PreparedStatement} and return its result as a stream.
    *
-   * @param rs the ResultSet to use
-   * @return a list of field maps
+   * @param ps to execute
+   * @param fetchSize used for {@link PreparedStatement#setFetchSize(int)}
+   * @return a stream of maps
+   * @throws SQLException if an exception occurred while executing the query
    */
-  public List<Map<String, String>> extractFieldMaps(ResultSet rs) {
-    List<Map<String, String>> fieldMaps = new ArrayList<>();
-    try {
-      while (rs.next()) {
-        fieldMaps.add(extractFieldMap(rs));
+  public Stream<Map<String, String>> toStream(PreparedStatement ps, int fetchSize)
+      throws SQLException {
+    ps.setFetchSize(fetchSize);
+    Iterator<Map<String, String>> sqlIterator = getSqlIterator(ps.executeQuery());
+
+    return StreamSupport.stream(
+        Spliterators.spliteratorUnknownSize(
+            sqlIterator, Spliterator.NONNULL | Spliterator.IMMUTABLE),
+        true);
+  }
+
+  /**
+   * Returns an {@link Iterator} for the given {@link ResultSet}.
+   *
+   * @param rs given result set
+   * @return an iterator
+   */
+  public Iterator<Map<String, String>> getSqlIterator(ResultSet rs) {
+    return new Iterator<>() {
+      @Override
+      public boolean hasNext() {
+        try {
+          return rs.next();
+        } catch (SQLException e) {
+          log.error("Exception at extracting next ResultSet: ", e);
+          throw new RuntimeException(e);
+        }
       }
-    } catch (SQLException e) {
-      log.error("Exception at extracting ResultSet: ", e);
-    }
-    return fieldMaps;
+
+      @Override
+      public Map<String, String> next() {
+        return extractFieldMap(rs);
+      }
+    };
   }
 
   /**

src/main/java/edu/ie3/datamodel/io/source/sql/SqlDataSource.java

@@ -125,7 +125,6 @@ public Optional<Set<String>> getSourceFields(String tableName) {
       ResultSet rs =
           connector.getConnection().getMetaData().getColumns(null, null, tableName, null);
       Set<String> columnNames = new HashSet<>();
-
       while (rs.next()) {
         String name = rs.getString("COLUMN_NAME");
         columnNames.add(StringUtils.snakeCaseToCamelCase(name));
@@ -177,11 +176,15 @@ protected Stream<Map<String, String>> buildStreamByTableName(String tableName) {
    * table name.
    */
   protected Stream<Map<String, String>> executeQuery(String query, AddParams addParams) {
-    try (PreparedStatement ps = connector.getConnection().prepareStatement(query)) {
+    try {
+      PreparedStatement ps = connector.getConnection().prepareStatement(query);
       addParams.addParams(ps);
 
-      ResultSet resultSet = ps.executeQuery();
-      return connector.extractFieldMaps(resultSet).stream();
+      Stream<Map<String, String>> stream = connector.toStream(ps, 1000);
+
+      // don't work with `try with resource`, but is necessary to get results
+      ps.closeOnCompletion();
+      return stream;
     } catch (SQLException e) {
       log.error("Error during execution of query {}", query, e);
     }

src/test/groovy/edu/ie3/datamodel/io/connectors/SqlConnectorIT.groovy

@@ -121,16 +121,15 @@ class SqlConnectorIT extends Specification implements TestContainerHelper {
   def "A SQL connector is able to extract all field to value maps from result set"() {
     given:
     def preparedStatement = connector.getConnection(false).prepareStatement("SELECT * FROM public.test;")
-    def resultSet = preparedStatement.executeQuery()
 
     when:
-    def actual = connector.extractFieldMaps(resultSet)
+    def actual = connector.toStream(preparedStatement, 1).toList()
 
     then:
     actual.size() == 2
 
     cleanup:
-    resultSet.close()
+    preparedStatement.close()
   }
 
   def "A SQL connector shuts down correctly, if no connection was opened"() {