Fix several design issues related to DataSourceBasedStatementExecutor…

… class.

Guarantee that connections are closed.
Remove side-effect from constructor.
Lazily get connection before call to execute.
Suppress on-close exceptions within DataSourceBasedStatementExecutor

release_notes.txt

@@ -1,3 +1,7 @@
+New in 1.8.2
+    Remove side-effect in constructor of `DataSourceBasedStatementExecutor`. Lazily initialize connection before call to `execute` method
+    FIXED possible connection leak in `DataSourceBasedStatementExecutor` class.
+
 New in 1.8.1
     Reintroduce and deprecate the API removed in 1.8.0 release
 

src/main/java/com/opower/persistence/jpile/jdbc/DataSourceBasedStatementExecutor.java

@@ -2,67 +2,137 @@
 
 import com.google.common.base.Preconditions;
 import com.google.common.base.Throwables;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 import javax.sql.DataSource;
+import java.sql.Connection;
 import java.sql.SQLException;
 
 /**
- * This implementation operates the {@link DataSource}. It creates a
- * {@link ConnectionBasedStatementExecutor} using {@link java.sql.Connection}
- * from {@link DataSource} and delegates all calls to it.
- * <p/>
- * This implementation has no mechanism to process failures from underlying
- * {@link ConnectionBasedStatementExecutor} and retry execution with another
- * {@link java.sql.Connection}. If you need such behavior then create a subclass.
- * <p/>
- * Be aware that this implementation is not thread-safe.
+ * <p>
+ *  This implementation operates the {@link DataSource}. It creates a
+ *  {@link ConnectionBasedStatementExecutor} using {@link java.sql.Connection}
+ *  from {@link DataSource} and delegates all calls to it.
+ * </p>
+ * <p>
+ *  Be aware that this implementation is not thread-safe.
+ * </p>
  *
  * @author ivan.german
  */
 public class DataSourceBasedStatementExecutor implements StatementExecutor {
+    private static final Logger LOGGER = LoggerFactory.getLogger(DataSourceBasedStatementExecutor.class);
 
     private ConnectionBasedStatementExecutor connectionBasedStatementExecutor;
     private final DataSource dataSource;
 
-    public DataSourceBasedStatementExecutor(DataSource dataSource) throws SQLException {
+    public DataSourceBasedStatementExecutor(DataSource dataSource) {
         Preconditions.checkNotNull(dataSource, "dataSource must not be null");
 
         this.dataSource = dataSource;
-        initNewConnection();
     }
 
     /**
-     * {@inheritDoc}
+     * <p>
+     *  Processes {@link StatementCallback} in context of used connection.
+     *  If this {@link StatementExecutor} don't own any open connection then new one
+     *  is requested from the data source.
+     * </p>
+     * <p>
+     *  Be aware that foreign key constraints are disabled for used connection.
+     * </p>
+     *
+     * @throws RuntimeException if any problem occur during execution
+     * @return result of {@link StatementCallback}
      */
     @Override
     public <T> T execute(StatementCallback<T> statementCallback) {
-        Preconditions.checkNotNull(
-                this.connectionBasedStatementExecutor, "underlying connectionBasedStatementExecutor is null");
+        try {
+            if (this.connectionBasedStatementExecutor == null) {
+                initNewConnection();
+            }
+
+            return this.connectionBasedStatementExecutor.execute(statementCallback);
+        }
+        /* Propagated exception */
+        catch (Exception e) {
+            try {
+                shutdown();
+            }
+            catch (Exception shutdownException) {
+                e.addSuppressed(shutdownException);
+            }
 
-        return this.connectionBasedStatementExecutor.execute(statementCallback);
+            throw Throwables.propagate(e);
+        }
     }
 
     /**
-     * {@inheritDoc}
+     * Re-enables foreign key constraints for owned connection and returns it
+     * back to connection pool.
+     *
+     * @throws RuntimeException if any error occur during execution
      */
     @Override
     public void shutdown() {
         if (this.connectionBasedStatementExecutor != null) {
-            this.connectionBasedStatementExecutor.shutdown();
-            this.connectionBasedStatementExecutor.closeConnection();
-            this.connectionBasedStatementExecutor = null;
+            try {
+                this.connectionBasedStatementExecutor.shutdown();
+            }
+            /* Propagated exception should be logged and rethrown */
+            catch (Exception e) {
+                LOGGER.warn("Fail to re-enable foreign key constraints", e);
+                throw Throwables.propagate(e);
+            }
+            finally {
+                this.connectionBasedStatementExecutor.closeConnection();
+                this.connectionBasedStatementExecutor = null;
+            }
         }
     }
 
     /**
-     * Ensure that previous connection is closed and open new one.
+     * <p>
+     *   Ensure that previous connection is closed and open new one.
+     * </p>
+     *
+     * @throws RuntimeException if any problem occur
      */
     public void initNewConnection() {
+        Connection connection = null;
+
         try {
-            shutdown();
-            this.connectionBasedStatementExecutor = new ConnectionBasedStatementExecutor(dataSource.getConnection());
+            try {
+                shutdown();
+            }
+            catch (RuntimeException ignoredException) {
+                /*
+                 * Exception from shutdown is already printed. Connection is returned to pool.
+                 * So basically nothing restricts us from opening new connection.
+                 */
+            }
+
+            connection = this.dataSource.getConnection();
+            this.connectionBasedStatementExecutor = new ConnectionBasedStatementExecutor(connection);
         }
-        catch (SQLException e) {
+        /**
+         * Exceptions propagated from {@link ConnectionBasedStatementExecutor}
+         * and exception from {@link DataSource#getConnection()}
+         */
+        catch (Exception e) {
+            LOGGER.error("Fail to init new connection", e);
+
+            try {
+                if (connection != null) {
+                    connection.close();
+                }
+            }
+            catch (SQLException closeException) {
+                LOGGER.error("Fail to close connection", e);
+                e.addSuppressed(closeException);
+            }
+
             throw Throwables.propagate(e);
         }
     }

src/test/java/com/opower/persistence/jpile/jdbc/IntTestDataSourceBasedStatementExecutor.java

@@ -1,5 +1,6 @@
 package com.opower.persistence.jpile.jdbc;
 
+import org.junit.After;
 import org.junit.Test;
 
 import javax.sql.DataSource;
@@ -10,6 +11,7 @@
 import static com.opower.persistence.jpile.util.JdbcTestUtil.openNewConnection;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.fail;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
@@ -20,34 +22,58 @@
  */
 public class IntTestDataSourceBasedStatementExecutor {
 
+    private DataSourceBasedStatementExecutor executor;
+
+    @After
+    public void tearDown() {
+        if (this.executor != null) {
+            this.executor.shutdown();
+        }
+    }
+
     /**
-     * Create {@link DataSourceBasedStatementExecutor}, execute query on it and shut it down.
+     * Execute query using {@link DataSourceBasedStatementExecutor}.
      */
     @Test
     public void testExecute() throws SQLException {
-        final String query = "SELECT 1";
+        String query = "SELECT 1";
 
         Connection connection = openNewConnection();
 
         DataSource dataSource = mock(DataSource.class);
         when(dataSource.getConnection()).thenReturn(connection);
 
-        DataSourceBasedStatementExecutor executor = new DataSourceBasedStatementExecutor(dataSource);
+        this.executor = new DataSourceBasedStatementExecutor(dataSource);
 
-        boolean result = executor.execute(new StatementCallback<Boolean>() {
-            @Override
-            public Boolean doInStatement(Statement statement) throws SQLException {
-                return statement.execute(query);
-            }
-        });
+        boolean result = this.executor.execute(statementCallbackForQuery(query));
 
         assertTrue(result);
+    }
+
+    /**
+     * Execute query and close {@link DataSourceBasedStatementExecutor}.
+     * Check that connection is closed.
+     */
+    @Test
+    public void testShutdownExecutor() throws SQLException {
+        String query = "SELECT 1";
+
+        Connection connection = openNewConnection();
+
+        DataSource dataSource = mock(DataSource.class);
+        when(dataSource.getConnection()).thenReturn(connection);
 
-        executor.shutdown();
+        this.executor = new DataSourceBasedStatementExecutor(dataSource);
+
+        this.executor.execute(statementCallbackForQuery(query));
+        this.executor.shutdown();
+
+        assertTrue(connection.isClosed());
     }
 
     /**
-     * Create {@link DataSourceBasedStatementExecutor}, shut it down and init new connection.
+     * Shut {@link DataSourceBasedStatementExecutor} down and init new connection.
+     * Check that first connection is closed and second one is opened.
      */
     @Test
     public void testShutdownAndInitialize() throws SQLException {
@@ -57,15 +83,46 @@ public void testShutdownAndInitialize() throws SQLException {
         DataSource dataSource = mock(DataSource.class);
         when(dataSource.getConnection()).thenReturn(connection1, connection2);
 
-        DataSourceBasedStatementExecutor executor = new DataSourceBasedStatementExecutor(dataSource);
-        executor.shutdown();
-        executor.initNewConnection();
+        this.executor = new DataSourceBasedStatementExecutor(dataSource);
+        this.executor.initNewConnection();
+        this.executor.shutdown();
+        this.executor.initNewConnection();
 
         assertTrue(connection1.isClosed());
         assertFalse(connection2.isClosed());
+    }
+
+    /**
+     * Execute query that fails and check that connection is closed automatically.
+     */
+    @Test
+    public void testExecuteQueryThatFails() throws Throwable {
+        Connection connection = openNewConnection();
 
-        executor.shutdown();
+        DataSource dataSource = mock(DataSource.class);
+        when(dataSource.getConnection()).thenReturn(connection);
+
+        this.executor = new DataSourceBasedStatementExecutor(dataSource);
+
+        try {
+            String query = "invalid query";
+            this.executor.execute(statementCallbackForQuery(query));
+            fail("Exception must be thrown in the previous line");
+        }
+        catch (Exception e) {
+            /* Ignore this exception. We know that syntax is invalid. */
+        }
+        finally {
+            assertTrue(connection.isClosed());
+        }
+    }
 
-        assertTrue(connection2.isClosed());
+    private StatementCallback<Boolean> statementCallbackForQuery(final String query) {
+        return new StatementCallback<Boolean>() {
+            @Override
+            public Boolean doInStatement(Statement statement) throws SQLException {
+                return statement.execute(query);
+            }
+        };
     }
 }