🐛 Fixed connection leak in StreamingJdbcDatabase (#20888)

* Fixed connection leak in StreamingJdbcDatabase * fixed checkstyle
airbytehq · Jan 22, 2023 · 126ce36 · 126ce36
1 parent 0548d5f
commit 126ce36
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 4 deletions.
diff --git a/airbyte-db/db-lib/src/main/java/io/airbyte/db/jdbc/JdbcDatabase.java b/airbyte-db/db-lib/src/main/java/io/airbyte/db/jdbc/JdbcDatabase.java
@@ -29,6 +29,8 @@
 public abstract class JdbcDatabase extends SqlDatabase {
 
   protected final JdbcCompatibleSourceOperations<?> sourceOperations;
+  protected Exception streamException;
+  protected boolean isStreamFailed;
 
   public JdbcDatabase(final JdbcCompatibleSourceOperations<?> sourceOperations) {
     this.sourceOperations = sourceOperations;

diff --git a/airbyte-db/db-lib/src/main/java/io/airbyte/db/jdbc/StreamingJdbcDatabase.java b/airbyte-db/db-lib/src/main/java/io/airbyte/db/jdbc/StreamingJdbcDatabase.java
@@ -71,6 +71,9 @@ public <T> Stream<T> unsafeQuery(final CheckedFunction<Connection, PreparedState
             try {
               connection.setAutoCommit(true);
               connection.close();
+              if (isStreamFailed) {
+                throw new RuntimeException(streamException);
+              }
             } catch (final SQLException e) {
               throw new RuntimeException(e);
             }
@@ -84,9 +87,9 @@ public <T> Stream<T> unsafeQuery(final CheckedFunction<Connection, PreparedState
    * This method differs from {@link DefaultJdbcDatabase#toUnsafeStream} in that it takes a streaming
    * config that adjusts the fetch size dynamically according to sampled row size.
    */
-  protected static <T> Stream<T> toUnsafeStream(final ResultSet resultSet,
-                                                final CheckedFunction<ResultSet, T, SQLException> mapper,
-                                                final JdbcStreamingQueryConfig streamingConfig) {
+  protected <T> Stream<T> toUnsafeStream(final ResultSet resultSet,
+                                         final CheckedFunction<ResultSet, T, SQLException> mapper,
+                                         final JdbcStreamingQueryConfig streamingConfig) {
     return StreamSupport.stream(new Spliterators.AbstractSpliterator<>(Long.MAX_VALUE, Spliterator.ORDERED) {
 
       @Override
@@ -102,7 +105,11 @@ public boolean tryAdvance(final Consumer<? super T> action) {
           return true;
         } catch (final SQLException e) {
           LOGGER.error("SQLState: {}, Message: {}", e.getSQLState(), e.getMessage());
-          throw new RuntimeException(e);
+          streamException = e;
+          isStreamFailed = true;
+          // throwing an exception in tryAdvance() method lead to the endless loop in Spliterator and stream
+          // will never close
+          return false;
         }
       }