From 2716c3a9e98ae4b4d3eea681bfafd39858ade832 Mon Sep 17 00:00:00 2001
From: Lukas Eder <lukas.eder@gmail.com>
Date: Mon, 5 May 2014 16:15:18 +0200
Subject: [PATCH] [#3191] ResultQuery.keepStatement(true) doesn't work
 correctly with DataSourceConnectionProvider

---
 .../jooq/test/_/testcases/StatementTests.java | 78 ++++++++++++++++---
 .../java/org/jooq/impl/AbstractQuery.java     |  8 +-
 .../org/jooq/impl/DefaultExecuteContext.java  | 22 ++++--
 3 files changed, 90 insertions(+), 18 deletions(-)

diff --git a/jOOQ-test/src/test/java/org/jooq/test/_/testcases/StatementTests.java b/jOOQ-test/src/test/java/org/jooq/test/_/testcases/StatementTests.java
index ec5edf145a..189162ebc4 100644
--- a/jOOQ-test/src/test/java/org/jooq/test/_/testcases/StatementTests.java
+++ b/jOOQ-test/src/test/java/org/jooq/test/_/testcases/StatementTests.java
@@ -65,6 +65,7 @@ import java.util.logging.Logger;
 import javax.sql.DataSource;
 
 import org.jooq.Configuration;
+import org.jooq.ConnectionProvider;
 import org.jooq.Cursor;
 import org.jooq.DSLContext;
 import org.jooq.ExecuteContext;
@@ -179,10 +180,9 @@ extends BaseTest<A, AP, B, S, B2S, BS, L, X, DATE, BOOL, D, T, U, UU, I, IPK, T7
 
         // "Normal" DataSource behaviour, without keeping an open statement
         CountingDataSource ds1 = new CountingDataSource(getConnection());
+        ConnectionProviderProxy p1 = new ConnectionProviderProxy(new DataSourceConnectionProvider(ds1));
 
-        Configuration c1 =
-            create().configuration().derive(new DataSourceConnectionProvider(ds1));
-
+        Configuration c1 = create().configuration().derive(p1);
         for (int i = 0; i < 5; i++) {
             ResultQuery<Record1<Integer>> query =
             DSL.using(c1)
@@ -190,13 +190,33 @@ extends BaseTest<A, AP, B, S, B2S, BS, L, X, DATE, BOOL, D, T, U, UU, I, IPK, T7
                .from(TBook())
                .where(TBook_ID().eq(param("p", 0)));
 
+            assertEquals(i * 2, p1.acquire);
+            assertEquals(i * 2, p1.release);
+            assertEquals(i * 2, ds1.open);
+            assertEquals(i * 2, ds1.close);
+
             query.bind("p", 1);
             assertEquals(1, (int) query.fetchOne().getValue(TBook_ID()));
+            assertEquals(i * 2 + 1, p1.acquire);
+            assertEquals(i * 2 + 1, p1.release);
+            assertEquals(i * 2 + 1, ds1.open);
+            assertEquals(i * 2 + 1, ds1.close);
+
             query.bind("p", 2);
             assertEquals(2, (int) query.fetchOne().getValue(TBook_ID()));
 
-            assertEquals((i + 1) * 2, ds1.open);
-            assertEquals((i + 1) * 2, ds1.close);
+            assertEquals(i * 2 + 2, p1.acquire);
+            assertEquals(i * 2 + 2, p1.release);
+            assertEquals(i * 2 + 2, ds1.open);
+            assertEquals(i * 2 + 2, ds1.close);
+
+            // Has no effect
+            query.close();
+
+            assertEquals(i * 2 + 2, p1.acquire);
+            assertEquals(i * 2 + 2, p1.release);
+            assertEquals(i * 2 + 2, ds1.open);
+            assertEquals(i * 2 + 2, ds1.close);
         }
 
         assertEquals(10, ds1.open);
@@ -204,10 +224,9 @@ extends BaseTest<A, AP, B, S, B2S, BS, L, X, DATE, BOOL, D, T, U, UU, I, IPK, T7
 
         // Keeping an open statement [#3191]
         CountingDataSource ds2 = new CountingDataSource(getConnection());
+        ConnectionProviderProxy p2 = new ConnectionProviderProxy(new DataSourceConnectionProvider(ds2));
 
-        Configuration c2 =
-            create().configuration().derive(new DataSourceConnectionProvider(ds2));
-
+        Configuration c2 = create().configuration().derive(p2);
         for (int i = 0; i < 5; i++) {
             ResultQuery<Record1<Integer>> query =
             DSL.using(c2)
@@ -216,22 +235,61 @@ extends BaseTest<A, AP, B, S, B2S, BS, L, X, DATE, BOOL, D, T, U, UU, I, IPK, T7
                .where(TBook_ID().eq(param("p", 0)))
                .keepStatement(true);
 
+            assertEquals(i, p2.acquire);
+            assertEquals(i, p2.release);
+            assertEquals(i, ds2.open);
+            assertEquals(i, ds2.close);
+
             query.bind("p", 1);
             assertEquals(1, (int) query.fetchOne().getValue(TBook_ID()));
+
+            assertEquals(i + 1, p2.acquire);
+            assertEquals(i    , p2.release);
+            assertEquals(i + 1, ds2.open);
+            assertEquals(i    , ds2.close);
+
             query.bind("p", 2);
             assertEquals(2, (int) query.fetchOne().getValue(TBook_ID()));
 
+            assertEquals(i + 1, p2.acquire);
+            assertEquals(i    , p2.release);
             assertEquals(i + 1, ds2.open);
             assertEquals(i    , ds2.close);
 
             query.close();
 
+            assertEquals(i + 1, p2.acquire);
+            assertEquals(i + 1, p2.release);
             assertEquals(i + 1, ds2.open);
             assertEquals(i + 1, ds2.close);
         }
 
-        assertEquals(5, ds1.open);
-        assertEquals(5, ds1.close);
+        assertEquals(5, ds2.open);
+        assertEquals(5, ds2.close);
+    }
+
+    private static class ConnectionProviderProxy implements ConnectionProvider {
+
+        final ConnectionProvider delegate;
+        int                      acquire;
+        int                      release;
+
+        ConnectionProviderProxy(ConnectionProvider delegate) {
+            this.delegate = delegate;
+        }
+
+        @Override
+        public Connection acquire() throws DataAccessException {
+            acquire++;
+            return delegate.acquire();
+        }
+
+        @Override
+        public void release(Connection connection) throws DataAccessException {
+            release++;
+            delegate.release(connection);
+        }
+
     }
 
     private static class CountingDataSource implements DataSource {
diff --git a/jOOQ/src/main/java/org/jooq/impl/AbstractQuery.java b/jOOQ/src/main/java/org/jooq/impl/AbstractQuery.java
index 169f400fa6..ddec3213a6 100644
--- a/jOOQ/src/main/java/org/jooq/impl/AbstractQuery.java
+++ b/jOOQ/src/main/java/org/jooq/impl/AbstractQuery.java
@@ -271,11 +271,17 @@ abstract class AbstractQuery extends AbstractQueryPart implements Query, Attacha
             //         This may be used to provide jOOQ with a JDBC connection,
             //         in case this Query / Configuration was previously
             //         deserialised
-            ExecuteContext ctx = new DefaultExecuteContext(c, this);
+            DefaultExecuteContext ctx = new DefaultExecuteContext(c, this);
             ExecuteListener listener = new ExecuteListeners(ctx);
 
             int result = 0;
             try {
+
+                // [#3191] Pre-initialise the ExecuteContext with a previous connection, if available.
+                if (statement != null) {
+                    ctx.connection(c.connectionProvider(), statement.getConnection());
+                }
+
                 if (ctx.connection() == null) {
                     throw new DetachedException("Cannot execute query. No Connection configured");
                 }
diff --git a/jOOQ/src/main/java/org/jooq/impl/DefaultExecuteContext.java b/jOOQ/src/main/java/org/jooq/impl/DefaultExecuteContext.java
index d1e88ba160..b641ce648a 100644
--- a/jOOQ/src/main/java/org/jooq/impl/DefaultExecuteContext.java
+++ b/jOOQ/src/main/java/org/jooq/impl/DefaultExecuteContext.java
@@ -423,19 +423,27 @@ class DefaultExecuteContext implements ExecuteContext {
         // Settings.getStatementType() correctly.
 
         ConnectionProvider provider = connectionProvider != null ? connectionProvider : configuration.connectionProvider();
-
         if (connection == null && provider != null) {
-            Connection c = provider.acquire();
-
-            if (c != null) {
-                LOCAL_CONNECTION.set(c);
-                connection = new SettingsEnabledConnection(new ProviderEnabledConnection(provider, c), configuration.settings());
-            }
+            connection(provider, provider.acquire());
         }
 
         return connection;
     }
 
+    /**
+     * Initialise this {@link DefaultExecuteContext} with a pre-existing
+     * {@link Connection}.
+     * <p>
+     * [#3191] This is needed, e.g. when using
+     * {@link Query#keepStatement(boolean)}.
+     */
+    final void connection(ConnectionProvider provider, Connection c) {
+        if (c != null) {
+            LOCAL_CONNECTION.set(c);
+            connection = new SettingsEnabledConnection(new ProviderEnabledConnection(provider, c), configuration.settings());
+        }
+    }
+
     @Override
     public final void record(Record r) {
         this.record = r;